cuwave 0.1.0__tar.gz → 0.2.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (37) hide show
  1. {cuwave-0.1.0 → cuwave-0.2.0}/PKG-INFO +3 -23
  2. {cuwave-0.1.0 → cuwave-0.2.0}/README.md +2 -22
  3. {cuwave-0.1.0 → cuwave-0.2.0}/cuwave/anisotropic.py +4 -11
  4. {cuwave-0.1.0 → cuwave-0.2.0}/cuwave/elastic.py +5 -6
  5. {cuwave-0.1.0 → cuwave-0.2.0}/cuwave/kernels/scalar.cu +107 -7
  6. {cuwave-0.1.0 → cuwave-0.2.0}/cuwave/maxwell.py +5 -6
  7. {cuwave-0.1.0 → cuwave-0.2.0}/cuwave/scalar.py +7 -15
  8. {cuwave-0.1.0 → cuwave-0.2.0}/cuwave/sensitivity.py +101 -86
  9. {cuwave-0.1.0 → cuwave-0.2.0}/cuwave/wave.py +256 -106
  10. {cuwave-0.1.0 → cuwave-0.2.0}/cuwave.egg-info/PKG-INFO +3 -23
  11. {cuwave-0.1.0 → cuwave-0.2.0}/pyproject.toml +1 -1
  12. {cuwave-0.1.0 → cuwave-0.2.0}/LICENSE +0 -0
  13. {cuwave-0.1.0 → cuwave-0.2.0}/cuwave/__init__.py +0 -0
  14. {cuwave-0.1.0 → cuwave-0.2.0}/cuwave/boundary.py +0 -0
  15. {cuwave-0.1.0 → cuwave-0.2.0}/cuwave/evals.py +0 -0
  16. {cuwave-0.1.0 → cuwave-0.2.0}/cuwave/geometry.py +0 -0
  17. {cuwave-0.1.0 → cuwave-0.2.0}/cuwave/kernels/__init__.py +0 -0
  18. {cuwave-0.1.0 → cuwave-0.2.0}/cuwave/kernels/anisotropic.cu +0 -0
  19. {cuwave-0.1.0 → cuwave-0.2.0}/cuwave/kernels/anisotropic_sensitivity.cu +0 -0
  20. {cuwave-0.1.0 → cuwave-0.2.0}/cuwave/kernels/common.cuh +0 -0
  21. {cuwave-0.1.0 → cuwave-0.2.0}/cuwave/kernels/elastic.cu +0 -0
  22. {cuwave-0.1.0 → cuwave-0.2.0}/cuwave/kernels/elastic_sensitivity.cu +0 -0
  23. {cuwave-0.1.0 → cuwave-0.2.0}/cuwave/kernels/maxwell.cu +0 -0
  24. {cuwave-0.1.0 → cuwave-0.2.0}/cuwave/kernels/maxwell_sensitivity.cu +0 -0
  25. {cuwave-0.1.0 → cuwave-0.2.0}/cuwave/kernels/scalar_sensitivity.cu +0 -0
  26. {cuwave-0.1.0 → cuwave-0.2.0}/cuwave/nn.py +0 -0
  27. {cuwave-0.1.0 → cuwave-0.2.0}/cuwave/optimization.py +0 -0
  28. {cuwave-0.1.0 → cuwave-0.2.0}/cuwave/postprocessing.py +0 -0
  29. {cuwave-0.1.0 → cuwave-0.2.0}/cuwave/regularization.py +0 -0
  30. {cuwave-0.1.0 → cuwave-0.2.0}/cuwave/signals.py +0 -0
  31. {cuwave-0.1.0 → cuwave-0.2.0}/cuwave/stencils.py +0 -0
  32. {cuwave-0.1.0 → cuwave-0.2.0}/cuwave/utils.py +0 -0
  33. {cuwave-0.1.0 → cuwave-0.2.0}/cuwave.egg-info/SOURCES.txt +0 -0
  34. {cuwave-0.1.0 → cuwave-0.2.0}/cuwave.egg-info/dependency_links.txt +0 -0
  35. {cuwave-0.1.0 → cuwave-0.2.0}/cuwave.egg-info/requires.txt +0 -0
  36. {cuwave-0.1.0 → cuwave-0.2.0}/cuwave.egg-info/top_level.txt +0 -0
  37. {cuwave-0.1.0 → cuwave-0.2.0}/setup.cfg +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: cuwave
3
- Version: 0.1.0
3
+ Version: 0.2.0
4
4
  Summary: GPU finite-difference wave solver with differentiable adjoints
5
5
  Author-email: Leon Herrmann <herrmann.leon@pm.me>
6
6
  License-Expression: MIT
@@ -93,31 +93,11 @@ Additional benefits of **CuWave** are
93
93
 
94
94
  ## Install
95
95
 
96
- Dependencies are kept **lightweight**. Only CuPy is required beyond standard Python library.
97
-
98
96
  ```bash
99
- pip install cupy-cuda12x # or cupy-cuda11x, to match your CUDA
100
- pip install cuwave # or `pip install -e .` from a checkout
97
+ pip install cuwave
101
98
  ```
102
99
 
103
- CuPy must be installed separately because the wheel depends on your CUDA toolkit;
104
- all remaining dependencies are declared in `pyproject.toml`.
105
-
106
- PyTorch is optional for the regularization via neural optimization; see [pytorch](https://pytorch.org/get-started/locally/) for the installation. Otherwise it is not needed.
107
-
108
- > [!NOTE]
109
- > Match PyTorch's CUDA version to CuPy's, or the two runtimes clash at the first kernel launch.
110
- > With `cupy-cuda12x`:
111
- > ```bash
112
- > pip install torch --index-url https://download.pytorch.org/whl/cu128
113
- > ```
114
-
115
- The tests under `tests/` are `unittest` classes, but `pytest` is the recommended runner:
116
-
117
- ```bash
118
- pip install pytest
119
- python -m pytest tests/ -q # ~20 s on a GPU, ~3 s without: CUDA and PyTorch tests skip when unavailable
120
- ```
100
+ Requires an NVIDIA GPU and [CuPy](https://docs.cupy.dev/en/stable/install.html) matching your CUDA toolkit (e.g. `pip install cupy-cuda12x`), which is not pulled in automatically. For the full installation, including the optional PyTorch and running the tests, see [docs/install.md](https://github.com/cmpmech/cuwave/blob/main/docs/install.md).
121
101
 
122
102
  ## References
123
103
 
@@ -67,31 +67,11 @@ Additional benefits of **CuWave** are
67
67
 
68
68
  ## Install
69
69
 
70
- Dependencies are kept **lightweight**. Only CuPy is required beyond standard Python library.
71
-
72
70
  ```bash
73
- pip install cupy-cuda12x # or cupy-cuda11x, to match your CUDA
74
- pip install cuwave # or `pip install -e .` from a checkout
71
+ pip install cuwave
75
72
  ```
76
73
 
77
- CuPy must be installed separately because the wheel depends on your CUDA toolkit;
78
- all remaining dependencies are declared in `pyproject.toml`.
79
-
80
- PyTorch is optional for the regularization via neural optimization; see [pytorch](https://pytorch.org/get-started/locally/) for the installation. Otherwise it is not needed.
81
-
82
- > [!NOTE]
83
- > Match PyTorch's CUDA version to CuPy's, or the two runtimes clash at the first kernel launch.
84
- > With `cupy-cuda12x`:
85
- > ```bash
86
- > pip install torch --index-url https://download.pytorch.org/whl/cu128
87
- > ```
88
-
89
- The tests under `tests/` are `unittest` classes, but `pytest` is the recommended runner:
90
-
91
- ```bash
92
- pip install pytest
93
- python -m pytest tests/ -q # ~20 s on a GPU, ~3 s without: CUDA and PyTorch tests skip when unavailable
94
- ```
74
+ Requires an NVIDIA GPU and [CuPy](https://docs.cupy.dev/en/stable/install.html) matching your CUDA toolkit (e.g. `pip install cupy-cuda12x`), which is not pulled in automatically. For the full installation, including the optional PyTorch and running the tests, see [docs/install.md](https://github.com/cmpmech/cuwave/blob/main/docs/install.md).
95
75
 
96
76
  ## References
97
77
 
@@ -25,7 +25,7 @@ import numpy.typing as npt
25
25
 
26
26
  from .boundary import Clamped, Traction, faces_with
27
27
  from .elastic import voigt
28
- from .wave import PAIRS, Simulation, apply_cell_weights, grid_block
28
+ from .wave import PAIRS, Simulation, apply_cell_weights, axis_geometry, grid_block
29
29
 
30
30
  KERNEL_PATH = Path(__file__).parent / "kernels" / "anisotropic.cu"
31
31
  SENSITIVITY_PATH = Path(__file__).parent / "kernels" / "anisotropic_sensitivity.cu"
@@ -253,13 +253,6 @@ class AnisotropicElasticWave(Simulation):
253
253
  """Source scaling, unscaled since rho0 is already folded into the lumped inertia."""
254
254
  return 1.0
255
255
 
256
- def axis_geometry(self) -> list:
257
- """Extents and previous-axis strides, the step kernel's tail without the factors."""
258
- geom = [self.Nx[0]]
259
- for d in range(1, self.ndim):
260
- geom += [self.Nx[d], self.strides[d - 1]]
261
- return geom
262
-
263
256
  def gradient_fields(self, mat: dict) -> dict[str, cpt.NDArray]:
264
257
  """Nodal accumulators plus the cell one the stiffness density lands in first."""
265
258
  grads = {
@@ -285,7 +278,7 @@ class AnisotropicElasticWave(Simulation):
285
278
  grads["material"],
286
279
  grads["design"],
287
280
  self.dtype(1.0 / 2.0**self.ndim),
288
- *self.axis_geometry(),
281
+ *axis_geometry(self),
289
282
  ],
290
283
  )
291
284
  return {"mass": grads["mass"], "stiff": grads["stiff"]}
@@ -302,7 +295,7 @@ class AnisotropicElasticWave(Simulation):
302
295
  mat["stencil"],
303
296
  mass_factor,
304
297
  np.int32(self.comp_stride),
305
- *self.axis_geometry(),
298
+ *axis_geometry(self),
306
299
  ]
307
300
 
308
301
  def gradient_step(u0, u1, u2, l1):
@@ -323,7 +316,7 @@ class AnisotropicElasticWave(Simulation):
323
316
  self.dtype(sign * self.density * volume / (2.0 * self.dt) ** 2),
324
317
  self.dtype(-sign),
325
318
  np.int32(self.comp_stride),
326
- *self.axis_geometry(),
319
+ *axis_geometry(self),
327
320
  ]
328
321
 
329
322
  def frechet_step(u0, u1, u2):
@@ -26,13 +26,12 @@ from .wave import (
26
26
  Simulation,
27
27
  apply_cell_weights,
28
28
  axis_geometry,
29
- component_weights,
30
29
  grid_block,
31
30
  pair_average,
32
31
  pair_average_adjoint,
33
- pair_weights,
34
32
  point_average,
35
33
  point_average_adjoint,
34
+ wall_weights,
36
35
  )
37
36
 
38
37
  KERNEL_PATH = Path(__file__).parent / "kernels" / "elastic.cu"
@@ -160,7 +159,7 @@ class ElasticWave(Simulation):
160
159
  for c in range(self.ncomp):
161
160
  mass = (
162
161
  self.density
163
- * component_weights(self, c)
162
+ * wall_weights(self, (c,))
164
163
  * point_average(self, gamma, c)
165
164
  )
166
165
  minv[c] = 1.0 / cp.maximum(mass, cp.finfo(self.dtype).tiny)
@@ -178,7 +177,7 @@ class ElasticWave(Simulation):
178
177
  if self.ndim > 1:
179
178
  gshear = cp.zeros((self.npairs, *self.Nx_padded), dtype=self.dtype)
180
179
  for p, axes in enumerate(PAIRS[self.ndim][self.ndim :]):
181
- gshear[p] = pair_weights(self, axes) * pair_average(self, gamma, axes)
180
+ gshear[p] = wall_weights(self, axes) * pair_average(self, gamma, axes)
182
181
  mat["gshear"] = cp.ascontiguousarray(gshear)
183
182
  if self.damping is not None:
184
183
  mat["damping"] = self.damping
@@ -269,12 +268,12 @@ class ElasticWave(Simulation):
269
268
  gamma = grads["design"]
270
269
  g_mass = cp.zeros(self.Nx_padded, dtype=self.dtype)
271
270
  for c in range(self.ncomp):
272
- density = grads["mass"][c] * component_weights(self, c) * self.density
271
+ density = grads["mass"][c] * wall_weights(self, (c,)) * self.density
273
272
  g_mass += point_average_adjoint(self, density, c)
274
273
  # the normal density sits on the nodes, so its chain rule is the weight alone
275
274
  g_stiff = apply_cell_weights(self, grads["normal"].copy())
276
275
  for p, axes in enumerate(PAIRS[self.ndim][self.ndim :]):
277
- density = grads["shear"][p] * pair_weights(self, axes)
276
+ density = grads["shear"][p] * wall_weights(self, axes)
278
277
  g_stiff += pair_average_adjoint(self, density, gamma, axes)
279
278
  return {"mass": g_mass, "stiff": g_stiff}
280
279
 
@@ -2,6 +2,7 @@
2
2
  // Compile-time configuration this file responds to:
3
3
  // NDIM = 1 | 2 | 3
4
4
  // USE_DAMPING
5
+ // USE_DOMAIN
5
6
 
6
7
  // spatial finite difference stencil for Laplacian
7
8
  __device__ __forceinline__ real_t flux_divergence_axis(
@@ -25,6 +26,81 @@ __device__ __forceinline__ real_t flux_divergence_axis(
25
26
  return factor * (Dp * gp - Dm * gm); // outer grad (incl. inner grad)
26
27
  }
27
28
 
29
+ // ----------------------------------- domain helpers
30
+ #ifdef USE_DOMAIN
31
+ #define DIRICHLET 7 // the radius code of a cell open onto a node held at zero
32
+
33
+ // the same with each cell at its own radius, zero for a cell leaving the domain
34
+ __device__ __forceinline__ real_t flux_divergence_cells(
35
+ const real_t *__restrict__ u1, const real_t *__restrict__ stiff,
36
+ const int idx, const int s, const real_t uc, const real_t sc,
37
+ const real_t factor, const int rp, const int rm) {
38
+ real_t Dp = 0, Dm = 0; // a closed cell carries no flux
39
+ if (rp == DIRICHLET)
40
+ Dp = -uc * sc; // onto a zero halfway across, at the node's own stiffness
41
+ else if (rp) {
42
+ const real_t sp = stiff[idx + s];
43
+ Dp = OP_W(rp, 1) * (u1[idx + s] - uc);
44
+ #pragma unroll
45
+ for (int k = 2; k <= STENCIL_RADIUS; ++k)
46
+ if (k <= rp)
47
+ Dp += OP_W(rp, k) * (u1[idx + k * s] - u1[idx - (k - 1) * s]);
48
+ Dp *= sc * sp / (sc + sp); // harmonic mean
49
+ }
50
+ if (rm == DIRICHLET)
51
+ Dm = uc * sc;
52
+ else if (rm) {
53
+ const real_t sm = stiff[idx - s];
54
+ Dm = OP_W(rm, 1) * (uc - u1[idx - s]);
55
+ #pragma unroll
56
+ for (int k = 2; k <= STENCIL_RADIUS; ++k)
57
+ if (k <= rm)
58
+ Dm += OP_W(rm, k) * (u1[idx + (k - 1) * s] - u1[idx - k * s]);
59
+ Dm *= sc * sm / (sc + sm); // harmonic mean
60
+ }
61
+ return factor * (Dp - Dm);
62
+ }
63
+ #endif
64
+
65
+ #ifdef USE_DOMAIN
66
+ #define BLOCK_X (tile & 1023) // one block per domain tile, 10 bits an axis
67
+ #define BLOCK_Y ((tile >> 10) & 1023)
68
+ #define BLOCK_Z ((tile >> 20) & 1023)
69
+ #else
70
+ #define BLOCK_X blockIdx.x
71
+ #define BLOCK_Y blockIdx.y
72
+ #define BLOCK_Z blockIdx.z
73
+ #endif
74
+
75
+ #ifdef USE_DOMAIN
76
+ #define CELLS(s, f, d) \
77
+ flux_divergence_cells(u1, stiff, idx, s, uc, sc, f, (code >> 6 * (d)) & 7, \
78
+ (code >> (6 * (d) + 3)) & 7)
79
+
80
+ // out of line, so the rare wall node costs the deep ones no registers
81
+ __device__ __noinline__ real_t
82
+ wall_laplacian(const real_t *__restrict__ u1, const real_t *__restrict__ stiff,
83
+ const int idx, const real_t uc, const real_t sc, const int code,
84
+ const real_t f0
85
+ #if NDIM >= 2
86
+ ,
87
+ const real_t f1, const int s0
88
+ #endif
89
+ #if NDIM >= 3
90
+ ,
91
+ const real_t f2, const int s1
92
+ #endif
93
+ ) {
94
+ #if NDIM == 1
95
+ return CELLS(1, f0, 0);
96
+ #elif NDIM == 2
97
+ return CELLS(s0, f0, 0) + CELLS(1, f1, 1);
98
+ #elif NDIM == 3
99
+ return CELLS(s0, f0, 0) + CELLS(s1, f1, 1) + CELLS(1, f2, 2);
100
+ #endif
101
+ }
102
+ #endif
103
+
28
104
  // ----------------------------- boundary condition helper
29
105
  #if NDIM == 1
30
106
  #define BC_PARAMS const int N0
@@ -80,6 +156,9 @@ fd_kernel(const real_t *__restrict__ u0, const real_t *__restrict__ u1,
80
156
  const real_t *__restrict__ minv, const int derive_inertia,
81
157
  #ifdef USE_DAMPING
82
158
  const real_t *__restrict__ damping, const real_t dt,
159
+ #endif
160
+ #ifdef USE_DOMAIN
161
+ const int *__restrict__ cells, const int *__restrict__ tiles,
83
162
  #endif
84
163
  const real_t f0, const int N0
85
164
  #if NDIM >= 2
@@ -91,23 +170,26 @@ fd_kernel(const real_t *__restrict__ u0, const real_t *__restrict__ u1,
91
170
  const real_t f2, const int N2, const int s1
92
171
  #endif
93
172
  ) {
173
+ #ifdef USE_DOMAIN
174
+ const int tile = tiles[blockIdx.x];
175
+ #endif
94
176
  #if NDIM == 1
95
- const int a0 = blockIdx.x * blockDim.x + threadIdx.x;
177
+ const int a0 = BLOCK_X * blockDim.x + threadIdx.x;
96
178
  if (!(a0 > 0 && a0 < N0 - 1))
97
179
  return;
98
180
  const int idx = a0;
99
181
  const int r0 = CLOSURE(a0, N0);
100
182
  #elif NDIM == 2
101
- const int a1 = blockIdx.x * blockDim.x + threadIdx.x;
102
- const int a0 = blockIdx.y * blockDim.y + threadIdx.y;
183
+ const int a1 = BLOCK_X * blockDim.x + threadIdx.x;
184
+ const int a0 = BLOCK_Y * blockDim.y + threadIdx.y;
103
185
  if (!(a0 > 0 && a0 < N0 - 1 && a1 > 0 && a1 < N1 - 1))
104
186
  return;
105
187
  const int idx = a0 * s0 + a1;
106
188
  const int r0 = CLOSURE(a0, N0), r1 = CLOSURE(a1, N1);
107
189
  #elif NDIM == 3
108
- const int a2 = blockIdx.x * blockDim.x + threadIdx.x;
109
- const int a1 = blockIdx.y * blockDim.y + threadIdx.y;
110
- const int a0 = blockIdx.z * blockDim.z + threadIdx.z;
190
+ const int a2 = BLOCK_X * blockDim.x + threadIdx.x;
191
+ const int a1 = BLOCK_Y * blockDim.y + threadIdx.y;
192
+ const int a0 = BLOCK_Z * blockDim.z + threadIdx.z;
111
193
  if (!(a0 > 0 && a0 < N0 - 1 && a1 > 0 && a1 < N1 - 1 && a2 > 0 &&
112
194
  a2 < N2 - 1))
113
195
  return;
@@ -117,7 +199,6 @@ fd_kernel(const real_t *__restrict__ u0, const real_t *__restrict__ u1,
117
199
 
118
200
  const real_t uc = u1[idx]; // load once
119
201
  const real_t sc = stiff[idx]; // load once
120
-
121
202
  #if NDIM == 1
122
203
  real_t laplacian = flux_divergence_axis(u1, stiff, idx, 1, uc, sc, f0, r0);
123
204
  #elif NDIM == 2
@@ -128,6 +209,25 @@ fd_kernel(const real_t *__restrict__ u0, const real_t *__restrict__ u1,
128
209
  flux_divergence_axis(u1, stiff, idx, s1, uc, sc, f1, r1) +
129
210
  flux_divergence_axis(u1, stiff, idx, 1, uc, sc, f2, r2);
130
211
  #endif
212
+ #ifdef USE_DOMAIN
213
+ // bit 30 marks a tile touching the wall, so the branch is uniform per block
214
+ if (tile >> 30) {
215
+ const int code = cells[idx]; // 0 outside, bit 30 deep, else the cell radii
216
+ if (!code)
217
+ return; // outside the domain nothing is stepped
218
+ #if NDIM == 1
219
+ if (!(code >> 30))
220
+ laplacian = wall_laplacian(u1, stiff, idx, uc, sc, code, f0);
221
+ #elif NDIM == 2
222
+ if (!(code >> 30))
223
+ laplacian = wall_laplacian(u1, stiff, idx, uc, sc, code, f0, f1, s0);
224
+ #elif NDIM == 3
225
+ if (!(code >> 30))
226
+ laplacian =
227
+ wall_laplacian(u1, stiff, idx, uc, sc, code, f0, f1, s0, f2, s1);
228
+ #endif
229
+ }
230
+ #endif
131
231
 
132
232
  const real_t mi = derive_inertia ? 1.f / sc : minv[idx];
133
233
  #ifdef USE_DAMPING
@@ -27,13 +27,12 @@ from .wave import (
27
27
  Simulation,
28
28
  apply_cell_weights,
29
29
  axis_geometry,
30
- component_weights,
31
30
  grid_block,
32
31
  pair_average,
33
32
  pair_average_adjoint,
34
- pair_weights,
35
33
  point_average,
36
34
  point_average_adjoint,
35
+ wall_weights,
37
36
  )
38
37
 
39
38
  KERNEL_PATH = Path(__file__).parent / "kernels" / "maxwell.cu"
@@ -217,7 +216,7 @@ class MaxwellWave(Simulation):
217
216
  permittivity = cp.ascontiguousarray(permittivity, dtype=self.dtype)
218
217
  minv = cp.zeros((self.ncomp, *self.Nx_padded), dtype=self.dtype)
219
218
  for c in range(self.ncomp):
220
- mass = component_weights(self, c) * point_average(self, permittivity, c)
219
+ mass = wall_weights(self, (c,)) * point_average(self, permittivity, c)
221
220
  minv[c] = 1.0 / cp.maximum(mass, cp.finfo(self.dtype).tiny)
222
221
  # a conductor holds the tangential field, which is what a zeroed inertia does
223
222
  for face in faces_with(self, Conductor):
@@ -231,7 +230,7 @@ class MaxwellWave(Simulation):
231
230
  nu = cp.ascontiguousarray(nu, dtype=self.dtype)
232
231
  pairs = cp.zeros((self.npairs, *self.Nx_padded), dtype=self.dtype)
233
232
  for p, axes in enumerate(PAIRS[self.ndim][self.ndim :]):
234
- pairs[p] = pair_weights(self, axes) * pair_average(self, nu, axes)
233
+ pairs[p] = wall_weights(self, axes) * pair_average(self, nu, axes)
235
234
  mat["nu_pair"] = cp.ascontiguousarray(pairs)
236
235
  mat["nu"] = nu
237
236
  if self.damping is not None:
@@ -310,13 +309,13 @@ class MaxwellWave(Simulation):
310
309
  """Chain the point densities through the averages onto the nodal design fields."""
311
310
  g_mass = cp.zeros(self.Nx_padded, dtype=self.dtype)
312
311
  for c in range(self.ncomp):
313
- density = grads["mass"][c] * component_weights(self, c)
312
+ density = grads["mass"][c] * wall_weights(self, (c,))
314
313
  g_mass += point_average_adjoint(self, density, c)
315
314
  # a non-magnetic medium has no stiffness design dependence: zeros say so
316
315
  g_stiff = cp.zeros(self.Nx_padded, dtype=self.dtype)
317
316
  if self.magnetic:
318
317
  for p, axes in enumerate(PAIRS[self.ndim][self.ndim :]):
319
- density = grads["nu"][p] * pair_weights(self, axes)
318
+ density = grads["nu"][p] * wall_weights(self, axes)
320
319
  g_stiff += pair_average_adjoint(self, density, grads["design"], axes)
321
320
  return {"mass": g_mass, "stiff": g_stiff}
322
321
 
@@ -21,6 +21,7 @@ from .boundary import Neumann
21
21
  from .wave import (
22
22
  Simulation,
23
23
  apply_cell_weights,
24
+ axis_geometry,
24
25
  grid_block,
25
26
  mirror_ghosts,
26
27
  sensor_cell_weights,
@@ -40,6 +41,7 @@ class PressureWave(Simulation):
40
41
  default_boundary = Neumann
41
42
 
42
43
  derive_inertia = False # set where m == k (rho scaling): minv derived from stiff
44
+ accepts_domain = True
43
45
 
44
46
  def build_materials(self, indicator: cpt.NDArray) -> dict:
45
47
  """Turn `indicator` into the kernel's material dict, `damping` included."""
@@ -92,14 +94,8 @@ class PressureWave(Simulation):
92
94
  grid, block = grid_block(self)
93
95
  # the operator without the dt^2 the step folds into it: L, not dt^2 L
94
96
  factors = [self.dtype(float(f) / self.dt**2) for f in self.step_factors()]
95
- geom = [factors[0], self.Nx[0]]
96
- for d in range(1, self.ndim):
97
- geom += [factors[d], self.Nx[d], self.strides[d - 1]]
98
- args = [grads["mass"], grads["stiff"], None, None, None, None] + [
99
- mat["stiff"],
100
- self.dtype(1.0 / self.dt**2),
101
- *geom,
102
- ]
97
+ args = [grads["mass"], grads["stiff"], None, None, None, None, mat["stiff"]]
98
+ args += [self.dtype(1.0 / self.dt**2), *axis_geometry(self, factors)]
103
99
 
104
100
  def gradient_step(u0, u1, u2, l1):
105
101
  args[2], args[3], args[4], args[5] = u0, u1, u2, l1
@@ -115,13 +111,9 @@ class PressureWave(Simulation):
115
111
  frechet_kernel = kernels.get_function("frechet_kernel")
116
112
  grid, block = grid_block(self)
117
113
  # the stiffness density enters negated, so the epilogue scales both alike
118
- geom = [self.dtype(sign / (2.0 * self.dt) ** 2)]
119
- for d in range(self.ndim):
120
- geom.append(self.dtype(-sign / (2.0 * self.dx[d]) ** 2))
121
- geom.append(self.Nx[d])
122
- if d:
123
- geom.append(self.strides[d - 1])
124
- args = [accs["mass"], accs["stiff"], None, None, None] + geom
114
+ factors = [self.dtype(-sign / (2.0 * d) ** 2) for d in self.dx]
115
+ args = [accs["mass"], accs["stiff"], None, None, None]
116
+ args += [self.dtype(sign / (2.0 * self.dt) ** 2), *axis_geometry(self, factors)]
125
117
 
126
118
  def frechet_step(u0, u1, u2):
127
119
  args[2], args[3], args[4] = u0, u1, u2