cuwave 0.1.0__py3-none-any.whl

cuwave/postprocessing.py

@@ -0,0 +1,181 @@
+"""Matplotlib helpers for the bare field figures the drivers produce
+
+`matplotlib` is imported here and nowhere else in the package, the way torch is
+confined to `nn.py`, so the solver keeps its cupy-only dependency. Everything works in
+node index coordinates: an axes is the grid at one pixel per node, and a marker is
+sized in nodes, so the same call gives the same figure at any resolution and any dpi
+"""
+
+import io
+
+import cupy.typing as cpt
+import matplotlib.pyplot as plt
+import numpy as np
+import numpy.typing as npt
+from matplotlib.axes import Axes
+from matplotlib.figure import Figure
+from PIL import Image
+
+
+# -------------------------------------- helpers --------------------------------------
+def _host(field: cpt.NDArray | npt.NDArray) -> npt.NDArray:
+    """`field` on the host, whichever array module it came from"""
+    return np.asarray(field.get() if hasattr(field, "get") else field)
+
+
+def _panel_width(ax: Axes) -> float:
+    """Width of `ax` in inches, without asking for a renderer"""
+    return ax.get_position().width * ax.figure.get_figwidth()
+
+
+# -------------------------------------- figures --------------------------------------
+def field_axes(
+    resolution: tuple[int, ...],
+    dpi: int = 100,
+    pad: float = 0.05,
+    panels: int = 1,
+    gap: float = 0.15,
+) -> tuple[Figure, Axes]:
+    """Borderless axes the size of the grid, padded so a marker on the edge survives.
+
+    Args:
+        resolution: (nx, ny) nodes of the plotted array, which sets the figure size.
+        dpi: dots per inch, the figure being `resolution` pixels wide at 100.
+        pad: margin around the grid, as a fraction of the longest axis.
+        panels: side-by-side panels, each one `resolution` wide.
+        gap: spacing between panels, as a fraction of one panel's width.
+
+    Returns:
+        (fig, ax), `ax` a single axes for one panel and a `panels`-long array beyond.
+    """
+    width = resolution[0] / 100 * (panels + gap * (panels - 1))
+    fig, axes = plt.subplots(1, panels, figsize=(width, resolution[1] / 100), dpi=dpi)
+    margin = pad * max(resolution)
+    for ax in np.atleast_1d(axes):
+        ax.set_xlim(-margin, resolution[0] + margin)
+        ax.set_ylim(-margin, resolution[1] + margin)
+        ax.set_aspect("equal")
+        ax.axis("off")
+    fig.subplots_adjust(left=0, right=1, top=1, bottom=0, wspace=gap)
+    return fig, axes
+
+
+def show(
+    ax: Axes,
+    field: cpt.NDArray | npt.NDArray | None = None,
+    indicator: cpt.NDArray | npt.NDArray | None = None,
+    cmap: str = "seismic",
+    saturation: float = 1.0,
+    scale: float | None = None,
+    indicator_cmap: str = "binary",
+    gray: float = 1.0,
+) -> Axes:
+    """Draw a wave field, a design, or the design over the field.
+
+    Args:
+        ax: the axes from `field_axes`, which fixes the node index coordinates.
+        field: (nx, ny) wave field, colored symmetrically about zero.
+        indicator: (nx, ny) design, drawn over `field` where it is solid and left
+            transparent elsewhere, or alone as a [0, 1] field in `cmap`.
+        cmap: colormap of `field`, or of `indicator` when that is drawn alone.
+        saturation: fraction of `scale` the colormap runs to, below 1 to clip the peak.
+        scale: the amplitude `field` saturates at, defaulting to its own peak.
+        indicator_cmap: colormap of the overlay, only reached together with `field`.
+        gray: the constant the overlay is drawn at, 1 black and 0 white in `binary`.
+    """
+    if field is not None:
+        values = _host(field)
+        scale = float(np.max(np.abs(values))) if scale is None else scale
+        limit = scale * saturation
+        ax.pcolormesh(values.T, cmap=cmap, vmin=-limit, vmax=limit)
+    if indicator is not None and field is None:
+        ax.pcolormesh(_host(indicator).T, cmap=cmap, vmin=0, vmax=1)
+    elif indicator is not None:
+        solid = np.where(_host(indicator), gray, np.nan)
+        ax.pcolormesh(solid.T, cmap=indicator_cmap, vmin=0, vmax=1)
+    return ax
+
+
+def save(fig: Figure, path: str) -> None:
+    """Write `fig` to `path` on a transparent background, palettized where it can be.
+
+    A colormapped field holds a few hundred distinct colors, so 32 bits per pixel is
+    mostly waste: a figure that came out fully opaque is written as an 8-bit palette
+    instead. One carrying transparency keeps RGBA, since a palette entry has a single
+    alpha and cannot hold the intermediate ones of an antialiased marker.
+    """
+    buffer = io.BytesIO()
+    fig.savefig(buffer, format="png", transparent=True)
+    image = Image.open(buffer)
+    if image.mode != "RGBA" or image.getchannel("A").getextrema()[0] == 255:
+        image = image.convert("RGB").quantize(
+            colors=256, method=Image.Quantize.MEDIANCUT, dither=Image.Dither.NONE
+        )
+    image.save(path, optimize=True)
+
+
+# ------------------------------------- annotation ------------------------------------
+def marker_points(ax: Axes, nodes: float = 6.0) -> float:
+    """Marker size in points spanning `nodes` grid nodes, which points alone do not"""
+    low, high = ax.get_xlim()
+    return nodes * 80.0 * _panel_width(ax) / (high - low)
+
+
+def markers(
+    ax: Axes,
+    positions: cpt.NDArray | npt.ArrayLike,
+    dx: tuple[float, ...] | None = None,
+    origin: int = 0,
+    nodes: float = 6.0,
+    color: str = "silver",
+    **style,
+) -> None:
+    """Dots at grid `positions`, sized in nodes so they match across figures.
+
+    Args:
+        ax: the axes from `field_axes`, whose limits fix the points-per-node scale.
+        positions: (ndim, count) grid node indices, or (count, ndim) physical
+            coordinates when `dx` is given.
+        dx: grid spacing, which turns physical coordinates into node indices.
+        origin: node index of the plotted array's first cell, 1 for an interior slice.
+        nodes: diameter of a dot in grid nodes, held the same figure to figure.
+        color: fill and edge color, gray reading over both a light and a dark field.
+        **style: passed to `plot`, e.g. `zorder` or `alpha`.
+    """
+    index = np.atleast_2d(_host(positions)).astype(float)
+    if dx is not None:
+        index = (index / np.asarray(dx)).T + 1.0
+    # the pcolormesh cell of node i is centered half a node past its corner
+    ax.plot(
+        index[0] - origin + 0.5,
+        index[1] - origin + 0.5,
+        "o",
+        markersize=marker_points(ax, nodes),
+        markerfacecolor=color,
+        markeredgecolor=color,
+        linestyle="none",
+        clip_on=False,
+        zorder=3,
+        **style,
+    )
+
+
+def outline(
+    ax: Axes,
+    low: tuple[int, ...],
+    high: tuple[int, ...],
+    origin: int = 0,
+    color: str = "silver",
+    linewidth: float = 1.5,
+) -> None:
+    """Rectangle around the nodes `low` to `high`, marking a target or design region"""
+    ax.add_patch(
+        plt.Rectangle(
+            (low[0] - origin, low[1] - origin),
+            high[0] - low[0] + 1,
+            high[1] - low[1] + 1,
+            fill=False,
+            edgecolor=color,
+            linewidth=linewidth,
+        )
+    )

cuwave/regularization.py

@@ -0,0 +1,243 @@
+"""Regularization and design-map penalties for gradient-based optimization.
+
+A `Regularization` is both a differentiable map (`__call__`) and its own adjoint
+(`grad`), so filters, projections, and penalties compose the same way the simulation's
+forward/adjoint pair does: chain the calls forward, chain `grad` backward. `set` lets a
+`continuation` schedule mutate a live instance's parameters between iterations without
+rebuilding the pipeline.
+"""
+
+from __future__ import annotations
+
+import abc
+import math
+from collections.abc import Sequence
+
+import cupy as cp
+import cupy.typing as cpt
+import cupyx.scipy.ndimage as ndi
+import numpy.typing as npt
+
+
+# -------------------------------------- helpers --------------------------------------
+def _axis_slice(ndim: int, axis: int, cut: int | slice) -> tuple:
+    """Build an ndim-length index tuple selecting `cut` along `axis`, `slice(None)` elsewhere."""
+    index = [slice(None)] * ndim
+    index[axis] = cut
+    return tuple(index)
+
+
+def _spatial_grad(x: cpt.NDArray, axis: int) -> cpt.NDArray:
+    """forward difference along `axis`, zero-padded at the far boundary"""
+    head = _axis_slice(x.ndim, axis, slice(None, -1))
+    tail = _axis_slice(x.ndim, axis, slice(1, None))
+    d = cp.zeros_like(x)
+    d[head] = x[tail] - x[head]
+    return d
+
+
+def _spatial_grad_adjoint(u: cpt.NDArray, axis: int) -> cpt.NDArray:
+    """transpose of `_spatial_grad`: the negative divergence of `u`"""
+    head = _axis_slice(u.ndim, axis, slice(None, -1))
+    tail = _axis_slice(u.ndim, axis, slice(1, None))
+    inner = u[head]
+    out = cp.zeros_like(u)
+    out[head] -= inner
+    out[tail] += inner
+    return out
+
+
+# ------------------------------- regularization classes ------------------------------
+class Regularization(abc.ABC):
+    """Forward pass in `__call__(x)`, backward pass in `grad(x, dy)`, re-tuned by `set`."""
+
+    @abc.abstractmethod
+    def __call__(self, x: cpt.NDArray) -> cpt.NDArray:
+        """Apply the map, or evaluate the penalty."""
+
+    @abc.abstractmethod
+    def grad(self, x: cpt.NDArray, dy: cpt.NDArray | float = 1.0) -> cpt.NDArray:
+        """Gradient with respect to `x`, given the gradient `dy` towards the output."""
+
+    def set(self, **params) -> Regularization:
+        """Update settings of the regularizer: useful in continuation schemes."""
+        for name, value in params.items():
+            if not hasattr(self, name):
+                raise AttributeError(f"{type(self).__name__} has no parameter {name!r}")
+            setattr(self, name, value)
+        return self
+
+
+# ------------------------------ design-map modification ------------------------------
+class DensityFilter(Regularization):
+    """Conic density filter on a structured grid of any dimension, with an OC variant.
+
+    Args:
+        rmin: filter radius in nodes, which sets the conic kernel's support.
+        shape: the design field's shape, which fixes the dimension of the kernel and
+            the normalization `Hs` it is built for.
+        dtype: kernel dtype, matched to the design field to keep the convolution
+            in one precision.
+    """
+
+    def __init__(
+        self, rmin: float, shape: tuple[int, ...], dtype: npt.DTypeLike | None = None
+    ) -> None:
+        ceil_r = int(math.ceil(rmin))
+        taps = cp.arange(-ceil_r, ceil_r + 1)
+        offsets = cp.meshgrid(*(taps,) * len(shape), indexing="ij")
+        self.kernel = cp.maximum(0.0, rmin - cp.sqrt(sum(k**2 for k in offsets)))
+        if dtype is not None:
+            self.kernel = self.kernel.astype(dtype)
+        self.Hs = ndi.convolve(
+            cp.ones(shape, self.kernel.dtype), self.kernel, mode="constant", cval=0.0
+        )
+
+    def __call__(self, x: cpt.NDArray) -> cpt.NDArray:
+        """Filtered design: convolve `x` with the conic kernel, normalized by `Hs`."""
+        return ndi.convolve(x, self.kernel, mode="constant", cval=0.0) / self.Hs
+
+    def grad(self, x: cpt.NDArray, dy: cpt.NDArray | float = 1.0) -> cpt.NDArray:
+        """Adjoint of the filter applied to `dy`; `x` is unused since the filter is linear."""
+        return ndi.convolve(dy / self.Hs, self.kernel, mode="constant", cval=0.0)
+
+    def sensitivity(self, rho: cpt.NDArray, dc: cpt.NDArray) -> cpt.NDArray:
+        """Filter sensitivities `dc` at density `rho`, the OC scheme in place of `grad`."""
+        num = ndi.convolve(rho * dc, self.kernel, mode="constant", cval=0.0)
+        return num / (cp.maximum(rho, 1e-3) * self.Hs)
+
+
+class Projection(Regularization):
+    """Smoothed Heaviside about threshold `eta`, with sharpness `beta`."""
+
+    def __init__(self, beta: float, eta: float = 0.5) -> None:
+        self.beta, self.eta = beta, eta
+
+    def _ends(self) -> tuple[float, float]:
+        """tanh at the two ends of [0, 1], used to rescale `__call__` and `grad`."""
+        return math.tanh(self.beta * self.eta), math.tanh(self.beta * (1.0 - self.eta))
+
+    def __call__(self, x: cpt.NDArray) -> cpt.NDArray:
+        """Projected design: rescaled `tanh(beta * (x - eta))`, mapped onto [0, 1]."""
+        a, b = self._ends()
+        return (a + cp.tanh(self.beta * (x - self.eta))) / (a + b)
+
+    def grad(self, x: cpt.NDArray, dy: cpt.NDArray | float = 1.0) -> cpt.NDArray:
+        """Adjoint of the projection: `dy` scaled by the local tanh derivative."""
+        a, b = self._ends()
+        t = cp.tanh(self.beta * (x - self.eta))
+        return dy * self.beta * (1.0 - t * t) / (a + b)
+
+
+class SIMP(Regularization):
+    """Power-law penalization making intermediate designs uneconomical (Bendsoe 1989).
+
+    See https://doi.org/10.1007/BF01650949
+    """
+
+    def __init__(self, p: float = 3.0, x_min: float = 0.0) -> None:
+        self.p, self.x_min = p, x_min
+
+    def __call__(self, x: cpt.NDArray) -> cpt.NDArray:
+        """Penalized design: `x_min + (1 - x_min) * x**p`."""
+        # x in [0, 1] is assumed, as a fractional p needs x >= 0
+        return self.x_min + (1.0 - self.x_min) * x**self.p
+
+    def grad(self, x: cpt.NDArray, dy: cpt.NDArray | float = 1.0) -> cpt.NDArray:
+        """Adjoint of the penalization: `dy` scaled by the power-law derivative."""
+        return dy * (1.0 - self.x_min) * self.p * x ** (self.p - 1.0)
+
+
+# ----------------------------- penalization in objective -----------------------------
+class Tikhonov(Regularization):
+    """L2 penalty, damping toward a reference or smoothing by penalizing the gradient.
+
+    Args:
+        alpha: penalty weight.
+        order: 0 penalizes deviation from `x_ref`, 1 penalizes the spatial gradient.
+        x_ref: the reference `order` 0 damps toward, or None for zero.
+    """
+
+    def __init__(
+        self, alpha: float, order: int = 0, x_ref: cpt.NDArray | None = None
+    ) -> None:
+        if order not in (0, 1):
+            raise ValueError("order must be 0 (damping) or 1 (smoothing)")
+        self.alpha, self.order, self.x_ref = alpha, order, x_ref
+
+    def _residual(self, x: cpt.NDArray) -> cpt.NDArray:
+        """`x` relative to `x_ref`, or `x` itself when no reference is set."""
+        return x if self.x_ref is None else x - self.x_ref
+
+    def __call__(self, x: cpt.NDArray) -> cpt.NDArray:
+        """Penalty value: `0.5 * alpha` times the squared residual, or its squared gradient."""
+        r = self._residual(x)
+        if self.order == 0:
+            return 0.5 * self.alpha * cp.sum(r * r)
+        squares = (cp.sum(_spatial_grad(r, axis) ** 2) for axis in range(r.ndim))
+        return 0.5 * self.alpha * sum(squares)
+
+    def grad(self, x: cpt.NDArray, dy: cpt.NDArray | float = 1.0) -> cpt.NDArray:
+        """Adjoint of the penalty: `dy * alpha` times the residual, or its Laplacian."""
+        r = self._residual(x)
+        if self.order == 0:
+            return (dy * self.alpha) * r
+        out = cp.zeros_like(r)
+        for axis in range(r.ndim):
+            out += _spatial_grad_adjoint(_spatial_grad(r, axis), axis)
+        return (dy * self.alpha) * out
+
+
+class TotalVariation(Regularization):
+    """Smoothed isotropic total variation, `alpha * sum sqrt(|D x|**2 + eps**2)`"""
+
+    def __init__(self, alpha: float, eps: float = 1e-3) -> None:
+        self.alpha, self.eps = alpha, eps
+
+    def _magnitude(self, diffs: Sequence[cpt.NDArray]) -> cpt.NDArray:
+        """Smoothed gradient magnitude `sqrt(sum(d**2) + eps**2)` over the differences `diffs`."""
+        total = self.eps**2
+        for d in diffs:
+            total = total + d * d
+        return cp.sqrt(total)
+
+    def __call__(self, x: cpt.NDArray) -> cpt.NDArray:
+        """Penalty value: `alpha` times the summed smoothed gradient magnitude of `x`."""
+        diffs = [_spatial_grad(x, axis) for axis in range(x.ndim)]
+        return self.alpha * cp.sum(self._magnitude(diffs))
+
+    def grad(self, x: cpt.NDArray, dy: cpt.NDArray | float = 1.0) -> cpt.NDArray:
+        """Adjoint of the penalty: divergence of the normalized gradient, scaled by `dy * alpha`."""
+        diffs = [_spatial_grad(x, axis) for axis in range(x.ndim)]
+        magnitude = self._magnitude(diffs)
+        out = cp.zeros_like(x)
+        for axis, d in enumerate(diffs):
+            out += _spatial_grad_adjoint(d / magnitude, axis)
+        return (dy * self.alpha) * out
+
+
+# ------------------------------ continuation schedules -------------------------------
+def continuation(
+    scheme: str, iters: int, start: float, stop: float, stages: int = 4
+) -> list[float]:
+    """`iters` values ramping `start` -> `stop`, a `Projection` sharpness schedule
+
+    * `"constant"`: `start` throughout, the reference
+    * `"linear"`: equal increments, so most of the run is already sharp
+    * `"exponential"`: equal factors, spending the early iterations near `start`
+    * `"staircase"`: `stages` levels held for `iters // stages` iterations
+      each, the classic continuation (Wang, Lazarov & Sigmund 2011,
+      https://doi.org/10.1007/s00158-010-0602-y)
+    """
+    last = max(iters - 1, 1)
+    if scheme == "constant":
+        return [start] * iters
+    if scheme == "linear":
+        return [start + (stop - start) * i / last for i in range(iters)]
+    if scheme == "exponential":
+        return [start * (stop / start) ** (i / last) for i in range(iters)]
+    if scheme == "staircase":  # the exponential ramp, held over `stages` levels
+        every, levels = max(iters // stages, 1), max(stages - 1, 1)
+        factors = (min(i // every, levels) / levels for i in range(iters))
+        return [start * (stop / start) ** f for f in factors]
+    raise ValueError(f"unknown continuation scheme {scheme!r}")

cuwave/scalar.py

@@ -0,0 +1,224 @@
+"""Scalar pressure wave equations, the nodal sibling of the vector elastic schemes.
+
+One unknown per node and a per-axis flux, so the operator is the cheapest of the three
+families and needs no component offsets. `PressureWave` holds the nodal `stiff`/`minv`
+fields and the adjoint hooks; `ScalarWave` scales both with one indicator gamma, which
+is why `minv` is never formed (`derive_inertia`), and `AcousticWave` interpolates
+inverse density and inverse bulk modulus between two phases for topology optimization.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from pathlib import Path
+
+import cupy as cp
+import cupy.typing as cpt
+import numpy as np
+
+from .boundary import Neumann
+from .wave import (
+    Simulation,
+    apply_cell_weights,
+    grid_block,
+    mirror_ghosts,
+    sensor_cell_weights,
+)
+
+KERNEL_PATH = Path(__file__).parent / "kernels" / "scalar.cu"
+SENSITIVITY_PATH = Path(__file__).parent / "kernels" / "scalar_sensitivity.cu"
+
+
+# -------------------------------- discretization setup -------------------------------
+@dataclass
+class PressureWave(Simulation):
+    """Scalar wave equation base: nodal `stiff`/`minv` material fields, optional damping."""
+
+    kernel_path = KERNEL_PATH
+    sensitivity_path = SENSITIVITY_PATH
+    default_boundary = Neumann
+
+    derive_inertia = False  # set where m == k (rho scaling): minv derived from stiff
+
+    def build_materials(self, indicator: cpt.NDArray) -> dict:
+        """Turn `indicator` into the kernel's material dict, `damping` included."""
+        stiff, minv = self.parametrization(indicator)
+        # mirrored in place, so a caller's ghost ring is normalised to Neumann
+        mat = {"stiff": mirror_ghosts(self, stiff)}
+        mat["minv"] = None if minv is None else mirror_ghosts(self, minv)
+        if self.damping is not None:
+            mat["damping"] = self.damping
+        return mat
+
+    def inverse_inertia(self, indicator: cpt.NDArray) -> cpt.NDArray:
+        """Nodal `1 / m` for `indicator`, derived from the stiffness where `m == k`."""
+        stiff, minv = self.parametrization(indicator)
+        return 1.0 / stiff if minv is None else minv
+
+    def step_kernel_args(self, mat: dict) -> tuple:
+        """Material and damping arguments for the finite-difference step kernel."""
+        minv = mat["stiff"] if self.derive_inertia else mat["minv"]
+        args = (mat["stiff"], minv, np.int32(self.derive_inertia))
+        if self.damping is not None:
+            args += (mat["damping"], self.dtype(self.dt))
+        return args
+
+    gradient_names = ("mass", "stiff")  # the fields the adjoint differentiates
+
+    def gradient_fields(self, mat: dict) -> dict[str, cpt.NDArray]:
+        """Zeroed accumulators the adjoint kernels add into, one per material field."""
+        return {
+            name: cp.zeros(self.Nx_padded, dtype=self.dtype)
+            for name in self.gradient_names
+        }
+
+    def adjoint_weights(self, sensors: cpt.NDArray[cp.int32]) -> cpt.NDArray:
+        """Divisor the adjoint source carries: the cell weights W over the source factor."""
+        return sensor_cell_weights(self, sensors) * self.source_factor()
+
+    def finalize_gradients(self, grads: dict, kernels: cp.RawModule) -> dict:
+        """Weight the accumulators by W once the time loop is done."""
+        for field in grads.values():
+            apply_cell_weights(self, field)
+        return grads
+
+    def define_gradient(
+        self, kernels: cp.RawModule, mat: dict, grads: dict
+    ) -> Callable:
+        """Closure accumulating both gradient densities from a forward triplet and `l1`."""
+        # one kernel for both gradients: a launch costs more host time than either body
+        gradient_kernel = kernels.get_function("gradient_kernel")
+        grid, block = grid_block(self)
+        # the operator without the dt^2 the step folds into it: L, not dt^2 L
+        factors = [self.dtype(float(f) / self.dt**2) for f in self.step_factors()]
+        geom = [factors[0], self.Nx[0]]
+        for d in range(1, self.ndim):
+            geom += [factors[d], self.Nx[d], self.strides[d - 1]]
+        args = [grads["mass"], grads["stiff"], None, None, None, None] + [
+            mat["stiff"],
+            self.dtype(1.0 / self.dt**2),
+            *geom,
+        ]
+
+        def gradient_step(u0, u1, u2, l1):
+            args[2], args[3], args[4], args[5] = u0, u1, u2, l1
+            gradient_kernel(grid, block, args)
+
+        return gradient_step
+
+    def define_frechet(
+        self, kernels: cp.RawModule, accs: dict, sign: float
+    ) -> Callable:
+        """Closure accumulating both Frechet densities of one field triplet, times `sign`."""
+        # sign is fixed per pass, so it is folded into the factors, not recomputed
+        frechet_kernel = kernels.get_function("frechet_kernel")
+        grid, block = grid_block(self)
+        # the stiffness density enters negated, so the epilogue scales both alike
+        geom = [self.dtype(sign / (2.0 * self.dt) ** 2)]
+        for d in range(self.ndim):
+            geom.append(self.dtype(-sign / (2.0 * self.dx[d]) ** 2))
+            geom.append(self.Nx[d])
+            if d:
+                geom.append(self.strides[d - 1])
+        args = [accs["mass"], accs["stiff"], None, None, None] + geom
+
+        def frechet_step(u0, u1, u2):
+            args[2], args[3], args[4] = u0, u1, u2
+            frechet_kernel(grid, block, args)
+
+        return frechet_step
+
+    def excitation_weights(
+        self, mat: dict, lin_index: cpt.NDArray[cp.int32]
+    ) -> cpt.NDArray:
+        """Source weights `dt**2 / inertia` at `lin_index`, never forming inertia on the grid."""
+        field = mat["stiff"] if self.derive_inertia else mat["minv"]
+        weight = field.ravel()[lin_index]
+        if self.derive_inertia:
+            weight = 1.0 / weight
+        if self.damping is not None:
+            # the damped update divides by 1 + beta, so the source has to share it
+            beta = 0.5 * weight * mat["damping"].ravel()[lin_index] * self.dt
+            weight = weight / (1.0 + beta)
+        return (self.dtype(self.dt**2 * self.source_factor()) * weight).astype(
+            self.dtype
+        )
+
+
+@dataclass
+class ScalarWave(PressureWave):
+    """Constant-speed scalar wave equation, parametrized by a density-scaling indicator gamma."""
+
+    wavespeed: float = None  # background wave speed c0
+    density: float = None  # background density rho0
+
+    derive_inertia = True  # gamma scales inertia and stiffness alike
+
+    def __post_init__(self) -> None:
+        """Validate `wavespeed` and `density` are set, on top of `Simulation.__post_init__`."""
+        super().__post_init__()
+        if self.wavespeed is None or self.density is None:
+            raise ValueError("ScalarWave requires wavespeed and density")
+
+    def parametrization(
+        self, indicator: cpt.NDArray
+    ) -> tuple[cpt.NDArray, cpt.NDArray | None]:
+        """`indicator` is the density-scaling field gamma; `minv` is left to be derived from it."""
+        return indicator, None
+
+    def parametrization_jacobian(self, indicator: cpt.NDArray) -> tuple:
+        """Both mass and stiffness coefficients are gamma itself, so both derivatives are 1."""
+        return 1.0, 1.0
+
+    def step_factors(self) -> list:
+        """Per-axis finite-difference step factors `2 * c0**2 * dt**2 / dx**2`."""
+        return [
+            self.dtype(2.0 * self.wavespeed**2 * self.dt**2 / dxk**2) for dxk in self.dx
+        ]
+
+    def source_factor(self) -> float:
+        """Source scaling `1 / rho0`."""
+        return 1.0 / self.density
+
+
+@dataclass
+class AcousticWave(PressureWave):
+    """Two-phase acoustic wave equation (TATO), gamma interpolating between air and solid."""
+
+    # TATO material constants (gamma = 0 -> air, gamma = 1 -> solid)
+    rho1: float = None
+    rho2: float = None
+    kappa1: float = None
+    kappa2: float = None
+
+    def __post_init__(self) -> None:
+        """Validate the four material constants are set, on top of `Simulation.__post_init__`."""
+        super().__post_init__()
+        if None in (self.rho1, self.rho2, self.kappa1, self.kappa2):
+            raise ValueError("AcousticWave requires rho1, rho2, kappa1, kappa2")
+
+    def parametrization(
+        self, indicator: cpt.NDArray
+    ) -> tuple[cpt.NDArray, cpt.NDArray]:
+        """`indicator` interpolates inverse density and inverse bulk modulus between the phases."""
+        # the kernel wants 1 / rho, so rho is never formed
+        rho_inv = 1 / self.rho1 + indicator * (1 / self.rho2 - 1 / self.rho1)
+        kappa_inv = 1 / self.kappa1 + indicator * (1 / self.kappa2 - 1 / self.kappa1)
+        return rho_inv, 1 / kappa_inv
+
+    def parametrization_jacobian(self, indicator: cpt.NDArray) -> tuple:
+        """Derivatives of (mass, stiff) with respect to gamma; both coefficients are affine in it."""
+        # affine in gamma as (mass, stiff) = (1 / kappa, 1 / rho), hence constant
+        return (
+            1 / self.kappa2 - 1 / self.kappa1,
+            1 / self.rho2 - 1 / self.rho1,
+        )
+
+    def step_factors(self) -> list:
+        """Per-axis finite-difference step factors `2 * dt**2 / dx**2`."""
+        return [self.dtype(2.0 * self.dt**2 / dxk**2) for dxk in self.dx]
+
+    def source_factor(self) -> float:
+        """Source scaling, unscaled since rho is already folded into `parametrization`."""
+        return 1.0