- numpy - matplotlib

# interactive setup adapted from https://github.com/pyscript/pyscript/blob/main/pyscriptjs/examples/numpy_canvas_fractals.html from pyodide import to_js, create_proxy import numpy as np from numpy import random from js import ( console, document, devicePixelRatio, ImageData, Uint8ClampedArray, CanvasRenderingContext2D as Context2d, requestAnimationFrame, ) ###### boilerplate from https://github.com/PWhiddy/tensor-canvas/blob/master/tensorcanvas/BaseShapes.py def make_uv(t): uvx = np.tile(np.expand_dims(np.arange(0.0, t.shape[1], 1), axis=0), (t.shape[0],1)) uvy = np.tile(np.expand_dims(np.arange(0.0, t.shape[0], 1), axis=0), (t.shape[1],1)).transpose() return uvx, uvy def dist_to_col(dist, color, blend, t): msk = np.clip((dist+blend) / (2.0*blend), 0.0, 1.0) msk = msk * msk * (3.0 - 2.0 * msk) msk = np.tile(np.expand_dims(msk, axis=2), (1,1,3)) col_t = np.tile(np.expand_dims(np.expand_dims(np.array(color), axis=0), axis=0), ((t.shape[0],t.shape[1],1))) return msk*t + (1.0-msk)*col_t def draw_circle(xp, yp, radius, color, t, blend=0.75): uvx, uvy = make_uv(t) dist = np.sqrt((xp-uvx)**2.0 + (yp-uvy)**2.0 + 1.0)-radius t = dist_to_col(dist, color, blend, t) return t ###### ##### from https://github.com/PWhiddy/jax-experiments/blob/main/nbody.ipynb def make_init_state(p_count): return random.rand(p_count, 2), random.rand(p_count, 2)-0.5 def compute_forces(pos, scale, eps=0.1): a, b = np.expand_dims(pos, 1), np.expand_dims(pos, 0) diff = a - b dist = (diff * diff).sum(axis=-1) ** 0.5 dist = np.expand_dims(dist, 2) force = diff / ((dist * scale) ** 3 + eps) return force.sum(0) def sim_update_force(parts_pos, parts_vel, t_delta=0.05, scale=5, repel_mag=0.1, center_mag=2.5, steps=10, damp=0.99): p_p = np.array(parts_pos) p_v = np.array(parts_vel) for _ in range(steps): p_p = p_p + t_delta * p_v force = compute_forces(p_p, scale) center_diff = p_p-0.5 centering_force = center_diff / ((center_diff ** 2).sum() ** 0.5) p_v = damp * p_v - t_delta * (force * repel_mag + centering_force * center_mag) return p_p, p_v def draw_sim(parts_pos, parts_vel, grid_r, opacity=1.0, p_size=4.0): canvas = np.zeros((grid_r, grid_r, 3)) + 253/255 # would be interesting to use jax.experimental.loops for these for part_p, part_v in zip(parts_pos, parts_vel): sp = part_p*grid_r # snipe only the pixels we're going to update (numpy version only) lx = (sp[0]-5) ux = (sp[0]+5) ly = (sp[1]-5) uy = (sp[1]+5) mag = 10*(np.abs(part_v[0]) + np.abs(part_v[1])) canvas[int(lx):int(ux), int(ly):int(uy), :] = draw_circle( sp[0]-lx, sp[1]-ly, p_size, np.array([mag,0.0,1-mag]), canvas[int(lx):int(ux), int(ly):int(uy), :] ) return np.clip(canvas, 0, 1) width, height = 384, 384 grid_res = 384 should_push = False def generate_step(width, height, state): global should_push if state == None: p_state, v_state = make_init_state(128) v_state *= 0 base_canvas = np.ones((width, height, 4), dtype=np.uint8) * 255 else: base_canvas, p_state, v_state = state center_mag = 0.0 if should_push else 0.5 for _ in range(5): p_state, v_state = sim_update_force(p_state, v_state, t_delta=0.05, scale=10, center_mag=center_mag, repel_mag=0.05, damp=0.996, steps=2) base_canvas[:, :, :3] = 255*draw_sim(p_state, v_state, grid_res, p_size=4.0) #base_canvas[:, :, 3] = 1 #base_canvas[:, :, :3] = np.random.randint(0,255,size=(width,height, 3), dtype=np.uint8) return base_canvas, p_state, v_state #return np.random.randint(0,255,size=(width,height, 4), dtype=np.uint8) ##### def prepare_canvas(width: int, height: int, canvas: Element) -> Context2d: ctx = canvas.getContext("2d") canvas.style.width = f"{width}px" canvas.style.height = f"{height}px" canvas.width = width canvas.height = height ctx.clearRect(0, 0, width, height) return ctx def draw_image(ctx: Context2d, image: np.array) -> None: data = Uint8ClampedArray.new(to_js(image.tobytes())) width, height, _ = image.shape image_data = ImageData.new(data, width, height) ctx.putImageData(image_data, 0, 0) current_image = None async def draw_particles(state) -> None: spinner = document.querySelector("#nbody .loading") canvas = document.querySelector("#nbody canvas") spinner.style.display = "" canvas.style.display = "none" ctx = prepare_canvas(width, height, canvas) console.log("Computing Nbody interactions ...") poly_in = document.querySelector("#poly") coef_in = document.querySelector("#coef") conv_in = document.querySelector("#conv") iter_in = document.querySelector("#iter") image, p_state, v_state = generate_step(width, height, state) global current_image current_image = image draw_image(ctx, image) spinner.style.display = "none" canvas.style.display = "block" requestAnimationFrame(create_proxy(lambda _event: draw_particles((image, p_state, v_state)))) canvas = document.querySelector("#nbody canvas") async def mousedown(event): global should_push should_push = True async def mouseup(event): global should_push should_push = False canvas.addEventListener("mousedown", create_proxy(mousedown)) canvas.addEventListener("mouseup", create_proxy(mouseup)) # mobile canvas.addEventListener("touchstart", create_proxy(mousedown)) canvas.addEventListener("touchend", create_proxy(mouseup)) import asyncio _ = await asyncio.gather( draw_particles(None) )

I’ve ported my simple jax particle sim experiment to pyscript/numpy! The code ended up quite ugly and inefficient, but it works! Overall I’m impressed with performance and usability of pyodide and pyscript. Python libraries like numpy that are implemented in C seem to be able to run at full speed via wasm, so out of the box many numerical operations will probably be faster than naive JS. The integration with the web environment is very slick, DOM manipulation and JS interaction is pretty straightforward and this feels like a great option for interactive and graphical python applications. Because it’s running the regular CPython interpreter, you get error messages just like you’d normally expect. This means debugging works right away and doesn’t require any extra tooling as one might want using compiled languages in wasm. In the case of this particular code however giving up jax’s jit and GPU acceleration has made this simulation run quite slow. It will be interesting to see if ML frameworks add support for pyodide or new frameworks entirely pop up to take advantage of this. This setup won’t be useful to anyone working with large amounts of data or compute obviously, but is a much better fit for lightweight scripts than something like google colab. This page is a single markdown file which embeds html which embeds javascript which embeds python via pyscript. Exciting!

Bonus


The inverse-square force between two particles is defined by:

\[\mathbf{F(r_{a},r_{b})} = k\frac{\mathbf{r_{a}} - \mathbf{r_{b}}}{|\mathbf{r_{a}} - \mathbf{r_{b}}|^3 + \epsilon}\]

Using array broadcasting, the force calculation for all pairs of particles can be vectorized without looping in python:

def compute_forces(pos, scale, eps=0.1):
  a, b = np.expand_dims(pos, 1), np.expand_dims(pos, 0)
  diff = a - b
  dist = (diff * diff).sum(axis=-1) ** 0.5
  dist = np.expand_dims(dist, 2)
  force = diff / ((dist * scale) ** 3 + eps)
  return force.sum(0)

See full code for more details