SweepModule: state built once, reused by every call

The standard way to make an expensive solver affordable is to run it on a coarse grid once and interpolate afterwards. That gives the sweep a piece of state: a lookup table that costs a lot to build, nothing to use, and must not be rebuilt per point.

A closure cannot hold it, because the process executor has to pickle the callable. A module global can, badly. SweepModule gives it a home, modelled on the torch.nn.Module idiom: __init__ builds the state, forward is the physics, __call__ does the sweep.

Run with:

pixi run -e dev python examples/08_module_with_state.py

from __future__ import annotations

import sys
import time
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import xarray as xr

import xsweep

# sphinx-gallery executes examples without `__file__` set, so the gallery
# directory is located from the already-installed xsweep package instead.
sys.path.insert(0, str(Path(xsweep.__file__).resolve().parents[2] / "examples"))

from _solvers import mc_reflectance  # noqa: E402
from xsweep import SweepModule, SweepPolicy  # noqa: E402

N_PHOTONS = 4000

The module

contract is a class attribute, so a malformed one raises when the module is imported rather than after twenty minutes of engine time. forward never mentions the sweep and can be called on its own.

class LayerLUT(SweepModule):
    """Monte-Carlo reflectance, tabulated once and interpolated afterwards."""

    contract = "loop(tau, ssa) -> reflectance()"

    def __init__(
        self,
        n_tau: int = 8,
        n_ssa: int = 6,
        policy: SweepPolicy | None = None,
    ) -> None:
        super().__init__(policy)
        self.tau_nodes = np.linspace(0.05, 3.0, n_tau)
        self.ssa_nodes = np.linspace(0.80, 1.0, n_ssa)
        self.table = np.array(
            [
                [
                    mc_reflectance(
                        float(tau),
                        float(ssa),
                        g=0.6,
                        mu0=0.5,
                        n_photons=N_PHOTONS,
                        seed=0,
                    )["reflectance"]
                    for ssa in self.ssa_nodes
                ]
                for tau in self.tau_nodes
            ]
        )

    def forward(self, tau: float, ssa: float) -> float:
        """Bilinearly interpolate the table at one point."""
        i = int(
            np.clip(
                np.searchsorted(self.tau_nodes, tau) - 1, 0, len(self.tau_nodes) - 2
            )
        )
        j = int(
            np.clip(
                np.searchsorted(self.ssa_nodes, ssa) - 1, 0, len(self.ssa_nodes) - 2
            )
        )
        u = (tau - self.tau_nodes[i]) / (self.tau_nodes[i + 1] - self.tau_nodes[i])
        v = (ssa - self.ssa_nodes[j]) / (self.ssa_nodes[j + 1] - self.ssa_nodes[j])
        return float(
            (1 - u) * (1 - v) * self.table[i, j]
            + u * (1 - v) * self.table[i + 1, j]
            + (1 - u) * v * self.table[i, j + 1]
            + u * v * self.table[i + 1, j + 1]
        )

Building the state

48 Monte-Carlo runs, paid once. Every call afterwards is an interpolation.

start = time.perf_counter()
module = LayerLUT()
build_time = time.perf_counter() - start

print(f"table: {module.table.shape} built in {build_time:.2f} s")
print(f"that is {module.table.size} Monte-Carlo runs at {N_PHOTONS} photons")

table: (8, 6) built in 1.05 s
that is 48 Monte-Carlo runs at 4000 photons

forward is testable without a sweep

No space, no policy, no store: just the physics, with the arguments the contract says it takes.

print(f"forward(1.0, 0.95) = {module.forward(1.0, 0.95):.4f}")

forward(1.0, 0.95) = 0.2888

The sweep

A grid far finer than the table, which is the entire point of building the table in the first place.

space = xr.Dataset(
    {
        "tau": ("tau", np.linspace(0.1, 2.9, 60)),
        "ssa": ("ssa", np.linspace(0.82, 0.99, 30)),
    }
)

print(module.explain(space))

+- LayerLUT (v0) ---------------------------------------------+
| loop(tau, ssa) -> reflectance()
+------------------------------------------------------------+

SPACE
  tau        axis                60
  ssa        axis                30
  points                       1800
  dedup                      disabled

CALLS
  to compute     1800
  cached            0
  skipped           0
  total          1800

ARGUMENTS
  tau          loop    float64
  ssa          loop    float64

RESULT
  reflectance (tau: 60, ssa: 30) float64

STORE      none (in-memory: no cache, no resume)
EXECUTOR   serial

start = time.perf_counter()
result = module(space)
print(
    f"{space.tau.size * space.ssa.size} points in {time.perf_counter() - start:.3f} s"
)

1800 points in 0.222 s

What the table costs in accuracy

Interpolating is not free of consequence, and a page that showed only the speed-up would be selling something.

Comparing against a single direct Monte-Carlo run would not settle anything, because that run has a one-sigma noise of its own of around 0.006 at this photon count, the same size as the effect being looked for. The reference is therefore averaged over twelve seeds, which divides its noise by about three and makes the remaining difference readable.

REFERENCE_SEEDS = 12
checks = [(0.4, 0.87), (1.2, 0.93), (2.2, 0.96), (2.7, 0.985)]

print(f"{'tau':>5} {'ssa':>6} {'table':>8} {'direct':>8} {'noise':>7} {'error':>8}")
for tau, ssa in checks:
    draws = np.array(
        [
            mc_reflectance(tau, ssa, g=0.6, mu0=0.5, n_photons=N_PHOTONS, seed=seed)[
                "reflectance"
            ]
            for seed in range(REFERENCE_SEEDS)
        ]
    )
    direct = draws.mean()
    noise = draws.std() / np.sqrt(REFERENCE_SEEDS)
    interpolated = module.forward(tau, ssa)
    print(
        f"{tau:5.2f} {ssa:6.3f} {interpolated:8.4f} {direct:8.4f} "
        f"{noise:7.4f} {interpolated - direct:+8.4f}"
    )

 tau    ssa    table   direct   noise    error
0.40  0.870   0.1211   0.1234  0.0017  -0.0022
1.20  0.930   0.2966   0.2863  0.0020  +0.0103
2.20  0.960   0.4194   0.4085  0.0023  +0.0110
2.70  0.985   0.5103   0.4946  0.0029  +0.0157

Three of the four sit four to five standard errors above the reference, which is what a linear interpolation of a convex function does: the chord lies above the curve. The fourth comes out slightly below, by little more than one standard error, so it is not distinguishable from noise and should not be read as anything.

The largest error is in the top corner, where the table is coarsest relative to the curvature. That is a statement about node spacing, so refining the table is the fix and more photons would buy nothing.

The instance travels into worker processes

This is the combination the class facade is really for: state built once, then handed to several workers. The table goes with the instance, and the answers do not change.

parallel = LayerLUT(policy=SweepPolicy(executor="process", max_workers=4))
parallel.table = module.table  # reuse the table rather than rebuild it here

np.testing.assert_array_equal(
    result.reflectance.values, parallel(space).reflectance.values
)
print("four workers, identical result")

four workers, identical result

A bare @sweep function would have been simpler here, and is the right choice whenever there is no per-instance setup. Reach for the class when __init__ earns its keep, which is exactly when something expensive must be built before the first point and kept after it.

fig, ax = plt.subplots(figsize=(6.5, 4.0))
image = ax.pcolormesh(
    result.ssa, result.tau, result.reflectance, cmap="viridis", shading="auto"
)
ax.plot(
    *np.meshgrid(module.ssa_nodes, module.tau_nodes),
    "w.",
    markersize=3,
    label="table nodes",
)
ax.set_xlabel(r"single-scattering albedo $\omega$")
ax.set_ylabel(r"optical thickness $\tau$")
ax.set_title("1800 interpolated points over a 48-node table")
ax.legend(fontsize="small", loc="upper left")
fig.colorbar(image, ax=ax, label="reflectance")
fig.tight_layout()
plt.show()