Hessam's blog

Ancestral sampling with variational relaxation

2026-07-16T00:00:00+00:00

This post is an interactive HTML document. Read it on Hessam's blog.

Exploration: Flow matching, NeuTraHMC and normalizing flow guides in NumPyro

2026-05-06T00:00:00+00:00

This post is an interactive HTML document. Read it on Hessam's blog.

First bee sighting of 2026

2026-02-25T00:00:00+00:00

First bee sighting of 2026

I was crossing the Waterloo bridge to go to a seminar when I saw my first bee of 2026. Most years this would be a big fuzzy bumble bee but this time it’s a honey bee. What a wonderful little friend to see; it made my day. :)

Nice strange attractor visualisation

2026-02-24T00:00:00+00:00

Nice strange attractor visualisation

Zensical is a new static site generator implemented in Python that I heard mentioned on the Talk Python to Me podcast. The website features a stunning mathematical visualisation that I was able to extract into a GitHub pages website, using Claude Opus 4.6 (low) via Pi. The parameters involved in the differential equation can be modified and the default flythrough effect can be interrupted to allow zooming/panning/rotating to explore the geometry.

Draft: Evolving interactions with AI in 2026

2026-01-19T00:00:00+00:00

Draft: Evolving interactions with AI in 2026

Living through the first couple of weeks of 2026 it is clear, whether I like it or not, that it will be yet another year dominated by AI. I have to be honest, looking back, this growing presence is something I’ve had mixed feelings about, having gone through the usual cycles of excitement, epiphany, frustration and despair that countless of us seem to have experienced. It’s been a technical and emotional rollercoaster like discovering, falling in love with, then good old times

Up #1: Before GitHub Copilot, there was GitHub Codex, the predictive autocomplete for VS Code. This was great for those obscure pandas incantations that never seemed to make it into muscle memory. A function signature, maybe a comment or the beginning of the line combined_df = and more likely than not the right completion would show up as ghost text.

Down #1: Over the Christmas break of 2023, I set myself two challenges: use SQL to reason about molecular assembly; and write a Raspberry Pi camera image acquisition GUI to make life easier for my students who had to use the command line to capture images and a text file for documents. With the camera challenge specially Codex seemed completely clueless. This was around the raspicam to libcamera transition and there just didn’t seem to be a way to prompt it to use the more modern stack (in fairness to Codex, there was probably next to nothing in its training data about this).

Up #2: In 2024, when GitHub Copilot’s Edit mode popped up in my editor one day, there was a feeling of immense power as I prompted Claude Sonnet 3.5 to create interactive web UIs that I could never dream of having the time to learn coding on my own. I even prototyped an IDE for probabilistic modelling, even though I have little frontend coding experience.

Down #2: Then the code (and dead code) started to grow, I lost track of how and where things were implemented, what was still in use vs left over from the last prompt. You had to be sure that the right source files were added to context or the model would ignore/reimplement existing functionality, blissfully unaware that it’s right there. Things worked, but not quite the way I wanted them to, and prompting the model to go from working to correct seemed almost as hard as implementing them myself, since I was no longer able to wrap my head around all the react code and API endpoints.

Up #3: The first taste of Claude Code and the incredible feeling of power, entire projects written from scratch, handling git features that I didn’t even know about (worktrees).

Down #3: With more difficult challenges Claude’s tendency to hack its way to success became annoyingly apparent. It would move goalposts, keep celebrating trivial partial successes, growing the number of celebratory emojis used to hype up what it’s achieved whilst in reality it had made no progress, running around making a mess until it would run out of tokens. My last experiment was autoclaude, a script that would call claude -p with a set of instructions that included how to keep track of progress and gave it the task of implementing a WordPress-like content management system with a Python rather than PHP backend. Sonnet 3.7 running some 80 “awakenings” made some good initial progress but it was hard to keep it on track as it kept marking things done that weren’t really done at all.

Draft: Thinking about aerosols probabilistically — Part 1: Evaporation

2026-01-16T00:00:00+00:00

This post is an interactive HTML document. Read it on Hessam's blog.

A little snippet for nicer plots using matplotlib

2026-01-09T00:00:00+00:00

A little snippet for nicer plots using matplotlib

If you’ve seen previous blog posts (or read any of our group papers) you may have noticed the plots have a consistent look.

import seaborn as sns
from matplotlib import pyplot as plt

The default matplotlib look

plt.plot([1,2,5,3])

[<matplotlib.lines.Line2D at 0x11530e960>]

Not bad but try the following now

from matplotlib_inline.backend_inline import set_matplotlib_formats
set_matplotlib_formats('svg')

sns.set_theme('talk', 'ticks', font='Arial', font_scale=1.0, rc={'svg.fonttype': 'none'})

plt.plot([1,2,5,3])

[<matplotlib.lines.Line2D at 0x1154e8ce0>]

An added perk is that the plot is now embedded in your notebook as SVG so if you export to Markdown/HTML they will stay nice and crisp.

Bonus: If you’re a marimo user (and if you’re reading this you have to try it), it’s even easier:

import seaborn as sns

sns.set_theme(
    context="notebook",
    style="ticks",
    font="Inter",
    rc={"svg.fonttype": "none", "savefig.format": "svg"},
)

This has the unintended consequence of making all figures fill up the available space. There’s a little trick to keep things just the right size:

import marimo as mo

def fig(f):
    return mo.as_html(f).style({"width": "max-content", "display": "block"})

fig(plt.gcf())

This is a Jupyter notebook so I can’t show you the output but trust me, it looks good.

3D-printed mathematical lampshades

2025-12-24T00:00:00+00:00

3D-printed mathematical lampshades

I’ve been playing with translucent PLA for lighting projects at home. Earlier this year, I built a diffuser/housing for our Lidl LED strip, since the light from the exposed LEDs was too cool and jarring. The design files are on OnShape, although they might need tweaking for the exact strip and required installation area.

A few viral videos have recently popped up in my YouTube feed showing bespoke 3D-printed lampshades. These ondulating membranes struck me as a surface that would arise from sweeping 2D profile as it’s raised and rotated through space. Making one from scratch seemed like a fun challenge (not to mention a nice Christmas present!).

2D profile

I started by graphing the following equation in Desmos (link will take you to the formula).

r = cos (5θ) + 0.2cos (9θ) + 0.05cos (200θ) + 4

Something like this with big gentle ondulations on top of which small jittery ridges are overlapped.

Trying (and failing) 3D with OnShape

Desmos helpfully allows export to SVG, which I then converted to DXF in Inkscape so it can be imported into OnShape. DXF files can be added to OnShape sketches directly. So far so good. Creating the loft object in OnShape proved challenging though. A loft is pretty much what I described earlier, the volume or surface created as one 2D profile is interpolated to another. In this case, the second profile is just a point placed above the initial one. Something like this:

The loft tool provide a few nice settings, specifically the initial and final angle of approach to each profile. I want the shape to rise straight from the 2D contour and form a tangent to the final point. This way the top will form a nice dome rather than a sharp spike.

This complex organic-looking profile, made up of 1047 points, doesn’t really seem to fit OnShape’s main usecase. The loft operation kept failing to process (although it looks nice) and froze a couple of times too.

At this point, I had another idea. Since the point coordinates are known, it shouldn’t be too hard to create the shape in Python. This is not something I had done before, so I asked ChatGPT for advice. It gave me a script, which I was able to modify (repo here) to create the transition to a single point at the top (essentially scaling the profile depending on height). The results is a massive 70+MB STL file, which is currently printing …

Update: I’ve converted the Jupyter notebook to marimo so it’s easier to visualise the cross section and profile of the generated lampshade. You don’t even need to install Python to run this code, it will run directly in the browser using WebAssembly (link).

Using `uv` projects on a network share

2025-12-17T00:00:00+00:00

Using `uv` projects on a network share

The wondeful uv has a strong preference for virtual environments that are placed within a project tree rather than kept centrally a la conda. One unfortunate consequence of this is for projects stored on network shares, for us usually data analysis scripts and Jupyter/marimo notebooks for visualisation and as dashboards. Simply uv adding numpy alone will dump heaps of tiny files into the folder and completely overwhelm our network share (normally OneDrive since we get 1 TB free as part of 365).

A little trick to avoid this situation, while still being able to use pyproject.toml to specify dependencies via uv. Here it goes.

uv init --bare to initialise your pyproject.toml if it doesn’t already exist.
uv add --no-sync <dependencies>, the key is --no-sync so uv doesn’t make a .venv folder and install packages there.
uvx --with-requirements pyproject.toml <program that you want to run>. Here are a few examples:
- uvx --with-requirements pyproject.toml marimo edit
- uvx --with-requirements pyproject.toml jupyter lab
- uvx --with-requirements pyproject.toml ipython
- uvx --with-requirements pyproject.toml python script.py

You might be tempted to run uv run script.py. Don’t! That’s going to create a virtual environment in the project folder. Either use the last example above or the nifty --isolated flag: uv run --isolated script.py.

And just like that, the most frustration-free Python workflow that I am aware of.

Reparameterizing distributions in numpyro

2025-11-26T00:00:00+00:00

Reparameterizing distributions in numpyro

Every single time I’ve tried to “chance” it with probabilistic methods I am once again reminded that you can’t treat them like a black box. The following is a self-contained example of a pathological case so simple it almost feels like it should just work, and yet it fails miserably without intervention (in this case by reparameterizing the distributions). Part of me wonders if this all just means we need better tools or whether we just have to accept the fundamental complexity and model defensively.

import numpy as np
import pandas as pd
import seaborn as sns
from jax import lax
from jax import numpy as jnp
from jax.random import PRNGKey
from matplotlib import pyplot as plt
from numpyro import deterministic, handlers, sample
from numpyro import distributions as dist
from numpyro.infer import MCMC, NUTS, SVI, Trace_ELBO, autoguide, reparam
from numpyro.optim import Adam
from matplotlib_inline.backend_inline import set_matplotlib_formats

set_matplotlib_formats('svg')

sns.set_theme('talk', 'ticks', font='Arial', font_scale=1.0, rc={'svg.fonttype': 'none'})

Problem setting

Consider a catalyzed reaction $A \xrightarrow{\mathrm{cat}} \cdots$ where the catalyst itself is slowly undergoing decomposition cat → ⋯. Initial concentrations [A]₀ and

$\frac{d [A]}{d t} = -k[A][\mathrm{cat}]$
$\frac{d [\mathrm{cat}]}{d t} = -k_d[\mathrm{cat}]$

A couple of simple helper functions to integrate these differential equations.

def solve(a_0, k, cats, dt):
    return lax.scan(lambda a_n, cat: (a_n - dt * k * cat * a_n, a_n - dt * k * cat * a_n), a_0, cats)[1]

def solve_cat(cat_0, k_d, dt, n):
    return lax.scan(lambda cat_n, _: (cat_n - dt * k_d * cat_n, cat_n - dt * k_d * cat_n), cat_0, length=n)[1]
    

plt.plot(solve(1.0, 1.0, jnp.ones((1000,)), 0.01))
plt.plot(solve_cat(1.0, 1.0, 0.02, 1000))

[<matplotlib.lines.Line2D at 0x1185b9f40>]

Now a simple model, first without worrying about measurement error. We’ll use both MCMC and SVI.

def model(a_0, cat_0, dt, n):
    k = sample('k', dist.LogNormal(0.0, 1.0))
    k_d = sample('k_d', dist.LogNormal(-2.0, 1.0))
    cats = deterministic('cat', solve_cat(cat_0, k_d, dt, n))
    a_n = deterministic('a_n', solve(a_0, k, cats, dt))

args = (1.0, 0.1, 0.01, 500)
mcmc = MCMC(NUTS(model), num_warmup=1000, num_samples=100)
mcmc.run(PRNGKey(0), *args)
mcmc_samples = mcmc.get_samples()

guide = autoguide.AutoNormal(model)
svi = SVI(model, guide, Adam(0.01), Trace_ELBO())
svi_result = svi.run(PRNGKey(1), 2000, *args)
svi_samples = guide.sample_posterior(PRNGKey(0), svi_result.params, sample_shape=(100,))

sample: 100%|██████████| 1100/1100 [00:00<00:00, 1791.63it/s, 3 steps of size 9.31e-01. acc. prob=0.92]
100%|██████████| 2000/2000 [00:00<00:00, 7085.99it/s, init loss: 7.5166, avg. loss [1901-2000]: 0.0118]

f, (a1, a2) = plt.subplots(ncols=2, sharey=True, figsize=(10, 4))
a1.set_title('MCMC')
a2.set_title('SVI')

a_line, *_ = a1.plot(mcmc_samples['a_n'].T, c='b', alpha=0.1, label='[A]')
a_line.set_alpha(1)
cat_line, *_ = a1.plot(mcmc_samples['cat'].T, c='r', alpha=0.1, label='[cat]')
cat_line.set_alpha(1)
a2.plot(svi_samples['a_n'].T, c='b', alpha=0.1)
a2.plot(svi_samples['cat'].T, c='r', alpha=0.1)
f.legend(handles=[a_line, cat_line])
a_line.set_alpha(0.1)
cat_line.set_alpha(0.1)

Very nice, we get a range of [A] trajectories based on possible k’s but also k_d’s. But now let’s just sample simulate adding measurement error as a LogNormal. The previous variables shouldn’t be affected because we are not making an observation. What’s notable about this a_draw variable is that it has a very narrow distribution (σ = 0.02).

def model(a_0, cat_0, dt, n, err):
    k = sample('k', dist.LogNormal(0.0, 1.0))
    k_d = sample('k_d', dist.LogNormal(-2.0, 1.0))
    cats = deterministic('cat', solve_cat(cat_0, k_d, dt, n))
    a_n = deterministic('a_n', solve(a_0, k, cats, dt))
    a_draw = sample('a_draw', dist.LogNormal(jnp.log(a_n), err))

args = (1.0, 0.1, 0.01, 500, 0.05)
mcmc = MCMC(NUTS(model), num_warmup=1000, num_samples=100)
mcmc.run(PRNGKey(0), *args)
mcmc_samples = mcmc.get_samples()

guide = autoguide.AutoNormal(model)
svi = SVI(model, guide, Adam(0.01), Trace_ELBO())
svi_result = svi.run(PRNGKey(1), 2000, *args)
svi_samples = guide.sample_posterior(PRNGKey(0), svi_result.params, sample_shape=(100,))

sample: 100%|██████████| 1100/1100 [00:02<00:00, 549.16it/s, 63 steps of size 4.88e-02. acc. prob=0.92]
100%|██████████| 2000/2000 [00:00<00:00, 3172.51it/s, init loss: 139124.0000, avg. loss [1901-2000]: 26.2503]

f, (a1, a2) = plt.subplots(ncols=2, sharey=True, figsize=(10, 4))
a1.set_title('MCMC')
a2.set_title('SVI')

a_line, *_ = a1.plot(mcmc_samples['a_n'].T, c='b', alpha=0.1, label='[A]')
a_line.set_alpha(1)
cat_line, *_ = a1.plot(mcmc_samples['cat'].T, c='r', alpha=0.1, label='[cat]')
cat_line.set_alpha(1)
meas_line, *_ = a1.plot(mcmc_samples['a_draw'].T, c='k', alpha=0.02, label='[A] measured')
meas_line.set_alpha(1)
a2.plot(svi_samples['a_n'].T, c='b', alpha=0.1)
a2.plot(svi_samples['cat'].T, c='r', alpha=0.1)
a2.plot(svi_samples['a_draw'].T, c='k', alpha=0.05)
f.legend(handles=[a_line, cat_line, meas_line])
a_line.set_alpha(0.1)
cat_line.set_alpha(0.1)
meas_line.set_alpha(0.02)

Yes, not looking good at all. We were expecting the exact sample plots! I tried increasing the number of samples, warmup and SVI steps and those didn’t help either. The culprit is clearly the narrow distribution of a_draw as increasing its std to 1.0 everything goes back to normal.

After a bit of head scratching, it turns out there is a way to fix this without rewriting the model and manually de-centering/scaling. This happens on the level of the Normal distributions being sampled with (μ = 0 and σ = 1.0) but the LogNormal distribution is a further exponential transform away from so that’s where TransformReparam comes in.

reparam_config = {
    'k': reparam.TransformReparam(),
    'k_d': reparam.TransformReparam(),
    'a_draw': reparam.TransformReparam(),
    'k_base': reparam.LocScaleReparam(0),
    'k_d_base': reparam.LocScaleReparam(0),
    'a_draw_base': reparam.LocScaleReparam(0),
}

def model(a_0, cat_0, dt, n, err):
    k = sample('k', dist.LogNormal(0.0, 1.0))
    k_d = sample('k_d', dist.LogNormal(-2.0, 1.0))
    cats = deterministic('cat', solve_cat(cat_0, k_d, dt, n))
    a_n = deterministic('a_n', solve(a_0, k, cats, dt))
    a_draw = sample('a_draw', dist.LogNormal(jnp.log(a_n), err))

model = handlers.reparam(model, reparam_config)

args = (1.0, 0.1, 0.01, 500, 0.05)
mcmc = MCMC(NUTS(model), num_warmup=1000, num_samples=100)
mcmc.run(PRNGKey(0), *args)
mcmc_samples = mcmc.get_samples()

guide = autoguide.AutoNormal(model)
svi = SVI(model, guide, Adam(0.01), Trace_ELBO())
svi_result = svi.run(PRNGKey(1), 2000, *args)
svi_samples = guide.sample_posterior(PRNGKey(0), svi_result.params, sample_shape=(100,))

sample: 100%|██████████| 1100/1100 [00:00<00:00, 1576.36it/s, 15 steps of size 3.03e-01. acc. prob=0.89]
100%|██████████| 2000/2000 [00:00<00:00, 4285.20it/s, init loss: 1233.4703, avg. loss [1901-2000]: 2.5142]

f, (a1, a2) = plt.subplots(ncols=2, sharey=True, figsize=(10, 4))
a1.set_title('MCMC')
a2.set_title('SVI')

a_line, *_ = a1.plot(mcmc_samples['a_n'].T, c='b', alpha=0.1, label='[A]')
a_line.set_alpha(1)
cat_line, *_ = a1.plot(mcmc_samples['cat'].T, c='r', alpha=0.1, label='[cat]')
cat_line.set_alpha(1)
meas_line, *_ = a1.plot(mcmc_samples['a_draw'].T, c='k', alpha=0.05, label='[A] measured')
meas_line.set_alpha(1)
a2.plot(svi_samples['a_n'].T, c='b', alpha=0.1)
a2.plot(svi_samples['cat'].T, c='r', alpha=0.1)
a2.plot(svi_samples['a_draw'].T, c='k', alpha=0.05)
f.legend(handles=[a_line, cat_line, meas_line])
a_line.set_alpha(0.1)
cat_line.set_alpha(0.1)
meas_line.set_alpha(0.02)

Beautiful! The latent concentrations are now exactly as before and we also have a nice extra variable showing measurement with error. How about a quick look at the sampled values of k and k_d, which should ideally be uncorrelated and have a lognormal marginal, i.e. each is lognormally distributed irrespective of the other one.

def posterior_df(samples):
    dfs = []

    for var, vals in samples.items():
        dfs.append(pd.DataFrame(samples[var]).melt().rename(columns={"value": var})[[var]])
    
    return pd.concat([dfs[0], *[df.iloc[:, 0] for df in dfs[1:]]], axis=1)

sns.jointplot(posterior_df(mcmc_samples), x='k', y='k_d', kind='kde', fill=True)
plt.gca().set(ylabel='$k_d$', xlabel='$k$')
plt.gca().set_xlim(left=0.0)
plt.gca().set_ylim(bottom=0.0);

Nice, normal as expected!

You will notice a quirk: I had to know that the name of the underlying Normal sites. And one downside of using this automatic reparameterisation is that TransformReparam doesn’t currently support observations. That’s fine though, you can split the sampled and observed parts (the observed part won’t need reparameterising as it won’t need to be sampled).

Adding support for observations

Visualising the prior predictive, i.e. what’s predicted solely from our priors without observing anything is nice but not that exciting. It took some trial and error to find the right combination of types and effect handlers to do this but the solution is to use a masked distribution. Here is how it works:

For any unobserved values in a_draw that need imputing (i.e. drawing from the prior predictive rather than observing), we can continue to use our reparamterized TransformedDistribution, but masked so the log probability of the any observed sites is ignored. Essentially, you are still sampling them but not accounting for their log probability, since we are later going to observe them.
For any observed values in a_draw_obs, reparameterizing is not necessary, since we are not going to sample them. We apply the opposite mask: observed values simply add to our overall joint log probability without having to be sampled; unobserved values are also observed from the imputed values but their log probs are ignored because we mask them. This is because we already calculated their log prob in the a_draw site.

I have to confess that this took a while to get right. Especially because some of the types don’t compose very well, so for example TransformedDistribution(MaskedDistribution(...)) is not supported but using a mask handler works.

reparam_config = {
    'k': reparam.TransformReparam(),
    'k_d': reparam.TransformReparam(),
    'a_draw': reparam.TransformReparam(),
    'k_base': reparam.LocScaleReparam(0),
    'k_d_base': reparam.LocScaleReparam(0),
    'a_draw_base': reparam.LocScaleReparam(0),
}

def model(a_0, cat_0, dt, n, err, obs):
    has_obs = ~jnp.isnan(obs)
    k = sample('k', dist.LogNormal(0.0, 1.0))
    k_d = sample('k_d', dist.LogNormal(-2.0, 1.0))
    cats = deterministic('cat', solve_cat(cat_0, k_d, dt, n))
    a_n = deterministic('a_n', solve(a_0, k, cats, dt))
    with handlers.mask(mask=~has_obs):
        d = dist.Normal(jnp.log(a_n), err)
    a_draw = sample('a_draw', dist.TransformedDistribution(d, dist.transforms.ExpTransform()))
    obs = jnp.where(has_obs, obs, a_draw)
    a_draw_combined = deterministic('a_draw_combined', obs)
    a_draw_obs = sample('a_draw_obs', dist.LogNormal(jnp.log(a_n), err).mask(has_obs), obs=obs)

def infer(model, args, guide=None):
    mcmc = MCMC(NUTS(model), num_warmup=2000, num_samples=100)
    mcmc.run(PRNGKey(0), *args)
    mcmc_samples = mcmc.get_samples()

    guide = guide or autoguide.AutoNormal(model)
    svi = SVI(model, guide, Adam(0.01), Trace_ELBO())
    svi_result = svi.run(PRNGKey(1), 10000, *args)
    svi_samples = guide.sample_posterior(PRNGKey(0), svi_result.params, sample_shape=(100,))
    return locals()


model = handlers.reparam(model, reparam_config)

obs = np.full((500,), np.nan)
args = (1.0, 0.1, 0.01, 500, 0.05, obs)

results = infer(model, args)

sample: 100%|██████████| 2100/2100 [00:01<00:00, 1359.92it/s, 15 steps of size 3.03e-01. acc. prob=0.85]
100%|██████████| 10000/10000 [00:01<00:00, 5709.20it/s, init loss: 1233.4703, avg. loss [9501-10000]: 2.4468]

Sampling from the prior predictive, i.e. when the observations are all np.nan.

def plot_samples(mcmc_samples, svi_samples, obs):
    f, (a1, a2) = plt.subplots(ncols=2, sharey=True, figsize=(12, 5))
    a1.set_title('MCMC')
    a2.set_title('SVI')
    has_obs = ~jnp.isnan(obs)


    a_line, *_ = a1.plot(mcmc_samples['a_n'].T, c='b', alpha=0.1, label='[A]')
    a_line.set_alpha(1)
    cat_line, *_ = a1.plot(mcmc_samples['cat'].T, c='r', alpha=0.1, label='[cat]')
    cat_line.set_alpha(1)
    meas_line, *_ = a1.plot(mcmc_samples['a_draw_combined'].T, c='y', alpha=0.05, label='[A] measured', zorder=-1)
    meas_line.set_alpha(1)
    a2.plot(svi_samples['a_n'].T, c='b', alpha=0.1)
    a2.plot(svi_samples['cat'].T, c='r', alpha=0.1)
    a2.plot(svi_samples['a_draw_combined'].T, c='y', alpha=0.05, zorder=-1)
    f.legend(handles=[a_line, cat_line, meas_line])
    a1.scatter(jnp.nonzero(has_obs)[0], obs[has_obs], c='k', zorder=10)
    a2.scatter(jnp.nonzero(has_obs)[0], obs[has_obs], c='k', zorder=10)
    a1.set(ylim=(-0.1,1.1))
    a_line.set_alpha(0.1)
    cat_line.set_alpha(0.1)
    meas_line.set_alpha(0.05)
    return f

plot_samples(results['mcmc_samples'], results['svi_samples'], obs)
plt.tight_layout()

Now let’s supply a single observation, specifically one that suggests the reaction is faster than predicted by most trajectories. In the posterior we would expect to see either a higher than anticipated k or a lower than anticipated k_d and the two should be positively correlated.

obs[100] = 0.6

args = (1.0, 0.1, 0.01, 500, 0.05, obs)

results = infer(model, args)

plot_samples(results['mcmc_samples'], results['svi_samples'], obs)
plt.tight_layout()

sample: 100%|██████████| 2100/2100 [00:01<00:00, 1255.81it/s, 15 steps of size 3.08e-01. acc. prob=0.85]
100%|██████████| 10000/10000 [00:01<00:00, 5971.33it/s, init loss: 1234.0802, avg. loss [9501-10000]: 3.8389]

all_data = pd.concat([
    posterior_df(results['mcmc_samples'])[['k', 'k_d']].assign(Method='MCMC'),
    posterior_df(results['svi_samples'])[['k', 'k_d']].assign(Method='SVI'),
    ], ignore_index=True)
sns.kdeplot(all_data, x='k', y='k_d', hue="Method", fill=True)
plt.gca().set(xlabel='$k$', ylabel='$k_d$')

[Text(0.5, 0, '$k$'), Text(0, 0.5, '$k_d$')]

And here we go, a very different picture now, and as expected the AutoNormal SVI guide fails to capture the correlation between k and k_d. A full rank multivariate normal should be able to handle this, albeit taking about 10x as long as MCMC. For some reason autoguide.AutoLowRankMultivariateNormal(rank=2) doesn’t seem to do the trick.

A more advanced option to get the best of both worlds would be creating a separate full-rank guide for k and k_d then using a simple AutoNormal for the rest (specifically a_draw) using AutoGuideList.

# "Easy" but slow and doesn't really scale
# results = infer(model, args, autoguide.AutoMultivariateNormal(model))

# More involved but very fast
guide = autoguide.AutoGuideList(model)
guide.append(autoguide.AutoNormal(handlers.block(model, expose=['a_draw_base_decentered'])))
guide.append(autoguide.AutoMultivariateNormal(handlers.block(model, expose=['k_base_decentered', 'k_d_base_decentered'])))
results = infer(model, args, guide)

sample: 100%|██████████| 2100/2100 [00:01<00:00, 1313.56it/s, 15 steps of size 3.08e-01. acc. prob=0.85]
100%|██████████| 10000/10000 [00:01<00:00, 5945.90it/s, init loss: 1272.4036, avg. loss [9501-10000]: 3.7564]

This is more involved and it seems like AutoGuideList doesn’t deal with deterministic sites, so we have derive them manually. I had to read a lot of numpyro’s source code to figure out how.

from numpyro.infer.util import soft_vmap

def predictive(sample):
    with handlers.seed(rng_seed=PRNGKey(0)):
        with handlers.substitute(data=sample):
            return {k: v['value'] for k, v in handlers.trace(model).get_trace(*args).items()}

results['svi_samples'].update(soft_vmap(predictive, results['svi_samples']))

all_data = pd.concat([
    posterior_df(results['mcmc_samples'])[['k', 'k_d']].assign(Method='MCMC'),
    posterior_df(results['svi_samples'])[['k', 'k_d']].assign(Method='SVI'),
    ], ignore_index=True)
sns.kdeplot(all_data, x='k', y='k_d', hue="Method", fill=True)
plt.gca().set(xlabel='$k$', ylabel='$k_d$')

[Text(0.5, 0, '$k$'), Text(0, 0.5, '$k_d$')]

And here we go!

plot_samples(results['mcmc_samples'], results['svi_samples'], obs)
plt.tight_layout()

I hope this was fun to explore together. Very happy to chat if you have any suggestions or questions. This experiment was motivated by my upcoming RSC book on digital chemistry, as part of which I am hoping to present probabilistic programming as a powerful data interpretation aid for experimentalists working with automated platforms that can generate a wealth of experimental data.

Live variational inference

2025-10-15T00:00:00+00:00

Live variational inference

This is mostly to show two possibilities:

The ability to update the data going into a model in the course of running variational inference. Imagine you are accumulating data as you go, or (this is more technical/hacky) you find that your multimodal distributions get locked into a specific mode that the optimiser cannot escape and are looking for a way to “ease into” the right mode by e.g. tightening your Dirichlet concentration factor gradually as you go.
Inspect model predictions and loss in real time.

The key is using the SVI.update call, which I have avoided so far for the convenience of SVI.run. A batch of steps can be take in a jitted loop for performance and the loss monitored for convergence every time a parameter/input data is changed.

With matplotlib you need to use an interactive backend like qt or osx and add a plt.pause(...) call in the loop to make sure there is a chance to re-render the figure. Super useful and something I didn’t know about (thanks ChatGPT).

A minimal example below, where the concentration parameter for a Dirichlet prior to a categorical observation is varied smoothly from 1.0 (i.e. flat) to 0.1 (pretty pointy at the extremes).

%matplotlib osx

UsageError: Invalid GUI request 'macosx', valid ones are:dict_keys(['inline', 'nbagg', 'webagg', 'notebook', 'ipympl', 'widget', None, 'qt', 'qt5', 'qt6', 'wx', 'tk', 'gtk', 'gtk3', 'osx', 'asyncio'])

import jax
import jax.numpy as jnp
import numpy as np
import numpyro
import numpyro.distributions as dist
import seaborn as sns
from matplotlib import pyplot as plt
from numpyro.infer import SVI, Trace_ELBO, autoguide
from numpyro.optim import Adam
from tqdm import tqdm

sns.set_theme(
    "talk", "ticks", font="Arial", font_scale=1.0, rc={"svg.fonttype": "none"}
)

num_steps = 100


def model(alpha_scale):
    alpha = jnp.ones(12) * alpha_scale
    theta = numpyro.sample("theta", dist.Dirichlet(alpha))
    numpyro.sample("obs1", dist.Categorical(theta), obs=jnp.array([2,2,2,2,2,5,5,5]))

guide = autoguide.AutoNormal(model)

optimizer = Adam(1e-3)
svi = SVI(model, guide, optimizer, loss=Trace_ELBO())

rng_key = jax.random.PRNGKey(0)
state = svi.init(
    rng_key,
    alpha_scale=1.0,
)


@jax.jit
def run_stage(state, init_loss, alpha_scale):
    def body_fn(i, val):
        return svi.update(val[0], alpha_scale)

    return jax.lax.fori_loop(0, 100, body_fn, (state, init_loss))


f, (a1, a2) = plt.subplots(nrows=2, sharex=True)
a1.set(ylabel="ELBO loss")
a2.set(xlabel="SVI batch", ylabel=r"$\theta$ (posterior)")
f.tight_layout()
for i in tqdm(range(num_steps)):
    alpha_scale = 0.02 ** (i / num_steps)
    a1.set_title(f"$\\alpha$: {alpha_scale:.2f}")
    while True:
        # Run SVI steps until convergence
        init_loss = svi.evaluate(state, alpha_scale)
        state, loss = run_stage(state, init_loss, alpha_scale)
        if jnp.abs(loss - init_loss) / np.abs(init_loss) < 0.02:
            break
    params = svi.get_params(state)
    a1.scatter([i], loss, c="k", s=5)
    posterior = guide.sample_posterior(rng_key, params)
    theta = posterior["theta"]
    a2.scatter(i * np.ones_like(theta), theta, c=np.arange(12), s=5)
    # This is really important to get live updates.
    plt.pause(0.01)

f.savefig("2025-10-15-live-svi_result.svg")
print("Posterior theta mean:", posterior["theta"])
plt.close()

100%|██████████| 100/100 [00:06<00:00, 14.36it/s]


Posterior theta mean: [1.26911415e-08 7.69821611e-08 6.65301085e-01 7.23654193e-06
 1.26773830e-06 2.86240101e-01 2.82157103e-10 9.47567692e-04
 2.23408958e-10 4.75026183e-02 6.67153069e-11 5.59603919e-09]

Not sure if the result will show up as expected when this notebook is converted to HTML but it looks very nice and it’s good to be able to see it converge.

Amortized probabilistic models for chemical microscopy

2025-07-21T00:00:00+00:00

Amortized probabilistic models for chemical microscopy

Can we used an amortized model to speed up inference in our droplet microscopy model?

My last go using a probabilistic model to analyze a microscope image seemed to work well enough, but I wanted to take a more flexible approach to modelling the appearance of droplets without having to roll out a more sophisticated physical model. Also, it seemed impractical to require a beefy GPU and minutes of compute for a single image.

I’ve been meaning to experiment with amortized inference a bit more recently. The idea is that instead of all latent variables being inferred, a small model (think a miniature multi-layer perceptron) is trained to predict a subset of these variables. We are interested in inferring droplet locations and compositions, not so much about rendering the droplets themselves. This approach can give the best of both worlds, a mechanistic interpretable model for the parts that we care about or are easy to reason about, and an data-driven, learned representation for the complex, inherently introspectable parts.

import jax
import jax.numpy as jnp
import jax.nn as jnn
import flax.linen as nn
import matplotlib.pyplot as plt
import numpy as np
import numpyro.distributions as dist
import seaborn as sns
from numpyro import deterministic, plate, sample
from numpyro.handlers import seed, trace, substitute
from numpyro.infer import SVI, Trace_ELBO, MCMC, NUTS
from numpyro.infer.autoguide import AutoNormal
from numpyro.optim import Adam
from PIL import Image

plt.rcParams['figure.dpi'] = 200

sns.set_theme(context='paper', style='ticks', font='Arial')

We’ll use the same delightful microscope image as last time, part of the experiments that went into our latest paper.

img = Image.open('data/example.jpg')
img = img.resize((img.width // 4, img.height // 4))
img

img = np.array(img) / 255.0

class DropletOpticsModel(nn.Module):
    hidden_dims: tuple = (32, 16, 8)
    
    @nn.compact
    def __call__(self, background, dx, dy, radius, composition):
        """
        Args:
            background: (batch, n_channels) - existing/background pixel values
            dx: (batch, 1) - normalized distance from droplet center in x-direction
            dy: (batch, 1) - normalized distance from droplet center in y-direction
            radius: (batch, 1) - droplet radius
            composition: (batch, n_composition_features) - droplet composition vector
        
        Returns:
            new_pixel_value: (batch, n_channels) - predicted new pixel values
        """
        # Concatenate all input features
        x = jnp.concatenate([background, dx, dy, radius, composition], axis=-1)

        for i, dim in enumerate(self.hidden_dims):
            x = nn.Dense(dim, name=f'hidden_{i}')(x)
            x = nn.LayerNorm(name=f'ln_{i}')(x)
            x = jnn.relu(x)
        
        # Output layer - predict change in pixel values
        delta = nn.Dense(background.shape[-1], name='output')(x)
        
        # Add residual connection and apply sigmoid to keep values in [0, 1]
        new_pixel_value = jnn.sigmoid(background + delta)
        
        return new_pixel_value

# Initialize model
model = DropletOpticsModel()

# Test with dummy data
key = jax.random.PRNGKey(42)
batch_size, n_channels, n_composition = 32, 3, 10

dummy_background = jax.random.uniform(key, (batch_size, n_channels))
dummy_dx = jax.random.uniform(key, (batch_size, 1), minval=-1.0, maxval=1.0)
dummy_dy = jax.random.uniform(key, (batch_size, 1), minval=-1.0, maxval=1.0)
dummy_radius = jax.random.uniform(key, (batch_size, 1)) * 0.1  # Small radii
dummy_composition = jax.random.uniform(key, (batch_size, n_composition))

# Initialize parameters
print(f"{dummy_background.shape=}, {dummy_dx.shape=}, {dummy_dy.shape=}, {dummy_radius.shape=}, {dummy_composition.shape=}")
params = model.init(key, dummy_background, dummy_dx, dummy_dy, dummy_radius, dummy_composition)

# Test forward pass
output = model.apply(params, dummy_background, dummy_dx, dummy_dy, dummy_radius, dummy_composition)
print(f"Input background shape: {dummy_background.shape}")
print(f"Output pixel values shape: {output.shape}")
print(f"Output range: [{output.min():.3f}, {output.max():.3f}]")

# Print model summary
print(f"\nModel summary:")
print(model.tabulate(key, dummy_background, dummy_dx, dummy_dy, dummy_radius, dummy_composition))

dummy_background.shape=(32, 3), dummy_dx.shape=(32, 1), dummy_dy.shape=(32, 1), dummy_radius.shape=(32, 1), dummy_composition.shape=(32, 10)
Input background shape: (32, 3)
Output pixel values shape: (32, 3)
Output range: [0.260, 0.897]

Model summary:

                           DropletOpticsModel Summary                           
┏━━━━━━━━━━┳━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━┓
┃ path     ┃ module         ┃ inputs         ┃ outputs        ┃ params         ┃
┡━━━━━━━━━━╇━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━┩
│          │ DropletOptics… │ -              │ float32[32,3]  │                │
│          │                │ float32[32,3]  │                │                │
│          │                │ -              │                │                │
│          │                │ float32[32,1]  │                │                │
│          │                │ -              │                │                │
│          │                │ float32[32,1]  │                │                │
│          │                │ -              │                │                │
│          │                │ float32[32,1]  │                │                │
│          │                │ -              │                │                │
│          │                │ float32[32,10] │                │                │
├──────────┼────────────────┼────────────────┼────────────────┼────────────────┤
│ hidden_0 │ Dense          │ float32[32,16] │ float32[32,32] │ bias:          │
│          │                │                │                │ float32[32]    │
│          │                │                │                │ kernel:        │
│          │                │                │                │ float32[16,32] │
│          │                │                │                │                │
│          │                │                │                │ 544 (2.2 KB)   │
├──────────┼────────────────┼────────────────┼────────────────┼────────────────┤
│ ln_0     │ LayerNorm      │ float32[32,32] │ float32[32,32] │ bias:          │
│          │                │                │                │ float32[32]    │
│          │                │                │                │ scale:         │
│          │                │                │                │ float32[32]    │
│          │                │                │                │                │
│          │                │                │                │ 64 (256 B)     │
├──────────┼────────────────┼────────────────┼────────────────┼────────────────┤
│ hidden_1 │ Dense          │ float32[32,32] │ float32[32,16] │ bias:          │
│          │                │                │                │ float32[16]    │
│          │                │                │                │ kernel:        │
│          │                │                │                │ float32[32,16] │
│          │                │                │                │                │
│          │                │                │                │ 528 (2.1 KB)   │
├──────────┼────────────────┼────────────────┼────────────────┼────────────────┤
│ ln_1     │ LayerNorm      │ float32[32,16] │ float32[32,16] │ bias:          │
│          │                │                │                │ float32[16]    │
│          │                │                │                │ scale:         │
│          │                │                │                │ float32[16]    │
│          │                │                │                │                │
│          │                │                │                │ 32 (128 B)     │
├──────────┼────────────────┼────────────────┼────────────────┼────────────────┤
│ hidden_2 │ Dense          │ float32[32,16] │ float32[32,8]  │ bias:          │
│          │                │                │                │ float32[8]     │
│          │                │                │                │ kernel:        │
│          │                │                │                │ float32[16,8]  │
│          │                │                │                │                │
│          │                │                │                │ 136 (544 B)    │
├──────────┼────────────────┼────────────────┼────────────────┼────────────────┤
│ ln_2     │ LayerNorm      │ float32[32,8]  │ float32[32,8]  │ bias:          │
│          │                │                │                │ float32[8]     │
│          │                │                │                │ scale:         │
│          │                │                │                │ float32[8]     │
│          │                │                │                │                │
│          │                │                │                │ 16 (64 B)      │
├──────────┼────────────────┼────────────────┼────────────────┼────────────────┤
│ output   │ Dense          │ float32[32,8]  │ float32[32,3]  │ bias:          │
│          │                │                │                │ float32[3]     │
│          │                │                │                │ kernel:        │
│          │                │                │                │ float32[8,3]   │
│          │                │                │                │                │
│          │                │                │                │ 27 (108 B)     │
├──────────┼────────────────┼────────────────┼────────────────┼────────────────┤
│          │                │                │          Total │ 1,347 (5.4 KB) │
└──────────┴────────────────┴────────────────┴────────────────┴────────────────┘
                                                                                
                        Total Parameters: 1,347 (5.4 KB)

What ended up working in the end is when each pixel refers to its closest droplet for colour. Not ideal but it keeps the memory requirement manageable.

def tree_to_dists(tree, path=''):
    if isinstance(tree, dict):
        return {k: tree_to_dists(v, path + '/' + k) for k, v in tree.items()}
    else:
        # print(f"Sampling {path} with shape {tree.shape}")
        return sample(path, dist.Normal().expand(tree.shape))

def model_with_nn(img, n_droplets, types=10):
    h, w, n_channels = img.shape
    
    # Create coordinate grids
    y_coords, x_coords = jnp.mgrid[:h, :w]
    
    # Sample background per channel
    with plate("channels", n_channels):
        bg = sample("bg", dist.Uniform(0, 1).expand((n_channels,)))

    # Sample droplet parameters
    with plate("droplets", n_droplets):
        x = sample("x", dist.Uniform(0, w))
        y = sample("y", dist.Uniform(0, h))
        r = sample("r", dist.LogNormal(0, 0.5))
        with plate("types", types):
            composition = sample("composition", dist.Uniform(0, 1)).T

    model = DropletOpticsModel()

    # Initialize background image
    prediction = jnp.broadcast_to(bg, (h, w, n_channels))
    nn_params = model.init(key, 
                           jnp.zeros((h * w, n_channels)), 
                           jnp.zeros((h * w, 1)),
                           jnp.zeros((h * w, 1)),
                           jnp.zeros((h * w, 1)),
                           jnp.zeros((h * w, types)))    
    nn_params = tree_to_dists(nn_params, path='nn_params')

    distance = ((x_coords[..., None] - x) / r)**2 + ((y_coords[..., None] - y) / r)**2
    nearest = jnp.argmin(distance, axis=-1)

    # Calculate relative distances from droplet center
    dx = (x_coords - x[nearest]) / r[nearest]  # Normalized by radius
    dy = (y_coords - y[nearest]) / r[nearest]  # Normalized by radius


    # Flatten spatial dimensions for neural network processing
    dx_flat = dx.flatten()[:, None]
    dy_flat = dy.flatten()[:, None]
    r_flat = r[nearest].flatten()[:, None]

    
    # Repeat background and composition for all pixels
    bg_flat = prediction.reshape(-1, n_channels)
    comp_flat = composition[nearest, :].reshape(-1, types)
    
    # Apply neural network to get new pixel values
    prediction = model.apply(nn_params, bg_flat, dx_flat, dy_flat, r_flat, comp_flat)
    
    # Reshape back to image dimensions
    prediction = prediction.reshape(h, w, n_channels)
    
    prediction = jnp.clip(prediction, 0, 1)
    prediction = deterministic('prediction', prediction)
    diff = deterministic('diff', img - prediction)
    sample('obs', dist.Normal(scale=0.05), obs=diff)
    # print(f"{x_coords.shape=}, {y_coords.shape=}, {x.shape=}, {y.shape=}, {r.shape=}, {composition.shape=}, {distance.shape=}, {nearest.shape=}, {dx.shape=}, {dy.shape=}, {dx_flat.shape=}, {dy_flat.shape=}, {r_flat.shape=}, {bg_flat.shape=}, {comp_flat.shape=}, {prediction.shape=}")
    return nn_params

tr = trace(seed(model_with_nn, 0)).get_trace(img, 1000, types=5)
nn_params = seed(model_with_nn, 0)(img, 1000, types=5)
{k: v['value'].shape for k, v in tr.items() if 'value' in v}

{'channels': (3,),
 'bg': (3,),
 'droplets': (1000,),
 'x': (1000,),
 'y': (1000,),
 'r': (1000,),
 'types': (5,),
 'composition': (5, 1000),
 'nn_params/params/hidden_0/kernel': (11, 32),
 'nn_params/params/hidden_0/bias': (32,),
 'nn_params/params/ln_0/scale': (32,),
 'nn_params/params/ln_0/bias': (32,),
 'nn_params/params/hidden_1/kernel': (32, 16),
 'nn_params/params/hidden_1/bias': (16,),
 'nn_params/params/ln_1/scale': (16,),
 'nn_params/params/ln_1/bias': (16,),
 'nn_params/params/hidden_2/kernel': (16, 8),
 'nn_params/params/hidden_2/bias': (8,),
 'nn_params/params/ln_2/scale': (8,),
 'nn_params/params/ln_2/bias': (8,),
 'nn_params/params/output/kernel': (8, 3),
 'nn_params/params/output/bias': (3,),
 'prediction': (380, 507, 3),
 'diff': (380, 507, 3),
 'obs': (380, 507, 3)}

guide = AutoNormal(model_with_nn)
svi = SVI(model_with_nn, guide, Adam(0.01), Trace_ELBO())

svi_result = svi.run(jax.random.PRNGKey(0), 100000, img, 800, types=3)
samples_svi = guide.sample_posterior(jax.random.PRNGKey(0), svi_result.params, sample_shape=(100,))
fig, ax = plt.subplots(figsize=(5, 2))
ax.plot(svi_result.losses)

100%|██████████| 100000/100000 [08:35<00:00, 193.87it/s, init loss: 6371727.5000, avg. loss [95001-100000]: -854501.1250]





[<matplotlib.lines.Line2D at 0x721bf0302960>]

plt.imshow(samples_svi['prediction'].mean(axis=0))

<matplotlib.image.AxesImage at 0x721cd022ff80>

Not a bad reconstruction, and this time we capture color as well.

samples_svi['composition'][:1].shape

(1, 3, 800)

plt.imshow(img)

plt.scatter(samples_svi['x'][0], samples_svi['y'][0], s=4, alpha=1.0, c=samples_svi['composition'][0].T, marker='x')

<matplotlib.collections.PathCollection at 0x721c10238f80>

Droplet composition seems nicely consistent with the image.

Other attempts

A couple of other approaches that didn’t quite work.

Fully flattened model with aggregation

Two issues: memory use and how to aggregate the results at the end.

def model_with_nn(img, n_droplets, types=10):
    h, w, n_channels = img.shape
    
    # Create coordinate grids
    y_coords, x_coords = jnp.mgrid[:h, :w]
    
    # Sample background per channel
    with plate("channels", n_channels):
        bg = sample("bg", dist.Uniform(0, 1).expand((n_channels,)))

    # Sample droplet parameters
    with plate("droplets", n_droplets):
        x = sample("x", dist.Uniform(0, w))
        y = sample("y", dist.Uniform(0, h))
        r = sample("r", dist.LogNormal(0, 0.5))
        with plate("types", types):
            composition = sample("composition", dist.Uniform(0, 1)).T

    model = DropletOpticsModel()

    # Initialize background image
    bg = jnp.broadcast_to(bg, (h, w, n_channels))
    nn_params = model.init(key, 
                           jnp.zeros((h * w, n_channels)), 
                           jnp.zeros((h * w, 1)),
                           jnp.zeros((h * w, 1)),
                           jnp.zeros((h * w, 1)),
                           jnp.zeros((h * w, types)))    
    nn_params = tree_to_dists(nn_params, path='nn_params')
    
    # Calculate relative distances from droplet center
    dx = (x_coords[..., None] - x) / r  # Normalized by radius
    dy = (y_coords[..., None] - y) / r  # Normalized by radius

    # Flatten spatial dimensions for neural network processing
    dx_flat = dx.flatten()[:, None]
    dy_flat = dy.flatten()[:, None]
    r_flat = jnp.broadcast_to(r, (h, w, n_droplets)).flatten()[:, None]
    
    # Repeat background and composition for all pixels
    bg_flat = jnp.broadcast_to(bg[:, :, None, :], (h, w, n_droplets, n_channels)).reshape(-1, n_channels)
    comp_flat = jnp.broadcast_to(composition, (h * w, n_droplets, types)).reshape(-1, types)
    
    # Apply neural network to get new pixel values
    new_pixels = model.apply(nn_params, bg_flat, dx_flat, dy_flat, r_flat, comp_flat)
    
    # Reshape back to image dimensions
    new_pixels = new_pixels.reshape(h, w, n_channels)
    
    # Update prediction (could be additive or replacement - using replacement here)
    prediction = new_pixels
    
    prediction = jnp.clip(prediction, 0, 1)
    prediction = deterministic('prediction', prediction)
    diff = deterministic('diff', img - prediction)
    sample('obs', dist.Normal(scale=0.05), obs=diff)

Iterative with `lax.scan`

def model_with_nn(img, n_droplets, types=10):
    h, w, n_channels = img.shape
    
    # Create coordinate grids
    y_coords, x_coords = jnp.mgrid[:h, :w]
    
    # Sample background per channel
    with plate("channels", n_channels):
        bg = sample("bg", dist.Uniform(0, 1).expand((n_channels,)))

    # Sample droplet parameters
    with plate("droplets", n_droplets):
        x = sample("x", dist.Uniform(0, w))
        y = sample("y", dist.Uniform(0, h))
        r = sample("r", dist.LogNormal(0, 0.5))
        with plate("types", types):
            composition = sample("composition", dist.Uniform(0, 1))

    model = DropletOpticsModel()

    # Initialize background image
    prediction = jnp.broadcast_to(bg, (h, w, n_channels))
    nn_params = model.init(key, 
                           jnp.zeros((h * w, n_channels)), 
                           jnp.zeros((h * w, 1)),
                           jnp.zeros((h * w, 1)),
                           jnp.zeros((h * w, 1)),
                           jnp.zeros((h * w, types)))    
    nn_params = tree_to_dists(nn_params, path='nn_params')
    # For each droplet, compute its effect using the neural network
    def apply_droplet(carry, droplet_params):
        current_prediction = carry
        x_i, y_i, r_i, comp_i = droplet_params
        
        # Calculate relative distances from droplet center
        dx = (x_coords - x_i) / r_i  # Normalized by radius
        dy = (y_coords - y_i) / r_i  # Normalized by radius
        
        # Flatten spatial dimensions for neural network processing
        dx_flat = dx.flatten()[:, None]
        dy_flat = dy.flatten()[:, None]
        r_flat = jnp.full((h * w, 1), r_i)
        
        # Repeat background and composition for all pixels
        bg_flat = current_prediction.reshape(-1, n_channels)
        comp_flat = jnp.broadcast_to(comp_i, (h * w, types))
        
        # Apply neural network to get new pixel values
        new_pixels = model.apply(nn_params, bg_flat, dx_flat, dy_flat, r_flat, comp_flat)
        
        # Reshape back to image dimensions
        new_pixels = new_pixels.reshape(h, w, n_channels)
        
        return new_pixels, None
    
    prediction, _ = jax.lax.scan(apply_droplet, prediction, (x, y, r, composition.T))
    prediction = jnp.clip(prediction, 0, 1)
    prediction = deterministic('prediction', prediction)
    diff = deterministic('diff', img - prediction)
    sample('obs', dist.Normal(scale=0.05), obs=diff)
    return nn_params

Trialing generative processes for chemical microscopy (part 2)

2025-02-23T00:00:00+00:00

Trialing generative processes for chemical microscopy (part 2)

Last time we used a very rigid generative model (droplets modeled as gaussians). This time we’ll use a learned representation of droplets.

import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np
import numpyro.distributions as dist
import seaborn as sns
from numpyro import deterministic, plate, sample
from numpyro.infer import SVI, Trace_ELBO
from numpyro.infer.autoguide import AutoNormal
from numpyro.optim import Adam
from PIL import Image

plt.rcParams['figure.dpi'] = 200

sns.set_theme(context='paper', style='ticks', font='Arial')

img = Image.open('data/example.jpg')
img = img.resize((img.width // 4, img.height // 4))
img

For simplicity, we’ll focus on modeling the H (hue) channel of the image.

img_hsv = np.array(img.convert('HSV')) / 255.0

plt.imshow(img_hsv[..., 0], cmap='gray')
plt.colorbar()

<matplotlib.colorbar.Colorbar at 0x7861b29ea210>

def model(w, h, n_droplets, channel, types=10, mask_shape=(15, 15)):
    # Sample droplet locations and sizes
    with plate("droplets", n_droplets):
        x = sample("x", dist.Uniform(0, 1)) * (w - mask_shape[1] - 1)
        y = sample("y", dist.Uniform(0, 1)) * (h - mask_shape[0] - 1)
        with plate("types", types):
            composition = sample("composition", dist.Uniform(0, 1))
    mask = sample('mask', dist.Uniform(0, 1).expand((types, *mask_shape)))

    # Instead of round, use floor and linear interpolation
    x_floor = jnp.floor(x)
    y_floor = jnp.floor(y)
    x_frac = x - x_floor
    y_frac = y - y_floor

    dx = jnp.arange(mask_shape[1])[:, None, None]
    dy = jnp.arange(mask_shape[0])[None, :, None]

    # Generate coordinates for bilinear interpolation
    x1 = x_floor[None, None, :].astype(int) + dx
    x2 = x1 + 1
    y1 = y_floor[None, None, :].astype(int) + dy
    y2 = y1 + 1

    # Calculate weights for bilinear interpolation
    wx2 = x_frac[None, None, :]
    wx1 = 1 - wx2
    wy2 = y_frac[None, None, :]
    wy1 = 1 - wy2

    # Calculate droplet masks
    droplet_masks = jnp.einsum('tn,thw->hwn', composition, mask)

    # Initialize background
    bg = sample("bg", dist.Uniform(0, 1))
    img = jnp.full((h, w), bg)

    # Add droplets using bilinear interpolation
    img = img.at[y1, x1].add(droplet_masks * (wx1 * wy1))
    img = img.at[y1, x2].add(droplet_masks * (wx2 * wy1))
    img = img.at[y2, x1].add(droplet_masks * (wx1 * wy2))
    img = img.at[y2, x2].add(droplet_masks * (wx2 * wy2))

    img = jnp.clip(img, 0, 1)
    img = deterministic('img', img)
    diff = deterministic('diff', channel - img)
    sample('obs', dist.Normal(scale=0.05), obs=diff)

guide = AutoNormal(model)
svi = SVI(model, guide, Adam(0.01), Trace_ELBO())

svi_result = svi.run(jax.random.PRNGKey(0), 100000, img.width, img.height, 2000, img_hsv[..., 0])
samples_svi = guide.sample_posterior(jax.random.PRNGKey(0), svi_result.params, sample_shape=(100,))
fig, ax = plt.subplots(figsize=(5, 2))
ax.plot(svi_result.losses)

100%|██████████| 100000/100000 [01:38<00:00, 1015.79it/s, init loss: 13266128.0000, avg. loss [95001-100000]: 447068.1562]





[<matplotlib.lines.Line2D at 0x78611a185820>]

plt.imshow(samples_svi['img'].mean(axis=0), cmap='gray')
plt.colorbar()

<matplotlib.colorbar.Colorbar at 0x78610cf92bd0>

Looks quite good!

plt.imshow(img_hsv[:, :, 0]/255.0, cmap='gray')
plt.colorbar()
plt.scatter(samples_svi['x'][:100] * img_hsv.shape[1], samples_svi['y'][:100] * img_hsv.shape[0], s=4, alpha=0.01, c='red', marker='x')

<matplotlib.collections.PathCollection at 0x786118d06030>

Most droplets are now detected — very nice!

Let’s have a look at the inferred droplet masks:

fig, axes = plt.subplots(1, samples_svi['mask'].shape[1], figsize=(samples_svi['mask'].shape[1], 1), sharey=True)
for i, ax in enumerate(axes):
    ax.imshow(samples_svi['mask'].mean(axis=0)[i], cmap='gray', vmin=0, vmax=1)

This model took about 90 seconds to fit on a rusty RTX 2080. Still there is much that could be improved. If new samples include the same droplet types, only at different locations, the masks could be “frozen”. Even better, we could use amortised inference for almost instant results. Definitely something to explore in the future.

Trialing generative processes for chemical microscopy (part 1)

2024-12-28T00:00:00+00:00

Trialing generative processes for chemical microscopy (part 1)

Is it possible to use a generative process to model microscope images like this (and is it worth the effort?)

import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np
import numpyro.distributions as dist
import seaborn as sns
from numpyro import deterministic, plate, sample
from numpyro.infer import MCMC, NUTS, SVI, Trace_ELBO
from numpyro.infer.autoguide import AutoNormal
from numpyro.optim import Adam
from PIL import Image

sns.set_theme('notebook', 'ticks', font='Arial')

plt.rcParams['figure.dpi'] = 200

img = Image.open('data/example.jpg')
img = img.resize((img.width // 2, img.height // 2))
img

Try a couple of different colour spaces in case something interesting stands out.

from skimage import color

img_array = np.array(img)
lab_img = color.rgb2lab(img_array)

fig, axes = plt.subplots(3, 1, figsize=(5, 10), sharex=True)

for i, (ax, title) in enumerate(zip(axes, ['L channel', 'a channel', 'b channel'])):
    im = ax.imshow(lab_img[:,:,i], cmap='gray')
    ax.set_title(title)
    fig.colorbar(im, ax=ax, fraction=0.03, pad=0.04)

plt.tight_layout()

hsv_img = color.rgb2hsv(img_array)

fig, axes = plt.subplots(3, 1, figsize=(5, 10), sharex=True)

for i, (ax, title) in enumerate(zip(axes, ['H channel', 'S channel', 'V channel'])):
    im = ax.imshow(hsv_img[:,:,i], cmap='gray')
    ax.set_title(title)
    fig.colorbar(im, ax=ax, fraction=0.03, pad=0.04)

plt.tight_layout()

First approach, fixed number of droplets; model centres and radii

def model1(w, h, n_droplets):
    with plate("droplets", n_droplets):
        x = sample("x", dist.Uniform(0, w))
        y = sample("y", dist.Uniform(0, h))
        r = sample("r", dist.LogNormal(1.5, 0.75))


mcmc = MCMC(NUTS(model1), num_warmup=1000, num_samples=100)
mcmc.run(jax.random.PRNGKey(0), w=img.width, h=img.height, n_droplets=100)
samples = mcmc.get_samples()

fig, ax = plt.subplots()
ax.imshow(np.ones_like(np.array(img)) * 255, cmap="gray")
for sample_no in range(5):
    for i in range(100):
        circle = plt.Circle(
            (samples["x"][sample_no][i], samples["y"][sample_no][i]),
            samples["r"][sample_no][i],
            color=plt.cm.tab10(sample_no),
        )

        ax.add_artist(circle)

sample: 100%|██████████| 1100/1100 [00:09<00:00, 112.09it/s, 15 steps of size 3.08e-01. acc. prob=0.84]

def model2(w, h, n_droplets):
    # Sample droplet locations and sizes
    with plate("droplets", n_droplets):
        x = sample("x", dist.Uniform(0, w))
        y = sample("y", dist.Uniform(0, h))
        r = sample("r", dist.LogNormal(1.5, 0.75))
        
        # Sample HSV values for each droplet
        h_val = sample("h", dist.Uniform(0, 1))
        s_val = sample("s", dist.Beta(2, 2))
        v_val = sample("v", dist.Beta(5, 2))  # Biased towards brighter values

mcmc = MCMC(NUTS(model2), num_warmup=1000, num_samples=100)
mcmc.run(jax.random.PRNGKey(0), w=img.width, h=img.height, n_droplets=100)
samples = mcmc.get_samples()
samples = {k: np.array(v) for k, v in samples.items()}

# Visualize with HSV colors
fig, ax = plt.subplots()
ax.imshow(np.ones_like(np.array(img)) * 255, cmap="gray")
for i in range(100):
    circle = plt.Circle(
        (samples["x"][0][i], samples["y"][0][i]),
        samples["r"][0][i],
        color=color.hsv2rgb(
            np.array(
                [
                    samples["h"][0][i],
                    samples["s"][0][i],
                    samples["v"][0][i],
                ]
            )
        ),
    )
    ax.add_artist(circle)

sample: 100%|██████████| 1100/1100 [00:14<00:00, 74.39it/s, 15 steps of size 2.52e-01. acc. prob=0.86]

def model2(w, h, n_droplets):
    # Sample droplet locations and sizes
    with plate("droplets", n_droplets):
        x = sample("x", dist.Uniform(0, w))
        y = sample("y", dist.Uniform(0, h))
        r = sample("r", dist.LogNormal(1.5, 0.75))
        
        with plate("pixels", w * h):
            x_dist = jnp.abs(x - jnp.arange(w)[:, None])
            y_dist = jnp.abs(y - jnp.arange(h)[:, None])
            distance = jnp.sqrt(x_dist ** 2 + y_dist[:, None] ** 2)
            val = deterministic('val', jnp.sum(jnp.exp(-distance ** 2 / (2 * r ** 2)), axis=-1))


mcmc = MCMC(NUTS(model2), num_warmup=500, num_samples=10)
mcmc.run(jax.random.PRNGKey(0), w=img.width, h=img.height, n_droplets=100)
samples = mcmc.get_samples()
samples = {k: np.array(v) for k, v in samples.items()}

# show the first 3 samples
fig, axes = plt.subplots(3, 1, figsize=(5, 10), sharex=True)
for i, ax in enumerate(axes):
    ax.imshow(1 - samples['val'][i], cmap='gray')

sample: 100%|██████████| 510/510 [00:06<00:00, 78.12it/s, 15 steps of size 3.10e-01. acc. prob=0.85]

Not a bad generative process to start with. Now let’s just fit the hue channel …

def model3(w, h, n_droplets, channel, error_scale):
    # Sample droplet locations and sizes
    bg = sample("bg", dist.Uniform(0, 1))
    with plate("droplets", n_droplets):
        x = sample("x", dist.Uniform(0, 1))*w
        y = sample("y", dist.Uniform(0, 1))*h
        r = sample("r", dist.LogNormal(1.5, 0.75))
        amplitude = sample("amplitude", dist.Uniform(0, 1))
        
        x_dist = jnp.abs(x - jnp.arange(w)[:, None])
        y_dist = jnp.abs(y - jnp.arange(h)[:, None])
        distance = jnp.sqrt(x_dist ** 2 + y_dist[:, None] ** 2)
        val = deterministic('val', bg + jnp.sum(amplitude[None, None, :] * jnp.exp(-distance ** 2 / (2 * r ** 2)), axis=-1))
        diff = deterministic('diff', val - channel)
    sample('obs', dist.Normal(0, error_scale), obs=val - channel)

Fitting, instead of sampling

guide = AutoNormal(model3)
svi = SVI(model3, guide, Adam(0.01), Trace_ELBO())
svi_result = svi.run(jax.random.PRNGKey(0), 20000, img.width, img.height, 500, hsv_img[:,:,0], 0.05)

100%|██████████| 20000/20000 [04:40<00:00, 71.43it/s, init loss: 19718028.0000, avg. loss [19001-20000]: 5279295.0000]

svi_result = svi.run(jax.random.PRNGKey(0), 50000, img.width, img.height, 500, hsv_img[:,:,0], 0.05, init_state=svi_result.state)

  0%|          | 0/50000 [00:00<?, ?it/s]

100%|██████████| 50000/50000 [11:36<00:00, 71.82it/s, init loss: 5101927.5000, avg. loss [47501-50000]: 5046932.5000]

samples = guide.sample_posterior(jax.random.PRNGKey(0), svi_result.params, sample_shape=(5,))

fig, axes = plt.subplots(3, 1, figsize=(5, 10), sharex=True)
fig.colorbar(axes[0].imshow(samples['val'][0], cmap='gray'), ax=axes[0], fraction=0.03, pad=0.04)
axes[0].set_title('Prediction')
fig.colorbar(axes[1].imshow(jnp.abs(samples['diff'][0]), cmap='gray'), ax=axes[1], fraction=0.03, pad=0.04)
axes[1].set_title('Difference')
fig.colorbar(axes[-1].imshow(hsv_img[:,:,0], cmap='gray'), ax=axes[-1], fraction=0.03, pad=0.04)
axes[-1].set_title('Ground truth')
fig.tight_layout()

Not a bad start. Most of the brightests spots have been fitted. I am surprised thought that a lot of the amplitudes are almost zero. At the end of the day, optimisations where pieces have to move into the right place first are tricky and I have no reason to believe that this is a global optimum, despite having spent a while trying to tease out a better outcome.

sns.histplot(samples['amplitude'].flatten(), bins=50);

Simple generation of locally constrained values in `numpyro`

2024-12-26T00:00:00+00:00

Simple generation of locally constrained values in `numpyro`

Just a simple experiment using numpyro to simulate values from a function with locally constrained values. This is often achieved using Gaussian processed but I thought it would be interesting to try something a bit more intuitive, plus priors on the points themselves and their interdependence can be anything, not just normal distributions.

import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np
import numpyro.distributions as dist
import seaborn as sns
from numpyro import sample
from numpyro.infer import MCMC, NUTS

sns.set_theme('notebook', 'ticks', font='Arial')

plt.rcParams['figure.dpi'] = 200

First, baseline: Independent draws from a normal distribution.

def model1():
    x = sample("x", dist.Normal().expand([100]))


mcmc = MCMC(NUTS(model1), num_warmup=1000, num_samples=100)
mcmc.run(jax.random.PRNGKey(0))
samples = mcmc.get_samples()

x_points = np.repeat(np.arange(100)[None, :], samples["x"].shape[0], axis=0)
plt.scatter(
    x_points.flatten(), samples["x"].flatten(), color="darkblue", alpha=0.01, s=10
)
plt.gca().set(xlabel="Spatial/temporal dimension", ylabel="Observable");

sample: 100%|██████████| 1100/1100 [00:00<00:00, 1656.79it/s, 7 steps of size 4.46e-01. acc. prob=0.86]

Adding some point observations …

obs_vals = jnp.array([3.3, -2.5, 4.9])
obs_idx = jnp.array([20, 40, 75])

def model2(obs_vals, obs_idx):
    x = sample('x', dist.Normal().expand([100]))
    sample('x_point_obs', dist.Normal(loc=obs_vals, scale=0.1), obs=x[obs_idx])


mcmc = MCMC(NUTS(model2), num_warmup=1000, num_samples=100)
mcmc.run(jax.random.PRNGKey(0), obs_vals, obs_idx)
samples = mcmc.get_samples()

plt.scatter(x_points.flatten(), samples['x'].flatten(), color='darkblue', alpha=0.01, s=10);
plt.scatter(obs_idx, obs_vals, label='Observed values', color='crimson', s=20)
plt.legend();

sample: 100%|██████████| 1100/1100 [00:00<00:00, 1624.28it/s, 15 steps of size 4.02e-01. acc. prob=0.90]

“Observing” each point to be the average of previous and following points. This would be equivalent to adjusting the log-likelihood or adding a potential in other packages.

def model3(obs_vals, obs_idx):
    x = sample('x', dist.Normal().expand([100]))
    sample('x_dependence', dist.Normal(loc=(x[:-2]+x[2:])/2.0, scale=0.1), obs=x[1:-1])
    sample('x_point_obs', dist.Normal(loc=obs_vals, scale=0.1), obs=x[obs_idx])


mcmc = MCMC(NUTS(model3), num_warmup=1000, num_samples=100)
mcmc.run(jax.random.PRNGKey(0), obs_vals, obs_idx)
samples = mcmc.get_samples()

plt.scatter(x_points.flatten(), samples['x'].flatten(), color='darkblue', alpha=0.01, s=10)
plt.scatter(obs_idx, obs_vals, label='Observed values', color='crimson', s=20)
plt.legend();

sample: 100%|██████████| 1100/1100 [00:01<00:00, 903.87it/s, 63 steps of size 8.62e-02. acc. prob=0.84]

Slightly fancier — weighted relation to next/previous 2 points. I don’t notice a dramatic change with this 1:2:2:1 weighting but it would be interesting to add asymmetric constraints, etc. I suppose this could be useful for probabilistic low-pass filtering with sinc weights.

def model4(obs_vals, obs_idx):
    x = sample('x', dist.Normal().expand([100]))
    sample('x_pre_obs', dist.Normal(loc=(x[:-4]+2.0*x[1:-3]+2.0*x[3:-1]+x[4:])/6.0, scale=0.1), obs=x[2:-2])
    sample('x_point_obs', dist.Normal(loc=obs_vals, scale=0.1), obs=x[obs_idx])


mcmc = MCMC(NUTS(model4), num_warmup=1000, num_samples=100)
mcmc.run(jax.random.PRNGKey(0), obs_vals, obs_idx)
samples = mcmc.get_samples()

plt.plot(samples['x'].T, color='darkblue', alpha=0.01)
plt.scatter(obs_idx, obs_vals, label='Observed values', color='crimson', s=20)
plt.legend();

sample: 100%|██████████| 1100/1100 [00:01<00:00, 1011.04it/s, 63 steps of size 9.88e-02. acc. prob=0.89]

Something a bit more interesting. Let’s change the “base” distribution to be bimodal and asymmetric. I have picked fairly broad humps so we have a fair shot at achieving good mixing.

mixing_distribution = dist.Categorical(jnp.array([0.3, 0.7]))
component_distribution = dist.Normal(
    loc=jnp.array([-3.0, 3.5]), scale=jnp.array([1.0, 1.5])
)

d = dist.MixtureSameFamily(mixing_distribution, component_distribution)

x = jnp.linspace(-10, 10, 200)
y = np.exp(d.log_prob(x))

plt.fill_between(x, y, alpha=0.3, color='darkblue')
plt.gca().set(xlabel='x', ylabel='PDF');

def model5(obs_vals, obs_idx):
    x = sample('x', d.expand([100]))

mcmc = MCMC(NUTS(model5), num_warmup=1000, num_samples=200)
mcmc.run(jax.random.PRNGKey(0), obs_vals, obs_idx)
samples = mcmc.get_samples()

x_points = np.repeat(np.arange(100)[None, :], samples['x'].shape[0], axis=0)
plt.scatter(x_points.flatten(), samples['x'].flatten(), color='darkblue', alpha=0.01, s=10);

sample: 100%|██████████| 1200/1200 [00:00<00:00, 1610.28it/s, 15 steps of size 1.68e-01. acc. prob=0.83]

Adding back the observation constraints without dependence.

def model6(obs_vals, obs_idx):
    x = sample('x', d.expand([100]))
    sample('x_point_obs', dist.Normal(loc=obs_vals, scale=0.1), obs=x[obs_idx])

mcmc = MCMC(NUTS(model6), num_warmup=1000, num_samples=200)
mcmc.run(jax.random.PRNGKey(0), obs_vals, obs_idx)
samples = mcmc.get_samples()

x_points = np.repeat(np.arange(100)[None, :], samples['x'].shape[0], axis=0)
plt.scatter(x_points.flatten(), samples['x'].flatten(), color='darkblue', alpha=0.01, s=10)
plt.scatter(obs_idx, obs_vals, label='Observed values', color='crimson', s=20)
plt.legend();

sample: 100%|██████████| 1200/1200 [00:00<00:00, 1463.46it/s, 15 steps of size 1.48e-01. acc. prob=0.87]

Finally, bringing back dependency on adjacent points …

def model6(obs_vals, obs_idx):
    x = sample('x', d.expand([100]))
    sample('x_dependence', dist.Normal(loc=(x[:-2]+x[2:])/2.0, scale=0.1), obs=x[1:-1])
    sample('x_point_obs', dist.Normal(loc=obs_vals, scale=0.1), obs=x[obs_idx])

mcmc = MCMC(NUTS(model6), num_warmup=1000, num_samples=200)
mcmc.run(jax.random.PRNGKey(0), obs_vals, obs_idx)
samples = mcmc.get_samples()

x_points = np.repeat(np.arange(100)[None, :], samples['x'].shape[0], axis=0)
# plt.scatter(x_points.flatten(), samples['x'].flatten(), color='darkblue', alpha=0.01, s=10)
plt.scatter(obs_idx, obs_vals, label='Observed values', color='crimson', s=20, zorder=10)
plt.plot(samples['x'].T, color='darkblue', alpha=0.01)
plt.legend();

sample: 100%|██████████| 1200/1200 [00:03<00:00, 356.29it/s, 255 steps of size 2.30e-02. acc. prob=0.88]

And now just for fun, let’s try a bimodal dependency distribution, here essentially saying that each point is likely to be larger than its neighbor (or linger in the same ballpark).

d_dep = dist.MixtureSameFamily(
    mixing_distribution,
    dist.Normal(loc=jnp.array([0, 0.5]), scale=0.1)
)

x = jnp.linspace(-1, 1, 200)
y = np.exp(d_dep.log_prob(x))

plt.fill_between(x, y, alpha=0.3, color='darkblue')
plt.gca().set(xlabel='x', ylabel='PDF', title='$x_i - x_{i-1}$ prior');

def model7(obs_vals, obs_idx):
    x = sample('x', d.expand([100]))
    sample('x_dependence', d_dep, obs=x[1:] - x[:-1])

mcmc = MCMC(NUTS(model7), num_warmup=1000, num_samples=200)
mcmc.run(jax.random.PRNGKey(0), obs_vals, obs_idx)
samples = mcmc.get_samples()

plt.scatter(x_points.flatten(), samples['x'].flatten(), color='darkblue', alpha=0.01, s=10)
plt.plot(x_points[0], samples['x'].mean(axis=0), color='crimson', lw=2, label='Mean')
plt.legend();

sample: 100%|██████████| 1200/1200 [00:03<00:00, 396.09it/s, 127 steps of size 6.49e-02. acc. prob=0.84]

With the observations added back in you can see how the function finds it easier to catch up with sudden rises than falls due to the prior on adjacent values.

def model8(obs_vals, obs_idx):
    x = sample('x', d.expand([100]))
    sample('x_dependence', d_dep, obs=x[1:] - x[:-1])
    sample('x_point_obs', dist.Normal(loc=obs_vals, scale=0.25), obs=x[obs_idx])

mcmc = MCMC(NUTS(model8), num_warmup=1000, num_samples=200)
mcmc.run(jax.random.PRNGKey(0), obs_vals, obs_idx)
samples = mcmc.get_samples()

plt.scatter(x_points.flatten(), samples['x'].flatten(), color='darkblue', alpha=0.01, s=10)
plt.plot(x_points[0], samples['x'].mean(axis=0), color='crimson', lw=2, label='Mean')
plt.scatter(obs_idx, obs_vals, label='Observed values', color='crimson', s=20)
plt.legend();

sample: 100%|██████████| 1200/1200 [00:02<00:00, 497.90it/s, 255 steps of size 5.76e-02. acc. prob=0.88]

Very cool and I think quite useful for modelling chromatograms.

A rant against macros

2020-06-14T00:00:00+00:00

A rant against macros

I used to be a huge fan of macros. I remember reading SICP and being amazed that you could use the language to generate and transform code. How cool is that? First a couple of examples: Clojure’s core.async library includes a go macro that lets you launch goroutine-like tasks without having to change the language.

; https://github.com/clojure/core.async/blob/master/examples/walkthrough.clj
(let [c1 (chan)
      c2 (chan)]
  (go (while true
        (let [[v ch] (alts! [c1 c2])]
          (println "Read" v "from" ch))))
  (go (>! c1 "hi"))
  (go (>! c2 "there")))

The Turing library lets you write probabilistic programs in Julia as if you’re using a dedicated probabilistic programming language (PPL):

@model gdemo(x, y) = begin
    # Assumptions
    σ ~ InverseGamma(2,3)
    μ ~ Normal(0,sqrt(σ))
    # Observations
    x ~ Normal(μ, sqrt(σ))
    y ~ Normal(μ, sqrt(σ))
end

Fast forward to 2018 when I sat down with Chris Rackauckas before JuliaCon and he mentioned he’d been in touch with the Turing developers. I thought he bring up their PPL syntax and how it’s so wonderful that Julia lets you mold the language, but when I prompted him he said the macros have gotten in the way of using Turing as a library. He said functions and types were the way forward if you want things to compose.

Since then, I’ve written a couple of macros of my own and, powerful as they are, I have come to the conclusion that the problems I used them for were better handled by i) new or more expressive data structures, ii) plain old functions, iii) accepting a small amount of extra verbosity. In return you get, i) better interoperability, ii) code that is more explicit and easier to undestand, iii) much easier debugging, iv) a more robust design, v) much better support from your tools (e.g. IDE, REPL).

Let’s look at a simpler model:

using Distributions

@model normal_model(x) = begin
    # just a simple transformation; z is still observed, just like x
    z = 2x
    # sample y
    y ~ Normal(0.0, 1.0)
    # observe z
    z ~ Normal(y, 1.0)
end

This is what it expands to

quote
    #= /home/group/.julia/packages/DynamicPPL/9OFG0/src/compiler.jl:348 =#
    function var"##evaluator#371"(_rng::Random.AbstractRNG, _model::DynamicPPL.Model, _varinfo::DynamicPPL.AbstractVarInfo, _sampler::AbstractMCMC.AbstractSampler, _context::DynamicPPL.AbstractContext)
        #= /home/group/.julia/packages/DynamicPPL/9OFG0/src/compiler.jl:355 =#
        begin
            x = (DynamicPPL.matchingvalue)(_sampler, _varinfo, _model.args.x)
        end
        #= /home/group/.julia/packages/DynamicPPL/9OFG0/src/compiler.jl:356 =#
        begin
            #= REPL[22]:1 =#
            #= REPL[22]:2 =#
            z = 2x
            #= REPL[22]:3 =#
            begin
                var"##tmpright#363" = Normal(0.0, 1.0)
                var"##tmpright#363" isa Union{Distribution, AbstractVector{<:Distribution}} || throw(ArgumentError("Right-hand side of a ~ must be subtype of Distribution or a vector of Distributions."))
                var"##vn#365" = y
                var"##inds#366" = ()
                y = (DynamicPPL.tilde_assume)(_rng, _context, _sampler, var"##tmpright#363", var"##vn#365", var"##inds#366", _varinfo)
            end
            #= REPL[22]:4 =#
            begin
                var"##tmpright#367" = Normal(y, 1.0)
                var"##tmpright#367" isa Union{Distribution, AbstractVector{<:Distribution}} || throw(ArgumentError("Right-hand side of a ~ must be subtype of Distribution or a vector of Distributions."))
                var"##vn#369" = z
                var"##inds#370" = ()
                z = (DynamicPPL.tilde_assume)(_rng, _context, _sampler, var"##tmpright#367", var"##vn#369", var"##inds#370", _varinfo)
            end
        end
    end
    #= /home/group/.julia/packages/DynamicPPL/9OFG0/src/compiler.jl:359 =#
    var"##generator#372"(x) = begin
            #= /home/group/.julia/packages/DynamicPPL/9OFG0/src/compiler.jl:359 =#
            (DynamicPPL.Model)(var"##evaluator#371", (DynamicPPL.namedtuple)(NamedTuple{(:x,), Tuple{Core.Typeof(x)}}, (x,)), (DynamicPPL.ModelGen){(:x,)}(var"##generator#372", NamedTuple()))
        end
    #= /home/group/.julia/packages/DynamicPPL/9OFG0/src/compiler.jl:360 =#
    var"##generator#372"(; x) = begin
            #= /home/group/.julia/packages/DynamicPPL/9OFG0/src/compiler.jl:344 =#
            var"##generator#372"(x)
        end
    #= /home/group/.julia/packages/DynamicPPL/9OFG0/src/compiler.jl:362 =#
    begin
        $(Expr(:meta, :doc))
        normal_model = (DynamicPPL.ModelGen){(:x,)}(var"##generator#372", NamedTuple())
    end
end

And now to sample it:

sample(normal_model(3.0), NUTS(), 1000)

# Summary Statistics
#   parameters    mean     std  naive_se    mcse       ess   r_hat
#   ──────────  ──────  ──────  ────────  ──────  ────────  ──────
#            y  0.0096  1.0146    0.0454  0.0978  168.5831  0.9986
#            z  0.0169  1.4692    0.0657  0.1204  158.7592  0.9992

And we see that Turing has sampled both y and z, where z should have been marked as deterministic and observed rather than sampled. Now, I’m sure this is well-documented somewhere but the point is that when you use a macro, your Julia code no longer functions the way you would expect. Worse, yet finding out why means being able to navigate the mess of generated symbols in the expanded version. And yes, the authors can fix this (if it’s actually a bug) but it doesn’t change the problem that the language inside that block is no longer Julia. You keep having to second guess yourself every time you reach for a new language feature.

Increasingly, macros, even nice hygienic ones remind me of the horrible mess that’s C/C++ macros: an untamed partial language with its own semantics that you need to learn and use, and how people have created whole programming languages in part to escape this ugly metalangauge problem. It’s true that homoiconic languages mostly get rid of the macro/preprocessor language, but the semantics of how language constucts behave within the macro and how they compose with other langauge features is still completely up to the programmer and, in my experience, quite hard to get right.

I see macros used in places that I find really troubling. I was writing a toy GTK application in Rust earlier today and learned that you need to use these weird macros to get memory management to play nicely with Rust.

use glib::clone;

let window = Rc::new(ApplicationWindow::new(app));

# moving a weak reference to `window` into the closure
butten.connect_activate(clone!(@weak window => move |_| {
    window.close(&button);
}));

I really don’t think introducing this metalanguage is a good idea at all. Also, how is this custom syntax supposed to be understood by the editor? Before rust-analyzer my editor (VSCode + RLS) would give up with the macro and I would have to guess my way out. Things are better now that we have rust-analyzer but I’m not even sure the Rust tooling is ever supposed to be able to make sense of this.

Bottom line (and I’m happy to be proven wrong): macros are an unsustainable convenience. They are never good enough to justify the readability/maintainability/tooling headaches.

The case for lazy computation and interactive optimization

2020-05-24T00:00:00+00:00

The case for lazy computation and interactive optimization

Since 2018, my colleague Dario and I have been working on a probabilistic model of chemical reactivity. In a nutshell, this model can take reactivity observations between a bunch of compounds and interpret them as manifestations of properties and their mutual reactivities. The Markov chain Monte carlo implementation of this model in PyMC3 was quite fast to start but over time we have expanded the model and now with all the bells and whistles sampling can take close to a day.

Given these performance issues it was only natural to look at alternative implementations in my high performance language of choice, Julia. So, over the past year I have done maybe a dozen experiments, from a more or less equivalent high-level description in Turing.jl, to encoding the log probability manually (on the CPU and the GPU) and using DynamicHMC.jl and friends to sample, to using the HMC implementation in AdvancedHMC.jl and hand coding the log probability + domain transformations. What has been surprising is that the performance gain hasn’t been nearly as large as I imagined it would be. In fact, my fastest implementation using AdvancedHMC is still only half as fast as the straightforward PyMC3 implementation. There may be differences in the formulation of NUTS between the different packages, of course, so I don’t think this should be taken as representative. Just that it got me thinking about performance.

Yesterday, I started thinking about PyMC3’s backend, Theano, and how it allows PyMC to be fast and expressive. Theano (and newer systems like TensorFlow) express computations as executable graphs of operations (ops for short). In principle, evaluation of these graphs can entail any of the following:

Simply running an interpreted implementation of each op.
Native code generation for all or part of the ops before execution.
Graph optimization, e.g. removing redundant ops or rewriting sequences of ops them with faster equivalents.

Libraries like Theano and TensorFlow have a fairly limited scope (numerical code) but I believe that borrowing certain of the above elements can make a great DSL for high performance computing in other domains. Specifically:

Computations described as graphs of ops.
Op-graph transformations, themselves described using #1. These transformations can be applied to a certain op or to all ops in a given scope (e.g. children of a certain op, or ops matching a certain pattern).
Interpretive execution of the op graph.

Using this design, optimizations will be part of libraries that can be imported and applied to existing code a la carte as opposed to hard-coded in the compiler/JIT. This model is somewhat similar to what Julia does, i.e. interpretation and JIT compilation of code based on inferred type, but goes beyond building everything around the type system. Moreover, it can be implemented as a library in Julia, Python, or any language with bindings to codegen backend being targetted (if any).

Update (2020-05-26): Relevant discussion on the Julia discourse. Forum users pointed out some of the promising developments in the area, e.g. Mjolnir and being able to customize the compilation pipeline through parameterized interpretation.

Fertile land at the confluence of staged programming and (typed) logic programming — Part 1

2019-04-07T00:00:00+00:00

Fertile land at the confluence of staged programming and (typed) logic programming — Part 1

I recently started sketching out a logic programming library called Logic.jl (it really is very much a sketch at the moment). As the name might imply, the library is implemented in Julia, which seems like an odd choice of implementation language. Choosing Julia and the twist on logic programming that I am aiming for are motivated by a particularly interesting confluence of ideas that I have become aware of in the past year and hope to describe in this blog post. I would love to hear your thoughts of course; just keep in mind that I am not a computer scientist.

Idea #1: Staged programming and languages with first-class JIT compilation machinery

There is no shortage of programming languages with just-in-time (JIT) accelerated virtual machines. Typical JITs are opaque: they step in at run time and speed up your code without you having to tell them anything. Increasingly, however, there are JITs of a different breed that, rather than stay hidden and opaque, work in conjunction with the program, giving rise to a class of programming languages where application code can inspect and influence the various stages of code lowering as a precise mechanism for on- demand code generation and behaviour adaptation. Using the notion of staged programming , the various components of the runtime monolith (type checker, GC, codegen) can also be used in an à la carte fashion, further blurring the line between compile and run time. I find this paradigm superior to the AST- level metaprogramming often encountered with LISP family languages. Where conventional JITs like HotSpot are often used to bridge the performance gap between bytecode-interpreted dynamically-typed languages like Python and statically-typed compiled languages like OCaml, staged programming, e.g. in Julia, also recovers some of the safety and expressive power of a sophisticated type system in an otherwise dynamically-typed environment.

The immediate trade-off is the overhead of including the full compiler toolchain with application code. The run time memory footprint would then resemble the equivalent of an AST/bytecode interpreter, libraries, and any compiler passes, including [possibly] the type system and codegen (e.g. LLVM). This is non-trivial and impractical in memory-constrained applications such as embedded systems. In principle at least, one can ahead-of-time compile the application and discard unused toolchain code as long as it can be inferred as unnecessary. That said, annotating/inferring dependencies between application code and the various pieces of compiler/runtime might require some effort.

Staged programming in the wild

Before we move on to Julia, let’s look at a few examples of programming langauges/environments that feature elements of staged programming. I suspect that there are many more such systems out there that I am not aware of.

MetaOCaml is described as an OCaml dialect for multi-stage programming. OCaml seems like a particularly good choice of language in my opinion, because a lot of the power afforded by its type system is challenged when dealing with data the structure of which is unknown until run time. Here for example is a (simplified) definition of the json type from the OCaml package Yojson.

type json = [
  | `Assoc of (string * json) list
  | `Bool of bool
  | `Float of float
  | `Int of int
  | `List of json list
  | `Null
  | `String of string
]

Here is the Yojson package in use

let doc =

Terra, essentially a metaprogramming system using Lua and LLVM aimed at low-level system programming.

-- Terra allows the Lua interpreter and LLVM to interact.
-- Lua code can invoke LLVM for code generation.
-- LLVM can also call Lua, here to partially evaluate an expression.
local a = 5
terra sin5()
    return [ math.sin(a) ]
end

-- output bitcode
sin5:printpretty() 
> output:
> sin50 = terra() : {double}
>    return -0.95892427466314
> end

-- example code from terralang.org

Staged programming in Julia

In my opinion Julia is the most successful implementation of the staged programming paradigm today. I say this based on the number of Julia users as well as how far it has taken the paradigm.

In Julia functions act as the basic unit of JIT compilation, with types guiding the process through a language feature called multiple dispatch. Take the following simple function, for instance.

""" sum(col)
Return the sum of the elements of collection `col`
"""
function sum(col)
    result = zero(eltype(col))
    for elem in col
        result += elem
    end
end

No code generation happens for this function until it is invoked, e.g. sum([1,2,3]), at which point the type of its argument Array{Int64, 1} recursively propagates through the body of the function. Much of the logic inside the function can be constant-folded given this concrete type. far from perfect.

The use of multiple dispatch as the primary driver of code generation in Julia seems to have worked out really well for the language. Still, I wonder if there are mechanisms that allow more expressive programming.

Julia was conceived as a programming language for high performance numerical calculations but, unlike other languages in its league like Matlab and Mathematica, its type system and compilation machinery appear to be applicable well beyond the realm of scientific computing, as demonstrated by projects like imp and DataKnots.jl and Automa.jl.

Staged-programming is also known as multi

First-class embedding of logic programming

Having learned the functional paradigm, many people find it hard to go back to their old imperitive tools, which feel verbose and error-prone in comparison. My brief exposure to logic programming has left me with a similar feeling towards functional programming, namely the sense that logic programming allows solving the problem in a more direct and natural way. Just as functional programming seems removed from the physical reality of the computer because of its pervasive use of the function as a layer of indirection, logic programming’s symbolic variables and predicates provide a further layer of abstraction that allow the solution to be described by its properties rather than its realization from a given set of inputs.

Many Prolog users are familiar with the less-than-ideal interoperability of typical Prolog implementations with real-world code written in a language like Python. In a follow-up blog post I will try to explain how Julia’s multiple dispatch and staged programming facilitate embedding logic programming to bring most of its expressive power to Julia.

Funny

2013-12-28T00:00:00+00:00

Funny

… Likewise, claiming JS is homoiconic because you can eval strings of code is nonsense. If that’s the only criteria for homoiconicity, then C is too, since you can treat an array of bytes as code and jump to it.

From here.

A new way of doing chemistry

2013-12-23T00:00:00+00:00

A new way of doing chemistry

I’ve been thinking about a new way of doing chemistry. Take for example the case of growing crystals. In the more difficult cases, this involves trial and error using a number of educated guesses about the solvents and conditions. What if we had an algorithm that would decide the best procedure based on a database of molecular metrics, e.g., dipole moment, molecular weight, melting point, decomposition temperature, functional groups, hydrogen bonding, etc. for solutes as well as solvents, trained on a series of successful/unsuccessful combinations, e.g., naphthalene from hot ethanol, macrocycle x from DMSO + ether by diffusion at room temperature, compound y from slow evaporation of hexanes solution?

This of course can be generalized to any chemical transformation. Publicly available databases like Org. Syn. can be mined for information that no one chemist can hope to memorize, significantly reducing the amount of trial and error in day-to-day chemistry.

Hessam's blog

Ancestral sampling with variational relaxation

Exploration: Flow matching, NeuTraHMC and normalizing flow guides in NumPyro

First bee sighting of 2026

First bee sighting of 2026

Nice strange attractor visualisation

Nice strange attractor visualisation

Draft: Evolving interactions with AI in 2026

Draft: Evolving interactions with AI in 2026

Draft: Thinking about aerosols probabilistically — Part 1: Evaporation

A little snippet for nicer plots using matplotlib

A little snippet for nicer plots using matplotlib

3D-printed mathematical lampshades

3D-printed mathematical lampshades

2D profile

Trying (and failing) 3D with OnShape

Using `uv` projects on a network share

Using uv projects on a network share

Reparameterizing distributions in numpyro

Reparameterizing distributions in numpyro

Problem setting

Adding support for observations

Live variational inference

Live variational inference

Amortized probabilistic models for chemical microscopy

Amortized probabilistic models for chemical microscopy

Can we used an amortized model to speed up inference in our droplet microscopy model?

Other attempts

Fully flattened model with aggregation

Iterative with lax.scan

Trialing generative processes for chemical microscopy (part 2)

Trialing generative processes for chemical microscopy (part 2)

Trialing generative processes for chemical microscopy (part 1)

Trialing generative processes for chemical microscopy (part 1)

Simple generation of locally constrained values in `numpyro`

Simple generation of locally constrained values in numpyro

A rant against macros

A rant against macros

The case for lazy computation and interactive optimization

The case for lazy computation and interactive optimization

Fertile land at the confluence of staged programming and (typed) logic programming — Part 1

Fertile land at the confluence of staged programming and (typed) logic programming — Part 1

Idea #1: Staged programming and languages with first-class JIT compilation machinery

Staged programming in the wild

Staged programming in Julia

First-class embedding of logic programming

Funny

Funny

A new way of doing chemistry

A new way of doing chemistry

Using `uv` projects on a network share

Iterative with `lax.scan`

Simple generation of locally constrained values in `numpyro`