That said, hooking up the OpenAI API to a simple back-and-forth conversational system is still the entry point for AI engineering.

Even though 99.9% of this system’s complexity is abstracted away behind the inference API, there are a few concepts that I see people take a while to grasp as they learn how to build systems on top of AI.

This short note discusses a few of them. I haven’t found good sources or references for these ideas elsewhere, so I hope this note is useful.

Efficient citation via clipboarding

LLM-driven search à la Perplexity may seem like a 2024 thing, but it is probably still one of the top three most valuable use cases for AI. In this workstream, being able to reference sources and efficiently show them to the user is critical.

Many starter kits for chatbots (e.g., Chainlit and Open WebUI) can easily render Markdown. Markdown allows you to write hyperlinks with a string like Please click me!, which is easy to write and simple to remember.

LLMs are perfectly capable of writing hyperlink strings like these and require only a few short prompt instructions to do so. However, asking the LLM to write the hyperlink comes with a few limitations:

• While the user is waiting for the text to finish streaming, they will see partial hyperlink strings like [Please click me!](https://www.yo... Once the string is complete, it will be rendered, causing the partial text to disappear and the hyperlink to appear in its place. This jump is irritating and creates a poor user experience.
• As system complexity and scope grow, you will likely begin citing sources with very long URLs that take a long time (>100 ms) to write token by token.
• This approach is token-inefficient because the same result can be produced more cheaply.

The solution is simple. As the reader, you have probably already implemented it or reinvented it on the spot.

Regardless, here it is: instruct the LLM to write short citation keys like @d3 or #a1. A postprocessor sitting on top of the text stream then captures these keys and deterministically replaces them with full hyperlinks. I personally call this the “clipboard” approach in all of my codebases.

Some caveats do apply. Namely, you need some mechanism to give the source a nice descriptor (like “Annual Shareholder Report, 2026”). The naive approach from earlier lets you prompt the LLM to create the hyperlinked text on the fly; this approach doesn’t let you do that out of the box.

This approach also separates citation styling from prompting logic. That separation is a good idea (IMO), but it can be a pain if you’d rather just write prompts instead of extending application logic. Visually, here’s what it looks like:

You can also generalize this approach to any string that the LLM might generate inefficiently. If your LLM is powering an agent with tools for accessing APIs or reading files in a structured format, users may ask questions whose answers require the model to regurgitate data nearly verbatim. If the source is a long CSV or JSON file, having the LLM copy its contents token by token is wasteful. Use the clipboard approach to copy and paste the output instead.

Your choice of placeholder token is mostly unimportant, so long as it is small and unlikely to show up in the generated text for unrelated reasons. Avoid using standard formats like [1] or [2], since hallucinating LLMs will typically generate false citations in those formats. For richer styling, you could probably use XML or a similar format.

Constructing follow-up questions or actions

Most domain-specific AI platforms (e.g., Microsoft Copilot, Perplexity, and Glean) now supply recommended or follow-up questions for users to click. This is partly a gimmick from product managers to juice usage numbers, but it can sometimes be genuinely useful.

The best follow-up questions (e.g., “Would you like me to also search for restaurants closer to you?”) know both 1) what the user originally asked and 2) how the system responded in the most recent turn. The standard approach I saw in a number of systems from 2025 and earlier was to run a specific prompt to make the LLM generate these questions.

Again, this is wasteful and inefficient. Nowadays, using structured generation plus streaming, you can have an LLM answer the user’s question and generate the follow-ups in a single prompt using a schema like the following:

(show a schema that uses the OpenAI API to answer and generate follow-ups)

The combined approach can be cost-neutral, as most of the instruction tokens from the follow-up prompt are moved into the main response prompt. In either case, the instructions should remain in the prompt cache 100% of the time, so the prefill cost of additional instructions is low relative to the decode cost. Since both approaches emit the same number of tokens, the overall cost of generation is mostly unchanged.

From a user-experience perspective, emitting the follow-up questions from the same stream as the main response probably saves at least 100 milliseconds by avoiding the overhead of a separate API call.

The primary weakness of a combined response/follow-up prompt is that it requires the LLM to follow more instructions per prompt, which generally leads to poorer adherence and effectiveness as the prompt grows. Open-weight models released since late 2025 generally do not have this issue in standard workflows. It may still be a problem if you have a very complex prompt with many instructions.

Using cheap LLMs to capture user attention

Every few months, a new high-speed LLM inference provider comes on the market, and the industry reacts by touting the advantages of zero-lag conversations, instant coding agents, and so on. I think the market for that is actually very small. When we have faster LLMs, we don’t simply make existing workflows faster; we generally use the extra latency headroom to make those workflows better. In other words, long-running AI tasks are going to be a fact of life.

Consequently, we need to resort to all sorts of tricks to keep the user on our page before they’re emotionally spent from waiting 800 milliseconds for the next dopamine hit.

From an AI workflow perspective, this means providing continuous state updates on long-running tasks. Perplexity was an early leader here and had a great design that ensured something interesting was always happening on screen, even if you had to wait for 15+ seconds.

My go-to method is to task a very small, cheap LLM, such as Llama 3.2 1B or Gemma 2B, with providing short commentary snippets using prompts like this:

Your role is to provide one or two sentences of brief commentary on the state of a running program. Write at a high level, use concise language, and use the present tense.

Original user query:
{query}

Recent actions:
{tool_calls}

The output should be something like Searched Salesforce cases from the last two weeks. This can be hugely beneficial for keeping users engaged.

Make sure to keep the tool calls ordered from oldest to newest so that you can leverage the prompt cache effectively here. Putting the newest calls first invalidates the prompt cache. If you run this every 3 seconds during a 60-second task, you can easily send this prompt 20+ times, making cache management important even within the scope of a single turn.

OpenAI did this quite a bit with ChatGPT, as you can see below:

If you have any more ideas to include here, let me know!

In this post, I’ll show a workflow for using GPU-based inference on Modal for a model which is very large by the standards of Bayesian statisticians \((\vert \theta\vert \gt 10^6)\), by deploying to a datacenter GPU and renting it only for a short time.

Model & data

We’ll use synthetic data for this example.

I’ve chosen a hierarchical logistic regression for this post since it has a non-conjugate likelihood, appears commonly in practice, and can easily be assigned more parameters by increasing the number of covariates and/or the number of groups.

Let \(i\), \(g\), and \(k\) denote the indices over observations, groups, and covariates. Furthermore, let \(x_i \in \mathbb{R}^{20}\) be the covariate vector for observation \(i\), and let \(g_i \in \{1,\ldots,100000\}\) identify its group. The data-generating process uses population-level slopes \(\beta_k\), group intercept deviations \(\alpha_g\), and group slope deviations \(\gamma_{gk}\). The binary outcome is generated from

Essentially, this is a logistic regression with random slopes for 20 covariates and a random intercept for each of 100,000 groups in the data.

We’ll use a non-centered parameterization for the group effects. The prior specification is

The code below produces a synthetic dataset; we can control the overall sparsity of the response with the value of α_true.

import numpy as np

RANDOM_SEED = 827
rng = np.random.default_rng(RANDOM_SEED)

N = 1_000_000 # Number of data points
G = 100_000   # Number of groups
K = 20      # Number of covariates / features

group_idx = rng.integers(0, G, size=N, dtype=np.int64)
X = rng.normal(size=(N, K)).astype(np.float32)

α_true = np.float32(-1.0)
β_true = rng.normal(0.0, 0.45, size=K).astype(np.float32)
σ_α_true = np.float32(0.80)
σ_γ_true = rng.uniform(0.15, 0.35, size=K).astype(np.float32)
α_group_true = rng.normal(0.0, σ_α_true, size=G).astype(np.float32)
γ_group_true = rng.normal(0.0, σ_γ_true, size=(G, K)).astype(np.float32)

η = α_true + α_group_true[group_idx] + np.sum(X * (β_true + γ_group_true[group_idx]), axis=1)
p = 1.0 / (1.0 + np.exp(-η))
y = rng.binomial(1, p).astype(np.int64)

print(f"Average of y: {np.mean(y):.2f}")

Average of y: 0.36

Here, we define our model in PyMC. This is a fairly standard model definition without much nuance or many tricks.

import pymc as pm
import pytensor
import pytensor.tensor as pt

pytensor.config.floatX = "float32"

def build_model(X, y, group_idx):

    X = np.asarray(X, dtype=np.float32)
    y = np.asarray(y, dtype=np.int64)
    group_idx = np.asarray(group_idx, dtype=np.int64)

    N, K = X.shape
    G = int(group_idx.max()) + 1
    coords = {
        "obs": np.arange(N),
        "group": np.arange(G),
        "covariate": np.arange(K),
    }

    with pm.Model(coords=coords) as model:
        X_data = pm.Data("X", X, dims=("obs", "covariate"))
        group_lookup = pt.as_tensor_variable(group_idx, name="group_lookup")

        α = pm.Normal("α", np.float32(0.0), np.float32(1.5))
        β = pm.Normal("β", np.float32(0.0), np.float32(1.0), dims="covariate")
        σ_α = pm.HalfNormal("σ_α", np.float32(1.0))
        σ_γ = pm.HalfNormal("σ_γ", np.float32(0.5), dims="covariate")

        z_α = pm.Normal("z_α", np.float32(0.0), np.float32(1.0), dims="group")
        z_γ = pm.Normal("z_γ", np.float32(0.0), np.float32(1.0), dims=("group", "covariate"))
        α_group = σ_α * z_α 
        γ_group = σ_γ * z_γ

        η = α + α_group[group_lookup] + pm.math.sum(X_data * (β + γ_group[group_lookup]), axis=1)
        pm.Bernoulli("Y", logit_p=η, observed=y, dims="obs")

    return model

To illustrate why this model could be challenging to fit using only a CPU, we profile the gradient evaluation time. The dlogp function represents \(\frac{\partial}{\partial \theta}[\log \tilde{p}]\) with \(\tilde{p}\) denoting the unnormalized posterior density; it may be calculated up to 1000 times per sample with the evaluation count increasing with the difficulty of the problem in terms of posterior geometry and/or nonlinearity. Hamiltonian Monte Carlo, the workhorse of most modern Bayesian programming frameworks, requires these gradient evaluations to draw Monte Carlo samples.

The code cell below runs the gradient a few times and records the time taken.

import time

def median_runtime(fn, n=5):
    fn()
    times = []
    for _ in range(n):
        start = time.perf_counter()
        fn()
        times.append(time.perf_counter() - start)
    return float(np.median(times)), times

local_model = build_model(X, y, group_idx)
local_point = local_model.initial_point()

with local_model:
    dlogp_fn = local_model.compile_dlogp()

local_grad_median, local_grad_times = median_runtime(lambda: dlogp_fn(local_point))
print(f"Median grad eval time is {local_grad_median:.2f} seconds")

Median grad eval time is 0.20 seconds

If we try to extrapolate to the time required for running a full chain of 1000 samples, we find that assuming ~250 grad evals per sample, we would need 50 seconds per sample and 50,000 seconds (~14 hours) to run a full chain! This simply won’t do.

Let’s deploy this remotely and see what we get!

Running remotely on Modal

For those of you unfamiliar with the magical world of serverless GPU, please take a look at Modal. I will probably never be wealthy enough to outright own a beautiful server rack of datacenter GPUs, so this is the best I can get. Basically, Modal provides APIs for spinning up jobs on GPUs with very little friction and little downtime.

To make this example run nicely on GPU, we’ll make a few adjustments. First, we’ll toggle it to use the Jax + NumPyro backend to work out-of-the-box with an NVIDIA GPU.

import modal

modal_image = (
    modal.Image.debian_slim(python_version="3.13")
    .uv_pip_install(
        "arviz==1.1.0",
        "jax[cuda12]==0.7.2",
        "numpy==2.3.5",
        "numpyro==0.19.0",
        "pandas==2.3.3",
        "pymc==6.0.1",
    )
)

app = modal.App("i-should-be-working-right-now", image=modal_image)

We’ll set a few environment flags for the float precision and the devices, define a helper to benchmark the execution time, and apply some Jax-isms to prep the compute graph and evaluate it.

def remote_model(X, y, group_idx):
    import os

    os.environ["JAX_PLATFORMS"] = "cuda"

    import jax
    import jax.numpy as jnp
    import numpy as np
    import pytensor

    pytensor.config.floatX = "float32"
    X = np.asarray(X, dtype=np.float32)
    y = np.asarray(y, dtype=np.int64)
    group_idx = np.asarray(group_idx, dtype=np.int64)

    gpu_devices = jax.devices("gpu")
    gpu_check = jax.device_put(jnp.ones((512, 512), dtype=jnp.float32), gpu_devices[0]).sum().block_until_ready()
    model = build_model(X, y, group_idx)

    info = {
        "N": X.shape[0],
        "G": int(group_idx.max()) + 1,
        "K": X.shape[1],
        "jax_backend": jax.default_backend(),
        "jax_gpu_devices": [str(device) for device in gpu_devices],
        "gpu_check_sum": float(gpu_check),
        "x_dtype": str(X.dtype),
        "value_var_dtypes": {var.name: var.dtype for var in model.value_vars},
    }
    print(f"JAX backend: {info['jax_backend']}; GPU devices: {info['jax_gpu_devices']}")
    return model, info


def median_runtime(fn, n=5, synchronize=None):
    import time

    import numpy as np

    if synchronize is None:
        synchronize = lambda result: None

    def timed_run():
        start = time.perf_counter()
        result = fn()
        synchronize(result)
        return time.perf_counter() - start

    synchronize(fn())
    times = [timed_run() for _ in range(n)]
    return float(np.median(times)), times


def wait_jax(result):
    import jax

    jax.tree_util.tree_map(lambda x: x.block_until_ready(), result)


@app.function(gpu="A100", timeout=2 * 60 * 60)
def profile_dlogp(X, y, group_idx, n_evals=5):

    import jax
    import jax.numpy as jnp
    from pymc.sampling.jax import get_jaxified_logp

    model, info = remote_model(X, y, group_idx)
    point = model.initial_point()

    with model:
        pytensor_dlogp = model.compile_dlogp()

    pytensor_median, pytensor_times = median_runtime(lambda: pytensor_dlogp(point), n_evals)

    values = [jnp.asarray(point[var.name]) for var in model.value_vars]
    jax_loss = get_jaxified_logp(model)
    jax_dlogp = jax.jit(jax.value_and_grad(jax_loss))

    wait_jax(jax_dlogp(values))
    jax_median, jax_times = median_runtime(lambda: jax_dlogp(values), n_evals, wait_jax)

    profile = {
        "median_pytensor_dlogp_eval_seconds": pytensor_median,
        "median_jax_dlogp_eval_seconds": jax_median,
    }
    print(f"Median dlogp eval: PyTensor={pytensor_median:.3f}s, JAX={jax_median:.3f}s")
    return {**info, **profile}


@app.function(gpu="A100", timeout=2 * 60 * 60)
def fit_hierarchical_logistic(X, y, group_idx, draws=100, tune=100, random_seed=RANDOM_SEED):
    import time
    import arviz as az
    import pymc as pm

    model, info = remote_model(X, y, group_idx)

    with model:
        start = time.perf_counter()
        idata = pm.sample(
            draws,
            tune=tune,
            chains=1,
            nuts_sampler="numpyro",
            random_seed=random_seed,
            var_names=["α", "β", "σ_α", "σ_γ"],
            progressbar=False,
        )
        elapsed_seconds = time.perf_counter() - start

    β_mean = idata.posterior["β"].mean(dim=("chain", "draw")).to_numpy()

    return {
        **info,
        "elapsed_seconds": elapsed_seconds,
        "summary": az.summary(idata, var_names=["α", "β", "σ_α", "σ_γ"]),
        "β_mean": β_mean.tolist(),
    }

With all of this helper logic written, we can finally deploy to the cloud! We will start by running a short job to just profile the logp gradient function.

print("Launching Modal GPU dlogp profile")
with app.run():
    profile_result = profile_dlogp.remote(X, y, group_idx)
print(f"Finished running; profile results: {profile_result}")

Launching Modal GPU dlogp profile
Finished running; profile results: {'N': 1000000, 'G': 100000, 'K': 20, 'jax_backend': 'gpu', 'jax_gpu_devices': ['cuda:0'], 'gpu_check_sum': 262144.0, 'x_dtype': 'float32', 'value_var_dtypes': {'α': 'float32', 'β': 'float32', 'σ_α_log__': 'float32', 'σ_γ_log__': 'float32', 'z_α': 'float32', 'z_γ': 'float32'}, 'median_pytensor_dlogp_eval_seconds': 0.2271286229999987, 'median_jax_dlogp_eval_seconds': 0.0013876599999989025}

Interesting - the gradient takes 0.001 seconds on GPU and 0.22 seconds on CPU. That is around a 200x speedup!

Next, we run the Markov chain to completion and retrieve the results.

print("Launching Modal GPU run")
with app.run():
    result = fit_hierarchical_logistic.remote(X, y, group_idx)

print(f"MCMC run finished in {result["elapsed_seconds"]:.2f} seconds")

MCMC run finished in 272.16 seconds

Modal helpfully lists their prices per second; an A100 runs for $0.000583 / second, meaning this run cost me around 15 cents from start to finish.

Parameter Recovery

We can see all the sampler diagnostics in the posterior summary. We’d need to run more chains to get the \(\hat{R}\) value for this model.

import pandas as pd

posterior_summary = pd.DataFrame(result["summary"])
posterior_summary.iloc[0:5]

  parameter   mean     sd  eti89_lb  eti89_ub  ess_bulk  ess_tail  r_hat  \
0         α -0.980  0.015    -0.997    -0.952     1.510     5.445    NaN   
1      β[0] -0.275  0.005    -0.282    -0.266     2.042     5.659    NaN   
2      β[1]  0.630  0.010     0.612     0.641     1.627     5.374    NaN   
3      β[2] -0.358  0.007    -0.366    -0.345     2.390     5.593    NaN   
4      β[3]  0.359  0.006     0.348     0.366     1.834     5.607    NaN   

   mcse_mean  mcse_sd  
0      0.012    0.008  
1      0.004    0.003  
2      0.008    0.005  
3      0.005    0.003  
4      0.005    0.003

A quick diagnostic is to compare the posterior mean of the fixed effects with the values used to simulate the data. With \(10^6\) observations, this model has more than enough data to estimate the fixed-effects.

import matplotlib.pyplot as plt

blog_colors = {
    "background": "#1c1c1d",
    "text": "#e8e8e8",
    "accent": "#2698ba",
    "point": "#ffffff",
}
β_posterior_mean = np.asarray(result["β_mean"])
limits = [
    min(β_true.min(), β_posterior_mean.min()) - 0.05,
    max(β_true.max(), β_posterior_mean.max()) + 0.05,
]

fig, ax = plt.subplots(figsize=(4.8, 3.6))
fig.patch.set_facecolor(blog_colors["background"])
ax.set_facecolor(blog_colors["background"])
ax.scatter(β_true, β_posterior_mean, s=36, color=blog_colors["point"], alpha=0.82)
ax.plot(limits, limits, color=blog_colors["accent"], linewidth=1.5)
ax.set_xlim(limits); ax.set_ylim(limits); ax.grid(False)
ax.set_xlabel("True fixed effect", color=blog_colors["text"]); ax.set_ylabel("Posterior mean fixed effect", color=blog_colors["text"])
ax.tick_params(colors=blog_colors["text"])
for spine in ax.spines.values():
    spine.set_color(blog_colors["text"])

Nice! The true value and posterior mean estimates line up perfectly. A job well done for MCMC.

I think that Bayes on GPU is tremendously undervalued. If this interests you or you have ideas to chat about, drop me a line.

To illustrate the problem better, we will use a dataset of uranium and vanadium point-referenced concentration measurements from Walker Lake. The data originate from Isaaks and Srivastava’s An Introduction to Applied Geostatistics and are distributed with the R package gstat.

Up to this point, you may have seen lots of examples of how Gaussian process models are use in robotics, spatial statistics, neuroscience, etc. Now, we work through a more exotic example which modifies the Gaussian process model to accommodate the case in which the actual location of our data points is not known precisely, and is only observed with substantial measurement noise. Spatial location error changes the covariance and prediction problem itself, a point emphasized in geostatistical work on location error and GP regression work on noisy spatial inputs by Cressie and Kornak and Cervone and Pillai.

This may seem like an unusual example at first glance, but we include it to give an example of how Bayesian modeling with appropriate priors lets us modify and change nearly any part of the model, given that we have some idea of how to represent our assumptions as part of the model process. Then, we use Monte Carlo methods to turn the inference crank and obtain reliable parameter estimates.

Introducing some more notation, let \(\tilde{\mathbf{s}}_i\) denote the recorded coordinate and \(\mathbf{s}_i\) the latent coordinate where the measurement actually occurred. We use \(\mathbf{s}_i = \tilde{\mathbf{s}}_i + \Delta_i\), with \(\Delta_i \sim \operatorname{Normal}(\mathbf{0}, \sigma_s^2 I_2)\), and evaluate the Gaussian process at \(\mathbf{s}_i\) rather than at \(\tilde{\mathbf{s}}_i\). Our choice of coordinate system here is somewhat arbitrary; we could also choose to work with polar coordinates and place priors over the magnitude and the angle of the location error. The scale \(\sigma_s\) is treated as known in this example so that the model represents different assumed levels of coordinate error.

This model is computationally harder than the fixed-location GP because the covariance matrix changes whenever the latent coordinates change. We use pm.gp.Marginal so that the latent GP values at the observations are integrated out.

We will construct datasets with increasing noise and examine how the model’s parameter estimates change. To begin, we perturb the original coordinates with increasing noise:

from pathlib import Path

import arviz as az
import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import pymc as pm
import seaborn as sns
from matplotlib import colors
from matplotlib.patches import Circle

from stsp import use_sepia_gunmetal

RANDOM_SEED = 8927
np.random.seed(RANDOM_SEED)

theme, sepia_cmap = use_sepia_gunmetal()

plot_bg_color = "#1c1c1d"
plot_text_color = "#e8e8e8"

plt.style.use("dark_background")
mpl.rcParams.update(
    {
        "figure.facecolor": plot_bg_color,
        "axes.facecolor": plot_bg_color,
        "savefig.facecolor": plot_bg_color,
        "savefig.edgecolor": plot_bg_color,
        "text.color": plot_text_color,
        "axes.labelcolor": plot_text_color,
        "xtick.color": plot_text_color,
        "ytick.color": plot_text_color,
        "axes.edgecolor": "#828282",
    }
)

theme["paper"] = plot_bg_color

df = pd.read_csv(Path("../../data/walker/cleaned.csv"))

X_walker = df[["x_m", "y_m"]].to_numpy()
y_walker = df["u_log10_p1"].to_numpy()

n_prediction_grid = 70
buffer_fraction = 0.1
x_range = df["x_m"].max() - df["x_m"].min()
y_range = df["y_m"].max() - df["y_m"].min()
x_limits = (df["x_m"].min() - buffer_fraction * x_range, df["x_m"].max() + buffer_fraction * x_range)
y_limits = (df["y_m"].min() - buffer_fraction * y_range, df["y_m"].max() + buffer_fraction * y_range)
x_new = np.linspace(*x_limits, n_prediction_grid)
y_new = np.linspace(*y_limits, n_prediction_grid)
x_new_mesh, y_new_mesh = np.meshgrid(x_new, y_new)
Xnew = np.column_stack([x_new_mesh.ravel(), y_new_mesh.ravel()])

eps = np.random.randn(*X_walker.shape)
multipliers = [12.0, 25.0, 40.0]
noisy_xs = [X_walker + m * eps for m in multipliers]

The perturbed coordinates are shown in Figure X. We keep the target variable unchanged in this analysis. Next, we construct the model, using the pm.Data object as a container to let us swap out coordinates easily.

location_coords = {"obs": np.arange(X_walker.shape[0]), "coord": ["x", "y"]}
selected_location_idx = np.sort(np.random.default_rng(RANDOM_SEED).choice(X_walker.shape[0], size=4, replace=False))
location_error_idatas = {}

with pm.Model(coords=location_coords) as location_error_gp_model:
    X_noisy = pm.Data("X_noisy", noisy_xs[0], dims=("obs", "coord"))
    y = pm.Data("y", y_walker, dims="obs")
    σ_s = pm.Data("σ_s", multipliers[0])

    Δs = pm.Normal("Δs", 0.0, σ_s, dims=("obs", "coord"))
    X_true = pm.Deterministic("X_true", X_noisy + Δs, dims=("obs", "coord"))

    μ = pm.Normal("μ", 2.0, 2.0)
    σ = pm.HalfNormal("σ", 1.0)
    ℓ = pm.HalfNormal("ℓ", 100.0)
    σ0 = pm.HalfNormal("σ0", 0.5)

    cov = σ**2 * pm.gp.cov.Matern52(2, ls=ℓ)
    gp_location = pm.gp.Marginal(mean_func=pm.gp.mean.Constant(μ), cov_func=cov)
    gp_location.marginal_likelihood("y_obs", X=X_true, y=y, sigma=σ0, dims="obs")

    for multiplier, X_noisy_value in zip(multipliers, noisy_xs):
        pm.set_data({"X_noisy": X_noisy_value, "σ_s": multiplier})
        location_error_idatas[multiplier] = pm.sample(
            chains=2,
            cores=2,
            target_accept=0.95,
            mp_ctx="spawn",
            random_seed=RANDOM_SEED,
            progressbar=False,
        )

location_error_parameter_summaries = {
    multiplier: az.summary(idata, var_names=["μ", "σ", "ℓ", "σ0"])
    for multiplier, idata in location_error_idatas.items()
}
location_error_diagnostics = pd.DataFrame(
    [
        {
            "Location error SD": multiplier,
            "Divergences": int(idata.sample_stats["diverging"].sum()),
            "Max $\\hat{R}$": location_error_parameter_summaries[multiplier]["r_hat"].max(),
            "Min bulk ESS": location_error_parameter_summaries[multiplier]["ess_bulk"].min(),
            "Mean displacement": np.sqrt((idata.posterior["Δs"] ** 2).sum("coord")).mean().item(),
        }
        for multiplier, idata in location_error_idatas.items()
    ]
)
location_error_diagnostics

Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (2 chains in 2 jobs)
NUTS: [Δs, μ, σ, ℓ, σ0]
Sampling 2 chains for 1_000 tune and 1_000 draw iterations (2_000 + 2_000 draws total) took 6340 seconds.
There were 24 divergences after tuning. Increase `target_accept` or reparameterize.
Chain 1 reached the maximum tree depth. Increase `max_treedepth`, increase `target_accept` or reparameterize.
We recommend running at least 4 chains for robust computation of convergence diagnostics
The rhat statistic is larger than 1.01 for some parameters. This indicates problems during sampling. See https://arxiv.org/abs/1903.08008 for details
The effective sample size per chain is smaller than 100 for some parameters.  A higher number is needed for reliable rhat and ess computation. See https://arxiv.org/abs/1903.08008 for details
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (2 chains in 2 jobs)
NUTS: [Δs, μ, σ, ℓ, σ0]
Sampling 2 chains for 1_000 tune and 1_000 draw iterations (2_000 + 2_000 draws total) took 6376 seconds.
There were 79 divergences after tuning. Increase `target_accept` or reparameterize.
Chain 0 reached the maximum tree depth. Increase `max_treedepth`, increase `target_accept` or reparameterize.
We recommend running at least 4 chains for robust computation of convergence diagnostics
The rhat statistic is larger than 1.01 for some parameters. This indicates problems during sampling. See https://arxiv.org/abs/1903.08008 for details
The effective sample size per chain is smaller than 100 for some parameters.  A higher number is needed for reliable rhat and ess computation. See https://arxiv.org/abs/1903.08008 for details
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (2 chains in 2 jobs)
NUTS: [Δs, μ, σ, ℓ, σ0]
Sampling 2 chains for 1_000 tune and 1_000 draw iterations (2_000 + 2_000 draws total) took 7756 seconds.
There were 52 divergences after tuning. Increase `target_accept` or reparameterize.
Chain 0 reached the maximum tree depth. Increase `max_treedepth`, increase `target_accept` or reparameterize.
Chain 1 reached the maximum tree depth. Increase `max_treedepth`, increase `target_accept` or reparameterize.
We recommend running at least 4 chains for robust computation of convergence diagnostics
The rhat statistic is larger than 1.01 for some parameters. This indicates problems during sampling. See https://arxiv.org/abs/1903.08008 for details
The effective sample size per chain is smaller than 100 for some parameters.  A higher number is needed for reliable rhat and ess computation. See https://arxiv.org/abs/1903.08008 for details

location_surface_grids = {}
with location_error_gp_model:
    for multiplier, X_noisy_value in zip(multipliers, noisy_xs):
        pm.set_data({"X_noisy": X_noisy_value, "σ_s": multiplier})
        posterior_mean_point = {
            name: location_error_idatas[multiplier].posterior[name].mean(("chain", "draw")).values
            for name in ["μ", "σ", "ℓ", "σ0", "Δs", "X_true"]
        }
        f_mean, _ = gp_location.predict(Xnew, point=posterior_mean_point, diag=True, pred_noise=False)
        location_surface_grids[multiplier] = f_mean.reshape(n_prediction_grid, n_prediction_grid)

naive_kde_grids = {}
kde_bandwidths = {
    multiplier: location_error_idatas[multiplier].posterior["ℓ"].mean(("chain", "draw")).item()
    for multiplier in multipliers
}
for multiplier, X_noisy_value in zip(multipliers, noisy_xs):
    squared_distance = (
        (Xnew[:, 0, None] - X_noisy_value[:, 0]) ** 2
        + (Xnew[:, 1, None] - X_noisy_value[:, 1]) ** 2
    )
    weights = np.exp(-0.5 * squared_distance / kde_bandwidths[multiplier]**2)
    naive_kde_grids[multiplier] = (weights @ y_walker / weights.sum(axis=1)).reshape(n_prediction_grid, n_prediction_grid)

location_norm = colors.Normalize(
    vmin=min(y_walker.min(), *(grid.min() for grid in location_surface_grids.values())),
    vmax=max(y_walker.max(), *(grid.max() for grid in location_surface_grids.values())),
)
point_colors = [theme["gunmetal"], theme["sepia"], theme["rust"], theme["steel"]]
arrow_head_width = 0.018 * max(x_range, y_range)

fig = plt.figure(figsize=(5, 11.2 * (5 / 7)))
gs = fig.add_gridspec(4, len(multipliers), hspace=0.08, wspace=0.08)
axes = np.array([[fig.add_subplot(gs[row, col]) for col in range(len(multipliers))] for row in range(4)])
panel_idx = 0
surface_meshes = []

for col, (multiplier, X_noisy_value) in enumerate(zip(multipliers, noisy_xs)):
    ax = axes[0, col]
    ax.scatter(X_walker[:, 0], X_walker[:, 1], facecolor=theme["paper"], edgecolor=theme["gunmetal"], s=14, linewidth=0.35, alpha=0.5)
    ax.scatter(X_noisy_value[:, 0], X_noisy_value[:, 1], c=y_walker, cmap=sepia_cmap, norm=location_norm, s=16, edgecolor=theme["paper"], linewidth=0.25, alpha=0.5)
    for point_idx, point_color in zip(selected_location_idx, point_colors):
        circle = Circle(X_noisy_value[point_idx], radius=multiplier, facecolor=colors.to_rgba(theme["steel"], 0.2), edgecolor=theme["gunmetal"], linewidth=0.6)
        ax.add_patch(circle)
        dx, dy = X_noisy_value[point_idx] - X_walker[point_idx]
        ax.arrow(X_walker[point_idx, 0], X_walker[point_idx, 1], dx, dy, color=plot_text_color, linewidth=0.8, length_includes_head=True, head_width=arrow_head_width, head_length=arrow_head_width)
    ax.scatter(X_walker[selected_location_idx, 0], X_walker[selected_location_idx, 1], facecolor=theme["paper"], edgecolor=plot_text_color, s=44, linewidth=0.8)
    ax.scatter(X_noisy_value[selected_location_idx, 0], X_noisy_value[selected_location_idx, 1], c=y_walker[selected_location_idx], cmap=sepia_cmap, norm=location_norm, s=44, edgecolor=plot_text_color, linewidth=0.8)
    ax.xaxis.set_label_position("top"); ax.set_xlabel(rf"$\sigma_s = {multiplier:.0f}$ m")

    ax = axes[2, col]
    mesh = ax.pcolormesh(x_new_mesh, y_new_mesh, location_surface_grids[multiplier], cmap=sepia_cmap, norm=location_norm, shading="auto")
    surface_meshes.append(mesh)

    ax = axes[3, col]
    ax.pcolormesh(x_new_mesh, y_new_mesh, naive_kde_grids[multiplier], cmap=sepia_cmap, norm=location_norm, shading="auto")

    ax = axes[1, col]
    ax.pcolormesh(x_new_mesh, y_new_mesh, location_surface_grids[multiplier], cmap=sepia_cmap, norm=location_norm, shading="auto", alpha=0.25)
    X_true_samples = location_error_idatas[multiplier].posterior["X_true"].stack(sample=("chain", "draw")).transpose("obs", "coord", "sample")
    for point_idx, point_color in zip(selected_location_idx, point_colors):
        x_draws = X_true_samples.isel(obs=point_idx).sel(coord="x").values
        y_draws = X_true_samples.isel(obs=point_idx).sel(coord="y").values
        sns.kdeplot(x=x_draws, y=y_draws, levels=3, color=point_color, linewidths=0.9, fill=False, ax=ax)
        ax.scatter(X_walker[point_idx, 0], X_walker[point_idx, 1], marker="x", color=point_color, s=34, linewidth=1.1)
        ax.scatter(X_noisy_value[point_idx, 0], X_noisy_value[point_idx, 1], marker="o", facecolor=theme["paper"], edgecolor=point_color, s=34, linewidth=1.1)

for row, row_label in enumerate(["Perturbed coordinates", "Posterior location density", r"Posterior mean of $f(s)$", "Naive smoothed KDE"]):
    axes[row, 0].set_ylabel(row_label)

for ax in axes.ravel():
    ax.set_xlim(*x_limits); ax.set_ylim(*y_limits); ax.set_aspect("equal"); ax.grid(False)
    ax.tick_params(axis="both", which="both", bottom=False, left=False, top=False, right=False, labelbottom=False, labelleft=False, labeltop=False, labelright=False)
    panel_idx += 1

cbar = fig.colorbar(surface_meshes[-1], ax=axes.ravel().tolist(), orientation="horizontal", fraction=0.035, pad=0.04)
cbar.set_label(r"Uranium, $\log_{10}(x + 1)$")

figure_dir = Path.home() / "ckrapu.github.io/images/2026-05-24-dont-know-where-your-data-is-from"
figure_dir.mkdir(parents=True, exist_ok=True)
fig.savefig(figure_dir / "error-in-location-grid.png", dpi=220, bbox_inches="tight", facecolor=fig.get_facecolor())

Figure 6 Data and posterior inference for a GP model with point coordinates which are only approximately observed. The top row shows the true, original coordinates in hollow circles while the noisily observed data points are shaded according to the uranium concentration at that point. The larger circles indicate the spatial scale of the coordinate error. The second row shows contours for the posterior high density region of the true location of four selected points in space; we see that as \(\sigma_s\) increases, these regions gracefully grow larger and larger.

As we examine Figure 6, we can see that some of the main features of the underlying surface are preserved even as the uncertainty in the coordinates grows. We see a few lighter features in the lower left corner, and a darker region in the top left and lower right. There is variation across these, but it is impressive that we can even work at all with this data given the severity of the perturbations! We include a comparison using a simpler approach, the Nadaraya-Watson Gaussian kernel smoother, with its bandwidth set to match the posterior mean estimate of the length scale under each scenario, and we find that it is unable to provide more than a rough average with little capacity to represent spatial variation.

There are some pretty sophisticated OCR tools out there but they tend to have usage limits or get expensive when you’re processing thousands of pages. DeepSeek recently released an open OCR model that handles mathematical notation well, and I figured I could run it myself if I had access to a GPU. Sadly, my daily driver is a decade-old Titan Xp which no longer supports the latest PyTorch versions and thus can’t run DeepSeek OCR.

I ended up using Modal for this.

What is Modal?

Modal is a serverless compute platform that lets you run Python code on cloud infrastructure without managing servers. The killer feature for machine learning work is that you can define a container image, attach a GPU, and pay only for the seconds your code is actually running.

import modal

image = modal.Image.from_registry(
    "nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04",
    add_python="3.11",
).pip_install("torch", "transformers", ...)

app = modal.App("my-gpu-app")

@app.function(image=image, gpu="A100")
def process_something():
    # This runs on an A100 with all your deps installed
    pass

The decorator pattern is what makes Modal pleasant to use. You write normal Python, sprinkle decorators on the functions that need special hardware, and Modal handles the rest: building the container, provisioning the GPU, routing your requests. For OCR, this is perfect.

The OCR script

The core idea is simple: deploy a FastAPI server on Modal that accepts images and returns markdown text. Let’s walk through the important pieces.

Defining the container image

First, we build a container with all the dependencies. DeepSeek’s OCR model needs PyTorch, transformers, and a few image processing libraries:

from pathlib import Path
import modal

APP_NAME = "deepseek-ocr-books-api-batch"
ROOT = Path(__file__).resolve().parents[1]
BOOKS_DIR = ROOT / "references" / "books"
PARSED_DIR = BOOKS_DIR / "parsed"

image = (
    modal.Image.from_registry(
        "nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04",
        add_python="3.11",
    )
    .apt_install("git", "libgl1", "libglib2.0-0")
    .pip_install(
        "torch==2.6.0",
        "torchvision==0.21.0",
        "transformers==4.46.3",
        "PyMuPDF",
        "Pillow",
        "numpy",
        extra_index_url="https://download.pytorch.org/whl/cu118",
    )
)

app = modal.App(APP_NAME)

The paths at the top let us find PDFs relative to the script location, keeping configuration close to where it’s used.

The FastAPI endpoint

Here’s where the magic happens. We wrap a FastAPI server in Modal’s @modal.asgi_app() decorator, which means Modal will handle spinning up GPU instances and routing HTTP requests to them:

@app.function(image=image, gpu="A100", timeout=60 * 60 *2) # timeout of 2 hours
@modal.asgi_app()
def fastapi_app():
    from fastapi import FastAPI, File, UploadFile
    from PIL import Image
    import torch
    from transformers import AutoModel, AutoTokenizer

    api = FastAPI()
    
    model_name = "deepseek-ai/DeepSeek-OCR"
    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
    model = AutoModel.from_pretrained(model_name, trust_remote_code=True)
    model = model.cuda().to(torch.bfloat16).eval()

The model loads once when the container starts. Subsequent requests reuse the same loaded model, which is crucial for throughput when processing hundreds of pages.

Handling batched inference

OCR is embarrassingly parallel since each page is independent. We can process multiple pages in a single forward pass through the model, which is faster than processing them one at a time:

@api.post("/ocr_batch")
async def ocr_batch(files: list[UploadFile] = File(...)) -> dict[str, list[str]]:
    images = []
    for file in files:
        image_bytes = await file.read()
        images.append(Image.open(io.BytesIO(image_bytes)).convert("RGB"))
    batch_items = [prepare_inputs(image) for image in images]
    texts = run_batch(batch_items)
    return {"texts": texts}

The run_batch function handles the actual model inference. It pads inputs to the same length, runs them through the model in one shot, and decodes the outputs:

def run_batch(batch_items):
    # Pad sequences to same length
    lengths = [item[0].size(0) for item in batch_items]
    max_len = max(lengths)
    
    input_ids = torch.full((len(batch_items), max_len), pad_id, dtype=torch.long)
    # ... (padding logic)
    
    with torch.autocast("cuda", dtype=torch.bfloat16):
        with torch.no_grad():
            output_ids = model.generate(
                input_ids.cuda(),
                images=images,
                max_new_tokens=8192,
                temperature=0.0,
            )
    
    # Decode outputs
    outputs = []
    for i, out_ids in enumerate(output_ids):
        token_ids = out_ids[lengths[i]:].tolist()
        text = tokenizer.decode(token_ids, skip_special_tokens=False)
        outputs.append(text.strip())
    return outputs

Setting temperature=0.0 makes the output deterministic, which helps the model generate results which are more reproducible.

The local client

With the server deployed on Modal, we need a client to feed it pages. The @app.local_entrypoint() decorator marks a function that runs on your local machine but can communicate with the Modal-deployed server:

@app.local_entrypoint()
def main(api_url: str, book: str = "", max_pages: int = None, batch_size: int = 1):
    import fitz  # PyMuPDF
    
    if book:
        pdf_paths = [BOOKS_DIR / book]
    else:
        pdf_paths = sorted(BOOKS_DIR.glob("*.pdf"))
    
    for pdf_path in pdf_paths:
        with fitz.open(pdf_path) as doc:
            batch_pages = []
            for page_index in range(doc.page_count):
                page = doc[page_index]
                pix = page.get_pixmap(matrix=fitz.Matrix(2, 2))  # 2x zoom
                batch_pages.append(pix.tobytes("png"))
                
                if len(batch_pages) >= batch_size:
                    # Send batch to server
                    response = requests.post(f"{api_url}/ocr_batch", files=files)
                    texts = response.json()["texts"]
                    # Save results...

The render-at-2x trick (fitz.Matrix(2, 2)) is important. Higher resolution input means the OCR model can read smaller text and mathematical subscripts more accurately.

Cleaning up the output

DeepSeek’s OCR model includes grounding tags, which are coordinates indicating where each piece of text appeared on the page. These can be useful for some applications, but I didn’t need them for searchable text:

tag_pattern = re.compile(
    r"<\|ref\|>(.*?)<\|/ref\|><\|det\|>.*?<\|/det\|>",
    flags=re.DOTALL,
)

for page_idx, text in zip(batch_page_indices, texts):
    text = tag_pattern.sub(r"\1", text)  # Keep text, drop coordinates
    page_path = pages_dir / f"page_{page_idx + 1:04d}.mmd"
    page_path.write_text(text, encoding="utf-8")

The .mmd extension stands for “multimodal markdown,” a convention for markdown that came from an OCR source and might have some artifacts.

Running it

To use this script, you first deploy the server:

modal deploy deepseek_ocr_modal.py

This gives you a URL like https://your-workspace--deepseek-ocr-books-api-batch-fastapi-app.modal.run. Then run the client:

modal run deepseek_ocr_modal.py --api-url "https://..." --book "Gelman - Bayesian Data Analysis.pdf"

For BDA’s ~600 pages, with a batch size of 4, processing takes about 45 minutes on an A100. The output is a directory of markdown files, one per page, plus a concatenated document.mmd with page markers. The whole thing cost maybe $2 bucks. I think it’s a great deal. I also now have a setup I can reuse for any PDF, including course notes, papers, and other textbooks.

The OCR quality on mathematical content is surprisingly good. Nearly all equations come through intact.

The real payoff comes from downstream use. I can now grep through BDA, paste sections into Claude and ask it to explain the notation, or build a proper search index. All of this from a PDF that was previously just a collection of images.

Here’s what the text looks like once it’s parsed and combined into a single file:

## Exponential model

The exponential distribution is commonly used to model 'waiting times' and other continuous, positive, real- valued random variables, often measured on a time scale. The sampling distribution of an outcome \(y\) , given parameter \(\theta\) , is

\[p(y|\theta) = \theta \exp (-y\theta),\mathrm{for} y  0,\]

and \(\theta = 1 / \mathrm{E}(y|\theta)\) is called the 'rate.' Mathematically, the exponential is a special case of the gamma distribution with the parameters \((\alpha ,\beta) = (1,\theta)\) . In this case, however, it is being used as a sampling distribution for an outcome \(y\) , not a prior distribution for a parameter \(\theta\) , as in the Poisson example.

The exponential distribution has a 'memoryless' property that makes it a natural model for survival or lifetime data; the probability that an object survives an additional length of time \(t\) is independent of the time elapsed to this point: \(\operatorname *{Pr}(y  t + s\mid y  s,\theta) = \operatorname *{Pr}(y  t\mid \theta)\) for any \(s,t\) . The conjugate prior distribution for the exponential parameter \(\theta\) , as for the Poisson mean, is \(\operatorname {Gamma}(\theta |\alpha ,\beta)\) with corresponding posterior distribution \(\operatorname {Gamma}(\theta |\alpha +1,\beta +y)\) . The sampling distribution of \(n\) independent exponential observations, \(y = (y_{1},\ldots ,y_{n})\) , with constant rate \(\theta\) is

\[p(y|\theta) = \theta^{n}\exp (-n\bar{y}\theta),\mathrm{for}\bar{y}\geq 0,\]

which when viewed as the likelihood of \(\theta\) , for fixed \(y\) , is proportional to a \(\operatorname {Gamma}(n + 1,n\bar{y})\) density. Thus the \(\operatorname {Gamma}(\alpha ,\beta)\) prior distribution for \(\theta\) can be viewed as \(\alpha - 1\) exponential observations with total waiting time \(\beta\) (see Exercise 2.19).
<--- Page Split --->
image
image_caption
<centerFigure 2.6 The counties of the United States with the highest \(10\%\) age-standardized death rates for cancer of kidney/ureter for U.S. white males, 1980-1989. Why are most of the shaded counties in the middle of the country? See Section 2.7 for discussion. </center

### 2.7 Example: informative prior distribution for cancer rates

At the end of Section 2.4, we considered the effect of the prior distribution on inference given a fixed quantity of data. Here, in contrast, we consider a large set of inferences, each based on different data but with a common prior distribution. In addition to illustrating the role of the prior distribution, this example introduces hierarchical modeling, to which we return in Chapter 5.

Compare that to the original passage from the textbook:

I think it did a pretty good job!

If you have a collection of scanned textbooks, this approach might be worth the 5 minutes it takes to set it up.

TL;DR: I’m in a non-engineering department and somehow became an extreme outlier for AI-assisted coding usage because I have to wear too many hats and I’m picky about our dependencies.

Last spring, my company onboarded onto Cursor and shortly after, it was gaining traction in virtually every department. I didn’t see myself as the target audience—my department doesn’t ship product code and, before Cursor, I generally had a low-tech approach: I had autocomplete turned off most of the time, and my AI assistance was mostly limited to ChatGPT buried 50 tabs deep in Chrome.

I didn’t know Cursor stats were visible to most managers, and when I found the department-level reports, I was genuinely worried at first. I was one of the more aggressive users of Claude 4 Opus at the time, and I thought I might be reprimanded for incurring too many API fees.

The usage dashboard showed that I was in the top 0.XX% of users for activity (as measured by Agent requests), but also that I had a very unusual activity pattern of roughly 4x as many Ask requests (LLM only explains, like ChatGPT) as compared to Agent requests (LLM actually writes code in the project directory structure). This didn’t surprise me much; I enjoy having a back-and-forth dialogue about all the technologies I don’t know yet. I find this almost as much fun as writing the code itself. Most users appeared to be showing a 0.3-0.5x ratio. They were less interested in having the AI explain it and simply wanted the code to be written. I will also note that this is all good-faith activity; I wasn’t churning out slop or using it for non-work purposes.

Despite being an early adopter of some AI technology (I started building around GPT-3 in early 2021) I genuinely enjoy typing, writing out functions, scaffolding, and the minutiae of writing code. To me, it’s the same feeling as watching TV or making small talk. Most of it doesn’t require a huge amount of thinking, and it feels pretty nice when you’re in the flow. I never thought of myself as an AI-first, 10x developer type of person.

Here’s what drove my usage:

1. Documentation gaps, rich tooling

We have an internal environment rich with tools, but light on documentation and frequently requiring workarounds. Documentation is lacking at many companies, and it’s super common for a service’s API to drift out of alignment with docs, especially when it’s just an outdated guide someone cobbled together.

I spent a lot of effort working around this. AI was a panacea for this problem. Lots of nominally-working but hard-to-use services were now perfectly usable, all of a sudden, because the LLM has the patience to poke around the API, try a dozen examples, read a few different conflicting sources, and figure out how to actually make it work.

2. No devops engineer (and that’s fine)

I don’t have a devops engineer, I will never get one, and that’s okay with me. Most of my projects are deployed via managed cloud services instead of k8s. Early on, I made the judgment call that, even in the best circumstances, our traffic was never going to exceed \(10^4\) active users and cloud costs wouldn’t be an issue, even unoptimized.

This means that I can largely control the entire apparatus via Terraform without needing an additional teammate to manage a k8s cluster.

3. The technical lead does analytics

In this case, that was me. Many teams operate by having BA or junior DS-type employees perform analytics and reporting. They’re cheaper, but less likely to automate their work given less coding experience—and I say this as someone with a DS background.

Since the reporting needs fell on my shoulders, it was definitely in my interest to vibe code as many SQL queries and dashboarding functions as possible to short-circuit the route from raw data to management decisions.

4. Not waiting for MCPs when a .cursorrules and some bash aliases would suffice

Pretty much all companies’ IT departments have been talking about MCP servers for the last year. Some are further along than others. I, being inordinately impatient, just wrote a few bash aliases to make API calls and, for many operational tasks (“run a report against our sales data”, “find all the newest companies we haven’t reached out to yet”, etc.), I could just tell Cursor that it had these command line tools for using Perplexity, Databricks, our company’s Powerpoint template, etc.

5. A great full-stack template

Many of the similar projects I’d seen at my company and elsewhere used software like Streamlit, Gradio, and OpenWebUI. These are fine, but they aren’t ready out-of-the-box for going into production.

I invested a lot of time in early 2024 figuring out how we’d deploy an internal chatbot and eventually landed on Chainlit. This was a stroke of luck as it was a solid template with batteries included for OAuth, telemetry, human-in-the-loop, and other features that many teams were authoring from scratch at that point. It had quirks (the default schema tanks after 100k conversations) but otherwise was great to work with. It also pushed us to work mostly in FastAPI + LlamaIndex + SQLAlchemy, all of which turned out to be rock-solid in production.

Ironically, this OSS project was a much better base than all of the paid consulting/dev shop-produced products that I’d seen.

By scaffolding around a quality codebase with sensible defaults, clean API structure, and a minimal but appealing UI, grounding on that foundation gave us much more effective LLM-produced code than if we had vibe coded from scratch.

These effects were compounding for me: it produced nice code which I enjoyed working on, which in turn made me want to do more of it instead of getting frustrated by opaque errors.

Since Gemini 3 and Google’s Antigravity IDE were released recently (and yes, I am aware it is basically Windsurf), I wanted to give it a try with an easy one - the Ising model of ferromagnetism.

Here’s what Antigravity with Gemini 3 Pro cooked up in an hour:

I’ve gotten decent results out of the last crop of OpenAI and Anthropic models, but the Chrome browser extension for retrieving the DOM really helped too. It’s a great feature, and I expect Cursor to have something similar soon. I think some of the other UI features like showing subtasks and intermediate steps were a little unnecessary. Overall, great work by the former Windsurf team and the other G staff!

This notebook provides a simple first example of using Warp to simulate a particle moving along a one-dimensional flow field. We’ll then optimize the direction and magnitude of the flow field to push the particle towards a pair of target locations specified at two distinct timesteps. It doesn’t assume much background in fluid mechanics or machine learning, and shows the basic workflow for manipulating the parameters of a simulation using Warp.

Simulating a particle’s movement in 1D

We start by importing packages and defining the simulation parameters. For the purposes of this project, Warp can be considered as a drop-in replacement for PyTorch or Jax - it includes both the tensor libraries and automatic differentiation logic.

from dataclasses import dataclass
from IPython.display import HTML
from matplotlib.animation import FuncAnimation
from tqdm import tqdm
from IPython.display import Image

import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
import warp as wp

plt.style.use('dark_background')

plot_bg_color = "#1c1c1d"

mpl.rcParams['figure.facecolor'] = plot_bg_color
mpl.rcParams['axes.facecolor'] = plot_bg_color

np.random.seed(8271991)

device = "cuda"

# Instantiate a variable to force Warp to start up
# and show its version + device
wp.array([0], dtype=float, device=device)

array(shape=(1,), dtype=float32)

Here, we specify that the domain will be one meter long and divided into 100 grid cells for numerical simulation. The entire simulation will run for 30 seconds in increments of 100 milliseconds.

@dataclass
class ParticleExampleConfig:
    length_m: float = 1.0
    nx: int = 100
    total_time_s: float = 30.
    dt: float = 0.1
    start_location_m = 0.1

particle_cfg = ParticleExampleConfig()

Next, we define the logic for tracking the location of the particle as it is moved around via the velocity field.

def simulate_particle(particle_cfg: ParticleExampleConfig, u0: wp.array) -> wp.array:
      '''
      Runs a simple 1D simulation, taking the flow field (a 1D vector in this case) and produces
      an array with the location of the particle at each timestep.
      '''
      n_steps = int(particle_cfg.total_time_s / particle_cfg.dt)
      positions = wp.zeros(n_steps + 1, dtype=wp.float32, requires_grad=True)

      @wp.kernel
      def integrate_trajectory(
          positions: wp.array(dtype=wp.float32),
          u0: wp.array(dtype=wp.float32),
          start_pos: float,
          dt: float,
          dx: float,
          nx: int,
          n_steps: int
      ):
          tid = wp.tid()

          if tid == 0:
              positions[0] = start_pos
              velocity = float(0.0)

              for step in range(n_steps):
                  current_pos = positions[step]

                  cell_idx = wp.int32(current_pos / dx)
                  cell_idx = wp.clamp(cell_idx, 0, nx - 1)

                  flow_velocity = u0[cell_idx]
                  force = flow_velocity - velocity

                  velocity = velocity + force * dt

                  new_pos = current_pos + velocity * dt
                  new_pos = wp.clamp(new_pos, 0.0, dx * float(nx))

                  positions[step + 1] = new_pos

      dx = particle_cfg.length_m / particle_cfg.nx

      wp.launch(
          integrate_trajectory,
          dim=1,
          inputs=[
              positions,
              u0,
              particle_cfg.start_location_m,
              particle_cfg.dt,
              dx,
              particle_cfg.nx,
              n_steps
          ]
      )

      return positions

Let’s go ahead and run it, with a velocity field which is randomly initialized but with a bias pointing towards the center of the domain.

u0_numpy = np.random.randn(particle_cfg.nx) + np.linspace(1, -1, particle_cfg.nx) * 2.5
u0 = wp.array(u0_numpy, dtype=float, device=device)
positions = simulate_particle(particle_cfg, u0)

plt.figure(figsize=(6, 2))
plt.xlabel("Time (s)"), plt.ylabel("Position (m)")
plt.title("Particle location over time")
plt.plot(positions.numpy())
plt.savefig("visuals/location_over_time.png")
plt.close()
Image("visuals/location_over_time.png")

We can also make an animation to provide a dynamic view of the particle’s behavior subject to the 1D velocity field.

fig, ax = plt.subplots(figsize=(8, 2))

ax.set_xlim(0, particle_cfg.length_m)
ax.set_ylim(-0.5, 0.5)
ax.set_xlabel("Position (m)")

# Turn off y-axis + spines
ax.axes.get_yaxis().set_visible(False)
for spine in ['top', 'right', 'left']:
    ax.spines[spine].set_visible(False)

# Add tight layout to prevent xlabel cutoff
plt.tight_layout()

particle = ax.scatter([], [], s=100, color='cyan')
time_text = ax.text(0.02, 0.85, '', transform=ax.transAxes, fontsize=10)

# Get position data
positions_np = positions.numpy()
n_frames = len(positions_np)

def init():
    particle.set_offsets(np.empty((0, 2)))
    time_text.set_text('')
    return particle, time_text

def animate(frame):
    particle.set_offsets([[positions_np[frame], 0]])
    time_text.set_text(f'Time: {frame * particle_cfg.dt:.1f}s')
    
    return particle, time_text

anim = FuncAnimation(fig, animate, init_func=init, 
                    frames=n_frames, interval=40, 
                    blit=True, repeat=True)
anim.save('visuals/01-position-animation.gif', writer='pillow', fps=25, dpi=80)
HTML(anim.to_jshtml())

Great! We’ve got a simple model for a particle moving subject to a velocity field. Let’s move onto the meat of this tutorial - manipulating the model parameters to do what we want.

Optimizing the velocity field via automatic differentiation

We’ve passed an initial test of basic functionality for our model; we assume a simple velocity field and let the particle be moved around subject to simple kinematics.

Next, we want to show how to optimize an arbitrary model quantity to perform an interesting task. Our new objective is to identify a setting of u0 that pushes the particle to a target location at a given timestep. Our optimization strategy is as follows:

For each iteration:

• We take the initial velocity field u0 which is just a 1D array for this problem, and run the forward simulation. We use the tape to record the execution trace
• Next, we use reverse-mode autodiff in Warp to get grad which represents \(\nabla_{u_0} L(x_1,...,x_T)\) where \(x_1,...,x_T\) denotes the particle position at each timestep.
• We apply an update to u0 per the update rule \(u_0^{(i+1)} = u_0^{(i)} - \eta ^{(i)} \cdot \nabla_{u_0^{i}} \left[ L(x_1^{(i)},...,x_T^{(i)}) \right]\) with \(u_0^{(i+1)}\) as the value of the velocity field at the \(i+1\)-th iteration and \(x_1^{(i-1)}\) denoting the value of the position of the particle at timestep 1 in optimization iteration \(i-1\). Our loss function \(L\) is simply the sum of the squared distances between the position locations and the target locations at the provided timesteps.
• We decrease the update size with \(\eta ^{(i+1)} = \eta ^{(i)} \times 0.99\)

'''
The configs for this version are mostly the same; we just want
to include a few validation checks.
'''

@dataclass
class ParticleOptimizationConfig(ParticleExampleConfig):
    target_times_s: tuple[float] = (1., 15.)
    target_locations_m: tuple[float] = (0.2, 0.75)
    step_size = 0.1
    n_iterations = 25
    decay_rate = 0.99 # Decreases the step size by this factor each iteration

    def __post_init__(self):
        if len(self.target_times_s) != len(self.target_locations_m):
            raise ValueError("target_times_s and target_locations_m must have the same length.")
        if any(t < 0 or t > self.total_time_s for t in self.target_times_s):
            raise ValueError("All target times must be within the total simulation time.")
        if any(l < 0 or l > self.length_m for l in self.target_locations_m):
            raise ValueError("All target locations must be within the simulation length.")
        
opt_cfg = ParticleOptimizationConfig()
    

Applying autodiff and gradient descent

The above configuration specifies that at 1 and 15 seconds in simulation time we would like the particle to be close to locations of 0.2 m and 0.75 m from the left-hand side of the domain. We’ll use Warp’s implementation of automatic differentiation with gradient descent and a simple decreasing stepsize schedule to find an optimal u0 such that the particle respects the previous conditions.

u0_history = []
u0 = wp.array(np.random.randn(opt_cfg.nx) * 0.1, dtype=wp.float32, device=device, requires_grad=True)

target_indices = [int(t / opt_cfg.dt) for t in opt_cfg.target_times_s]
target_positions = wp.array(opt_cfg.target_locations_m, dtype=wp.float32, device=device)

loss_history = []
grad_norm_history = []

pbar = tqdm(range(opt_cfg.n_iterations), desc="")
step_size = opt_cfg.step_size

for iteration in pbar:
    u0_history.append(u0.numpy().copy())
    
    tape = wp.Tape()
    
    with tape:
        positions = simulate_particle(opt_cfg, u0)
        
        loss = wp.zeros(1, dtype=wp.float32, device=device, requires_grad=True)
        
        @wp.kernel
        def compute_l2_loss(
            positions: wp.array(dtype=wp.float32),
            target_positions: wp.array(dtype=wp.float32),
            target_indices: wp.array(dtype=wp.int32),
            loss: wp.array(dtype=wp.float32)
        ):
            tid = wp.tid()
            if tid <= target_indices.shape[0]:
                idx = target_indices[tid]
                diff = positions[idx] - target_positions[tid]
                wp.atomic_add(loss, 0, diff * diff)
        
        target_indices_array = wp.array(target_indices, dtype=wp.int32, device=device)
        wp.launch(compute_l2_loss, dim=len(target_indices), 
                 inputs=[positions, target_positions, target_indices_array, loss])
    
    tape.backward(loss)
    
    loss_val = loss.numpy()[0]
    grad = tape.gradients[u0].numpy()
    grad_norm = np.linalg.norm(grad)
    
    loss_history.append(loss_val)
    grad_norm_history.append(grad_norm)
    
    pbar.set_description(f"Iter {iteration}: Loss = {loss_val:.6f}, Grad norm = {grad_norm:.6f}")
    
    u0_np = u0.numpy()
    u0_np -= opt_cfg.step_size * grad
    u0 = wp.array(u0_np, dtype=wp.float32, device=device, requires_grad=True)
    
    step_size *= opt_cfg.decay_rate

Iter 24: Loss = 0.000754, Grad norm = 0.028800: 100%|██████████| 25/25 [00:00<00:00, 156.56it/s]

Module __main__ 331110c load on device 'cuda:0' took 0.20 ms  (cached)

Animating the optimized trajectory

Our optimization converges to near-zero loss quite quickly. Let’s see what the solution looks like:

optimized_positions = simulate_particle(opt_cfg, u0)

fig, ax = plt.subplots(figsize=(8, 2))

ax.set_xlim(0, opt_cfg.length_m)
ax.set_ylim(-0.5, 0.5)
ax.set_xlabel("Position (m)")

ax.axes.get_yaxis().set_visible(False)
for spine in ['top', 'right', 'left']:
    ax.spines[spine].set_visible(False)

plt.tight_layout()

particle = ax.scatter([], [], s=100, color='cyan')
time_text = ax.text(0.02, 0.85, '', transform=ax.transAxes, fontsize=10)

for t, loc in zip(opt_cfg.target_times_s, opt_cfg.target_locations_m):
    ax.axvline(x=loc, color='cyan', linestyle='--', alpha=0.5)
    ax.text(loc, 0.3, f't={t}s', ha='center', fontsize=8, color='cyan')

positions_np = optimized_positions.numpy()
n_frames = len(positions_np)

def init():
    particle.set_offsets(np.empty((0, 2)))
    time_text.set_text('')
    return particle, time_text

def animate(frame):
    particle.set_offsets([[positions_np[frame], 0]])
    time_text.set_text(f'Time: {frame * opt_cfg.dt:.1f}s')
    
    return particle, time_text
    
anim = FuncAnimation(fig, animate, init_func=init, 
                    frames=n_frames, interval=40, 
                    blit=True, repeat=True)
anim.save('visuals/01-optimize-animation.gif', writer='pillow', fps=25, dpi=80)
HTML(anim.to_jshtml())

To understand how the solution evolves over iterations, let’s animate the particle movement for several values of u0 obtained at the start, middle, and end of optimization.

n_rows, n_cols = 2, 6

n_subplots = n_rows * n_cols
step = max(1, opt_cfg.n_iterations // n_subplots)
selected_iterations = list(range(0, opt_cfg.n_iterations, step))[:n_subplots]

fig, axes = plt.subplots(n_rows, n_cols, figsize=(16, n_rows*1.5))
axes = axes.flatten()

for i, ax in enumerate(axes):
    ax.set_xlim(0, opt_cfg.length_m)
    ax.set_ylim(-0.5, 0.5)
    ax.set_xticks([0, 0.5, 1.0])
    ax.tick_params(labelsize=8)

    ax.axis('off')

plt.tight_layout()

particles = []
titles = []

for i, (ax, iter_idx) in enumerate(zip(axes, selected_iterations)):
    particle = ax.scatter([], [], s=50, color='cyan', edgecolor='w')
    particles.append(particle)
    
    
    title = ax.text(0.05, 0.9, f'Iter. {iter_idx}', transform=ax.transAxes, 
                   ha='center', fontsize=11, color='white')
    titles.append(title)
    
    for t, loc in zip(opt_cfg.target_times_s, opt_cfg.target_locations_m):
        ax.axvline(x=loc, color='cyan', linestyle='--', alpha=0.2, linewidth=3.0)

def init():
    for particle in particles:
        particle.set_offsets(np.empty((0, 2)))
    return particles

def animate(frame):
    frame = frame * 2
    
    for i, (particle, iter_idx) in enumerate(zip(particles, selected_iterations)):
        u0_iter = wp.array(u0_history[iter_idx], dtype=wp.float32, device=device)
        positions = simulate_particle(opt_cfg, u0_iter)
        positions_np = positions.numpy()
        
        if frame < len(positions_np):
            particle.set_offsets([[positions_np[frame], 0]])
    
    return particles

n_frames = int(opt_cfg.total_time_s / opt_cfg.dt) + 1
anim = FuncAnimation(fig, animate, init_func=init, 
                    frames=n_frames//2, interval=80, 
                    blit=True, repeat=True)
anim.save('visuals/01-history-animation.gif', writer='pillow', fps=25, dpi=80)
HTML(anim.to_jshtml())

Language models have no reason to flip coins. They are (broadly speaking) not capable of generating random numbers nor should they be able to do so - there’s no step in the data or training pipelines for LLMs that encourages this behavior.

But, what if they could? What would it take to produce an LLM capable of drawing a dozen numbers from a simple distribution with statistically reliable results?

The point isn’t to see whether \(\mathcal{M}_{transformers}\) is sufficiently expressive to implement a decent random number generator (surprise, it is) but rather to see if our best computational model for intelligence can be pushed to do it, and if it offers any benefits.

Towards this end, I wanted to see what happened if I used a simple prompt for simulating a coin flip with as many models as possible.

I acknowledge this is a dumb question. Here is a chance for you to cease reading and go about your day without any ill feelings from me.

Flipping coins across all of OpenRouter

OpenRouter supports over 200 large language models’ inference endpoints. Most of them support text and are tuned for instruction-following and chat. There’s a delightful range of models including the tried-and-true gpt-4-turbo, the roleplaying models, the foreign-language finetunes like scb10x/llama-3-typhoon-v1.5x-70b-instruct, and even cutting-edge diffusion models from like mercury/inception.

I wrote a Python notebook to run API calls against each of them using a simple prompt borrowed from a paper on prompting for rejection sampling for flipping coins:

prompt = '''
  ** Descriptions: ** Your task is to generate a sample consisting of a single draw from the set {0, 1},
  where both 0 and 1 have a probability of 0.5. Please give your output strictly in the following format:

  <format>
  Explanations: [Your step-by-step analyses and results;
  You DO NOT have access to a computer or a random number generator]
  Output: [Your output MUST be an integer in {0, 1}]
  </format>

  Please ONLY reply according to this format, don't give me any other words.
  Write the output EXACTLY like "Output: 0" or "Output: 1". Make sure to include the "Output: " prefix.
  '''

Now, there are some problems with this prompt. Ideally, we’d replace 0 and 1 with randomly generated strings and marginalize over both orderings. This would be more work, so I skipped it.

We would also use structured generation to force outputs into {0, 1}. I skipped the latter because not all endpoints support structured generation.

If you were to run a similar analysis, here’s a list of models I would skip:

models_to_skip = (
    'eleutherai/llemma_7b',
    'google/gemini-2.5-pro-preview-05-06',
    'meta-llama/llama-3.1-405b', # Had some weird errors
    'meta-llama/llama-guard-2-8b',
    'meta-llama/llama-guard-3-8b',
    'meta-llama/llama-guard-4-12b',
    'mistralai/mistral-7b-instruct-v0.2',
    'morph/morph-v3',
    'morph/morph-v3-fast',
    'morph/morph-v3-large',
    'openai/gpt-4o-mini-search-preview-2025-03-11',
    'openai/gpt-5',  # Requires a key to access it via OpenRouter
    'openai/o1-mini',
    'openai/o1-pro',
    'openai/o3-pro',
    'openrouter/auto',
    'perplexity/sonar-deep-research',
    'qwen/qwq-32b-preview',
    'switchpoint/router',
    'x-ai/grok-vision-beta',
)

Some of these models cost too much to run (sonar-deep-research), others don’t support a standard chat API (morph-v3) while others are aggressively rate-limited or require additional login info (google and openai, respectively).

This leaves you with roughly 240 models. A quick random sample of their IDs shows frontier labs, big corpos like meta, as well as startups like liquid, ai21 and kimi. All told, I spent about 80 bucks with a few small test runs. A lot of that was due to me forgetting that sonar-deep-research was enabled and that is a very expensive API to run.

Here are some of the most expensive models to run:

I was pleasantly surprised that I could get \(n\approx 50\) draws from each model for less than $40. I’d certainly like to do more cross-model experiments using OpenRouter.

Which models are biased the most?

I did a bunch of work on creating predictor features for each model using Perplexity’s API and a prompt like “Categorize model X as dense / mixture of experts and also determine if it uses reasoning…”. I used these as the covariates for a Bayesian binomial regression to see if there was a strong effect discernable from the data for aspects like model architecture, size, the organization that released it, etc.

Unfortunately, there weren’t any interesting findings. I tried with different parameterizations and likelihoods. The only effect I repeatedly found was that larger models tended to be less biased.

The actual results

I had a much better experience by just listing the models as ordered by their bias, defined as the absolute value of the average of their draws (confined to {0,1}) minus 0.5. In other words, it’s the distance of the sample mean from a fair coin flip’s \(p(heads)\). Smaller is better. Here are the 20 least biased models:

Something that stood out to me immediately was the multiple roleplaying finetunes, like infermatic/mn-inferor-12b, arliai/qwq-32b-arliai-rpr-v1 , and sao10k/l3.3-euryale-70b among others.

None of the RP models show up on the top 20 most-biased:

Here, we see many of the heavy-hitters from the last 18 months. The best Claude models are all represented here, as is deepseek-chat-v3-0324 and qwen3-235b-a22b-thinking-2507. The coding-specialized models fron Mercury also both show up here.

This hints at a connection between sampling reliability and creativity. The models which are clearly most tuned for coding/problem solving like Claude and Deepseek are also the worst at giving a balanced sample of coin flips. Conversely, models finetuned for creative and interesting dialogue tend to show less bias.

This notebook aims to demonstrate how to perform a series of simplified models for understanding and visualizing these waves. We will start with basic analytical solutions and progress towards more complex simulations to explore their behavior and characteristics.

We’ll start by setting up the key parameters for our simulation. These parameters are fundamental to defining the environment and conditions under which the Rossby waves will be modeled. They include spatial dimensions such as the size of the domain in both the x and y directions, which will represent a portion of the Earth’s surface. We also define the grid resolution by specifying the number of grid cells in each direction, nx and ny, which directly impacts the level of detail captured in the simulation and the computational cost.

The time parameters, such as max_time_s and the number of iterations, determine the duration of the simulation and the frequency at which we observe the wave’s evolution. We set the angular velocity, omega_rad_s, which is the physical property driving the Coriolis effect, a fundamental force behind the generation and behavior of Rossby waves.

import numpy as np
import tqdm
import scipy
import matplotlib.pyplot as plt
import warnings

from matplotlib.animation import FuncAnimation
from IPython.display import HTML, clear_output
import matplotlib as mpl
from dataclasses import dataclass

plt.style.use('dark_background')

# Numpy complains when we use .real on a complex-valued array, so
# we suppress that warning to keep the notebook short.
warnings.filterwarnings('ignore', category=np.exceptions.ComplexWarning)

Our first version of this will be simplified to use a Cartesian coordinate system roughly matching Earth’s parameters, with a rectangular domain as wide and tall as half the Earth’s circumference, and with an angular velocity roughly matching our own.

Note that the role of the omega_rad_s is only to impart accelerations due to the Coriolis effect. The root cause of planetary waves is the differing strength of the Coriolis affect across latitudes.

@dataclass
class SimulationConfig():
  '''
  Values and parameters which won't be changing during the simulation runs.
  '''
  omega_rad_s = 7.29e-5
  earth_radius_m = 6_371_000
  U = 50 # Average zonal velocity, in meters / s
  x_size_m = 40_075_000 / 2
  y_size_m = 40_075_000 / 2
  nx = 100 # number of grid cells in horizontal direction
  ny = 100 # number of grid cells vertically
  dx = x_size_m / nx
  dy = y_size_m / ny
  max_time_s = 3600*24*30 # in terms of seconds, hours, and days
  iterations = 300
  timesteps = np.linspace(0, max_time_s, iterations)

  # separate 2d arrays where each element contains the x- or y-location of
  # that grid cell
  x_arr, y_arr = np.meshgrid(
      np.linspace(0, x_size_m, nx),
      np.linspace(0, y_size_m, ny)
  )

  # We also need the latitude values for each grid cell,
  # in units of radians
  latitude_arr  = y_arr / y_size_m * np.pi / 2
  beta_arr = 2 * omega_rad_s * np.cos(latitude_arr) / earth_radius_m

cfg = SimulationConfig()

Run Rossby wave simulation

The explanations in this section are based on the Wikipedia page for planetary waves where you can learn more about their derivation. Note that this section mixes both Cartesian (\(x,y\)) and polar (\(\phi\)) coordinates. To reconcile this, we only use the polar c

Our first version of this will be simplified to use a Cartesian coordinate system, which is a two-dimensional flat grid, as opposed to a spherical coordinate system that would more accurately represent the Earth’s curved surface. The rectangular domain used here roughly matches the scale of half the Earth’s circumference in both width and height. The angular velocity parameter is set to a value that approximates Earth’s rotational rate.

Note that the role of the omega_rad_s is primarily to impart accelerations due to the Coriolis effect, which is a fictitious force observed in a rotating frame of reference that deflects moving objects. The root cause of planetary waves, however, is not the Coriolis effect itself, but rather the differing strength of the Coriolis effect across latitudes. This variation is known as the beta effect, and it is this latitudinal gradient that provides the restoring force for Rossby waves, causing them to propagate westward relative to the mean flow. In a Cartesian system, the beta effect is often introduced as a linear variation of the Coriolis parameter with the y-coordinate (representing latitude). The resulting PDE is:

where \(U\) denotes the average zonal velocity, also described as the mean westerly flow. Note that \(\beta = \frac{2\omega \cos \phi}{a}\), the Rossby parameter, captures the dependence of the system on latitude, as well as the size of the Earth and the speed of Earth’s rotation.

Our first version will use an analytical solution, specified in terms of Cartesian coordinates \(x\), \(y\), wavenumbers \(k\) and \(l\), and angular velocity \(\omega\). The analytical solution provides a direct mathematical expression for the streamfunction \(\psi(x,y,t)\) as a function of spatial position \((x, y)\) and time \(t\).

Here, \(\psi_0\) represents the amplitude of the wave, and \(k\) and \(l\) are the zonal and meridional wavenumbers, respectively. Wavenumbers describe the spatial frequency of the wave in the x and y directions. A larger wavenumber corresponds to a shorter wavelength and finer spatial structure. The term \(\omega t\) represents the temporal evolution of the wave, with \(\omega\) being the angular frequency of the wave. Note that a realistic solution for atmospheric or oceanic flows would typically involve the superposition of many such terms with different wavenumbers and amplitudes, representing a spectrum of waves. This simulation, for simplicity, uses a combination of a few selected wavenumbers to illustrate the basic wave behavior. Our simulation logic tracks the values of the streamfunction \(\psi\) over time and space. The streamfunction itself is a scalar field, but its gradient \(\nabla \psi\) is physically meaningful as it represents the velocity field of the underlying fluid. Specifically, in this simplified model, the velocity components \((u, v)\) are given by \(u = -\partial\psi / \partial y\) and \(v = \partial\psi / \partial x\). Therefore, we also calculate and store the gradient of the streamfunction in both the x and y directions throughout the simulation.

A realistic solution would involve the superposition of terms from multiple wavenumbers. This simulation just uses a few of them. Our simulation logic tracks the values of the streamfunction \(\psi\) as well as its gradient, as the velocity of the underlying fluid field is specified in terms of the gradient \(\nabla \psi\).

@dataclass
class SimulationState():
  '''
  Contains the state variables for the simulation.
  '''
  ψ: np.ndarray
  u: np.ndarray
  v: np.ndarray

state = SimulationState(
  ψ = np.zeros((cfg.iterations, cfg.ny, cfg.nx)),
  u = np.zeros((cfg.iterations, cfg.ny, cfg.nx)),
  v = np.zeros((cfg.iterations, cfg.ny, cfg.nx))
)

def rossby_wave(x, y, ω, t, ψ_0, k, l):
  '''
  Copied from https://en.wikipedia.org/wiki/Rossby_wave#Free_barotropic_Rossby_waves_under_a_zonal_flow_with_linearized_vorticity_equation
  '''

  return ψ_0 * np.exp(1j * (k * x + l * y - ω * t))

def omega(k, l):
  return cfg.U*k - cfg.beta_arr * k / (k**2 + l**2)

# Wavenumbers for the Rossby waves
# Define wavenumbers in physical units (rad/m)
k_physical = 2 * np.pi / (10000e3)  # 10,000 km wavelength
l_physical = 2 * np.pi / (8000e3)   # 8,000 km wavelength

# Convert to grid units
ks = [k_physical * n for n in range(1, 3)] # zonal
ls = [l_physical] # meriodonal

# Note that only the gradient of the streamfunction is physically meaningful
# so this parameter is unimportant for the simulation.
ψ_0 = 1.0


with tqdm.tqdm(total=cfg.iterations, desc="Simulating Rossby waves") as pbar:
  for i, t in enumerate(cfg.timesteps):
    current_ψ = ψ_0 * sum(
        rossby_wave(cfg.x_arr, cfg.y_arr, omega(k, l), t, ψ_0, k, l) for k in ks for l in ls
    )
    state.ψ[i] = current_ψ
    grady, gradx = np.gradient(ψ[i], cfg.dy, cfg.dx)
    state.v[i] = gradx   # -∂ψ/∂y (eastward velocity)
    state.u[i] = - grady # ∂ψ/∂x (northward velocity)

    # Calculate RMS average of the streamfunction
    rms_psi = np.sqrt(np.mean(np.abs(current_ψ)**2))
    pbar.set_description("RMS(ψ): {:.2f}".format(rms_psi))
    pbar.update(1)