Skip to content

Marimo: The Antidote to Jupyter Notebooks

Don't use a notebook to do an HPC job-script's job, unless your notebook is also a job-script. Marimo, an alternative to Jupyter, enables this, and also helps mitigate some other common criticisms of notebooks. Through the magic of web-assembly, it even allows your Python code to live rent-free in someone else's web browser. Using Lyapunov fractals as a somewhat psychedelic example, we might as well experiment with multiprocessing and shared memory along the way.

OnDemand is an approachable way to access the HPC cluster. R Studio and Jupyter are the most popular applications it provides. Despite what is to follow, they will be supported as long as they are maintained and wanted. These applications allow experimentation in an interactive environment, but they consume HPC resources regardless of whether anything computationally intensive is being done. Conversely, these resources are only reserved as long as the user's web browser remains open. Being unable to close one's laptop and go home until one's code has run negates the point of using HPC in the first place. To quote a forum-post from an OnDemand developer:

JupyterLab Notebooks through OnDemand are for interactive jobs, i.e, jobs you’re actively interacting with. We often tell customers to run these programs in batch jobs if they require a long time to run.

Job scripts are generally recommended for long-running workflows, indeed OnDemand sessions are now restricted to 24 hours. Here, we will investigate ways of using notebooks as job scripts. However, beyond HPC, notebooks have other issues.

What's Wrong With Notebooks?

Starting up a Jupyter notebook server is arguably one of the easiest ways to provide neophyte programmers with access to a Python interpreter, if they have insufficient experience or admin permissions to do so for themselves. Indeed, mybinder.org will launch one based on an existing notebook in a GitHub repository on modest hardware at someone else's expense. They allow rich text and mathematical formatting to be combined with executable code, interactive user interfaces, images and charts. They are a good way of using data and code to construct a step-by-step argument. Who but the author would be sufficiently mean of spirit to find fault with them?

Joel Grus's talk entitled "I don't like notebooks" at JupyterCon back in 2018 remains highly relevant today. The RSE team would still agree with most of the points made. (It is also a refreshing example of a difference of opinion within tech being explored with polite magnanimity.) Given that you are already reading this blog-post, it's quite a time-commitment, but it should be required viewing if you are using notebooks in teaching, and are seeking to mitigate their shortcomings. Their main drawback is demonstrated by the unfortunate ease with which we can use them to assert that black is white:

Jupyter asserts that black is white
How has hidden state in Jupyter allowed this to happen?

Reassuringly, this doesn't happen in the Ipython console:

In [1]: BLACK = 'black'
In [2]: WHITE = 'white'
In [3]: assert BLACK == WHITE
--------------------------------------------------
AssertionError   Traceback (most recent call last)

Even running the cells in order down the page, one can execute __WHITE__ = 'black' in the 2nd cell, correct it without running it again, and still produce the unfortunate result in the 3rd. One could also run the cells out of order, or delete an intermediate cell after execution to produce similar results. Jupyter not only allows, but encourages out of order execution of code. In a long, complex analysis, it is all to easy to re-factor code and re-run cells up and down a notebook without realising that a mistake has been made. Reasoning about hidden state is hard for novice (or busy, or tired) programmers, and they tend to produce more stateful imperative and procedural code anyway, which makes matters worse. (More experienced programmers aren't better at thinking about hidden state, they're just better at avoiding it through vectorization, separation of concerns, or functional programming.)

Besides encouraging bad practice, notebooks can make adopting good practice hard. Viewing the .ipynb file of the previous the previous example in a text editor shows that each cell of Python code is embedded in an element in a JSON array:

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "cfaa94e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "__BLACK__ = 'black'"
   ]
  }
  ...
 ]
}

You are using version control when writing code for your research, aren't you? Suppose you want to explore making some complex changes to apparently working code. When you want to explore new approaches or features you create branches, secure in the knowledge that you can easily revert to the last working state of your code. Perhaps a colleague will work on one of the features in parallel. If you are working in notebooks, when you try merge the now-working branches into one and are reviewing changes in a pull request, differences between different versions of Python code embedded in different versions of JSON data will be almost impossible to make sense of. Woe betide you if you encounter a merge-conflict. You will not be able to resolve the changes easily into a valid notebook. Linters like Flake8 that enforce the PEP 8 Python style guide can be adapted for use in notebooks, but it is a little awkward. Notebooks can hinder effective collaboration, even if the collaborator is just you in a fortnight's time!

Notebooks are fine for constructing an argument, but they are hard for others to use as a software tool, except on their own terms as notebooks. You might retort that you are in the business of doing research, not writing software, and that software development best practice doesn't apply. We would strongly disagree. The papers you will publish are an account of what you did. The code you wrote is part of what you did. If others can easily read and run your code to verify that it works, it aids reproducibility. Good code is just good science; notebooks don't always help.

Forewarned of the deficiencies of Jupyter, we can attempt to work around and mitigate them. However, there is an alternative, Marimo. Marimo enforces the execution of cells in order, marking them as stale if their prerequisites change. Marimo notebook files are just normal .py files, so they work fine with version control. That takes care of our biggest concerns with Jupyter notebooks. At the time of writing, Marimo is barely two years old, and has some catching up to do with the Jupyter ecosystem. Let us explore a brief, but computationally demanding and hopefully aesthetic example.

Introducing the Lyapunov Fractal

The Lyapunov Fractal is calculated from repeated iterations of the Logistic Map:

\[ \Large{x_{n+1} = r_{n}x_{n}(1-x_{n})} \]

where at each iteration \(r_{n}\) takes values from some repeated sequence, for example \(AABAB\). For each point in an image \(A\) and \(B\) take the values of the X and Y coordinates. For a large number of iterations \(N\), the Lyapunov exponent \(\lambda\) is found for each point, and coloured accordingly:

\[ \Large{\lambda = \frac{1}{N}\sum_{n=1}^{N}\log|r_{n}(1-2x_{n})|} \]

On a local Linux terminal, we can set up and launch a virtualenv, install some packages, and start playing with Marimo:

mkdir lyapunov_venv
python3 -m venv lyapunov_venv
source ./lyapunov_venv/bin/activate
pip install numpy pillow marimo
marimo edit

Hopefully, a browser will open on a fresh Marimo session and you can create a new notebook. Otherwise, open the URL displayed in the terminal. If you hover the mouse over the right-hand side of a cell, you can opt to view it as markdown and \(\LaTeX\) via MathJax, just like in Jupyter, or indeed, this blog-post:

# Lyapunov fractals

The [Lyapunov Fractal](https://en.wikipedia.org/wiki/Lyapunov_fractal)
is calculated from repeated iterations of the
[Logistic Map](https://en.wikipedia.org/wiki/Logistic_map):

$$ \Large{x_{n+1} = r_{n}x_{n}(x_{n}-1)} $$

We shan't need to copy and paste in much code before we can experiment with Marimo's features. You can paste in each block of code into a new notebook cell as we go, or you can get the notebook from its repository. It's best to have the first cell containing imports as a set-up cell which can be created from the "hamburger" menu at the top-right.

from datetime import datetime
import io
import itertools
from functools import partial
try:
    from multiprocessing import Pool, shared_memory
except ModuleNotFoundError:
    # No shared memory in WASM
    pass
import os
import subprocess as sp
import sys
from typing import Generator, List, Tuple

from matplotlib import colormaps
import marimo as mo
try:
    import jax.numpy as np
    GOT_JAX = True
except ModuleNotFoundError:
    import numpy as np
    GOT_JAX = False
from numpy import typing as npt
from PIL import Image

We'll need to turn strings like \(AABAB\) into a Numpy array of 0s and 1s like \(\left[0, 0, 1, 0, 1\right]\). To do this, coerce the string to uppercase, and map the ord built-in function to it to obtain its ASCII values. Since the uppercase ASCII alphabet starts at "A" with a value of 65, turn our map object into a Numpy array via a list, then subtract 65 from it.

def seq_vector(seq: str) -> npt.ArrayLike:
    """Take a string of letters and return an array of ints."""
    assert str.isalpha(seq)
    return np.array(list(map(ord, seq.upper())), dtype=np.int32) - 65

Computing the points of the image is fairly straightforward. We pass in the sequence from the previous function, and the coordinates of the points as numpy.meshgrid-s. meshgrid takes 1D arrays of the X and Y coordinate ranges, and returns two 2D arrays. The star notation is used so that we can pass in multiple arrays of coefficients, not just those corresponding to the \(A\) and \(B\) values. We'll see why soon. We construct a cuboid array of coefficients we are going to use at each point as we cycle through the sequence by using numpy.stack and referencing it with our sequence array. The correct layer in the stack is found by taking the iteration number modulo the length of the repeating sequence. Next, we create an array that initialises the values of our iterations of the logistic map to 0.5. A for loop builds up the sum of the iterates one by one. We could allow the coefficients and iterates to build up into cuboid arrays as thick as the number of iterations and then compute the sum in a single vectorized operation along the appropriate axis, but this seems like a rather profligate use of RAM, especially when we come to generate many images in parallel later on. Numpy's += operator is perfectly capable of doing the summation in-place.

def lyapunov(seq: npt.ArrayLike, n_its: int, *points: List[npt.NDArray]) -> npt.NDArray:
    """Compute a Lyapunov fractal.

    seq: Coefficient sequence as a Numpy array of ints in [0..N-1].
    its: Number of iterations.
    points: List of N arrays, one for each coefficient, giving
    the coefficent value at each point in the image.
    """
    coeffs = np.stack(points)[seq]
    seq_len = len(seq)
    img_shape = coeffs.shape[1:3]
    prev = 0.5 * np.ones(img_shape, dtype=np.float32)
    img = np.zeros(img_shape)
    for i in range(1, n_its):
        r = coeffs[i % seq_len]
        nxt = r * prev * (1.0 - prev)
        img += np.log(np.abs(r * (1.0 - 2 * nxt)))
        prev = nxt
    return img / (n_its - 1)

We need to colour the image arrays before we can display them. Traditionally, positive and negative values are given different colours. matplotlib's diverging colour maps work well for this. The hyperbolic tangent coerces the large positive and negative values in the arrays to between zero and one. Scaling to unsigned integers from 0 to 255 allows us to return an image using the Pillow library.

def sigmoid(x: npt.NDArray) -> npt.NDArray:
    """Hyperboloic tangent normalized from 0 to 1."""
    return 0.5 * (1 + np.tanh(x))

def render_image(img: npt.NDArray, palette: str='Spectral') -> Image.Image:
    """Colour a Numpy array according to a Matplotlib palette."""
    colours = colormaps[palette]
    return Image.fromarray(
        (255 * colours(sigmoid(img))).astype(np.uint8)
    )

Next, we create a set of user-interface (UI) components to get the desired coefficient sequence, number of iterations, image bounds, and colour palette from the user:

seq_box = mo.ui.text(value='AABAB', label='coefficient sequence')
its_box = mo.ui.number(start=20, stop=400, step=20, value=100,
                       label="number of iterations")
x_img_slider = mo.ui.range_slider(start=2.0, stop=4.0, step=0.1, value=[2.0, 4.0],
                               label='x range')
y_img_slider = mo.ui.range_slider(start=2.0, stop=4.0, step=0.1, value=[2.0, 4.0],
                               label='y range')
cmap_names = ['seismic', 'vanimo', 'managua', 'berlin', 'Spectral',
    'twilight', 'twilight_shifted', 'ocean', 'cubehelix', 'turbo', 'plasma',
    'magma', 'PiYG']

# (Some of these might be missing in the WASM version.)
palettes = [name for name in cmap_names if name in colormaps]

colour_box = mo.ui.dropdown(palettes, value='twilight', label='palette')

Marimo insists that any reference to the values of UI elements must occur in subsequent cells. By manipulating the controls, we can mutate global state. In fact, that's the only way that Marimo allows you to manipulate global state. It tracks which cells need updating by wrapping the contents of each one in an outer function, itself wrapped by the @app.cell decorator. The function takes the cell's dependencies as inputs, and returns the contents as a tuple. This will have consequences later. For example, the cell containing the seq_vector function looks like:

@app.function
def seq_vector(seq: str) -> npt.ArrayLike:
    """Take a string of letters and return an array of ints."""
    assert str.isalpha(seq)
    return np.array(list(map(ord, seq.upper())), dtype=np.int32) - 65

Therefore, Marimo complains if you try and redefine a function or modify a variable across cells. Like it or not, you are forced away from an imperative, procedural style of programming and towards a more functional one.

Marimo won't let us mutate variables.
Marimo won't let you re-define or mutate variables or functions across cells.

We construct the meshgrid-s for the image coordinates and perform some very basic input sanitization on the sequence of coefficients. If the user is running the notebook as a script, they won't want to waste time generating an image for the UI, so we check if we're in a notebook.

IMG_SIZE = 400
SMALL = 0.000001

img_x_min, img_x_max = x_img_slider.value
img_y_min, img_y_max = y_img_slider.value

img_x_points, img_y_points = np.meshgrid(
    np.linspace(img_x_min + SMALL, img_x_max, IMG_SIZE),
    np.linspace(img_y_max, img_y_min + SMALL, IMG_SIZE),
    indexing='xy'
)

img_seq = seq_vector(
    ''.join(filter(lambda char: char in 'AB', seq_box.value.upper())),
)

if mo.running_in_notebook():
    img = render_image(lyapunov(
        img_seq, its_box.value, img_x_points, img_y_points),
        palette=colour_box.value
    )
else:
    img = None

Finally, we can display the UI by arranging the controls in a pleasing way with marimo.vstack and marimo.hstack:

mo.vstack([
    img,
    mo.hstack([x_img_box, y_img_box]),
    seq_box, its_box, colour_box
], align='center')

The user can adjust the controls and see the image react accordingly.

a basic UI for generating fractals
We can use Marimo to generate a basic reactive user-interface.

Sharing is Caring

Jupyter notebooks can be run as scripts with only a little extra friction, but since Marimo notebooks are regular Python files, they can be run directly. You can even add a standard Python shebang line to make them executable. This makes it easy to go from a local exploratory proof-of-concept to a usable HPC job. We will demonstrate Marimo doing this shortly. It is somewhat awkward and involved to import Jupyter notebooks as modules, but Marimo makes this rather easy. If a cell contains a single class or function, and has only dependencies from the set-up cell, it will have the label "reusable" in its bottom-right corner. (Reusable cells can also have other reusable cells as dependencies.) We can launch an Ipython console and import reusable functions directly. Here, the lyapunov_img function conveniently wraps the three functions needed to generate an image:

from lyapunov import lyapunov_img
lyapunov_img('BBBBBBAAAAAA',
    x_min=2.5, x_max=3.4, y_min=3.4, y_max=4.0, its=400, width=900, height=600,
palette='managua').save('zircon_zity.png')

If we don't want the constraint of reusable functions, we can always export notebooks to regular .py scripts.

marimo export script lyapunov.py -o lyap.py

Through the magic of WebAssembly, a Python interpreter and much of the PyData stack has been compiled for use within the web browser by the Pyodide project. Numpy, Pandas, Scipy, Scikit-learn and the Python Image library are all available. While Jupyter has been made available entirely in the browser with Pyodide, Marimo makes it very easy to publish existing notebooks in this way. Firstly, under the "hamburger" menu at the top right of the notebook is a "share" option that provides a Create WebAssembly link. This will produce a very long URL that encodes a compressed version of the entire notebook. You might be advised to use a URL shortening service to generate a link to it. https://marimo.app/ will decode and run the notebook entirely in the browser. You can also export your notebooks to static HTML, and host them somewhere, preferably for free on the likes GitHub Pages. For example, to publish our notebook when GitHub Pages has been configured to build from a directory called docs in the branch gh-pages:

git checkout gh-pages
git merge main
marimo export html-wasm lyapunov.py -o docs --mode run
git add docs
git commit -m "Pithy commit message..."
git push -u origin gh-pages

Alternatively, the whole process can be automated with GitHub Actions. The notebook can be viewed at https://augeas.github.io/lyapunov/ Here, the notebook was made available as a web-app, with only the rich text and user-interface visible by default, but it could also be made available as a fully-functioning notebook with the code editable. If your code can run using the sub-set of libraries available within Pyodide on the modest resources available from a reasonably modern browser, you can use WebAssembly notebooks to produce compelling demonstrations of your research. If a required package is missing, Marimo will attempt to install it via micropip. It's even possible to include data. You can also use the drag-and-drop marimo.ui.file component in both regular notebooks and WebAssembly to allow users to upload their own data. As yet, Marimo doesn't work directly with BinderHub, like Jupyter does, the most ubiquitous deployment being MyBinder.org. This allows the installation of arbitrary Python packages in a requirements.txt file, and even binary dependencies in apt.txt files. It can take a while for MyBinder.org to spin-up the containers it needs to run notebooks, and when it does, the latency can be considerable. You would not want to trust a workshop or laboratory class to this. However, if you do want to use these features of Jupyter notebooks, Marimo does allow exporting to Jupyter. In order to demonstrate our Marimo notebook running as a HPC job script, we must do a little more work to have something computationally worthy.

Making Movies

The Lyapunov fractal Wikipedia article mentions adding a third coefficient \(C\) that varies in time. More pleasingly, we can add two more coefficients, \(C, D\) whose values take those of coordinates of points on a circle. That way, the moving image can return to where it started and loop. We generate sets of extra coefficients with a stack of rotation matrices:

def rot_coeffs(x: float, y: float, radius: float, n: int) -> npt.NDArray:
    """Return an array of n pairs of coefficients centred at (x, y) with radius r."""
    theta = np.linspace(-np.pi, np.pi, n, dtype=np.float32)
    cos_theta = np.cos(theta)
    sin_theta = np.sin(theta)
    rot = np.array([
        [cos_theta, -sin_theta],
        [sin_theta, cos_theta]
    ]).T.reshape((n, 2, 2))
    point = np.array([[[0, radius]]], dtype=np.float32)
    return (
        np.array([x, y]).reshape((1, 1, 2)) + point @ rot
    ).reshape(n, 2)

For a given pair of coefficients, we generate the large arrays needed to pass to the lyapunov function:

def extra_coeffs(point: npt.ArrayLike, shape: Tuple[int, int]) -> Tuple[npt.NDArray, npt.NDArray]:
    """Return two arrays of coefficients with the given shape"""
    c, d = point
    c_coeff = c * np.ones(shape, dtype=np.float32)
    d_coeff = d * np.ones(shape, dtype=np.float32)
    return (c_coeff, d_coeff)

We can construct another set of UI elements to allow the user to place the centre of the circle \(CD\) and set its radius, carefully, so that no coefficients exceed 4, causing the logistic map to diverge. (For brevity, we shall omit that here.) By choosing the position around the circle, the user can watch the image morph and change. Since the generation of each frame doesn't depend on any other, it's "embarrassingly parallel". This gives us the excuse to think about speeding our code up by using multiple cores to generate videos.

Scream If You Want To Go Faster

To create a video, we can try and generate multiple frames at once, and stream them to FFmpeg. From a previous blog post, we had established that using multiple cores via multiprocessing doesn't always lead to a reduction in execution time, especially if we are passing large numpy arrays to and from our functions. If we naively pass large arrays of coordinates and coefficients to the lyapunov function, with multiprocessing.Pool and receive large arrays back again, we can keep all the cores in the pool at 100% usage, with the overall process taking almost as long as for a single core. The CPUs will be doing "busy work", serialising and de-serialising large arrays with the somewhat infamous pickle module. There is plenty of code "in the wild" that uses already parallel routines from numpy.linalg within a Pool, to little advantage as large arrays are being passed. Getting a useful performance improvement will take a little thought.

One complication when using Marimo with multiprocessing is that if a function isn't "reusable", it will not be visible to Pool.map, Pool.imap or Pool.starmap as it will be wrapped with an outer function, and hence not serializable with the rather brittle pickle module from the standard library. Imported functions not defined in the notebook work just fine. It's easy enough to ensure that all parallel functions are reusable, but if not, the errors are rather disconcerting to new programmers. An alternative is to use the third-party multiprocess library which uses the more robust dill library for serialisation, with unknown consequences for performance.

The book "High Performance Python" has some useful advice in its chapter on multiprocessing. We can create SharedMemory buffers attached to Numpy arrays, that can be written to in the process that created them, but accessed in other processes by name. The following function creates or recovers Numpy arrays of the required size and type:

def get_shared_np(shape: Tuple[int, ...], dtype: str='float32',
    name: str=None) -> Tuple[shared_memory.SharedMemory, npt.ArrayLike]:
    """Return a SharedMemory instance, and a numpy array of the given
    shape and dtype that points to it. If name is given, retrieve an
    existing SharedMemory object."""
    dtype = np.dtype(dtype)
    size = dtype.itemsize * np.prod(np.array(shape))
    if name is None:
        buff = shared_memory.SharedMemory(create=True, size=size)
    else:
        buff = shared_memory.SharedMemory(name=name, create=False, size=size)
    arr = np.ndarray(shape, dtype, buffer=buff.buf)
    return buff, arr

We can also create SharedMemory buffers from existing Numpy arrays:

def array_to_shared(arr: npt.ArrayLike) -> shared_memory.SharedMemory:
    """Return a SharedMemory instance that points to a given numpy array"""
    buff = shared_memory.SharedMemory(create=True, size=arr.nbytes)
    buff.buf[:] = arr.tobytes()
    return buff

Next we need a wrapper around the original lyapunov function to access the shared arrays of \(A\) and \(B\) coefficients, and return the images as references to SharedMemory arrays.

def lyapunov_mp(cd: Tuple[float, float], shape: Tuple[int, int],
    seq, its: int, x_name: str, y_name: str) -> str:
    """Compute a Lyapunov fractal using arrays backed by shared memory.

    cd_out: Tuple that allows the function to be called by Pool.imap containing:
        cd: Tuple of floats giving the C, D coefficients constant across the image.
        out_name: Name of a SharedMemory buffer to hold the Lyapunov exponents.
    shape: Tuple of ints giving the shape of the A, B coefficient and output arrays.
    seq:
    its: Number of iterations.
    x_name, y_name: Names of SharedMemory buffers pointing to the A and B coefficient arrays.
    """
    x_buff, x_coeff = get_shared_np(shape, name=x_name)
    y_buff, y_coeff = get_shared_np(shape, name=y_name)
    c_coeff, d_coeff = extra_coeffs(cd, shape)
    # Don't use the "render_image" function, keep the sigmoid function inside the Pool:
    out_buff = array_to_shared(sigmoid(
        lyapunov(seq, its, x_coeff, y_coeff, c_coeff, d_coeff)
    ).astype(np.float32))
    for buff in (x_buff, y_buff, out_buff):
        buff.close()
    return out_buff.name

The function video_seq_mp acquires SharedMemory buffers and a Pool. Pool.imap is used to apply lyapunov_mp to the \(CD\) coefficients, and the image arrays are recovered from the buffers before being converted into .pngs and yielded:

def video_seq_mp(seq: str, x_mi: float, x_mx: float, y_mi: float, y_mx: float,
    x: float, y: float, r: float, n: int, cores: int, its: int=100,
    pal: str='managua', w: int=512, h: int=512) -> Generator[npt.NDArray, None, None]:
    """Yield an animated sequence of Lyapunov images as Numpy arrays using
    multiprocessing and SharedMemory.

    seq: String representing the repeating coefficient sequence, e.g: "AACBABD".
    x_mi, y_mi, x_mx, y_mx: Floats describing the boundaries of the images, the ranges
    of the A and B coefficients.
    x, y, r: Floats giving the centre and radius of a circle on which the C, D coefficients lie.
    n: Number of points around the circle, the number of frames.
    cores: Number of cores to use in the Pool, not including the calling process which
    compresses the image and the ffmpeg process that consumes them.
    its: Number of iterations for each image.
    pal: Name of the Matplotlib palette to use.
    w, h: Image width and height.
    """

    img_shape = (h, w)
    chunk_size = 8 * cores

    # Reserve SharedMemory for the A, B coefficients:
    x_buff, y_buff = map(array_to_shared, np.meshgrid(
        np.linspace(x_mi, x_mx, w, dtype=np.float32),
        np.linspace(y_mx, y_mi, h, dtype=np.float32),
    indexing='xy'))

    seq_vec = seq_vector(seq)
    cd_coeff = rot_coeffs(x, y, r, n)

    """All the arguments of lyapunov_mp remain constant, except for the C, D coefficients
    Thus, no large numpy arrays will be serialized when the function is passed to the Pool."""
    lyap = partial(lyapunov_mp,
        shape=img_shape, seq=seq_vec, its=its,
        x_name=x_buff.name, y_name=y_buff.name,
    )

    with Pool(cores) as pool:
        for out_name in pool.imap(lyap, map(tuple, cd_coeff), chunksize=chunk_size):
            out_buff, out = get_shared_np(img_shape, name=out_name)
            buff = io.BytesIO()
            """The job of turning the returned arrays to images is left to the calling process.
            There's no point compressing the images when ffmpeg would have to uncompress
            them afterwards."""
            Image.fromarray(
                (255 * colours(out)).astype(np.uint8)
            ).save(buff, format='PNG', compress_level=0)
            yield buff.getvalue()
            out_buff.close()
            out_buff.unlink()

    for shm in (x_buff, y_buff):
        shm.close()
        shm.unlink()

    # Need to yield something so that the SharedMemory is freed.
    yield None

We need a function to launch ffmpeg in a subprocess and pipe in the sequence of images:

def render_video(fname: str, im_seq: Generator[npt.NDArray, None, None],
    fps: int=30, quiet: bool=True) -> None:
    """Stream a sequence of .png images to ffmpeg, and turn them into an .mp4 video.
    fname: Filename for the video.
    fps: Frames per second, defaults to 30.
    quiet: Whether to suppress ffmpeg's rather verbose output.
    """
    ffmpeg_cmd = [
        'ffmpeg', '-threads', '1', '-f', 'image2pipe', '-vcodec', 'png', '-r', str(fps),
        '-i', '-', '-vcodec', 'libx264', '-q:a', '0', fname
    ]
    if quiet:
        ffmpeg_out = sp.DEVNULL
    else:
        ffmpeg_out = None
    with sp.Popen(ffmpeg_cmd, stdin=sp.PIPE, stdout=ffmpeg_out,
        stderr=ffmpeg_out) as proc:
        # Filter the sequence to omit the final "None".
        for im in filter(None, im_seq):
            proc.stdin.write(im)
        proc.stdin.close()
        proc.wait()

For the sake of brevity, we shall omit the UI components needed to enable video generation in a notebook. However, it's worth knowing that we can prevent them appearing in WebAssembly notebooks by testing pyodide not in sys.modules. To generate videos from a job-script, we first need to be sure how many cores we have and define some default command-line arguments:

TOTAL_CORES = int(os.environ.get('OMP_NUM_THREADS', 1))
if TOTAL_CORES > 2:
    MAX_VID_SEQ_CORES = TOTAL_CORES - 2
else:
    MAX_VID_SEQ_CORES = 1

# Cores refers to the number of cores used to create the image sequence.
# Don't forget one for ffmpeg, and one to pass the images to it.

DEFAULT_ARGS = {
    'seq': 'AACBABD',
    'xmin': 2.0, 'xmax': 4.0, 'ymin': 2.0, 'ymax': 4.0,
    'xc': 3.0, 'yc': 3.0, 'rad': 0.2, 'its': 100,
    'width': 512, 'height': 512, 'dur': 60, 'fps': 30,
    'cores': MAX_VID_SEQ_CORES, 'pal': 'managua'
}

Marimo provides the convenience function marimo.cli_args(), but we could also use argparse from the standard library:

# If this is being run as a script:
if not mo.running_in_notebook():
    args = mo.cli_args()

    def get_arg(arg):
        return args.get(arg, DEFAULT_ARGS.get(arg))

    sq = args.get('seq')
    fname = args.get('fname')

    fps = get_arg('fps')
    n_frames = fps * get_arg('dur')

    xmin, xmax, ymin, ymax, xc, yc, rad = map(
        get_arg, ['xmin', 'xmax', 'ymin', 'ymax', 'xc', 'yc', 'rad'])

    its, cores, pal, width, height = map(get_arg, ['its', 'cores', 'pal', 'width', 'height'])

    if cores > 1 and not GOT_JAX:
        img_sq = video_seq_mp(sq, xmin, xmax, ymin, ymax, xc, yc, rad, n_frames, cores,
            its=its, pal=pal, w=width, h=height)
    else:
        img_sq = video_seq(sq, xmin, xmax, ymin, ymax, xc, yc, rad, n_frames,
            its=its, pal=pal, w=width, h=height)

    right_now = datetime.now()
    render_video(fname, img_sq, fps=fps, quiet=False)
    elapsed = (datetime.now()-right_now).total_seconds()
    print(f'Wrote {fname} in {elapsed}s.')

Now we can finally use our notebook as a script to generate a 1080x1080 video:

python lyapunov.py --seq=ACDBCD --fname=ACDBCD_1080.mp4 --width=1080  --height=1080 --xc=2.95 --yc=2.95 --rad=0.25 --pal=twilight --cores=4

Didn't We Do Well?

It is possible to write Jupyter notebooks in such a way that out-of-order execution of cells either causes no issues, or raises an error rather than quietly returning different results. However, how would you know that you'd got it right? How likely is a neophyte developer to get it right? Instead of a workflow where you make repeated edits to notebooks to generate results which are not preserved, you could have a chronological record of what you have done, worthy of the name "notebook". Marimo clearly helps enforce this. In our Lyapunov fractal example, we have used notebook user-interface elements to create a tool for our own experimentation. The ability to export to WebAssembly allows others to explore it with very little effort. As the notebook is a pure Python file, we can use it as a module or a script with ease, and so can others. We can engage in best-practices like version-control, testing and linting without any friction. Indeed, the somewhat pedantic Pylint rates the notebook at nearly 8/10.

Note

At the time of writing, Marimo is not yet available on our deployment of Open OnDemand, but a development implementation is being considered. In the meantime, you can start using Marimo locally in such a way that your notebooks can be used as job-scripts with ease.

Coda: Did We Do Well?

We chose to experiment with the Lyapunov fractal as a vehicle to evaluate Marimo. Our first attempt at parallelising our code brought the issues with multiprocessing to light. We might as well see how effective it was. On a single i7 core, a 1080x1080 video with 1800 frames takes around 45 minutes to render. In this case, a simplified function that generates the frames directly was used, rather than using SharedMemory to pass them from a Pool with only a single core. Using multiprocessing with 2 cores produces a modest 20% improvement. The law of diminishing returns kicks in rather quickly, there is not much benefit to using 4 cores. Experimenting with the chunksize parameter of Pool.imap had little effect, although breaking iterables into large chunks is often reckoned to make multiprocessing more efficient. You might have noticed that we attempt to import Jax at the top of the notebook, reverting to Numpy if we can't, as will be the case with WebAssembly. We can use the two libraries interchangeably if the Numpy code is written such that no arrays are mutated. Jax can work with multiprocessing, but there are some caveats, the fork strategy doesn't work, and there were issues with the SharedMemory code. If Jax is detected, multiprocessing is bypassed by the script. Splitting the task of video generation into frames seems like a more natural division of labour than splitting individual images across cores and having to re-assemble them. However, without setting any environment variables, Jax used 350% of CPU without being prompted. This is rather concerning on HPC! As others have found, limiting Jax's CPU usage can take some persuasion and experimentation. In the end, setting the environment variable NPROC seemed to work. On a single core, Jax ran twice as fast as Numpy. Setting NPROC=2 yielded a 15% improvement. Whether the likes of Numba or PyTorch could scale better will have to be another story.

Jax seems to be more performant than multiprocessing/Numpy
Jax on a single core out-performs Numpy with multiprocessing on two.