cuwave 0.2.0__tar.gz → 0.3.0__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cuwave
-Version: 0.2.0
+Version: 0.3.0
 Summary: GPU finite-difference wave solver with differentiable adjoints
 Author-email: Leon Herrmann <herrmann.leon@pm.me>
 License-Expression: MIT

{cuwave-0.2.0 → cuwave-0.3.0}/cuwave/kernels/common.cuh

@@ -69,6 +69,22 @@ excitation_kernel(real_t *__restrict__ u, const real_t *__restrict__ source,
   }
 }
 
+// ------------------------------------------------------------------------------------
+__global__ void adjoint_excitation_kernel(
+    real_t *__restrict__ l2, const real_t *__restrict__ signal,
+    const int offset, const int *__restrict__ lin_index, const int num_sensors,
+    const real_t *__restrict__ weight, real_t *__restrict__ g_mass,
+    const real_t *__restrict__ u1, const real_t mf) {
+  const int idx = blockIdx.x * blockDim.x + threadIdx.x;
+  if (idx < num_sensors) {
+    const int n = lin_index[idx];
+    const real_t load = weight[idx] * signal[offset + idx];
+    atomicAdd(&l2[n], load);
+    // the injected load is part of the second time difference of lambda
+    atomicAdd(&g_mass[n], -mf * u1[n] * load);
+  }
+}
+
 // ------------------------------------------------------------------------------------
 __global__ void get_signal_kernel(const real_t *__restrict__ u,
                                   real_t *__restrict__ um, const int offset,

cuwave-0.3.0/cuwave/kernels/scalar_sensitivity.cu

@@ -0,0 +1,302 @@
+// Prepended by wave.compile_kernels: stencils.preamble, then common.cuh.
+// Compile-time configuration this file responds to:
+//   NDIM = 1 | 2 | 3
+//   USE_DAMPING
+
+// ------------------------------- interior guard macros
+#if NDIM == 1
+#define INTERIOR_OR_RETURN                                                     \
+  const int a0 = blockIdx.x * blockDim.x + threadIdx.x;                        \
+  if (!(a0 > 0 && a0 < N0 - 1))                                                \
+    return;                                                                    \
+  const int idx = a0
+#define AXIS_RADII const int r0 = CLOSURE(a0, N0)
+#define AXIS_OFFSETS const int o0 = 1
+#elif NDIM == 2
+#define INTERIOR_OR_RETURN                                                     \
+  const int a1 = blockIdx.x * blockDim.x + threadIdx.x;                        \
+  const int a0 = blockIdx.y * blockDim.y + threadIdx.y;                        \
+  if (!(a0 > 0 && a0 < N0 - 1 && a1 > 0 && a1 < N1 - 1))                       \
+    return;                                                                    \
+  const int idx = a0 * s0 + a1
+#define AXIS_RADII const int r0 = CLOSURE(a0, N0), r1 = CLOSURE(a1, N1)
+#define AXIS_OFFSETS const int o0 = s0, o1 = 1
+#elif NDIM == 3
+#define INTERIOR_OR_RETURN                                                     \
+  const int a2 = blockIdx.x * blockDim.x + threadIdx.x;                        \
+  const int a1 = blockIdx.y * blockDim.y + threadIdx.y;                        \
+  const int a0 = blockIdx.z * blockDim.z + threadIdx.z;                        \
+  if (!(a0 > 0 && a0 < N0 - 1 && a1 > 0 && a1 < N1 - 1 && a2 > 0 &&            \
+        a2 < N2 - 1))                                                          \
+    return;                                                                    \
+  const int idx = a0 * s0 + a1 * s1 + a2
+#define AXIS_RADII                                                             \
+  const int r0 = CLOSURE(a0, N0), r1 = CLOSURE(a1, N1), r2 = CLOSURE(a2, N2)
+#define AXIS_OFFSETS const int o0 = s0, o1 = s1, o2 = 1
+#endif
+
+// ---------------------------------- gradient helpers
+__device__ __forceinline__ real_t stiffness_gradient_axis(
+    const real_t *__restrict__ u1, const real_t *__restrict__ l1,
+    const real_t *__restrict__ stiff, const int idx, const int s,
+    const real_t uc, const real_t lc, const real_t sc, const real_t factor,
+    const int r) {
+  const real_t sp = stiff[idx + s];                     // plus of sc
+  const real_t sm = stiff[idx - s];                     // minus of sc
+  const real_t dgp = sp * sp / ((sc + sp) * (sc + sp)); // d(harmonic mean)/dsc
+  const real_t dgm = sm * sm / ((sc + sm) * (sc + sm)); // d(harmonic mean)/dsc
+  real_t Dp = OP_W(r, 1) * (u1[idx + s] - uc);          // inner grad
+  real_t Dm = OP_W(r, 1) * (uc - u1[idx - s]);          // inner grad
+#pragma unroll
+  for (int k = 2; k <= STENCIL_RADIUS; ++k)
+    if (k <= r) {
+      Dp += OP_W(r, k) * (u1[idx + k * s] - u1[idx - (k - 1) * s]);
+      Dm += OP_W(r, k) * (u1[idx + (k - 1) * s] - u1[idx - k * s]);
+    }
+  return factor * (dgp * Dp * (l1[idx + s] - lc) +
+                   dgm * Dm * (lc - l1[idx - s])); // both cells of the node
+}
+
+// the adjoint step's flux divergence and the stiffness gradient of one axis,
+// sharing the cell stiffnesses
+__device__ __forceinline__ void adjoint_gradient_axis(
+    const real_t *__restrict__ u1, const real_t *__restrict__ l1,
+    const real_t *__restrict__ stiff, const int idx, const int s,
+    const real_t uc, const real_t lc, const real_t sc, const real_t factor,
+    const int r, real_t &div_l, real_t &g) {
+  const real_t sp = stiff[idx + s];                     // plus of sc
+  const real_t sm = stiff[idx - s];                     // minus of sc
+  const real_t gp = sc * sp / (sc + sp);                // harmonic mean
+  const real_t gm = sc * sm / (sc + sm);                // harmonic mean
+  const real_t dgp = sp * sp / ((sc + sp) * (sc + sp)); // d(harmonic mean)/dsc
+  const real_t dgm = sm * sm / ((sc + sm) * (sc + sm)); // d(harmonic mean)/dsc
+  const real_t lp = l1[idx + s];
+  const real_t lm = l1[idx - s];
+  real_t Lp = OP_W(r, 1) * (lp - lc);          // inner grad of lambda
+  real_t Lm = OP_W(r, 1) * (lc - lm);          // inner grad of lambda
+  real_t Dp = OP_W(r, 1) * (u1[idx + s] - uc); // inner grad of u
+  real_t Dm = OP_W(r, 1) * (uc - u1[idx - s]); // inner grad of u
+#pragma unroll
+  for (int k = 2; k <= STENCIL_RADIUS; ++k)
+    if (k <= r) {
+      Lp += OP_W(r, k) * (l1[idx + k * s] - l1[idx - (k - 1) * s]);
+      Lm += OP_W(r, k) * (l1[idx + (k - 1) * s] - l1[idx - k * s]);
+      Dp += OP_W(r, k) * (u1[idx + k * s] - u1[idx - (k - 1) * s]);
+      Dm += OP_W(r, k) * (u1[idx + (k - 1) * s] - u1[idx - k * s]);
+    }
+  div_l += factor * (Lp * gp - Lm * gm); // the step's flux divergence
+  g += factor * (dgp * Dp * (lp - lc) + dgm * Dm * (lc - lm));
+}
+
+// byte-identical to scalar.cu: each file is its own compilation unit
+__device__ __forceinline__ real_t flux_divergence_axis(
+    const real_t *__restrict__ u1, const real_t *__restrict__ stiff,
+    const int idx, const int s, const real_t uc, const real_t sc,
+    const real_t factor, const int r) {
+  const real_t sp = stiff[idx + s];            // plus of sc
+  const real_t sm = stiff[idx - s];            // minus of sc
+  const real_t gp = sc * sp / (sc + sp);       // harmonic mean
+  const real_t gm = sc * sm / (sc + sm);       // harmonic mean
+  real_t Dp = OP_W(r, 1) * (u1[idx + s] - uc); // initialization: inner grad
+  real_t Dm = OP_W(r, 1) * (uc - u1[idx - s]); // initialization: inner grad
+#pragma unroll
+  for (int k = 2; k <= STENCIL_RADIUS; ++k)
+    if (k <= r) {
+      Dp +=
+          OP_W(r, k) * (u1[idx + k * s] - u1[idx - (k - 1) * s]); // inner grad
+      Dm +=
+          OP_W(r, k) * (u1[idx + (k - 1) * s] - u1[idx - k * s]); // inner grad
+    }
+  return factor * (Dp * gp - Dm * gm); // outer grad (incl. inner grad)
+}
+
+// -------------------------------------- kernels
+extern "C" {
+
+// ------------------------------------------------------------------------------------
+__global__ void
+gradient_kernel(real_t *__restrict__ g_mass, real_t *__restrict__ g_stiff,
+                const real_t *__restrict__ u0, const real_t *__restrict__ u1,
+                const real_t *__restrict__ u2, const real_t *__restrict__ l1,
+                const real_t *__restrict__ stiff, const real_t inv_dt2,
+                const real_t F0, const int N0
+#if NDIM >= 2
+                ,
+                const real_t F1, const int N1, const int s0
+#endif
+#if NDIM >= 3
+                ,
+                const real_t F2, const int N2, const int s1
+#endif
+) {
+  INTERIOR_OR_RETURN;
+  AXIS_RADII;
+
+  const real_t uc = u1[idx];
+  const real_t lc = l1[idx];
+
+  // dJ/dmass: no neighbour and no material load
+  g_mass[idx] -= inv_dt2 * lc * (u2[idx] - 2.f * uc + u0[idx]);
+
+  // dJ/dstiff: scalar.cu's harmonic cell mean differentiated in place
+  const real_t sc = stiff[idx];
+#if NDIM == 1
+  const real_t g =
+      stiffness_gradient_axis(u1, l1, stiff, idx, 1, uc, lc, sc, F0, r0);
+#elif NDIM == 2
+  const real_t g =
+      stiffness_gradient_axis(u1, l1, stiff, idx, s0, uc, lc, sc, F0, r0) +
+      stiffness_gradient_axis(u1, l1, stiff, idx, 1, uc, lc, sc, F1, r1);
+#elif NDIM == 3
+  const real_t g =
+      stiffness_gradient_axis(u1, l1, stiff, idx, s0, uc, lc, sc, F0, r0) +
+      stiffness_gradient_axis(u1, l1, stiff, idx, s1, uc, lc, sc, F1, r1) +
+      stiffness_gradient_axis(u1, l1, stiff, idx, 1, uc, lc, sc, F2, r2);
+#endif
+  g_stiff[idx] -= g;
+}
+
+// ------------------------------------------------------------------------------------
+__global__ void frechet_kernel(real_t *__restrict__ acc_mass,
+                               real_t *__restrict__ acc_stiff,
+                               const real_t *__restrict__ u0,
+                               const real_t *__restrict__ u1,
+                               const real_t *__restrict__ u2, const real_t ft,
+                               const real_t f0, const int N0
+#if NDIM >= 2
+                               ,
+                               const real_t f1, const int N1, const int s0
+#endif
+#if NDIM >= 3
+                               ,
+                               const real_t f2, const int N2, const int s1
+#endif
+) {
+  INTERIOR_OR_RETURN;
+  AXIS_OFFSETS;
+
+  const real_t dudt = u2[idx] - u0[idx];
+  acc_mass[idx] += ft * dudt * dudt;
+
+  const real_t g0 = u1[idx + o0] - u1[idx - o0];
+  real_t sum = f0 * g0 * g0;
+#if NDIM >= 2
+  const real_t g1 = u1[idx + o1] - u1[idx - o1];
+  sum += f1 * g1 * g1;
+#endif
+#if NDIM >= 3
+  const real_t g2 = u1[idx + o2] - u1[idx - o2];
+  sum += f2 * g2 * g2;
+#endif
+  acc_stiff[idx] += sum;
+}
+
+// ------------------------------------------------------------------------------------
+__global__ void adjoint_gradient_kernel(
+    real_t *__restrict__ l0, const real_t *__restrict__ l1,
+    real_t *__restrict__ g_mass, real_t *__restrict__ g_stiff,
+    const real_t *__restrict__ u1, const real_t *__restrict__ stiff,
+    const real_t *__restrict__ minv, const int derive_inertia,
+#ifdef USE_DAMPING
+    const real_t *__restrict__ damping, const real_t dt,
+#endif
+    const real_t mf, const real_t f0, const int N0
+#if NDIM >= 2
+    ,
+    const real_t f1, const int N1, const int s0
+#endif
+#if NDIM >= 3
+    ,
+    const real_t f2, const int N2, const int s1
+#endif
+) {
+  INTERIOR_OR_RETURN;
+  AXIS_RADII;
+
+  const real_t uc = u1[idx];
+  const real_t lc = l1[idx];
+  const real_t sc = stiff[idx];
+  real_t div_l = 0, g = 0;
+#if NDIM == 1
+  adjoint_gradient_axis(u1, l1, stiff, idx, 1, uc, lc, sc, f0, r0, div_l, g);
+#elif NDIM == 2
+  adjoint_gradient_axis(u1, l1, stiff, idx, s0, uc, lc, sc, f0, r0, div_l, g);
+  adjoint_gradient_axis(u1, l1, stiff, idx, 1, uc, lc, sc, f1, r1, div_l, g);
+#elif NDIM == 3
+  adjoint_gradient_axis(u1, l1, stiff, idx, s0, uc, lc, sc, f0, r0, div_l, g);
+  adjoint_gradient_axis(u1, l1, stiff, idx, s1, uc, lc, sc, f1, r1, div_l, g);
+  adjoint_gradient_axis(u1, l1, stiff, idx, 1, uc, lc, sc, f2, r2, div_l, g);
+#endif
+
+  const real_t mi = derive_inertia ? 1.f / sc : minv[idx];
+  const real_t lo = l0[idx];
+#ifdef USE_DAMPING
+  const real_t beta = 0.5f * mi * damping[idx] * dt;
+  const real_t ln = (2.f * lc - lo * (1.f - beta) + mi * div_l) / (1.f + beta);
+#else
+  const real_t ln = -lo + 2.f * lc + mi * div_l;
+#endif
+  l0[idx] = ln;
+  // dJ/dmass by parts in time: u against the second difference of lambda
+  g_mass[idx] -= mf * uc * (ln - 2.f * lc + lo);
+  g_stiff[idx] -= mf * g;
+}
+
+// ------------------------------------------------------------------------------------
+__global__ void superposed_kernel(
+    const real_t *__restrict__ u0, const real_t *__restrict__ u1,
+    real_t *__restrict__ u2, real_t *__restrict__ acc_mass,
+    real_t *__restrict__ acc_stiff, const real_t *__restrict__ stiff,
+    const real_t *__restrict__ minv, const int derive_inertia, const real_t ft,
+    const real_t fs, const real_t f0, const int N0
+#if NDIM >= 2
+    ,
+    const real_t f1, const int N1, const int s0
+#endif
+#if NDIM >= 3
+    ,
+    const real_t f2, const int N2, const int s1
+#endif
+) {
+  INTERIOR_OR_RETURN;
+  AXIS_RADII;
+  AXIS_OFFSETS;
+
+  const real_t uc = u1[idx];
+  const real_t sc = stiff[idx];
+#if NDIM == 1
+  const real_t laplacian =
+      flux_divergence_axis(u1, stiff, idx, 1, uc, sc, f0, r0);
+#elif NDIM == 2
+  const real_t laplacian =
+      flux_divergence_axis(u1, stiff, idx, s0, uc, sc, f0, r0) +
+      flux_divergence_axis(u1, stiff, idx, 1, uc, sc, f1, r1);
+#elif NDIM == 3
+  const real_t laplacian =
+      flux_divergence_axis(u1, stiff, idx, s0, uc, sc, f0, r0) +
+      flux_divergence_axis(u1, stiff, idx, s1, uc, sc, f1, r1) +
+      flux_divergence_axis(u1, stiff, idx, 1, uc, sc, f2, r2);
+#endif
+
+  // the mass density of the triplet before, whose oldest slot u2 still holds
+  const real_t dudt = uc - u2[idx];
+  acc_mass[idx] += ft * dudt * dudt;
+
+  // the stiffness density of this triplet, its middle slot being u1
+  const real_t g0 = u1[idx + o0] - u1[idx - o0];
+  real_t sum = f0 * g0 * g0;
+#if NDIM >= 2
+  const real_t g1 = u1[idx + o1] - u1[idx - o1];
+  sum += f1 * g1 * g1;
+#endif
+#if NDIM >= 3
+  const real_t g2 = u1[idx + o2] - u1[idx - o2];
+  sum += f2 * g2 * g2;
+#endif
+  acc_stiff[idx] += fs * sum;
+
+  const real_t mi = derive_inertia ? 1.f / sc : minv[idx];
+  u2[idx] = -u0[idx] + 2.f * uc + mi * laplacian;
+}
+
+} // extern "C"

{cuwave-0.2.0 → cuwave-0.3.0}/cuwave/optimization.py

@@ -68,24 +68,26 @@ class Lbfgs:
             project: applied to every trial, a constraint on the design variables.
 
         Returns:
-            (design, cost, alpha, trials) of the accepted trial, which is the last one
-            attempted even when `alpha` bottomed out at `min_alpha`.
+            (design, cost, alpha, trials) of the accepted trial. A search that reaches
+            `min_alpha` without a decrease returns `x` and `cost` unchanged and drops
+            the secant pairs, so the next call restarts from the scaled gradient.
         """
         step = self.step(x, grad) - x
         if not self.pairs:
             step = step * (self.first_step / float(abs(step).max()))
-        slope = float(grad.ravel() @ step.ravel())
 
         alpha, trials = 1.0, 0
         while True:
             trial = project(x + alpha * step)
             trial_cost = f(trial)
             trials += 1
-            if (
-                trial_cost <= cost + self.armijo * alpha * slope
-                or alpha <= self.min_alpha
-            ):
+            # projected displacement: a clipped trial moves less than alpha * step
+            slope = float(grad.ravel() @ (trial - x).ravel())
+            if trial_cost <= cost + self.armijo * slope:
                 return trial, trial_cost, alpha, trials
+            if alpha <= self.min_alpha:
+                self.pairs.clear()
+                return x, cost, alpha, trials
             alpha *= self.shrink
 
     def step(self, x: cpt.NDArray, grad: cpt.NDArray) -> cpt.NDArray:

{cuwave-0.2.0 → cuwave-0.3.0}/cuwave/scalar.py

@@ -67,6 +67,7 @@ class PressureWave(Simulation):
         return args
 
     gradient_names = ("mass", "stiff")  # the fields the adjoint differentiates
+    fuses_adjoint = True  # one kernel steps the adjoint and accumulates both gradients
 
     def gradient_fields(self, mat: dict) -> dict[str, cpt.NDArray]:
         """Zeroed accumulators the adjoint kernels add into, one per material field."""
@@ -103,15 +104,59 @@ class PressureWave(Simulation):
 
         return gradient_step
 
+    def define_adjoint_gradient(
+        self, kernels: cp.RawModule, mat: dict, grads: dict
+    ) -> Callable:
+        """Closure stepping the adjoint `l0` over `l1`, accumulating both gradients of `l1`."""
+        adjoint_gradient_kernel = kernels.get_function("adjoint_gradient_kernel")
+        grid, block = grid_block(self)
+        args = [None, None, grads["mass"], grads["stiff"], None]
+        args += [*self.step_kernel_args(mat), self.dtype(self.adjoint_mass_factor())]
+        args += axis_geometry(self, self.step_factors())
+
+        def adjoint_gradient_step(l0, l1, u1):
+            args[0], args[1], args[4] = l0, l1, u1
+            adjoint_gradient_kernel(grid, block, args)
+            return l0
+
+        return adjoint_gradient_step
+
+    def define_superposed(
+        self, kernels: cp.RawModule, mat: dict, accs: dict, sign: float
+    ) -> Callable:
+        """Closure stepping (u0, u1, u2), adding `sign` times the stiffness density of this triplet and the mass density of the one before."""
+        superposed_kernel = kernels.get_function("superposed_kernel")
+        grid, block = grid_block(self)
+        factors = self.step_factors()
+        # both per-axis factors go as 1 / dx**2, so their ratio is one number
+        ratios = [-sign / (2.0 * d) ** 2 / float(f) for d, f in zip(self.dx, factors)]
+        if not np.allclose(ratios, ratios[0]):
+            raise ValueError(
+                f"step factors {factors} are not proportional to 1 / dx**2"
+            )
+        args = [None, None, None, accs["mass"], accs["stiff"]]
+        args += [*self.step_kernel_args(mat), self.dtype(sign / (2.0 * self.dt) ** 2)]
+        args += [self.dtype(ratios[0]), *axis_geometry(self, factors)]
+        massless = list(args)
+        massless[8] = self.dtype(0.0)
+
+        def superposed_step(u0, u1, u2, mass=True):
+            launch = args if mass else massless
+            launch[0], launch[1], launch[2] = u0, u1, u2
+            superposed_kernel(grid, block, launch)
+            return u2
+
+        return superposed_step
+
     def define_frechet(
-        self, kernels: cp.RawModule, accs: dict, sign: float
+        self, kernels: cp.RawModule, accs: dict, sign: float, stiffness: bool = True
     ) -> 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
-        factors = [self.dtype(-sign / (2.0 * d) ** 2) for d in self.dx]
+        factors = [self.dtype(-sign * stiffness / (2.0 * d) ** 2) for d in self.dx]
         args = [accs["mass"], accs["stiff"], None, None, None]
         args += [self.dtype(sign / (2.0 * self.dt) ** 2), *axis_geometry(self, factors)]
 

{cuwave-0.2.0 → cuwave-0.3.0}/cuwave/sensitivity.py

@@ -25,6 +25,7 @@ from .wave import (
     Source,
     Window,
     compile_kernels,
+    define_adjoint_excitation,
     define_excitation,
     define_get_signal,
     define_set_signal,
@@ -218,7 +219,6 @@ def sensitivity(
     bc_step = define_boundary(sim, kernels)
     excitation_step = define_excitation(sim, source.position, kernels, mat)
     grads = sim.gradient_fields(mat)
-    gradient_step = sim.define_gradient(sens_kernels, mat, grads)
 
 # ------------------------------------ forward pass -----------------------------------
     # stepped straight into the history, so the leading zeros are u^-2 / u^-1
@@ -238,10 +238,24 @@ def sensitivity(
 
 # --------------------------------- adjoint excitation --------------------------------
     signal = adjoint_signal(sim, dphi, sensors)
-    adjoint_excitation = define_excitation(sim, sensors, kernels, mat)
 
 # ----------------------------------- backward pass -----------------------------------
     P = cp.zeros((2, *sim.field_shape), dtype=sim.dtype)
+    if sim.fuses_adjoint:
+        adjoint_step = sim.define_adjoint_gradient(sens_kernels, mat, grads)
+        adjoint_excitation = define_adjoint_excitation(sim, sensors, kernels, mat, grads)
+
+        # reversed, so row m is the middle slot both gradients of lambda^(n + 1) read
+        def backward(m, signal, history):
+            p = adjoint_step(P[m % 2], P[1 - m % 2], history[m])  # lambda^n
+            adjoint_excitation(p, signal, history[m], m)
+            bc_step(p)
+
+        march(sim.N, backward, [Window(signal), Window(V[::-1])])
+        return cost, sim.finalize_gradients(grads, sens_kernels), um, {}
+
+    gradient_step = sim.define_gradient(sens_kernels, mat, grads)
+    adjoint_excitation = define_excitation(sim, sensors, kernels, mat)
 
     # reversed, so row m + 2 - i is u^(n + i) for n = N - 1 - m
     def backward(m, signal, history):
@@ -418,13 +432,18 @@ def superposition_sensitivity(
     kernels = compile_kernels(sim)
     sens_kernels = compile_kernels(sim, sim.sensitivity_path)
 
-    fd_step = sim.define_step(kernels, mat)
     bc_step = define_boundary(sim, kernels)
     excitation_step = define_excitation(sim, source.position, kernels, mat)
     get_signal = define_get_signal(sim, sensors, kernels)
     accs = sim.gradient_fields(mat)
-    subtract_step = sim.define_frechet(sens_kernels, accs, -1.0)
-    add_step = sim.define_frechet(sens_kernels, accs, 1.0)
+    fused = sim.fuses_adjoint
+    if fused:
+        subtract_step = sim.define_superposed(sens_kernels, mat, accs, -1.0)
+        add_step = sim.define_superposed(sens_kernels, mat, accs, 1.0)
+    else:
+        fd_step = sim.define_step(kernels, mat)
+        subtract_step = sim.define_frechet(sens_kernels, accs, -1.0)
+        add_step = sim.define_frechet(sens_kernels, accs, 1.0)
 
     U = cp.zeros((3, *sim.field_shape), dtype=sim.dtype)
     um = cp.zeros((sim.N, sensors.shape[1]), dtype=sim.dtype)
@@ -433,13 +452,21 @@ def superposition_sensitivity(
     # records the traces and subtracts the forward diagonal B(u, u)
     def forward(t, signal, um):
         u0, u1, u2 = U[t % 3], U[(t + 1) % 3], U[(t + 2) % 3]
-        fd_step(u0, u1, u2)
+        if fused:
+            subtract_step(u0, u1, u2)  # the mass density one triplet late
+        else:
+            fd_step(u0, u1, u2)
         excitation_step(u2, signal, t)
         bc_step(u2)
         get_signal(u2, um, t)
-        subtract_step(u0, u1, u2)
+        if not fused:
+            subtract_step(u0, u1, u2)
 
     march(sim.N, forward, [Window(source.signal), Window(um, store=True)])
+    if fused:
+        t = sim.N - 1
+        last = (U[t % 3], U[(t + 1) % 3], U[(t + 2) % 3])
+        sim.define_frechet(sens_kernels, accs, -1.0, stiffness=False)(*last)
 
     cost, dphi = objective(um)
 
@@ -461,14 +488,26 @@ def superposition_sensitivity(
 
 # ----------------------------------- backward pass -----------------------------------
     # starting from u^(N - 1) over u^(N - 2), so the one array carries u + k lambda
-    def backward(t, signal):
+    def backward(t, signal, row, mass=True):
         u0, u1, u2 = (U[(sim.N + i - t) % 3] for i in (1, 0, 2))
-        fd_step(u0, u1, u2)
-        backward_excitation(u2, signal, t)
+        if fused:
+            add_step(u0, u1, u2, mass)
+        else:
+            fd_step(u0, u1, u2)
+        backward_excitation(u2, signal, row)
         bc_step(u2)
-        add_step(u0, u1, u2)
-
-    march(sim.N, backward, [Window(backward_signal)])
+        if not fused:
+            add_step(u0, u1, u2)
+
+    def shifted(t, signal):
+        backward(t + 1, signal, t)
+
+    # outside the graph, as the first fused step's triplet before is the forward's last
+    backward(0, backward_signal, 0, mass=False)
+    march(sim.N - 1, shifted, [Window(backward_signal[1:])])
+    if fused:
+        last = (U[(1 + i) % 3] for i in (1, 0, 2))
+        sim.define_frechet(sens_kernels, accs, 1.0, stiffness=False)(*last)
 
     after = _accumulated(accs)
     cancellation = before / after if after > 0.0 else float("inf")

{cuwave-0.2.0 → cuwave-0.3.0}/cuwave/wave.py

@@ -271,6 +271,7 @@ class Simulation:
     dirichlet: cpt.NDArray | None = None  # nodal mask held at zero, its cells left open
 
     accepts_domain = False  # whether the step kernel honours -DUSE_DOMAIN
+    fuses_adjoint = False  # whether the gradients ride in the adjoint step kernels
 
     @property
     def compile_flags(self) -> tuple[str, ...]:
@@ -331,6 +332,10 @@ class Simulation:
                 f"index; coarsen the grid"
             )
 
+    def adjoint_mass_factor(self) -> float:
+        """Scale of the inertia gradient density: `1 / dt**2`, the residual's own."""
+        return 1.0 / self.dt**2
+
     def define_step(self, kernels: cp.RawModule, mat: dict) -> Callable:
         """Closure launching the finite-difference step kernel over (u0, u1, u2)."""
         fd_kernel = kernels.get_function("fd_kernel")
@@ -595,6 +600,32 @@ def define_excitation(
     return _define_transfer(sim, "excitation_kernel", position, kernels, weight)
 
 
+def define_adjoint_excitation(
+    sim: Simulation,
+    sensors: cpt.NDArray[cp.int32],
+    kernels: cp.RawModule,
+    mat: dict,
+    grads: dict,
+) -> Callable:
+    """Closure adding row `t_index` of the adjoint signal into `l2`, and its share of the inertia gradient against `u1`."""
+    lin_index = flatten_indices(sim, sensors)
+    kernel = kernels.get_function("adjoint_excitation_kernel")
+    threads = 256
+    num = sensors.shape[1]
+    blocks = ((num + threads - 1) // threads,)
+    args = [None, None, np.int32(0), lin_index, np.int32(num)]
+    args += [sim.excitation_weights(mat, lin_index), grads["mass"], None]
+    args.append(sim.dtype(sim.adjoint_mass_factor()))
+
+    def adjoint_excitation_step(l2, signal, u1, t_index):
+        args[0], args[1], args[7] = l2, signal, u1
+        args[2] = np.int32(t_index * num)
+        kernel(blocks, (threads,), args)
+        return l2
+
+    return adjoint_excitation_step
+
+
 def define_get_signal(
     sim: Simulation, sensors: cpt.NDArray[cp.int32], kernels: cp.RawModule
 ) -> Callable:

{cuwave-0.2.0 → cuwave-0.3.0}/cuwave.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cuwave
-Version: 0.2.0
+Version: 0.3.0
 Summary: GPU finite-difference wave solver with differentiable adjoints
 Author-email: Leon Herrmann <herrmann.leon@pm.me>
 License-Expression: MIT

{cuwave-0.2.0 → cuwave-0.3.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "cuwave"
-version = "0.2.0"
+version = "0.3.0"
 description = "GPU finite-difference wave solver with differentiable adjoints"
 readme = "README.md"
 license = "MIT"

cuwave-0.2.0/cuwave/kernels/scalar_sensitivity.cu

@@ -1,140 +0,0 @@
-// Prepended by wave.compile_kernels: stencils.preamble, then common.cuh.
-// Compile-time configuration this file responds to:
-//   NDIM = 1 | 2 | 3
-
-// ------------------------------- interior guard macros
-#if NDIM == 1
-#define INTERIOR_OR_RETURN                                                     \
-  const int a0 = blockIdx.x * blockDim.x + threadIdx.x;                        \
-  if (!(a0 > 0 && a0 < N0 - 1))                                                \
-    return;                                                                    \
-  const int idx = a0
-#define AXIS_RADII const int r0 = CLOSURE(a0, N0)
-#define AXIS_OFFSETS const int o0 = 1
-#elif NDIM == 2
-#define INTERIOR_OR_RETURN                                                     \
-  const int a1 = blockIdx.x * blockDim.x + threadIdx.x;                        \
-  const int a0 = blockIdx.y * blockDim.y + threadIdx.y;                        \
-  if (!(a0 > 0 && a0 < N0 - 1 && a1 > 0 && a1 < N1 - 1))                       \
-    return;                                                                    \
-  const int idx = a0 * s0 + a1
-#define AXIS_RADII const int r0 = CLOSURE(a0, N0), r1 = CLOSURE(a1, N1)
-#define AXIS_OFFSETS const int o0 = s0, o1 = 1
-#elif NDIM == 3
-#define INTERIOR_OR_RETURN                                                     \
-  const int a2 = blockIdx.x * blockDim.x + threadIdx.x;                        \
-  const int a1 = blockIdx.y * blockDim.y + threadIdx.y;                        \
-  const int a0 = blockIdx.z * blockDim.z + threadIdx.z;                        \
-  if (!(a0 > 0 && a0 < N0 - 1 && a1 > 0 && a1 < N1 - 1 && a2 > 0 &&            \
-        a2 < N2 - 1))                                                          \
-    return;                                                                    \
-  const int idx = a0 * s0 + a1 * s1 + a2
-#define AXIS_RADII                                                             \
-  const int r0 = CLOSURE(a0, N0), r1 = CLOSURE(a1, N1), r2 = CLOSURE(a2, N2)
-#define AXIS_OFFSETS const int o0 = s0, o1 = s1, o2 = 1
-#endif
-
-// ---------------------------------- gradient helpers
-__device__ __forceinline__ real_t stiffness_gradient_axis(
-    const real_t *__restrict__ u1, const real_t *__restrict__ l1,
-    const real_t *__restrict__ stiff, const int idx, const int s,
-    const real_t uc, const real_t lc, const real_t sc, const real_t factor,
-    const int r) {
-  const real_t sp = stiff[idx + s];                     // plus of sc
-  const real_t sm = stiff[idx - s];                     // minus of sc
-  const real_t dgp = sp * sp / ((sc + sp) * (sc + sp)); // d(harmonic mean)/dsc
-  const real_t dgm = sm * sm / ((sc + sm) * (sc + sm)); // d(harmonic mean)/dsc
-  real_t Dp = OP_W(r, 1) * (u1[idx + s] - uc);          // inner grad
-  real_t Dm = OP_W(r, 1) * (uc - u1[idx - s]);          // inner grad
-#pragma unroll
-  for (int k = 2; k <= STENCIL_RADIUS; ++k)
-    if (k <= r) {
-      Dp += OP_W(r, k) * (u1[idx + k * s] - u1[idx - (k - 1) * s]);
-      Dm += OP_W(r, k) * (u1[idx + (k - 1) * s] - u1[idx - k * s]);
-    }
-  return factor * (dgp * Dp * (l1[idx + s] - lc) +
-                   dgm * Dm * (lc - l1[idx - s])); // both cells of the node
-}
-
-// -------------------------------------- kernels
-extern "C" {
-
-// ------------------------------------------------------------------------------------
-__global__ void
-gradient_kernel(real_t *__restrict__ g_mass, real_t *__restrict__ g_stiff,
-                const real_t *__restrict__ u0, const real_t *__restrict__ u1,
-                const real_t *__restrict__ u2, const real_t *__restrict__ l1,
-                const real_t *__restrict__ stiff, const real_t inv_dt2,
-                const real_t F0, const int N0
-#if NDIM >= 2
-                ,
-                const real_t F1, const int N1, const int s0
-#endif
-#if NDIM >= 3
-                ,
-                const real_t F2, const int N2, const int s1
-#endif
-) {
-  INTERIOR_OR_RETURN;
-  AXIS_RADII;
-
-  const real_t uc = u1[idx];
-  const real_t lc = l1[idx];
-
-  // dJ/dmass: no neighbour and no material load
-  g_mass[idx] -= inv_dt2 * lc * (u2[idx] - 2.f * uc + u0[idx]);
-
-  // dJ/dstiff: scalar.cu's harmonic cell mean differentiated in place
-  const real_t sc = stiff[idx];
-#if NDIM == 1
-  const real_t g =
-      stiffness_gradient_axis(u1, l1, stiff, idx, 1, uc, lc, sc, F0, r0);
-#elif NDIM == 2
-  const real_t g =
-      stiffness_gradient_axis(u1, l1, stiff, idx, s0, uc, lc, sc, F0, r0) +
-      stiffness_gradient_axis(u1, l1, stiff, idx, 1, uc, lc, sc, F1, r1);
-#elif NDIM == 3
-  const real_t g =
-      stiffness_gradient_axis(u1, l1, stiff, idx, s0, uc, lc, sc, F0, r0) +
-      stiffness_gradient_axis(u1, l1, stiff, idx, s1, uc, lc, sc, F1, r1) +
-      stiffness_gradient_axis(u1, l1, stiff, idx, 1, uc, lc, sc, F2, r2);
-#endif
-  g_stiff[idx] -= g;
-}
-
-// ------------------------------------------------------------------------------------
-__global__ void frechet_kernel(real_t *__restrict__ acc_mass,
-                               real_t *__restrict__ acc_stiff,
-                               const real_t *__restrict__ u0,
-                               const real_t *__restrict__ u1,
-                               const real_t *__restrict__ u2, const real_t ft,
-                               const real_t f0, const int N0
-#if NDIM >= 2
-                               ,
-                               const real_t f1, const int N1, const int s0
-#endif
-#if NDIM >= 3
-                               ,
-                               const real_t f2, const int N2, const int s1
-#endif
-) {
-  INTERIOR_OR_RETURN;
-  AXIS_OFFSETS;
-
-  const real_t dudt = u2[idx] - u0[idx];
-  acc_mass[idx] += ft * dudt * dudt;
-
-  const real_t g0 = u1[idx + o0] - u1[idx - o0];
-  real_t sum = f0 * g0 * g0;
-#if NDIM >= 2
-  const real_t g1 = u1[idx + o1] - u1[idx - o1];
-  sum += f1 * g1 * g1;
-#endif
-#if NDIM >= 3
-  const real_t g2 = u1[idx + o2] - u1[idx - o2];
-  sum += f2 * g2 * g2;
-#endif
-  acc_stiff[idx] += sum;
-}
-
-} // extern "C"