cuwave 0.2.0__tar.gz → 0.6.0__tar.gz

{cuwave-0.2.0 → cuwave-0.6.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cuwave
-Version: 0.2.0
+Version: 0.6.0
 Summary: GPU finite-difference wave solver with differentiable adjoints
 Author-email: Leon Herrmann <herrmann.leon@pm.me>
 License-Expression: MIT
@@ -31,7 +31,7 @@ Dynamic: license-file
   <img src="https://raw.githubusercontent.com/cmpmech/cuwave/main/.assets/readme-light.webp" width="100%">
 </picture>
 
-**CuWave** is a single-GPU, differentiable finite difference wave propagation code.
+**CuWave** is a GPU-accelerated (single or multi), differentiable finite difference wave propagation code.
 Possible applications include
 
 <table width="100%">
@@ -67,13 +67,16 @@ Possible applications include
 - see the [documentation](https://github.com/cmpmech/cuwave/blob/main/docs/Home.md) for how the code works (AI-assisted; verify with sources for critical details)
 - see [examples](https://github.com/cmpmech/cuwave/tree/main/examples) for how to apply the code
 
-## Development status
-
-**Scalar** & **acoustic wave** equations have been developed over the last 2 years and are thoroughly validated.
+## Multiple GPUs
 
-> [!IMPORTANT]
-> **Elastic** & **electromagnetic wave** equations were developed with AI assistance (Claude) and have undergone less validation. The elastic wave equation is currently being validated against experimental results.
+Several GPUs take one argument, and the driver otherwise stays as it is:
+- `Simulation(..., devices=(0, 1))` splits **one grid** along its first axis over the listed GPUs, for grids too large or too slow for one. `simulate`, `sensitivity`, `reconstruction_sensitivity` and `superposition_sensitivity` then run split and return the same result as on one GPU, bitwise. Try it with the `DEVICES` setting of
+	- [examples/forward/scalarND.py](https://github.com/cmpmech/cuwave/blob/main/examples/forward/scalarND.py), the forward problem in 2D or 3D
+	- [examples/sensitivity/scalar2D_sensitivity.py](https://github.com/cmpmech/cuwave/blob/main/examples/sensitivity/scalar2D_sensitivity.py), the gradient
+- `misfit_gradient(..., devices=[0, 1])` (and `measure`, `misfit`) deals the **shots** of an inversion out over the GPUs, which needs no communication and is the faster choice whenever one shot fits one GPU.
+## Development status
 
+**Scalar** & **acoustic wave** equations have been developed over the last 2 years and are thoroughly validated. The **elastic wave** and **electromagnetic wave** equations were developed with AI assistance (Claude). The first is validated on experimental data, while the latter are cross-validated against the independent FDTD code [FDTDX](https://github.com/ymahlau/fdtdx). 
 ## Performance
 
 CuWave's runtime for identical discretizations is comparable to that of other established wave propagation finite difference codes. Speedups in 2D with reflecting boundaries (**a ratio above 1 means CuWave is that many times faster**) compared to the following frameworks:
@@ -105,6 +108,9 @@ If you use our code for your scientific research, please acknowledge this by ref
 
 _Herrmann, L., Bürchner, T., Kudela, L., Kollmannsberger, S., 2026, **A memory-efficient adjoint method to enable billion parameter optimization on a single GPU in dynamic problems**, Structural and Multidisciplinary Optimization, Volume 69, 52 (2026), DOI: [10.1007/s00158-025-04237-y](https://doi.org/10.1007/s00158-025-04237-y)_
 
+## Acknowledgements
+
+Ideas adapted from [FDTDX](https://github.com/ymahlau/fdtdx) ([Mahlau, Schubert, Berg & Rosenhahn 2026](https://doi.org/10.21105/joss.08912)), which also serves as the reference for validating the electromagnetic solvers.
 ## Contact
 
 For questions, bug reports, or collaboration inquiries, please don't hesitate to contact Leon Herrmann at [herrmann.leon@pm.me](mailto:herrmann.leon@pm.me).

{cuwave-0.2.0 → cuwave-0.6.0}/README.md

@@ -5,7 +5,7 @@
   <img src="https://raw.githubusercontent.com/cmpmech/cuwave/main/.assets/readme-light.webp" width="100%">
 </picture>
 
-**CuWave** is a single-GPU, differentiable finite difference wave propagation code.
+**CuWave** is a GPU-accelerated (single or multi), differentiable finite difference wave propagation code.
 Possible applications include
 
 <table width="100%">
@@ -41,13 +41,16 @@ Possible applications include
 - see the [documentation](https://github.com/cmpmech/cuwave/blob/main/docs/Home.md) for how the code works (AI-assisted; verify with sources for critical details)
 - see [examples](https://github.com/cmpmech/cuwave/tree/main/examples) for how to apply the code
 
-## Development status
-
-**Scalar** & **acoustic wave** equations have been developed over the last 2 years and are thoroughly validated.
+## Multiple GPUs
 
-> [!IMPORTANT]
-> **Elastic** & **electromagnetic wave** equations were developed with AI assistance (Claude) and have undergone less validation. The elastic wave equation is currently being validated against experimental results.
+Several GPUs take one argument, and the driver otherwise stays as it is:
+- `Simulation(..., devices=(0, 1))` splits **one grid** along its first axis over the listed GPUs, for grids too large or too slow for one. `simulate`, `sensitivity`, `reconstruction_sensitivity` and `superposition_sensitivity` then run split and return the same result as on one GPU, bitwise. Try it with the `DEVICES` setting of
+	- [examples/forward/scalarND.py](https://github.com/cmpmech/cuwave/blob/main/examples/forward/scalarND.py), the forward problem in 2D or 3D
+	- [examples/sensitivity/scalar2D_sensitivity.py](https://github.com/cmpmech/cuwave/blob/main/examples/sensitivity/scalar2D_sensitivity.py), the gradient
+- `misfit_gradient(..., devices=[0, 1])` (and `measure`, `misfit`) deals the **shots** of an inversion out over the GPUs, which needs no communication and is the faster choice whenever one shot fits one GPU.
+## Development status
 
+**Scalar** & **acoustic wave** equations have been developed over the last 2 years and are thoroughly validated. The **elastic wave** and **electromagnetic wave** equations were developed with AI assistance (Claude). The first is validated on experimental data, while the latter are cross-validated against the independent FDTD code [FDTDX](https://github.com/ymahlau/fdtdx). 
 ## Performance
 
 CuWave's runtime for identical discretizations is comparable to that of other established wave propagation finite difference codes. Speedups in 2D with reflecting boundaries (**a ratio above 1 means CuWave is that many times faster**) compared to the following frameworks:
@@ -79,6 +82,9 @@ If you use our code for your scientific research, please acknowledge this by ref
 
 _Herrmann, L., Bürchner, T., Kudela, L., Kollmannsberger, S., 2026, **A memory-efficient adjoint method to enable billion parameter optimization on a single GPU in dynamic problems**, Structural and Multidisciplinary Optimization, Volume 69, 52 (2026), DOI: [10.1007/s00158-025-04237-y](https://doi.org/10.1007/s00158-025-04237-y)_
 
+## Acknowledgements
+
+Ideas adapted from [FDTDX](https://github.com/ymahlau/fdtdx) ([Mahlau, Schubert, Berg & Rosenhahn 2026](https://doi.org/10.21105/joss.08912)), which also serves as the reference for validating the electromagnetic solvers.
 ## Contact
 
 For questions, bug reports, or collaboration inquiries, please don't hesitate to contact Leon Herrmann at [herrmann.leon@pm.me](mailto:herrmann.leon@pm.me).

{cuwave-0.2.0 → cuwave-0.6.0}/cuwave/anisotropic.py

@@ -23,9 +23,17 @@ import cupy.typing as cpt
 import numpy as np
 import numpy.typing as npt
 
-from .boundary import Clamped, Traction, faces_with
+from .boundary import Clamped, Traction, faces_with, pml_widths
 from .elastic import voigt
-from .wave import PAIRS, Simulation, apply_cell_weights, axis_geometry, grid_block
+from .wave import (
+    PAIRS,
+    PML_THREADS,
+    Simulation,
+    apply_cell_weights,
+    axis_geometry,
+    grid_block,
+    slab_grid,
+)
 
 KERNEL_PATH = Path(__file__).parent / "kernels" / "anisotropic.cu"
 SENSITIVITY_PATH = Path(__file__).parent / "kernels" / "anisotropic_sensitivity.cu"
@@ -96,6 +104,21 @@ def cell_stencil(
     return K
 
 
+def backward_axes(C: npt.NDArray, ndim: int, samples: int = 4096) -> list[int]:
+    """Axes along which a wave of the Voigt `C` runs against its phase, sampled over directions."""
+    row = np.zeros((ndim, ndim), dtype=int)
+    for r, (k, m) in enumerate(PAIRS[ndim]):
+        row[k, m] = row[m, k] = r
+    T = C[row[:, :, None, None], row[None, None]]
+    n = np.random.default_rng(0).standard_normal((samples, ndim))
+    n /= np.linalg.norm(n, axis=1, keepdims=True)
+    _, v = np.linalg.eigh(np.einsum("ikjl,nk,nl->nij", T, n, n))
+    # the group velocity of each mode, up to its positive 1 / (2 omega)
+    group = np.einsum("nip,imjl,nl,njp->npm", v, T + T.transpose(0, 3, 2, 1), n, v)
+    group /= np.linalg.norm(group, axis=2, keepdims=True)
+    return [d for d in range(ndim) if (n[:, None, d] * group[..., d]).min() < -1e-6]
+
+
 def cell_average(sim: Simulation, field: cpt.NDArray) -> cpt.NDArray:
     """Harmonic mean of `field` over the `2**ndim` corners of a cell, at its low corner.
 
@@ -137,6 +160,9 @@ class AnisotropicElasticWave(Simulation):
     sensitivity_path = SENSITIVITY_PATH
     default_boundary = Traction
     gradient_names = ("mass", "stiff", "cell")
+    accepts_pml = True
+    pml_shifted = True
+    pml_walls = (Clamped,)
 
     @property
     def ncomp(self) -> int:
@@ -172,6 +198,14 @@ class AnisotropicElasticWave(Simulation):
                 raise ValueError(f"C must be ({nvoigt}, {nvoigt}): {self.C.shape}")
             if not np.allclose(self.C, self.C.T):
                 raise ValueError("C must be symmetric")
+            widths = pml_widths(self) if self.pml else []
+            lined = [d for d, w in enumerate(widths) if sum(w)]
+            backward = [d for d in backward_axes(self.C, self.ndim) if d in lined]
+            if backward:
+                raise ValueError(
+                    f"C carries a wave against its phase along axes {backward}, which "
+                    f"a pml grows (Becache et al. 2003): use sponge"
+                )
 
     @property
     def lame(self) -> float:
@@ -183,14 +217,17 @@ class AnisotropicElasticWave(Simulation):
         """Second Lame parameter `rho0 * c_s**2`."""
         return self.density * self.wavespeed_s**2
 
-    def stencil(self) -> cpt.NDArray:
+    def stiffness_matrix(self) -> npt.NDArray:
+        """The Voigt stiffness at `gamma = 1`: `C`, or the isotropic one."""
+        if self.C is not None:
+            return self.C
+        return voigt(self.ndim, self.lame, self.shear, self.plane)
+
+    def stencil_table(self) -> cpt.NDArray:
         """Stencil table at `gamma = 1`: one cell's `-B^T C B`, flattened for the kernel."""
         if self._stencil is not None:
             return self._stencil
-        C = self.C
-        if C is None:
-            C = voigt(self.ndim, self.lame, self.shear, self.plane)
-        K = cell_stencil(self.ndim, self.dx, C)
+        K = cell_stencil(self.ndim, self.dx, self.stiffness_matrix())
         self._stencil = cp.asarray(K.ravel(), dtype=self.dtype)
         return self._stencil
 
@@ -215,7 +252,7 @@ class AnisotropicElasticWave(Simulation):
             "minv": cp.ascontiguousarray(minv, dtype=self.dtype),
             "gamma": cp.ascontiguousarray(indicator, dtype=self.dtype),
             "cell": cell_average(self, indicator),
-            "stencil": self.stencil(),
+            "stencil": self.stencil_table(),
         }
         if self.damping is not None:
             mat["damping"] = self.damping
@@ -245,6 +282,63 @@ class AnisotropicElasticWave(Simulation):
         """Gamma scales inertia and stiffness alike, so both derivatives are 1."""
         return 1.0, 1.0
 
+    def define_step(self, kernels: cp.RawModule, mat: dict) -> Callable:
+        """The cell step, followed by its pml terms if there are any."""
+        passes = ("pml_cell", "pml_force")
+        return self._define_pml(kernels, mat, passes)
+
+    def define_adjoint_step(
+        self, kernels: cp.RawModule, sens_kernels: cp.RawModule, mat: dict
+    ) -> Callable:
+        """The same step with the pml terms transposed, which breaks its symmetry."""
+        passes = ("pml_adjoint_memory", "pml_adjoint_cell", "pml_adjoint_force")
+        return self._define_pml(kernels, mat, passes)
+
+    def _define_pml(self, kernels: cp.RawModule, mat: dict, passes: tuple) -> Callable:
+        """The cell step, then each pass over every lined axis, all axes in one memory."""
+        fd_step = super().define_step(kernels, mat)
+        if self.pml is None:
+            return fd_step
+        slabs = self.pml_slabs()
+        sizes = [2 * self.ndim * (lo + hi) * plane for _, lo, hi, _, plane in slabs]
+        starts = np.cumsum([0, *sizes])
+        offsets = np.cumsum([0, *(slab[3].size for slab in slabs)])
+        # per axis: rows lo and hi, then where its memory and profile start
+        table = np.zeros((self.ndim, 4), dtype=np.int32)
+        for k, (d, lo, hi, _, _) in enumerate(slabs):
+            table[d] = lo, hi, starts[k], offsets[k]
+        cv = float(np.prod(self.dx)) * self.stiffness_matrix().ravel()
+        common = [
+            cp.zeros(int(starts[-1]), dtype=self.dtype),
+            mat["minv"],
+            mat["cell"],
+            cp.asarray(cv, dtype=self.dtype),
+            self.dtype(self.dt**2),
+            np.int32(self.comp_stride),
+            cp.concatenate([slab[3].ravel() for slab in slabs]),
+            cp.asarray(table.ravel()),
+        ]
+        geometry = axis_geometry(self, [self.dtype(1.0 / d) for d in self.dx])
+        launches = [
+            (
+                kernels.get_function(f"{name}_kernel"),
+                slab_grid(lo, hi, plane),
+                [None, None, *common, np.int32(d), *geometry],
+            )
+            for name in passes
+            for d, lo, hi, _, plane in slabs
+        ]
+
+        def pml_step(u0, u1, u2):
+            fd_step(u0, u1, u2)
+            for kernel, grid, args in launches:
+                args[0], args[1] = u1, u2
+                kernel(grid, (PML_THREADS,), args)
+            return u2
+
+        pml_step.memory = [common[0]]
+        return pml_step
+
     def step_factors(self) -> list:
         """`dt**2`; the grid spacing already sits in the stencil table."""
         return [self.dtype(self.dt**2)] * self.ndim
@@ -312,7 +406,7 @@ class AnisotropicElasticWave(Simulation):
         grid, block = grid_block(self)
         volume = float(np.prod(self.dx))
         args = [accs["mass"], accs["cell"], None, None, None] + [
-            self.stencil(),
+            self.stencil_table(),
             self.dtype(sign * self.density * volume / (2.0 * self.dt) ** 2),
             self.dtype(-sign),
             np.int32(self.comp_stride),

{cuwave-0.2.0 → cuwave-0.6.0}/cuwave/boundary.py

@@ -5,8 +5,9 @@ setup can place one per face in `Simulation.boundary` long before the module is
 compiled. `define_boundary` turns the markers into the launch closures once
 `compile_kernels` has run, which is why `simulate` needs no separate preparation phase.
 
-`sponge` opens the domain from the other side: it dissipates in the material behind a
-face rather than acting on the ghost ring, so it needs no kernel and no marker of its own.
+`sponge` and `Pml` open the domain from the other side: they absorb in the material
+behind a face rather than acting on the ghost ring. The sponge is a damping field; the
+PML is declarative like a condition, its profiles built once the step is defined.
 """
 
 from __future__ import annotations
@@ -253,3 +254,110 @@ def sponge(
     return cp.ascontiguousarray(
         (2.0 * beta * taper**2 / (sim.dt * minv)).astype(sim.dtype)
     )
+
+
+@dataclass(frozen=True)
+class Pml:
+    """A convolutional perfectly matched layer behind `faces`, for `Simulation.pml`.
+
+    Every derivative along an axis carries a memory variable inside the `width` nodes
+    behind that axis' lined faces, its graded profile set by the target `reflection`
+    at the fastest `wavespeed` (Komatitsch & Martin 2007,
+    https://doi.org/10.1190/1.2757586).
+    """
+
+    width: int  # layer depth in nodes, the wall node included
+    wavespeed: float  # fastest wave speed in the layer, which sizes the damping
+    reflection: float = 1e-6  # target reflection at normal incidence
+    frequency: float = 0.0  # CFS shift pi * frequency through the layer, 0 for none
+    faces: tuple[int, ...] | None = None  # the `2 * axis + side` codes, None for all
+
+
+def validate_pml(sim: Simulation) -> None:
+    """Raise if `sim.pml` cannot be stepped: unsupported, combined, or too wide."""
+    layer = sim.pml
+    if not sim.accepts_pml:
+        raise NotImplementedError(f"{type(sim).__name__} has no pml; use sponge")
+    if sim.damping is not None or sim.masked:
+        raise ValueError("a pml takes neither a damping field nor a domain mask")
+    least = 2 + sim.pml_margin
+    if layer.width < least or layer.wavespeed <= 0.0 or not 0.0 < layer.reflection < 1:
+        raise ValueError(
+            f"a pml needs width >= {least}, a speed and 0 < reflection < 1: {layer}"
+        )
+    if sim.pml_shifted and layer.frequency <= 0.0:
+        raise ValueError(
+            f"{type(sim).__name__} keeps a static mode in an unshifted pml: set "
+            f"frequency, near the low end of the band"
+        )
+    # the transposed taps reach `reach` nodes past the layer, which must not meet
+    faces = _validate_layer(sim, layer.width + sim.reach, layer.faces)
+    if sim.pml_walls is not None:
+        for face in faces:
+            wall = sim.boundary[face // 2][face % 2]
+            if wall not in sim.pml_walls:
+                raise ValueError(
+                    f"face {face} ends the pml on {wall}, which traps a growing wave: "
+                    f"make it {' or '.join(map(repr, sim.pml_walls))}"
+                )
+
+
+def pml_widths(sim: Simulation) -> list[tuple[int, int]]:
+    """Per axis, the (low, high) layer depths in nodes, 0 for a face left unlined."""
+    faces = _canonical_faces(sim.ndim, sim.pml.faces)
+    return [
+        tuple(sim.pml.width if 2 * d + side in faces else 0 for side in (0, 1))
+        for d in range(sim.ndim)
+    ]
+
+
+def pml_profiles(sim: Simulation) -> list[cpt.NDArray | None]:
+    """Per lined axis, the (4, Nx[d]) recursion coefficients `a, b` at nodes and half points.
+
+    Row 0 and 1 are the node values, rows 2 and 3 those of the half point `a + 1/2`,
+    graded from 0 at the front to the wall and mirrored across it, so the ghost cell
+    sees what its mirror does. An axis with no lined face gets None.
+    """
+    layer = sim.pml
+    # the front node is undamped, and so are the margin nodes in front of it
+    L = layer.width - 1 - sim.pml_margin
+    profiles = []
+    for d, (lo, hi) in enumerate(pml_widths(sim)):
+        if not lo + hi:
+            profiles.append(None)
+            continue
+        n = sim.Nx[d]
+        damping = (
+            -3.0 * layer.wavespeed * np.log(layer.reflection) / (2.0 * L * sim.dx[d])
+        )
+        rows = []
+        for position in (np.arange(n) + shift for shift in (0.0, 0.5)):
+            distance = np.full(n, np.inf)
+            if lo:
+                distance = np.minimum(distance, np.abs(position - 1))
+            if hi:
+                distance = np.minimum(distance, np.abs(n - 2 - position))
+            depth = np.clip((L - distance) / L, 0.0, 1.0)
+            d0 = damping * depth**2
+            alpha = np.pi * layer.frequency
+            inside = depth > 0.0
+            b = np.where(inside, np.exp(-(d0 + alpha) * sim.dt), 1.0)
+            a = np.where(inside, d0 / np.maximum(d0 + alpha, 1e-300) * (b - 1.0), 0.0)
+            rows += [a, b]
+        profiles.append(cp.asarray(np.stack(rows), dtype=sim.dtype))
+    return profiles
+
+
+def pml_nodes(sim: Simulation) -> cpt.NDArray:
+    """Interior nodes inside the layer, where the step carries memory terms."""
+    inside = cp.zeros(sim.Nx_padded, dtype=bool)
+    if sim.pml is None:
+        return inside
+    for d, (lo, hi) in enumerate(pml_widths(sim)):
+        index = [slice(None)] * sim.ndim
+        for rows in (slice(1, 1 + lo), slice(sim.Nx[d] - 1 - hi, sim.Nx[d] - 1)):
+            index[d] = rows
+            inside[tuple(index)] = True
+    interior = cp.zeros_like(inside)
+    interior[tuple(slice(1, n - 1) for n in sim.Nx)] = True
+    return inside & interior