Mandelbrot orbits with Nx and Kino.Live.JS
Mix.install(
[
{:nx, "~> 0.9.1"},
{:exla, "~> 0.9.1"},
{:kino_vega_lite, "~> 0.1.11"},
],
config: [nx: [default_backend: EXLA.Backend]]
)
Nx.Defn.global_default_options(compiler: EXLA, client: :host)
Orbit number
We return a tensor consisting of each iterates (points) of a given point $c$ under the map:
$p: z \mapsto z^2 +c$
starting at $p(0)=c$.
We allocate a tensor $t_0$ of length $n$, the max number of calculated iterates with $t_0=[c,0,\dots]$
At each step:
- while the squared norm of the iterate is less than 4 and max iterates is not reached, we add the iterate to the tensor,
- else we stop.
defmodule Orbit do
import Nx.Defn
defn p(z,c), do: z*z + c
defn sq_norm(z), do: Nx.real(z * Nx.conjugate(z))
defn calc(c, opts) do
n = opts[:n]
t0 = Nx.broadcast(0, {n}) |> Nx.put_slice([0], Nx.reshape(c, {1}))
while {i=0, t=t0, c, n}, Nx.less(i,n) and Nx.less(sq_norm(t[i]), 4) do
t_i_plus_1 = Nx.indexed_put(t, Nx.stack([i+1]), p(t[i], c))
{i+1, t_i_plus_1, c, n}
end
end
end
Plotting orbits
Given a point $c$, we return a tuple {interation_number, [[x_0,y_0], [x_1,y_1],...}
for the Kino.Live.JS
to send it to the Javascript to plot it.
defmodule Plot do
import Nx.Defn
import Nx.Constants, only: [i: 0]
defn point(cx,cy), do: cx + i() * cy
def to_js(cx, cy, imax) do
c = point(cx,cy)
{t_nb, t, _, _} = Orbit.calc(c, n: imax)
nb = Nx.to_number(t_nb) - 1
data =
t[0..nb]
|> Nx.to_list()
|> Enum.map(fn z -> [Complex.real(z), Complex.imag(z)] end)
{nb, data}
end
end
The module below is a Kino.Js.Live
module to interact between the browser and the Livebook as a server.
When you click on the 2D plan, it will plot teh orbit of this point.
More precisely, when you click on the plan:
-
the browser sends the coordinates to
Kino.Live.JS
(modulo a transformation canvas to 2D-plan), -
the
Kino.Live.JS
server calls thePlot.point
module that calculates the orbit (viaOrbit.calc
) -
then
Kino.Live.JS
sends the reuslt to the browser to plot it with a little animation.
The client server communication uses the primitives handle_event
and broadcast
server side, and pushEvent
and handleEvent
clent side.
defmodule LiveOrbit do
use Kino.JS
use Kino.JS.Live
def canvas() do
"""
Clicked point: ,   Iteration number:
"""
end
def run(), do: Kino.JS.Live.new(__MODULE__, canvas())
@impl true
def init(html, ctx) do
ctx = assign(ctx, %{imax: 150})
{:ok, assign(ctx, html: html)}
end
@impl true
def handle_connect(ctx) do
{:ok, ctx.assigns.html, ctx}
end
#-------------- received from client
@impl true
def handle_event("clicked", %{"x"=> x, "y" => y} = _evt, ctx) do
%{assigns: %{imax: imax}} = ctx
{nb, data} = Plot.to_js(x, y, imax)
# Nx.to_list(data) |> dbg()
# send data to the client
broadcast_event(ctx, "points", %{data: data})
broadcast_event(ctx, "number", %{nb: nb})
{:noreply, ctx}
end
asset "main.js" do
"""
export function init(ctx, html) {
ctx.root.innerHTML = html;
const canvas = document.getElementById("myCanvas");
const canvasCtx = canvas.getContext("2d");
const coordsElement = document.getElementById("coordinates");
const iterationNumber = document.getElementById("number")
// Canvas dimensions
const canvasWidth = canvas.width;
const canvasHeight = canvas.height;
// Coordinate plane bounds
const minX = -1.5;
const maxX = 1;
const minY = -1;
const maxY = 1;
// Function to convert canvas coordinate to complex coordinate
function toComplexCoord(canvasX, canvasY) {
const x = minX + ((canvasX / canvasWidth) * (maxX - minX));
const y = maxY - ((canvasY / canvasHeight) * (maxY - minY));
return { x, y };
}
// Function to convert compelx coordinate to canvas coordinate
function toCanvasCoord(x, y) {
const canvasX = ((x - minX) / (maxX - minX)) * canvasWidth;
const canvasY = canvasHeight - ((y - minY) / (maxY - minY)) * canvasHeight;
return { x: canvasX, y: canvasY };
}
// Draw x-axis
canvasCtx.beginPath();
canvasCtx.moveTo(0, canvasHeight * (Math.abs(maxY) / (Math.abs(maxY) + Math.abs(minY))));
canvasCtx.lineTo(canvasWidth, canvasHeight * (Math.abs(maxY) / (Math.abs(maxY) + Math.abs(minY))));
canvasCtx.strokeStyle = 'black';
canvasCtx.stroke();
// Draw y-axis
canvasCtx.beginPath();
canvasCtx.moveTo(canvasWidth * (Math.abs(minX) / (Math.abs(minX) + Math.abs(maxX))), 0);
canvasCtx.lineTo(canvasWidth * (Math.abs(minX) / (Math.abs(minX) + Math.abs(maxX))), canvasHeight);
canvasCtx.strokeStyle = 'black';
canvasCtx.stroke();
// Draw grid dynamically based on min/max X/Y
const gridStep = 0.2;
canvasCtx.strokeStyle = '#ccc';
canvasCtx.setLineDash([2, 4]);
for (let x = minX; x <= maxX; x += gridStep) {
const { x: canvasX } = toCanvasCoord(x, 0);
canvasCtx.beginPath();
canvasCtx.moveTo(canvasX, 0);
canvasCtx.lineTo(canvasX, canvasHeight);
canvasCtx.stroke();
}
for (let y = minY; y <= maxY; y += gridStep) {
const { y: canvasY } = toCanvasCoord(0, y);
canvasCtx.beginPath();
canvasCtx.moveTo(0, canvasY);
canvasCtx.lineTo(canvasWidth, canvasY);
canvasCtx.stroke();
}
canvasCtx.setLineDash([]);
//---------- client event: send click position --------------------
canvas.addEventListener("click", (e) => {
const rect = canvas.getBoundingClientRect();
const x = e.clientX - rect.left;
const y = e.clientY - rect.top;
const { x: complexX, y: complexY } = toComplexCoord(x, y);
coordsElement.textContent = `${complexX.toFixed(2)}, ${complexY.toFixed(2)}`;
//-- send data to the server -->
ctx.pushEvent("clicked", {x: Number(complexX), y: Number(complexY)})
});
//----------- callbacks from server after computations--------------
ctx.handleEvent("number", ({nb}) => {iterationNumber.textContent = nb});
ctx.handleEvent("points", drawPointsWithAnimation)
// Draw each point with animation
function drawPointsWithAnimation({data: points}) {
console.log(points)
let index = 0;
// loop with "setInterval"
const intervalId = setInterval(() => {
if (index >= points.length) {
clearInterval(intervalId); // Stop the animation when all points are drawn
return;
}
animatePoint(points, index)
index++;
}, 100);
}
function animatePoint(points, index) {
const currPoint = toCanvasCoord(points[index][0], points[index][1])
// Draw the point as a small circle
canvasCtx.beginPath();
const radius = index === 0 ? 6 : 3;
canvasCtx.arc(currPoint.x, currPoint.y, radius, 0, 2 * Math.PI);
canvasCtx.fillStyle = index === 0 ? "green" : 'red';
canvasCtx.fill();
// Draw a dotted line connecting to the previous point, if exists
if (index > 0) {
const prevPoint = toCanvasCoord(points[index - 1][0], points[index - 1][1]);
canvasCtx.beginPath();
canvasCtx.moveTo(prevPoint.x, prevPoint.y);
canvasCtx.lineTo(currPoint.x, currPoint.y);
canvasCtx.strokeStyle = 'blue';
canvasCtx.setLineDash([5, 5]); // Creates a dotted line
canvasCtx.lineWidth = 1;
canvasCtx.stroke();
canvasCtx.setLineDash([]); // Reset line dash
}
}
}
"""
end
end
You can visualize the orbits of any point by clicking into the plot below.
LiveOrbit.run()