Deduplication: a scene has fewer atmospheres than pixels

A retrieval product does not carry a different atmospheric state in every pixel. It comes out of a classifier or off a coarse grid, so neighbouring pixels repeat, and a sweep over the image calls the engine again and again with arguments it has already seen.

dedup=True collapses identical input rows before calling anything, computes each distinct one once, and expands the result back onto every pixel that shares it. It never changes a value or a shape, only how many times the engine runs to produce them.

Run with:

pixi run -e dev python examples/04_dedup_on_a_scene.py

from __future__ import annotations

import shutil
import sys
import tempfile
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 _scene import atmosphere_scene  # noqa: E402
from _solvers import mc_reflectance  # noqa: E402
from xsweep import SweepPolicy, sweep  # noqa: E402

N_PHOTONS = 1500
CALLS = {"n": 0}


@sweep("loop(tau, ssa) -> reflectance(), transmittance()")
def layer(tau: float, ssa: float) -> dict[str, float]:
    """Reflectance and transmittance of a scattering layer, by Monte-Carlo."""
    CALLS["n"] += 1
    return mc_reflectance(tau, ssa, g=0.6, mu0=0.5, n_photons=N_PHOTONS, seed=0)


def counted() -> int:
    """Return the calls made since the last check, and reset the counter."""
    made, CALLS["n"] = CALLS["n"], 0
    return made


tau_map, ssa_map = atmosphere_scene(24, block=4)
scene = xr.Dataset({"tau": (("y", "x"), tau_map), "ssa": (("y", "x"), ssa_map)})

rows = np.stack([tau_map.ravel(), ssa_map.ravel()]).T
print(f"pixels:               {tau_map.size}")
print(f"distinct atmospheres: {len(np.unique(rows, axis=0))}")

pixels:               576
distinct atmospheres: 6

The plan counts the distinct rows for you

Turning dedup on adds a line to the SPACE block: how many unique rows the space holds, and therefore how many calls a run will really make. Nothing has been computed yet.

print(layer.explain(scene, policy=SweepPolicy(dedup=True)))

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

SPACE
  y          tau, ssa            24
  x          tau, ssa            24
  points                        576
  dedup                      6 unique (570 duplicates)

CALLS
  to compute        6
  cached            0
  skipped           0
  total             6

ARGUMENTS
  tau          loop    float64
  ssa          loop    float64

RESULT
  reflectance (y: 24, x: 24) float64
  transmittance (y: 24, x: 24) float64

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

Without deduplication

One call per pixel, most of them recomputing an answer the engine has already given.

start = time.perf_counter()
plain = layer(scene)
plain_time = time.perf_counter() - start
plain_calls = counted()
print(f"{plain_calls} calls in {plain_time:.2f} s")

576 calls in 2.32 s

With deduplication

start = time.perf_counter()
deduped = layer(scene, policy=SweepPolicy(dedup=True))
dedup_time = time.perf_counter() - start
dedup_calls = counted()
print(f"{dedup_calls} calls in {dedup_time:.2f} s")

print(f"\ncalls:  {plain_calls / dedup_calls:.0f}x fewer")
print(f"time:   {plain_time / dedup_time:.0f}x faster")

6 calls in 0.05 s

calls:  96x fewer
time:   50x faster

Same numbers, to the last bit

Deduplication is a cost decision, not a numerical one. The engine is stochastic and the two runs still agree exactly, because the duplicated pixels were copied rather than recomputed.

np.testing.assert_array_equal(plain.reflectance.values, deduped.reflectance.values)
np.testing.assert_array_equal(plain.transmittance.values, deduped.transmittance.values)
print("identical on every pixel")

identical on every pixel

What it does not buy

The duplicate positions still have to be filled. In memory that is a gather; into a store it is a batched write pass. Deduplication removes engine calls, which is nearly the whole story when a call is expensive and almost none of it when a call is free.

run_dir = Path(tempfile.mkdtemp())
try:
    start = time.perf_counter()
    persisted = layer(
        scene, policy=SweepPolicy(dedup=True, store=run_dir / "scene.zarr")
    )
    print(f"{counted()} calls, persisted, in {time.perf_counter() - start:.2f} s")
    np.testing.assert_array_equal(
        plain.reflectance.values, persisted.reflectance.values
    )
    print("still identical once written to a store")
finally:
    shutil.rmtree(run_dir)

6 calls, persisted, in 0.26 s
still identical once written to a store

The scaling that matters

The number of distinct rows is a property of the scene, not of its size. Doubling the image doubles the pixels and leaves the call count where it was, so the saving grows with the image.

print(f"{'pixels':>8}  {'distinct':>9}  {'calls saved':>12}")
for size in (24, 48, 96, 192):
    t_map, s_map = atmosphere_scene(size, block=size // 6)
    unique = len(np.unique(np.stack([t_map.ravel(), s_map.ravel()]).T, axis=0))
    print(f"{t_map.size:>8}  {unique:>9}  {t_map.size - unique:>12}")

pixels   distinct   calls saved
   576          6           570
  2304          6          2298
  9216          6          9210
 36864          6         36858

The result is still an image, with the pixels the duplicates were expanded onto in their proper places.

fig, axes = plt.subplots(1, 2, figsize=(8.0, 3.4))
for ax, (data, title) in zip(
    axes,
    ((scene.tau, r"input $\tau$"), (deduped.reflectance, "reflectance, deduplicated")),
    strict=True,
):
    image = ax.imshow(data, origin="upper", cmap="viridis")
    ax.set_title(title)
    ax.set_xlabel("x")
    ax.set_ylabel("y")
    fig.colorbar(image, ax=ax, shrink=0.85)
fig.tight_layout()
plt.show()