cuwave 0.1.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/LICENSE +21 -0
  2. cuwave-0.1.0/PKG-INFO +134 -0
  3. cuwave-0.1.0/README.md +108 -0
  4. cuwave-0.1.0/cuwave/__init__.py +8 -0
  5. cuwave-0.1.0/cuwave/anisotropic.py +337 -0
  6. cuwave-0.1.0/cuwave/boundary.py +255 -0
  7. cuwave-0.1.0/cuwave/elastic.py +342 -0
  8. cuwave-0.1.0/cuwave/evals.py +130 -0
  9. cuwave-0.1.0/cuwave/geometry.py +226 -0
  10. cuwave-0.1.0/cuwave/kernels/__init__.py +0 -0
  11. cuwave-0.1.0/cuwave/kernels/anisotropic.cu +174 -0
  12. cuwave-0.1.0/cuwave/kernels/anisotropic_sensitivity.cu +226 -0
  13. cuwave-0.1.0/cuwave/kernels/common.cuh +95 -0
  14. cuwave-0.1.0/cuwave/kernels/elastic.cu +225 -0
  15. cuwave-0.1.0/cuwave/kernels/elastic_sensitivity.cu +217 -0
  16. cuwave-0.1.0/cuwave/kernels/maxwell.cu +154 -0
  17. cuwave-0.1.0/cuwave/kernels/maxwell_sensitivity.cu +139 -0
  18. cuwave-0.1.0/cuwave/kernels/scalar.cu +164 -0
  19. cuwave-0.1.0/cuwave/kernels/scalar_sensitivity.cu +140 -0
  20. cuwave-0.1.0/cuwave/maxwell.py +416 -0
  21. cuwave-0.1.0/cuwave/nn.py +99 -0
  22. cuwave-0.1.0/cuwave/optimization.py +123 -0
  23. cuwave-0.1.0/cuwave/postprocessing.py +181 -0
  24. cuwave-0.1.0/cuwave/regularization.py +243 -0
  25. cuwave-0.1.0/cuwave/scalar.py +224 -0
  26. cuwave-0.1.0/cuwave/sensitivity.py +535 -0
  27. cuwave-0.1.0/cuwave/signals.py +71 -0
  28. cuwave-0.1.0/cuwave/stencils.py +48 -0
  29. cuwave-0.1.0/cuwave/utils.py +472 -0
  30. cuwave-0.1.0/cuwave/wave.py +518 -0
  31. cuwave-0.1.0/cuwave.egg-info/PKG-INFO +134 -0
  32. cuwave-0.1.0/cuwave.egg-info/SOURCES.txt +35 -0
  33. cuwave-0.1.0/cuwave.egg-info/dependency_links.txt +1 -0
  34. cuwave-0.1.0/cuwave.egg-info/requires.txt +8 -0
  35. cuwave-0.1.0/cuwave.egg-info/top_level.txt +1 -0
  36. cuwave-0.1.0/pyproject.toml +42 -0
  37. cuwave-0.1.0/setup.cfg +4 -0
cuwave-0.1.0/LICENSE ADDED
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Leon Herrmann
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
cuwave-0.1.0/PKG-INFO ADDED
@@ -0,0 +1,134 @@
1
+ Metadata-Version: 2.4
2
+ Name: cuwave
3
+ Version: 0.1.0
4
+ Summary: GPU finite-difference wave solver with differentiable adjoints
5
+ Author-email: Leon Herrmann <herrmann.leon@pm.me>
6
+ License-Expression: MIT
7
+ Project-URL: Homepage, https://github.com/cmpmech/cuwave
8
+ Project-URL: Documentation, https://github.com/cmpmech/cuwave/blob/main/docs/Home.md
9
+ Project-URL: Changelog, https://github.com/cmpmech/cuwave/releases
10
+ Keywords: wave-equation,finite-difference,adjoint,full-waveform-inversion,topology-optimization,cuda,cupy
11
+ Classifier: Development Status :: 4 - Beta
12
+ Classifier: Environment :: GPU :: NVIDIA CUDA
13
+ Classifier: Intended Audience :: Science/Research
14
+ Classifier: Programming Language :: Python :: 3
15
+ Classifier: Topic :: Scientific/Engineering :: Physics
16
+ Requires-Python: >=3.10
17
+ Description-Content-Type: text/markdown
18
+ License-File: LICENSE
19
+ Requires-Dist: numpy>=1.24
20
+ Provides-Extra: examples
21
+ Requires-Dist: matplotlib>=3.7; extra == "examples"
22
+ Requires-Dist: pillow>=9.1; extra == "examples"
23
+ Provides-Extra: nn
24
+ Requires-Dist: torch>=2.0; extra == "nn"
25
+ Dynamic: license-file
26
+
27
+ # CuWave
28
+
29
+ <picture>
30
+ <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/cmpmech/cuwave/main/.assets/readme-dark.webp">
31
+ <img src="https://raw.githubusercontent.com/cmpmech/cuwave/main/.assets/readme-light.webp" width="100%">
32
+ </picture>
33
+
34
+ **CuWave** is a single-GPU, differentiable finite difference wave propagation code.
35
+ Possible applications include
36
+
37
+ <table width="100%">
38
+ <tr>
39
+ <td valign="middle"><a href="https://www.sciencedirect.com/science/article/pii/S0045782523000166"><strong>nondestructive testing via full waveform inversion</strong></a></td>
40
+ <td width="60%" align="right" valign="middle"><picture>
41
+ <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/cmpmech/cuwave/main/.assets/fwi-dark.png">
42
+ <img width="100%" src="https://raw.githubusercontent.com/cmpmech/cuwave/main/.assets/fwi-light.png" alt="nondestructive testing via full waveform inversion">
43
+ </picture></td>
44
+ </tr>
45
+ <tr>
46
+ <td valign="middle"><a href="https://doi.org/10.1007/s00158-025-04237-y"><strong>transient acoustic topology optimization</strong></a></td>
47
+ <td width="60%" align="center" valign="middle"><picture>
48
+ <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/cmpmech/cuwave/main/.assets/tato-dark.webp">
49
+ <img width="46.512%" src="https://raw.githubusercontent.com/cmpmech/cuwave/main/.assets/tato-light.webp" alt="transient acoustic topology optimization">
50
+ </picture></td>
51
+ </tr>
52
+ <tr>
53
+ <td valign="middle"><a href="https://www.science.org/doi/10.1126/sciadv.aay6946"><strong>analog neural networks</strong></a></td>
54
+ <td width="60%" align="right" valign="middle"></td>
55
+ </tr>
56
+ <tr>
57
+ <td valign="middle"><a href="https://opg.optica.org/josab/fulltext.cfm?uri=josab-38-2-496"><strong>transient photonic topology optimization</strong></a></td>
58
+ <td width="60%" align="right" valign="middle"><picture>
59
+ <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/cmpmech/cuwave/main/.assets/tpto-dark.webp">
60
+ <img width="100%" src="https://raw.githubusercontent.com/cmpmech/cuwave/main/.assets/tpto-light.webp" alt="transient photonic topology optimization">
61
+ </picture></td>
62
+ </tr>
63
+ </table>
64
+
65
+ ## Documentation
66
+
67
+ - 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)
68
+ - see [examples](https://github.com/cmpmech/cuwave/tree/main/examples) for how to apply the code
69
+
70
+ ## Development status
71
+
72
+ **Scalar** & **acoustic wave** equations have been developed over the last 2 years and are thoroughly validated.
73
+
74
+ > [!IMPORTANT]
75
+ > **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.
76
+
77
+ ## Performance
78
+
79
+ 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:
80
+ - **scalar wave equation**
81
+ - [Deepwave](https://github.com/ar4/deepwave) (forward: ~1.1x, sensitivity: ~1x)
82
+ - [NVIDIA Warp](https://github.com/NVIDIA/warp) (forward: ~2.1x, sensitivity: ~1.8x)
83
+ - [SeismicWaves.jl](https://github.com/GinvLab/SeismicWaves.jl) (forward: ~4.4x, sensitivity: ~4x)
84
+ - [Devito](https://github.com/devitocodes/devito) on CPU (forward: ~3.6x, sensitivity: ~3.8x)
85
+ - **elastic wave equation**
86
+ - Deepwave (forward: ~1.1x, sensitivity: ~1.3x)
87
+ - SeismicWaves.jl (forward: ~4x, sensitivity: ~5x)
88
+ Tested on one NVIDIA RTX PRO 500 Blackwell laptop GPU (6 GB) on the largest possible grids with CuWave's `superposition_sensitivity` as reference for the sensitivities. The specific numbers need to be taken with a grain of salt, as they are subject to specific hardware and simulation setup. All implementations operate on the same order of magnitude.
89
+
90
+ Additional benefits of **CuWave** are
91
+ - the built-in **higher order finite difference** schemes, allowing for fewer grid points
92
+ - a sensitivity analysis whose **memory is independent of the number of timesteps**, allowing for orders of magnitude larger grids
93
+
94
+ ## Install
95
+
96
+ Dependencies are kept **lightweight**. Only CuPy is required beyond standard Python library.
97
+
98
+ ```bash
99
+ pip install cupy-cuda12x # or cupy-cuda11x, to match your CUDA
100
+ pip install cuwave # or `pip install -e .` from a checkout
101
+ ```
102
+
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
+ ```
121
+
122
+ ## References
123
+
124
+ If you use our code for your scientific research, please acknowledge this by referring to the following publication:
125
+
126
+ _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)_
127
+
128
+ ## Contact
129
+
130
+ 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).
131
+
132
+ ## License
133
+
134
+ MIT; see [LICENSE](https://github.com/cmpmech/cuwave/blob/main/LICENSE).
cuwave-0.1.0/README.md ADDED
@@ -0,0 +1,108 @@
1
+ # CuWave
2
+
3
+ <picture>
4
+ <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/cmpmech/cuwave/main/.assets/readme-dark.webp">
5
+ <img src="https://raw.githubusercontent.com/cmpmech/cuwave/main/.assets/readme-light.webp" width="100%">
6
+ </picture>
7
+
8
+ **CuWave** is a single-GPU, differentiable finite difference wave propagation code.
9
+ Possible applications include
10
+
11
+ <table width="100%">
12
+ <tr>
13
+ <td valign="middle"><a href="https://www.sciencedirect.com/science/article/pii/S0045782523000166"><strong>nondestructive testing via full waveform inversion</strong></a></td>
14
+ <td width="60%" align="right" valign="middle"><picture>
15
+ <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/cmpmech/cuwave/main/.assets/fwi-dark.png">
16
+ <img width="100%" src="https://raw.githubusercontent.com/cmpmech/cuwave/main/.assets/fwi-light.png" alt="nondestructive testing via full waveform inversion">
17
+ </picture></td>
18
+ </tr>
19
+ <tr>
20
+ <td valign="middle"><a href="https://doi.org/10.1007/s00158-025-04237-y"><strong>transient acoustic topology optimization</strong></a></td>
21
+ <td width="60%" align="center" valign="middle"><picture>
22
+ <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/cmpmech/cuwave/main/.assets/tato-dark.webp">
23
+ <img width="46.512%" src="https://raw.githubusercontent.com/cmpmech/cuwave/main/.assets/tato-light.webp" alt="transient acoustic topology optimization">
24
+ </picture></td>
25
+ </tr>
26
+ <tr>
27
+ <td valign="middle"><a href="https://www.science.org/doi/10.1126/sciadv.aay6946"><strong>analog neural networks</strong></a></td>
28
+ <td width="60%" align="right" valign="middle"></td>
29
+ </tr>
30
+ <tr>
31
+ <td valign="middle"><a href="https://opg.optica.org/josab/fulltext.cfm?uri=josab-38-2-496"><strong>transient photonic topology optimization</strong></a></td>
32
+ <td width="60%" align="right" valign="middle"><picture>
33
+ <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/cmpmech/cuwave/main/.assets/tpto-dark.webp">
34
+ <img width="100%" src="https://raw.githubusercontent.com/cmpmech/cuwave/main/.assets/tpto-light.webp" alt="transient photonic topology optimization">
35
+ </picture></td>
36
+ </tr>
37
+ </table>
38
+
39
+ ## Documentation
40
+
41
+ - 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)
42
+ - see [examples](https://github.com/cmpmech/cuwave/tree/main/examples) for how to apply the code
43
+
44
+ ## Development status
45
+
46
+ **Scalar** & **acoustic wave** equations have been developed over the last 2 years and are thoroughly validated.
47
+
48
+ > [!IMPORTANT]
49
+ > **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.
50
+
51
+ ## Performance
52
+
53
+ 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:
54
+ - **scalar wave equation**
55
+ - [Deepwave](https://github.com/ar4/deepwave) (forward: ~1.1x, sensitivity: ~1x)
56
+ - [NVIDIA Warp](https://github.com/NVIDIA/warp) (forward: ~2.1x, sensitivity: ~1.8x)
57
+ - [SeismicWaves.jl](https://github.com/GinvLab/SeismicWaves.jl) (forward: ~4.4x, sensitivity: ~4x)
58
+ - [Devito](https://github.com/devitocodes/devito) on CPU (forward: ~3.6x, sensitivity: ~3.8x)
59
+ - **elastic wave equation**
60
+ - Deepwave (forward: ~1.1x, sensitivity: ~1.3x)
61
+ - SeismicWaves.jl (forward: ~4x, sensitivity: ~5x)
62
+ Tested on one NVIDIA RTX PRO 500 Blackwell laptop GPU (6 GB) on the largest possible grids with CuWave's `superposition_sensitivity` as reference for the sensitivities. The specific numbers need to be taken with a grain of salt, as they are subject to specific hardware and simulation setup. All implementations operate on the same order of magnitude.
63
+
64
+ Additional benefits of **CuWave** are
65
+ - the built-in **higher order finite difference** schemes, allowing for fewer grid points
66
+ - a sensitivity analysis whose **memory is independent of the number of timesteps**, allowing for orders of magnitude larger grids
67
+
68
+ ## Install
69
+
70
+ Dependencies are kept **lightweight**. Only CuPy is required beyond standard Python library.
71
+
72
+ ```bash
73
+ pip install cupy-cuda12x # or cupy-cuda11x, to match your CUDA
74
+ pip install cuwave # or `pip install -e .` from a checkout
75
+ ```
76
+
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
+ ```
95
+
96
+ ## References
97
+
98
+ If you use our code for your scientific research, please acknowledge this by referring to the following publication:
99
+
100
+ _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)_
101
+
102
+ ## Contact
103
+
104
+ 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).
105
+
106
+ ## License
107
+
108
+ MIT; see [LICENSE](https://github.com/cmpmech/cuwave/blob/main/LICENSE).
@@ -0,0 +1,8 @@
1
+ """CuWave: a single-GPU, differentiable finite-difference wave solver"""
2
+
3
+ from importlib.metadata import PackageNotFoundError, version
4
+
5
+ try:
6
+ __version__ = version("cuwave")
7
+ except PackageNotFoundError: # a checkout on sys.path without an install
8
+ __version__ = "0.0.0"
@@ -0,0 +1,337 @@
1
+ """Cell-assembled elasticity, the collocated sibling of the staggered `ElasticWave`.
2
+
3
+ The stencil is a 9-point one in 2D and 27-point in 3D, gathered cell by cell rather
4
+ than axis by axis, since a variable-coefficient elastic operator couples the components
5
+ through mixed derivatives that a per-axis flux cannot carry. Its coefficients are
6
+ derived as `-B^T C B` with `C` on the cell, which is what makes the operator exactly
7
+ symmetric for a varying material, never differentiates the material, leaves the
8
+ leapfrog reversible, and makes a traction-free surface the natural condition of the
9
+ interior-cell sum. `C` is any symmetric Voigt matrix, which is what earns this scheme
10
+ its keep next to the staggered one: order 2 only, but collocated components and a
11
+ general anisotropy.
12
+ """
13
+
14
+ from __future__ import annotations
15
+
16
+ import itertools
17
+ from collections.abc import Callable
18
+ from dataclasses import dataclass
19
+ from pathlib import Path
20
+
21
+ import cupy as cp
22
+ import cupy.typing as cpt
23
+ import numpy as np
24
+ import numpy.typing as npt
25
+
26
+ from .boundary import Clamped, Traction, faces_with
27
+ from .elastic import voigt
28
+ from .wave import PAIRS, Simulation, apply_cell_weights, grid_block
29
+
30
+ KERNEL_PATH = Path(__file__).parent / "kernels" / "anisotropic.cu"
31
+ SENSITIVITY_PATH = Path(__file__).parent / "kernels" / "anisotropic_sensitivity.cu"
32
+
33
+
34
+ # -------------------------------------- helpers --------------------------------------
35
+ def corner_bits(ndim: int) -> list[tuple[int, ...]]:
36
+ """The `2**ndim` nodes of a cell, node `c` having bit `d` set on the plus side of axis `d`."""
37
+ return [tuple((c >> d) & 1 for d in range(ndim)) for c in range(2**ndim)]
38
+
39
+
40
+ def cell_stencil(
41
+ ndim: int, dx: tuple[float, ...], C: npt.NDArray
42
+ ) -> npt.NDArray[np.float64]:
43
+ """One cell's contribution to the nodal stencil, ordered (node, component).
44
+
45
+ Built as `sum_q w_q B_q^T C B_q` over the cell, `B` the discrete symmetric gradient
46
+ on the `2**ndim` corners, so the assembled operator is symmetric by construction
47
+ rather than by inspection. The entries collapse to the 9-point (27-point in 3D)
48
+ finite difference stencil
49
+
50
+ f_x = (lame + 2 mu) [M_y * D2_x] u_x + mu [M_x * D2_y] u_x
51
+ + (lame + mu) [D_x * D_y] u_y
52
+
53
+ with the second difference `D2` = (1, -2, 1), the centred first difference `D` and
54
+ the transverse average `M` = (1, 4, 1) / 6. That average is the only departure from
55
+ the textbook elastic stencil, and it is what removes the checkerboard from the null
56
+ space: one quadrature point per axis would give (1, 2, 1) / 4 and leave it there as
57
+ an undamped hourglass mode.
58
+
59
+ Args:
60
+ ndim: dimensionality of the cell.
61
+ dx: cell side per axis.
62
+ C: (voigt, voigt) symmetric stiffness matrix, from `elastic.voigt` or the
63
+ caller.
64
+
65
+ Returns:
66
+ the (n * ndim, n * ndim) matrix for `n = 2**ndim`, node `c` and component `i`
67
+ occupying row `c * ndim + i`, `c` running over the corners in `corner_bits`
68
+ order.
69
+ """
70
+ corners = corner_bits(ndim)
71
+ nloc = len(corners) * ndim
72
+ nodes, weights = np.polynomial.legendre.leggauss(2)
73
+ nodes = 0.5 * (nodes + 1.0)
74
+ weights = 0.5 * weights * float(np.prod(dx)) ** (1.0 / ndim)
75
+ K = np.zeros((nloc, nloc))
76
+ for point in itertools.product(range(2), repeat=ndim):
77
+ weight = float(np.prod([weights[q] for q in point]))
78
+ value = [(1.0 - nodes[q], nodes[q]) for q in point]
79
+ dN = np.zeros((len(corners), ndim))
80
+ for m, entry in enumerate(corners):
81
+ for k in range(ndim):
82
+ term = (-1.0, 1.0)[entry[k]] / dx[k]
83
+ for axis in range(ndim):
84
+ if axis != k:
85
+ term *= value[axis][entry[axis]]
86
+ dN[m, k] = term
87
+ B = np.zeros((len(PAIRS[ndim]), nloc))
88
+ for r, (axis_k, axis_l) in enumerate(PAIRS[ndim]):
89
+ for m in range(len(corners)):
90
+ if axis_k == axis_l:
91
+ B[r, m * ndim + axis_k] += dN[m, axis_k]
92
+ else:
93
+ B[r, m * ndim + axis_l] += dN[m, axis_k] # engineering shear
94
+ B[r, m * ndim + axis_k] += dN[m, axis_l]
95
+ K += weight * (B.T @ C @ B)
96
+ return K
97
+
98
+
99
+ def cell_average(sim: Simulation, field: cpt.NDArray) -> cpt.NDArray:
100
+ """Harmonic mean of `field` over the `2**ndim` corners of a cell, at its low corner.
101
+
102
+ Harmonic and not arithmetic for the reason the scalar flux uses it: it keeps the
103
+ cell stiffness single valued across a material jump, and it is what holds the
104
+ stable timestep together at a high contrast, where a light node borders a stiff
105
+ cell.
106
+ """
107
+ out = cp.zeros(sim.Nx_padded, dtype=sim.dtype)
108
+ safe = cp.maximum(field, cp.finfo(sim.dtype).tiny)
109
+ inner = tuple(slice(0, n - 1) for n in sim.Nx)
110
+ for corner in corner_bits(sim.ndim):
111
+ shifted = tuple(slice(b, n - 1 + b) for b, n in zip(corner, sim.Nx))
112
+ out[inner] += 1.0 / safe[shifted]
113
+ out[inner] = 2.0**sim.ndim / out[inner]
114
+ return out
115
+
116
+
117
+ # ------------------------------- discretization setup --------------------------------
118
+ @dataclass
119
+ class AnisotropicElasticWave(Simulation):
120
+ """Cell-assembled elasticity, parametrized by a density-scaling indicator gamma.
121
+
122
+ Both wave speeds are held fixed and gamma scales the density, so
123
+ `C = gamma * rho0 * C` scales inertia and stiffness alike and the stable timestep
124
+ does not move with the design. `C` overrides the isotropic Voigt matrix with an
125
+ anisotropic one at gamma = 1, which the cell assembly carries where the staggered
126
+ scheme cannot. Order 2 only: the cell gather costs `(2r)**(2 ndim)` per node, so a
127
+ wide stencil belongs to `ElasticWave`.
128
+ """
129
+
130
+ density: float = None # background density rho0
131
+ wavespeed_p: float = None # pressure wave speed
132
+ wavespeed_s: float = None # shear wave speed
133
+ plane: str = "strain" # "strain" or "stress", 2D only
134
+ C: npt.NDArray | None = None # (voigt, voigt) stiffness at gamma = 1, or isotropic
135
+
136
+ kernel_path = KERNEL_PATH
137
+ sensitivity_path = SENSITIVITY_PATH
138
+ default_boundary = Traction
139
+ gradient_names = ("mass", "stiff", "cell")
140
+
141
+ @property
142
+ def ncomp(self) -> int:
143
+ """One displacement component per axis."""
144
+ return self.ndim
145
+
146
+ def __post_init__(self) -> None:
147
+ """Validate the material, the 2D plane assumption and the order, then derive the grid."""
148
+ super().__post_init__()
149
+ if None in (self.density, self.wavespeed_p, self.wavespeed_s):
150
+ raise ValueError(
151
+ "AnisotropicElasticWave requires density, wavespeed_p, wavespeed_s"
152
+ )
153
+ if self.plane not in ("strain", "stress"):
154
+ raise ValueError(f"plane must be strain or stress: {self.plane}")
155
+ if self.ndim == 3 and self.plane != "strain":
156
+ raise ValueError("plane stress is a 2D reduction, not a 3D one")
157
+ if self.space_order != 2:
158
+ raise ValueError(
159
+ f"the cell gather costs (2r)**(2 ndim) per node, so order "
160
+ f"{self.space_order} belongs to the staggered ElasticWave"
161
+ )
162
+ self._stencil = None # built once and reused across material rebuilds
163
+ if self.wavespeed_s >= self.wavespeed_p:
164
+ raise ValueError(
165
+ f"wavespeed_s must be below wavespeed_p: "
166
+ f"{self.wavespeed_s} >= {self.wavespeed_p}"
167
+ )
168
+ nvoigt = len(PAIRS[self.ndim])
169
+ if self.C is not None:
170
+ self.C = np.asarray(self.C, dtype=float)
171
+ if self.C.shape != (nvoigt, nvoigt):
172
+ raise ValueError(f"C must be ({nvoigt}, {nvoigt}): {self.C.shape}")
173
+ if not np.allclose(self.C, self.C.T):
174
+ raise ValueError("C must be symmetric")
175
+
176
+ @property
177
+ def lame(self) -> float:
178
+ """First Lame parameter `rho0 * (c_p**2 - 2 c_s**2)`."""
179
+ return self.density * (self.wavespeed_p**2 - 2.0 * self.wavespeed_s**2)
180
+
181
+ @property
182
+ def shear(self) -> float:
183
+ """Second Lame parameter `rho0 * c_s**2`."""
184
+ return self.density * self.wavespeed_s**2
185
+
186
+ def stencil(self) -> cpt.NDArray:
187
+ """Stencil table at `gamma = 1`: one cell's `-B^T C B`, flattened for the kernel."""
188
+ if self._stencil is not None:
189
+ return self._stencil
190
+ C = self.C
191
+ if C is None:
192
+ C = voigt(self.ndim, self.lame, self.shear, self.plane)
193
+ K = cell_stencil(self.ndim, self.dx, C)
194
+ self._stencil = cp.asarray(K.ravel(), dtype=self.dtype)
195
+ return self._stencil
196
+
197
+ def cell_weights(self) -> cpt.NDArray:
198
+ """Nodal cell weights W, halved once per wall the node sits on."""
199
+ return apply_cell_weights(self, cp.ones(self.Nx_padded, dtype=self.dtype))
200
+
201
+ def inverse_inertia(self, indicator: cpt.NDArray) -> cpt.NDArray:
202
+ """Lumped `1 / (gamma rho0 V W)`, the mass the interior-cell assembly implies."""
203
+ volume = float(np.prod(self.dx))
204
+ mass = indicator * (self.density * volume) * self.cell_weights()
205
+ return 1.0 / cp.maximum(mass, cp.finfo(self.dtype).tiny)
206
+
207
+ def build_materials(self, indicator: cpt.NDArray) -> dict:
208
+ """Lumped inverse inertia, the cell design field, and the stencil table."""
209
+ minv = self.inverse_inertia(indicator)
210
+ for face in faces_with(self, Clamped):
211
+ wall = [slice(None)] * self.ndim
212
+ wall[face // 2] = 1 if face % 2 == 0 else self.Nx[face // 2] - 2
213
+ minv[tuple(wall)] = 0.0
214
+ mat = {
215
+ "minv": cp.ascontiguousarray(minv, dtype=self.dtype),
216
+ "gamma": cp.ascontiguousarray(indicator, dtype=self.dtype),
217
+ "cell": cell_average(self, indicator),
218
+ "stencil": self.stencil(),
219
+ }
220
+ if self.damping is not None:
221
+ mat["damping"] = self.damping
222
+ return mat
223
+
224
+ def step_kernel_args(self, mat: dict) -> tuple:
225
+ """Material, stencil table and component stride for the step kernel."""
226
+ args = (mat["minv"], mat["cell"], mat["stencil"])
227
+ if self.damping is not None:
228
+ args += (mat["damping"], self.dtype(self.dt))
229
+ return args + (np.int32(self.comp_stride),)
230
+
231
+ def excitation_weights(
232
+ self, mat: dict, lin_index: cpt.NDArray[cp.int32]
233
+ ) -> cpt.NDArray:
234
+ """Source weights `dt**2 / inertia`, the spatial node read off the folded index."""
235
+ node = lin_index % np.int32(self.comp_stride)
236
+ weight = mat["minv"].ravel()[node]
237
+ if self.damping is not None:
238
+ beta = 0.5 * weight * mat["damping"].ravel()[node] * self.dt
239
+ weight = weight / (1.0 + beta)
240
+ return (self.dtype(self.dt**2 * self.source_factor()) * weight).astype(
241
+ self.dtype
242
+ )
243
+
244
+ def parametrization_jacobian(self, indicator: cpt.NDArray) -> tuple:
245
+ """Gamma scales inertia and stiffness alike, so both derivatives are 1."""
246
+ return 1.0, 1.0
247
+
248
+ def step_factors(self) -> list:
249
+ """`dt**2`; the grid spacing already sits in the stencil table."""
250
+ return [self.dtype(self.dt**2)] * self.ndim
251
+
252
+ def source_factor(self) -> float:
253
+ """Source scaling, unscaled since rho0 is already folded into the lumped inertia."""
254
+ return 1.0
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
+ def gradient_fields(self, mat: dict) -> dict[str, cpt.NDArray]:
264
+ """Nodal accumulators plus the cell one the stiffness density lands in first."""
265
+ grads = {
266
+ name: cp.zeros(self.Nx_padded, dtype=self.dtype)
267
+ for name in self.gradient_names
268
+ }
269
+ # the harmonic cell mean has a design dependent chain rule, so keep both fields
270
+ grads["material"] = mat["cell"]
271
+ grads["design"] = mat["gamma"]
272
+ return grads
273
+
274
+ def finalize_gradients(self, grads: dict, kernels: cp.RawModule) -> dict:
275
+ """Weight the mass density by W and spread the cell density onto its corners."""
276
+ apply_cell_weights(self, grads["mass"])
277
+ cell_to_node = kernels.get_function("cell_to_node_kernel")
278
+ grid, block = grid_block(self)
279
+ cell_to_node(
280
+ grid,
281
+ block,
282
+ [
283
+ grads["stiff"],
284
+ grads["cell"],
285
+ grads["material"],
286
+ grads["design"],
287
+ self.dtype(1.0 / 2.0**self.ndim),
288
+ *self.axis_geometry(),
289
+ ],
290
+ )
291
+ return {"mass": grads["mass"], "stiff": grads["stiff"]}
292
+
293
+ def define_gradient(
294
+ self, kernels: cp.RawModule, mat: dict, grads: dict
295
+ ) -> Callable:
296
+ """Closure accumulating the nodal mass and cell stiffness densities."""
297
+ gradient_kernel = kernels.get_function("gradient_kernel")
298
+ grid, block = grid_block(self)
299
+ # d(mass)/d(gamma) is rho0 V, the W of the lumping supplied by the epilogue
300
+ mass_factor = self.dtype(self.density * float(np.prod(self.dx)) / self.dt**2)
301
+ args = [grads["mass"], grads["cell"], None, None, None, None] + [
302
+ mat["stencil"],
303
+ mass_factor,
304
+ np.int32(self.comp_stride),
305
+ *self.axis_geometry(),
306
+ ]
307
+
308
+ def gradient_step(u0, u1, u2, l1):
309
+ args[2], args[3], args[4], args[5] = u0, u1, u2, l1
310
+ gradient_kernel(grid, block, args)
311
+
312
+ return gradient_step
313
+
314
+ def define_frechet(
315
+ self, kernels: cp.RawModule, accs: dict, sign: float
316
+ ) -> Callable:
317
+ """Closure accumulating both quadratic densities of one field triplet, times `sign`."""
318
+ frechet_kernel = kernels.get_function("frechet_kernel")
319
+ grid, block = grid_block(self)
320
+ volume = float(np.prod(self.dx))
321
+ args = [accs["mass"], accs["cell"], None, None, None] + [
322
+ self.stencil(),
323
+ self.dtype(sign * self.density * volume / (2.0 * self.dt) ** 2),
324
+ self.dtype(-sign),
325
+ np.int32(self.comp_stride),
326
+ *self.axis_geometry(),
327
+ ]
328
+
329
+ def frechet_step(u0, u1, u2):
330
+ args[2], args[3], args[4] = u0, u1, u2
331
+ frechet_kernel(grid, block, args)
332
+
333
+ return frechet_step
334
+
335
+ def adjoint_weights(self, sensors: cpt.NDArray[cp.int32]) -> cpt.NDArray:
336
+ """Ones: the lumped inertia already carries W, and dJ/du is nodal, not a density."""
337
+ return cp.ones(sensors.shape[1], dtype=self.dtype)