The code below is a Livebook.
We use the Nx library with the EXLA backend to speed up the computations.
We also propose to run the equivalent code in Zig in this Livebook if you want extra speed. This happens thanks to the
Input: image dimensions (eg w x h of 1500 x 1000), max iteration (eg 100)
Iterate over each pixel (i,j):
- map it into the 2D plane: compute its "complex coordinates"
- compute the iteration number
- compute a colour
- Sum-up and draw from the final tensor with
Kino.
Pixel to complex plan mapping
This module transforms a couple (i,j) into a complex number.
› Notice that once you are in a numerical function, the arguments becomes "tensors", and a tensor can be of type complex "c64". It natively understands complex numbers.
defmodule Pixel do
import Nx.Defn
defn map(index, {h,w}, {top_left_x, top_left_y, bottom_right_x,bottom_right_y}) do
scale_x = Nx.divide(bottom_right_x-top_left_x, w-1)
scale_y = Nx.divide(bottom_right_y-top_left_y, h-1)
# building a complex typed tensor
Ncx.new(
top_left_x + Nx.dot(index[1],scale_x),
top_left_y + Nx.dot(index[0], scale_y)
)
end
end
Orbit and iteration number
This module computes the iteration number for a given input c.
If |c|>2, then this point is unstable. Otherwise, we have to compute for each point whether it stays bounded or not.
If it is bounded, we get max_iter, otherwise a lower value.
It is also using numerical functions via defn.
We cannot use the recursion form we did earlier because numerical functions don't accept several headers as plain Elixir. Instead we run a specil while loop. Note how we use the Nx versions of cond, and also the double condition managed byNx.logical_and, and also the Nx version of cond. Also, true is 1.
defmodule Orbit do
import Nx.Defn
defn poly(z,c) do
z*z + c
end
defn number(c,max_iter) do
condition = (Nx.real(c) +1) ** 2 + (Nx.imag(c)**2)
cond do
# points in first cardioid are all stable. Save on iterations
Nx.less(condition, 0.0625) ->
max_iter
# these points are unbounded whenever the norm is > 2
Nx.greater(Ncx.sq_norm(c), 4) ->
0
# we have to evaluate each point as it can be or not bounded in the disk 2
1 ->
{_, _, j} =
while {z=c, c, i=max_iter}, Nx.logical_and(Nx.greater(i,1), Nx.less(Ncx.sq_norm(z), 4)) do
{poly(z,c), c,i-1}
end
max_iter - j
end
end
end
Examples:
st = Ncx.new(0.2, 0.2)
dv1 = Ncx.new(0.4, 0.4)
dv2 = Ncx.new(0.3, 0.6)
dv3 = Ncx.new(2,2)
iter_max = 100
iter_dv1 = Orbit.number(dv1, iter_max) #<- we should find 8 iterations before z_n escapes from the disk 2
iter_dv2 = Orbit.number(dv2, iter_max) #<- we should find 14 iterations before z_n escapes from the disk 2
iter_dv3 = Orbit.number(dv3, iter_max)
iter_st = Orbit.number(st, iter_max) #<- this point is stable and the loop reaches n interations.
%{
"unstable/2: #{Nx.to_number(dv2)}" => iter_dv2 |> Nx.to_number(),
"unstable/1: #{Nx.to_number(dv1)}" => iter_dv1 |> Nx.to_number(),
"out_of_disk2: #{Nx.to_number(dv3)}" => iter_dv3 |> Nx.to_number(),
"stable: #{Nx.to_number(st)}" => iter_st |> Nx.to_number(),
}
%{
"out_of_disk2: 2.0+2.0i" => 0,
"stable: 0.20000000298023224+0.20000000298023224i" => 99,
"unstable/1: 0.4000000059604645+0.4000000059604645i" => 8,
"unstable/2: 0.30000001192092896+0.6000000238418579i" => 14
}
A Colour palette
Each iteration number is an integer n. We want to associate a colour [r(n),g(n),b(n)].
This will help us to visualise which point of the complex plane is stable, and if not how fast it escapes.
The choice below is just an example. Other choices can be made.
defmodule Colour do
import Nx.Defn
defn normalize(n, max_iter) do
n / max_iter
end
defn rgb(n) do
cond do
Nx.equal(n, 0) ->
Nx.stack([255, 255, 0]) |> Nx.as_type(:u8)
Nx.less(n, 0.5) ->
scaled = n * 2
r = 255 * (1 - scaled)
g = 255 * (1 - scaled/2)
b = 127 * scaled
Nx.stack([r, g, b]) |> Nx.as_type(:u8)
true ->
scaled = (n - 0.5) * 2;
r = 255*(1+scaled/2)
g = 128 * (1+scaled/2)
b = 255 * (1 - scaled)
Nx.stack([r, g, b]) |> Nx.as_type(:u8)
end
end
end
Computing the Mandelbrot set
We will know reassemble our modules.
Firstly, an example.
dim = {500,500}; iter_max = 100
p = Nx.tensor([30,40])
c_i_j = Pixel.map(p,dim, defining_points)
n_i_j = Orbit.number(c_i_j, iter_max)
nm_i_j = Colour.normalize(n_i_j, iter_max)
{Nx.to_number(n_i_j), Colour.rgb(nm_i_j)} |> dbg()
p = Nx.tensor([40,70])
c_i_j = Pixel.map(p,dim, defining_points)
n_i_j = Orbit.number(c_i_j, iter_max)
nm_i_j = Colour.normalize(n_i_j, iter_max)
{Nx.to_number(n_i_j), Colour.rgb(nm_i_j)} |> dbg()
We found that this pixel reached a point in the complex plan that escapes rather quickly from the disk 2. It get stamped with some colour.
{4,
#Nx.Tensor<
u8[3]
EXLA.Backend<host:0, 0.3807096825.1655832596.50891>
[234, 244, 10]
>}
The final module
We then reassemble the tensor into the desired format for Kino to consume it and display.
Note that you want to pass arguments into a
defnfunction that you don't want to be treated as tensors, you need to use a keyword list or a map.
defmodule Mandelbrot do
import Nx.Defn
defn compute(opts) do
top_left_x = -2; top_left_y = 1.2; bottom_right_x = 0.6; bottom_right_y = - 1.2;
defining_points = {top_left_x, top_left_y, bottom_right_x, bottom_right_y}
h = opts[:h]
w = opts[:w]
max_iter = opts[:max_iter]
# build the tensor [[0,0],, ...[0,m], [1,1]...[n,m]]. Thks to PValente
iota_rows = Nx.iota({h}, type: :u16) |> Nx.vectorize(:rows)
iota_cols = Nx.iota({w}, type: :u16) |> Nx.vectorize(:cols)
cross_product = Nx.stack([iota_rows, iota_cols])
Pixel.map(cross_product,{h,w}, defining_points)
|> Orbit.number(max_iter)
|> Colour.normalize(max_iter)
|> Colour.rgb()
|> Nx.devectorize()
|> Nx.reshape({h, w, 3})
|> Nx.as_type(:u8)
end
end
Depending on your machine, the computation below can be lengthy.
If you want to simply evaluate, set h = w = 400.
h = w = 400;
Mandelbrot.compute(h: h, w: w, max_iter: 100)
|> Kino.Image.new()
Parallelise it with async_stream
When the resolution of the image increases, it is interesting to parallelise the computations.
We divide the image in horizontal bands, as many as the number of CPU cores on the machine.
When you use async_stream, the BEAM - the VM that runs this code - parallelises the running code on the cores.
This is worth only if the size of the image is large enough as this comes with non negligible overhead.
We also set ordered: true as we need to sum-up the results in an ordered manner.
Another possible optimisation is to remark that the image is symmetric. You can compute half of the image (redefine
hto beh-rem(h, cpus*2)but you would need to be able to reverse a tensor.
defmodule StreamMandelbrot do
import Nx.Defn
@doc"""
Example: 42 rows, 8 cpus
42 rows = 8cpus * 5 + 2
We run 8 threads consuming 5 rows each
We just ignore the last 2 rows.
"""
def run(%{h: h, w: w} = opts) do
cpus = :erlang.system_info(:logical_processors_available)
# we eliminate a few rows from the final image, 8 at most.
h = h - rem(h,cpus)
rows_per_cpu = div(h, cpus)
Task.async_stream(0..cpus-1, fn cpu_count ->
# we shift the start index by the number of rows already consummed
iota_rows = Nx.iota({rows_per_cpu}, type: :u16) |> Nx.add(cpu_count * rows_per_cpu)|> Nx.vectorize(:rows)
# full width
iota_cols = Nx.iota({w}, type: :u16) |> Nx.vectorize(:cols)
cross_product = Nx.stack([iota_rows, iota_cols])
Nx.Defn.jit_apply(fn t ->
compute(t, opts) end, [cross_product])
end,
timeout: :infinity, ordered: true
)
|> Enum.map(fn {:ok, t} -> t end) #&elem(&1, 1)
|> Nx.stack()
|> Nx.reshape({h,w,3})
end
defn compute(cross_product, %{h: h, w: w, max_iter: max_iter}) do
top_left_x = -2; top_left_y = 1.2; bottom_right_x = 0.6; bottom_right_y = -1.2;
defining_points = {top_left_x, top_left_y, bottom_right_x, bottom_right_y}
Pixel.map(cross_product,{h,w}, defining_points)
|> Orbit.number(max_iter)
|> Colour.normalize(max_iter)
|> Colour.rgb()
|> Nx.devectorize()
|> Nx.as_type(:u8)
end
end
When we run the code, we have much faster results. On my machine, it took 44s to draw a 1M pixels image. We get the expected performance boost.
h= w = 400;
StreamMandelbrot.run( %{h: h, w: w, max_iter: 200})
|> Kino.Image.new()
Run embedded Zig code
If we still need or want extra speed, we can also embed Zig code in Elixir within a Livebook.
Zigler offers a remarkable documentation.
❗ You may to have Zig installed on your machine.
In the Livebook, we add the dependencies (in the first cell):
Mix.install([{:zigler, "~> 0.13.3"},{:zig_get, "~> 0.13.1"},])
With the Zigler, we can even inline Zig code.
The code below runs the same algorithm and runs OS threads for concurrency.
we use the
beammemory allocator from the library.
the slice is returned as a binary - typed as
beam.term- to be easily consumed byNxand thenKino.
defmodule Zigit do
use Zig, otp_app: :zigler,
nifs: [..., generate_mandelbrot: [:threaded]]
# release_mode: :fast
~Z"""
const beam = @import("beam");
const std = @import("std");
const Cx = std.math.Complex(f64);
const topLeft = Cx{ .re = -2.1, .im = 1.2 };
const bottomRight = Cx{ .re = 0.6, .im = -1.2 };
const w = bottomRight.re - topLeft.re;
const h = bottomRight.im - topLeft.im;
const Context = struct {res_x: usize, res_y: usize, imax: usize};
/// nif: generate_mandelbrot/3 Threaded
pub fn generate_mandelbrot(res_x: usize, res_y: usize, max_iter: usize) !beam.term {
const pixels = try beam.allocator.alloc(u8, res_x * res_y * 3);
defer beam.allocator.free(pixels);
const resolution = Context{ .res_x = res_x, .res_y = res_y, .imax = max_iter };
const res = try createBands(pixels, resolution);
return beam.make(res, .{ .as = .binary });
}
// <--- threaded version
fn createBands(pixels: []u8, ctx: Context) ![]u8 {
const cpus = try std.Thread.getCpuCount();
var threads = try beam.allocator.alloc(std.Thread, cpus);
defer beam.allocator.free(threads);
// half of the total rows
const rows_to_process = ctx.res_y / 2 + ctx.res_y % 2;
// one band is one count of cpus
// const nb_rows_per_band = rows_to_process / cpus + rows_to_process % cpus;
const rows_per_band = (rows_to_process + cpus - 1) / cpus;
for (0..cpus) |cpu_count| {
const start_row = cpu_count * rows_per_band;
// Stop if there are no rows to process
if (start_row >= rows_to_process) break;
const end_row = @min(start_row + rows_per_band, rows_to_process);
const args = .{ ctx, pixels, start_row, end_row };
threads[cpu_count] = try std.Thread.spawn(.{}, processRows, args);
}
for (threads[0..cpus]) |thread| {
thread.join();
}
return pixels;
}
fn processRows(ctx: Context, pixels: []u8, start_row: usize, end_row: usize) void {
for (start_row..end_row) |current_row| {
processRow(ctx, pixels, current_row);
}
}
fn processRow(ctx: Context, pixels: []u8, row_id: usize) void {
// Calculate the symmetric row
const sym_row_id = ctx.res_y - 1 - row_id;
if (row_id <= sym_row_id) {
// loop over columns
for (0..ctx.res_x) |col_id| {
const c = mapPixel(.{ @as(usize, @intCast(row_id)), @as(usize, @intCast(col_id)) }, ctx);
const iter = iterationNumber(c, ctx.imax);
const colour = createRgb(iter, ctx.imax);
const p_idx = (row_id * ctx.res_x + col_id) * 3;
pixels[p_idx + 0] = colour[0];
pixels[p_idx + 1] = colour[1];
pixels[p_idx + 2] = colour[2];
// Process the symmetric row (if it's different from current row)
if (row_id != sym_row_id) {
const sym_p_idx = (sym_row_id * ctx.res_x + col_id) * 3;
pixels[sym_p_idx + 0] = colour[0];
pixels[sym_p_idx + 1] = colour[1];
pixels[sym_p_idx + 2] = colour[2];
}
}
}
}
fn mapPixel(pixel: [2]usize, ctx: Context) Cx {
const px_width = ctx.res_x - 1;
const px_height = ctx.res_y - 1;
const scale_x = w / @as(f64, @floatFromInt(px_width));
const scale_y = h / @as(f64, @floatFromInt(px_height));
const re = topLeft.re + scale_x * @as(f64, @floatFromInt(pixel[1]));
const im = topLeft.im + scale_y * @as(f64, @floatFromInt(pixel[0]));
return Cx{ .re = re, .im = im };
}
fn iterationNumber(c: Cx, imax: usize) ?usize {
if (c.re > 0.6 or c.re < -2.1) return 0;
if (c.im > 1.2 or c.im < -1.2) return 0;
// first cardiod
if ((c.re + 1) * (c.re + 1) + c.im * c.im < 0.0625) return null;
var z = Cx{ .re = 0.0, .im = 0.0 };
for (0..imax) |j| {
if (sqnorm(z) > 4) return j;
z = Cx.mul(z, z).add(c);
}
return null;
}
fn sqnorm(z: Cx) f64 {
return z.re * z.re + z.im * z.im;
}
fn createRgb(iter: ?usize, imax: usize) [3]u8 {
// If it didn't escape, return black
if (iter == null) return [_]u8{ 0, 0, 0 };
// Normalize time to [0,1[ now that we know it isn't "null"
const normalized = @as(f64, @floatFromInt(iter.?)) / @as(f64, @floatFromInt(imax));
if (normalized < 0.5) {
const scaled = normalized * 2;
return [_]u8{ @as(u8, @intFromFloat(255 * (1 - scaled))), @as(u8, @intFromFloat(255.0 * (1 - scaled / 2))), @as(u8, @intFromFloat(127 + 128 * scaled)) };
} else {
const scaled = (normalized - 0.5) * 2.0;
return [_]u8{ 0, @as(u8, @intFromFloat(127 * (1 - scaled / 2))), @as(u8, @intFromFloat(255 * (1 - scaled))) };
}
}
"""
end
We run the Zig code. It returns a binary that we are able to consume with Nx and display the image.
To draw an image of 1M pixels, it takes a few milliseconds. Feels like magic.
h = w = 5_000
max_iter = 300;
Zigit.generate_mandelbrot(h, w, max_iter)
|> Nx.from_binary(:u8)
|> Nx.reshape({h, w, 3})
|> Kino.Image.new()
SOCIAL SHARE CARD GENERATOR