diff --git a/src/sweep/csrc/cuda/equations/visco_acoustic2d/backward.cu b/src/sweep/csrc/cuda/equations/visco_acoustic2d/backward.cu index eff1dc0e..e90788dd 100644 --- a/src/sweep/csrc/cuda/equations/visco_acoustic2d/backward.cu +++ b/src/sweep/csrc/cuda/equations/visco_acoustic2d/backward.cu @@ -224,10 +224,10 @@ void adjoint_damping_extra(AcousticWavefieldTensor& adj, const ViscoSpectral& d, if (d.disp) { const Buf X = visco_acoustic2d_real_alias(ws.C1); visco_ops::binary(X.data_ptr(), d.Gd1.data_ptr(), 1, u_now, 1, n); // Gd1 * u_now - visco_acoustic2d_lop_into(X.data_ptr(), d.Dk2.data_ptr(), ws); // C1.re = L_Dk2 + visco_acoustic2d_lop_into(X.data_ptr(), d.Dk2, ws); // C1.re = L_Dk2 visco_ops::real_copy(ws.R1.data_ptr(), visco_acoustic2d_complex_ptr(ws.C1), n); // R1 = L_Dk2 visco_ops::binary(X.data_ptr(), d.Gd2.data_ptr(), 1, u_now, 1, n); // Gd2 * u_now (L_Dk2 is in R1) - visco_acoustic2d_lop_into(X.data_ptr(), d.Dfrac.data_ptr(), ws); // C1.re = L_Dfrac + visco_acoustic2d_lop_into(X.data_ptr(), d.Dfrac, ws); // C1.re = L_Dfrac visco_ops::inplace(ws.R1.data_ptr(), visco_acoustic2d_real_ptr(ws.C1), VISCO_REAL_STRIDE, n); // R1 = L_Dk2 - L_Dfrac = e m = ws.R1.data_ptr(); m_stride = 1; } @@ -235,7 +235,7 @@ void adjoint_damping_extra(AcousticWavefieldTensor& adj, const ViscoSpectral& d, const Buf X = visco_acoustic2d_real_alias(ws.C1); visco_ops::binary(X.data_ptr(), u_prev, 1, u_now, 1, n); // u_prev - u_now visco_ops::inplace(X.data_ptr(), d.Gp.data_ptr(), 1, n); // == Gp * (u_prev - u_now) - visco_acoustic2d_lop_into(X.data_ptr(), d.kmul.data_ptr(), ws); // C1.re = L + visco_acoustic2d_lop_into(X.data_ptr(), d.kmul, ws); // C1.re = L if (d.disp) visco_ops::inplace(ws.R1.data_ptr(), visco_acoustic2d_real_ptr(ws.C1), VISCO_REAL_STRIDE, n); // e + m == m + e else { @@ -272,7 +272,7 @@ void accumulate_grad_A(float* grad_A, { if (!d.active) return; const int64_t n = ws.C1.numel() / 2; - visco_acoustic2d_lop_into(du, d.kmul.data_ptr(), ws); // C1.re = L(du) + visco_acoustic2d_lop_into(du, d.kmul, ws); // C1.re = L(du) const Buf P = visco_acoustic2d_real_alias(ws.C0); visco_ops::binary(P.data_ptr(), lam, 1, visco_acoustic2d_real_ptr(ws.C1), VISCO_REAL_STRIDE, n); @@ -294,10 +294,10 @@ void accumulate_grad_disp(float* grad_B1, float* grad_B2, // add_(P, Scalar(dt2)): the double dt*dt, loaded as the float alpha const float dt2 = static_cast(static_cast(dt) * static_cast(dt)); const Buf P = visco_acoustic2d_real_alias(ws.C0); - visco_acoustic2d_lop_into(u_it, d.Dk2.data_ptr(), ws); + visco_acoustic2d_lop_into(u_it, d.Dk2, ws); visco_ops::binary(P.data_ptr(), lam, 1, visco_acoustic2d_real_ptr(ws.C1), VISCO_REAL_STRIDE, n); visco_ops::axpy(grad_B1, P.data_ptr(), 1, dt2, n); - visco_acoustic2d_lop_into(u_it, d.Dfrac.data_ptr(), ws); + visco_acoustic2d_lop_into(u_it, d.Dfrac, ws); visco_ops::binary(P.data_ptr(), lam, 1, visco_acoustic2d_real_ptr(ws.C1), VISCO_REAL_STRIDE, n); visco_ops::axpy(grad_B2, P.data_ptr(), 1, -dt2, n); } diff --git a/src/sweep/csrc/cuda/equations/visco_acoustic2d/kernels.cuh b/src/sweep/csrc/cuda/equations/visco_acoustic2d/kernels.cuh index 3e1b348c..0c2f0e63 100644 --- a/src/sweep/csrc/cuda/equations/visco_acoustic2d/kernels.cuh +++ b/src/sweep/csrc/cuda/equations/visco_acoustic2d/kernels.cuh @@ -112,20 +112,25 @@ struct ViscoSpectral { Buf Gd1, Gd2; // dt^2 * B1, dt^2 * B2 }; +// A grid is shared by the batch, (nz, nx), or one per batch element, +// (B, nz, nx): per-shot models whose frozen exponents differ (the filter +// strides it per batch index, see visco_acoustic2d_lop_into). inline void visco_acoustic2d_check_grid( - const Buf& g, int nz, int nx, const char* name) + const Buf& g, int64_t batch, int nz, int nx, const char* name) { - SWEEP_CHECK(g.dim() == 2 && g.size(0) == nz && g.size(1) == nx, + const bool shared = g.dim() == 2 && g.size(0) == nz && g.size(1) == nx; + const bool per_batch = g.dim() == 3 && g.size(0) == batch && g.size(1) == nz && g.size(2) == nx; + SWEEP_CHECK(shared || per_batch, "visco eq_aux grid ", name, " must be (nz_runtime, nx_runtime) = (", nz, ", ", nx, - "), got ", g.sizes()); + ") or (B, nz_runtime, nx_runtime) with B = ", batch, ", got ", g.sizes()); SWEEP_CHECK(g.is_cuda() && g.scalar_type() == BoundaryDtype::FP32, "visco eq_aux grid ", name, " must be a float32 CUDA tensor"); } // The term switches and filter grids from the eq_aux composition (no tables). inline ViscoSpectral visco_acoustic2d_spectral_grids( - const BufList& eq_aux, int nz, int nx) + const BufList& eq_aux, int64_t batch, int nz, int nx) { ViscoSpectral s; const size_t n = eq_aux.size(); @@ -134,13 +139,13 @@ inline ViscoSpectral visco_acoustic2d_spectral_grids( s.disp = (n >= 2); if (s.active) { s.kmul = eq_aux[0]; - visco_acoustic2d_check_grid(s.kmul, nz, nx, "D_loss"); + visco_acoustic2d_check_grid(s.kmul, batch, nz, nx, "D_loss"); } if (s.disp) { s.Dk2 = eq_aux[n - 2]; s.Dfrac = eq_aux[n - 1]; - visco_acoustic2d_check_grid(s.Dk2, nz, nx, "D_k2"); - visco_acoustic2d_check_grid(s.Dfrac, nz, nx, "D_frac"); + visco_acoustic2d_check_grid(s.Dk2, batch, nz, nx, "D_k2"); + visco_acoustic2d_check_grid(s.Dfrac, batch, nz, nx, "D_frac"); } return s; } @@ -161,7 +166,7 @@ inline ViscoSpectral visco_acoustic2d_make_spectral_from( { SWEEP_CHECK(models.size() == 4, what, ": visco_acoustic2d expects the prepared models " "(vp_step, B1, B2, A); got ", models.size()); - ViscoSpectral s = visco_acoustic2d_spectral_grids(eq_aux, nz, nx); + ViscoSpectral s = visco_acoustic2d_spectral_grids(eq_aux, models[0].size(0), nz, nx); const float dt2 = dt * dt; // float on the host: the scalar of the replaced ``model * (dt * dt)`` products const auto c = derived::visco_coefficients( derived_models, models[1], models[2], models[3], s.active, s.disp, mode, dt, dt2, what); @@ -195,7 +200,7 @@ inline ViscoSpectral visco_acoustic2d_make_spectral_from( // builds the same CuFFTParams and issues the same three calls, with the // work area taken from the pool. The same plan on the same data gives the // same bits. -// * kmul * F, a float (nz, nx) grid times the complex (B, 1, nz, nx) +// * kmul * F, a float (nz, nx) [or (B, nz, nx)] grid times the complex (B, 1, nz, nx) // spectrum: mul_kernel_cuda on the common dtype complex64, kmul loaded as // (k, 0) and c10::complex operator* evaluated as (a*c - b*d, a*d + b*c) // with (a, b) = (k, 0). Here: at::mul_out(C0, kmul, F) -- functional mul @@ -461,16 +466,21 @@ inline ViscoScratch visco_acoustic2d_bind_scratch( // is consumed by the promote before C1 is overwritten) but NOT in C0's. // --------------------------------------------------------------------------- // ``x``: a contiguous float grid of the slot's shape (a wavefield, a real -// alias, a history row); ``kmul``: the (nz, nx) filter. The result is the +// alias, a history row); ``kmul``: the (nz, nx) filter, or (B, nz, nx) with +// one filter per batch element (cmul_grid indexes it modulo its own size, so +// the shared grid is read exactly as before). The result is the // real part of C1 -- visco_acoustic2d_real_ptr(ws.C1) at VISCO_REAL_STRIDE -- // valid until C1 is next written. -inline void visco_acoustic2d_lop_into(const float* x, const float* kmul, ViscoScratch& ws) +inline void visco_acoustic2d_lop_into(const float* x, const Buf& kmul, ViscoScratch& ws) { const int64_t n = ws.C1.numel() / 2; - const int64_t ng = ws.C1.size(-3) * ws.C1.size(-2); // (nz, nx) of the (..., nz, nx, 2) slot + const int64_t ng = kmul.numel(); + SWEEP_CHECK(ng == ws.C1.size(-3) * ws.C1.size(-2) || ng == n, + "visco filter grid of ", ng, " cells matches neither one (nz, nx) plane nor the " + "whole ", n, "-cell batch of the spectral slot"); visco_ops::promote(visco_acoustic2d_complex_ptr(ws.C0), x, n); ws.fft->forward(ws.C0, ws.C1, ws.fft_ws); - visco_ops::cmul_grid(visco_acoustic2d_complex_ptr(ws.C0), kmul, visco_acoustic2d_complex_ptr(ws.C1), n, ng); + visco_ops::cmul_grid(visco_acoustic2d_complex_ptr(ws.C0), kmul.data_ptr(), visco_acoustic2d_complex_ptr(ws.C1), n, ng); ws.fft->inverse(ws.C0, ws.C1, ws.fft_ws); } @@ -485,7 +495,7 @@ inline void visco_acoustic2d_lop_into(const float* x, const float* kmul, ViscoSc // subtracted in place, as before. inline void visco_acoustic2d_apply_damping_into( AcousticWavefieldTensor& wf, - const Buf& kmul, // (nz, nx) D_loss grid + const Buf& kmul, // (nz, nx) or (B, nz, nx) D_loss grid const Buf& dt2A, // (B, 1, nz, nx) dt^2 * A (derived table) float dt, ViscoScratch& ws) { @@ -495,7 +505,7 @@ inline void visco_acoustic2d_apply_damping_into( wf.u_now_t.template data_ptr(), 1, wf.u_prev_t.template data_ptr(), 1, n); visco_ops::scale(dudt.data_ptr(), 1.0f / dt, n); // div_(dt): times the reciprocal - visco_acoustic2d_lop_into(dudt.data_ptr(), kmul.data_ptr(), ws); // C1.re = L; dudt is gone + visco_acoustic2d_lop_into(dudt.data_ptr(), kmul, ws); // C1.re = L; dudt is gone const Buf R = visco_acoustic2d_real_alias(ws.C0); visco_ops::binary(R.data_ptr(), dt2A.data_ptr(), 1, visco_acoustic2d_real_ptr(ws.C1), VISCO_REAL_STRIDE, n); @@ -527,17 +537,16 @@ inline void visco_acoustic2d_apply_spectral_into( SWEEP_CHECK(ws.C2.defined(), "visco_acoustic2d: the dispersion pipeline needs the C2 slot " "(bind the scratch with dispersion on, in forward / ckpt / recursive mode)"); const int64_t n = ws.C1.numel() / 2; - const int64_t ng = ws.C1.size(-3) * ws.C1.size(-2); float2* C0 = visco_acoustic2d_complex_ptr(ws.C0); visco_ops::promote(C0, wf.u_now_t.template data_ptr(), n); ws.fft->forward(ws.C0, ws.C1, ws.fft_ws); // C1 = F, alive across both products - visco_ops::cmul_grid(C0, s.Dk2.data_ptr(), visco_acoustic2d_complex_ptr(ws.C1), n, ng); // C0 = Dk2 * F + visco_ops::cmul_grid(C0, s.Dk2.data_ptr(), visco_acoustic2d_complex_ptr(ws.C1), n, s.Dk2.numel()); // C0 = Dk2 * F ws.fft->inverse(ws.C0, ws.C2, ws.fft_ws); // C2 = IFFT2(Dk2 * F) const Buf R = visco_acoustic2d_real_alias(ws.C0); // the product in C0 is dead visco_ops::binary(R.data_ptr(), s.Gd1.data_ptr(), 1, visco_acoustic2d_real_ptr(ws.C2), VISCO_REAL_STRIDE, n); visco_ops::inplace(wf.u_next_t.template data_ptr(), R.data_ptr(), 1, n); - visco_ops::cmul_grid(C0, s.Dfrac.data_ptr(), visco_acoustic2d_complex_ptr(ws.C1), n, ng); // C0 = Dfrac * F (R's bits are dead) + visco_ops::cmul_grid(C0, s.Dfrac.data_ptr(), visco_acoustic2d_complex_ptr(ws.C1), n, s.Dfrac.numel()); // C0 = Dfrac * F (R's bits are dead) ws.fft->inverse(ws.C0, ws.C2, ws.fft_ws); // C2 = IFFT2(Dfrac * F) visco_ops::binary(R.data_ptr(), s.Gd2.data_ptr(), 1, visco_acoustic2d_real_ptr(ws.C2), VISCO_REAL_STRIDE, n); diff --git a/src/sweep/equations/visco_acoustic.py b/src/sweep/equations/visco_acoustic.py index d742d16b..d318e6d2 100644 --- a/src/sweep/equations/visco_acoustic.py +++ b/src/sweep/equations/visco_acoustic.py @@ -271,7 +271,7 @@ def prepare_models(self, models): vp, Q, omega = models op = self.op gamma = op.arctan(1.0 / Q) / np.pi - self._gbar = op.mean(gamma) + self._gbar = self._average_exponent(gamma) if self.phase_shift: csq_fac = op.cos(np.pi * gamma / 2) vp_step = vp * csq_fac @@ -290,10 +290,29 @@ def prepare_models(self, models): A = op.zeros_like(vp) return (vp_step, B1, B2, A) + def _average_exponent(self, gamma): + """The frozen exponent ``gbar``: the mean of ``gamma`` over the model. + + A shared model ``(nz, nx)`` gets one scalar. Per-shot batched models + (``(B, 1, nz, nx)``: the eager layout, and every impl='c' model) get + one mean per batch element, shaped ``(B, 1, 1, 1)``, so a shot's + filter exponent depends on its own Q only -- each mean is the same + reduction a single-shot run takes over that model. When all the + per-element means agree (a shared model the C path repeated across + the batch), the historical whole-batch scalar is kept bit-for-bit.""" + op = self.op + if gamma.ndim <= 2 or gamma.shape[0] == 1: + return op.mean(gamma) + g = op.stack([op.mean(gamma[b]) for b in range(gamma.shape[0])]) + if self.backend == 'torch' and bool((g == g[0]).all()): + return op.mean(gamma) + return g.reshape((gamma.shape[0],) + (1,) * (gamma.ndim - 1)) + def _spectral_grids(self, k): """(D_k2, D_frac, D_loss) fractional-power grids on the runtime |k| grid, using the average exponent ``self._gbar`` from the latest - :meth:`prepare_models` call. Cached per (gbar tensor, k shape): one + :meth:`prepare_models` call (one per batch element for per-shot + models, giving ``(B, 1, nz, nx)`` fractional grids). Cached per (gbar tensor, k shape): one rebuild per forward on torch (prepare runs once outside the loop), one per step inside a jax trace (where it stays trace-local). ``k = 0`` rows go through a safe branch so the pow's exponent-gradient stays @@ -352,8 +371,11 @@ def c_eq_aux(self, prop): from .base import init_wavenumbers shape = tuple(int(v) for v in prop.shape_cuda)[-2:] h = float(np.asarray(prop._grid_spacing).reshape(-1)[0]) - gbar = float(self._gbar.detach().cpu()) - key = (shape, h, str(prop.dev), round(gbar, 9), + # One exponent for a shared model; per-shot models with differing Q + # carry one per batch element, and the fractional grids become + # (B, nz, nx) stacks (the CUDA filter strides them per batch index). + gbars = [float(g) for g in self._gbar.detach().cpu().reshape(-1)] + key = (shape, h, str(prop.dev), tuple(round(g, 9) for g in gbars), self.phase_shift, self.amplitude_damping) cached = getattr(self, "_c_kmul_cache", None) if cached is None or cached[0] != key: @@ -361,14 +383,16 @@ def c_eq_aux(self, prop): def t(a): return torch.as_tensor( a.astype(np.float32), device=prop.dev).contiguous() + def frac(p): + g = [np.where(k_np > 0, np.power(k_np, 2.0*gbar + p, where=k_np > 0), 0.0) + for gbar in gbars] + return t(g[0] if len(g) == 1 else np.stack(g)) grids = [] if self.amplitude_damping: - grids.append(t(np.where( - k_np > 0, np.power(k_np, 2.0*gbar + 1.0, where=k_np > 0), 0.0))) + grids.append(frac(1.0)) if self.phase_shift: grids.append(t(k_np * k_np)) - grids.append(t(np.where( - k_np > 0, np.power(k_np, 2.0*gbar + 2.0, where=k_np > 0), 0.0))) + grids.append(frac(2.0)) self._c_kmul_cache = (key, tuple(grids)) return self._c_kmul_cache[1] diff --git a/test/test_visco_acoustic_batched_models.py b/test/test_visco_acoustic_batched_models.py new file mode 100644 index 00000000..322d0e12 --- /dev/null +++ b/test/test_visco_acoustic_batched_models.py @@ -0,0 +1,142 @@ +"""Per-shot batched models (B, nz, nx) for ViscoAcoustic: no cross-shot coupling. + +The NCQ step freezes the fractional exponent at the model average ``gbar``. +With per-shot models that average must be taken per batch element: a +whole-batch mean made every shot's |k| filter depend on the other shots' Q +(issue #123: ~1.7e-3 data / ~8e-4 gradient error vs running each model alone +when Q differs across the batch). Covered here: + +* the batched forward and the vp / Q gradients equal looping the shots one at + a time, on impl='eager' and impl='c', with Q DIFFERENT per shot; +* ``prepare_models`` keeps one scalar exponent for a shared model (and for a + shared model the C path repeated across the batch), so the shared path is + bit-for-bit the historical one. +""" +import numpy as np +import pytest +import torch + +from sweep.equations import ViscoAcoustic +from sweep.propagator.torch import PropTorch +from sweep.propagator.options import CUDAOptions, MemoryOptions, CkptOptions +from sweep.signal import ricker + +HAS_CUDA = torch.cuda.is_available() +DEV = "cuda" if HAS_CUDA else "cpu" +NZ, NX, DH, DT, NT, SO, ABCN = 48, 64, 8.0, 1.0e-3, 400, 4, 16 +F0 = 12.0 +B = 3 + +needs_cuda = pytest.mark.skipif(not HAS_CUDA, reason="impl='c' requires CUDA") +# impl='c' in each gradient memory mode (every backward reads the filter grids). +IMPLS = [pytest.param(m, marks=needs_cuda) for m in ("c", "c_ckpt", "c_rec")] + ["eager"] +CKPT = {"c_ckpt": CkptOptions(mode="chunk", chunks=10), + "c_rec": CkptOptions(mode="recursive", count=8)} + + +def _wavelet(): + t = np.arange(NT, dtype=np.float32) * DT - 0.06 + return torch.tensor((1e3 * ricker(t, f=F0)).astype(np.float32), device=DEV) + + +def _geometry(): + sx = np.array([16, 32, 48], dtype=np.int64) + sources = np.stack([sx, np.full(B, 4, np.int64)], -1) # (B, [x, z]) + rx = np.arange(4, NX - 4, 2, dtype=np.int64) + rec = np.stack([rx, np.full_like(rx, 2)], -1) + return sources, np.broadcast_to(rec, (B, *rec.shape)).copy() # (B, nrec, 2) + + +def _models(): + """B models, heterogeneous in vp and Q, with a different Q level per shot + (the case the whole-batch exponent coupled).""" + rng = np.random.default_rng(0) + d = np.linspace(0, 1, NZ, dtype=np.float32)[:, None] + vp = np.stack([2100 + 150 * b + 600 * d + rng.uniform(-40, 40, (NZ, NX)) + for b in range(B)]).astype(np.float32) + Q = np.stack([(15 + 25 * b) * (1 + 0.5 * d) + rng.uniform(-3, 3, (NZ, NX)) + for b in range(B)]).astype(np.float32) + om = np.full_like(vp, 2 * np.pi * F0) + return [torch.tensor(m, device=DEV) for m in (vp, Q, om)] + + +def _prop(impl): + eq = ViscoAcoustic(spatial_order=SO, device=DEV, backend="torch") + if impl == "eager": + kw = dict(impl="eager", use_compile=False, use_ckpt=False) + elif impl in CKPT: + kw = dict(impl="c", cuda_options=CUDAOptions(memory=MemoryOptions( + strategy="ckpt", ckpt=CKPT[impl]))) + else: + kw = dict(impl="c", use_ckpt=False) + return PropTorch(eq, shape=(NZ, NX), abcn=ABCN, dh=DH, dt=DT, nt=NT, **kw) + + +def _run(impl, models, sources, receivers): + ms = [m.clone().requires_grad_(i < 2) for i, m in enumerate(models)] + d = _prop(impl)(_wavelet(), sources, receivers, models=ms) + (d ** 2).sum().backward() + return d.detach().cpu(), [m.grad.detach().cpu() for m in ms[:2]] + + +def _rel(a, b): + return ((a - b).norm() / b.norm()).item() + + +@pytest.mark.parametrize("impl", IMPLS) +def test_visco_batched_matches_loop(impl): + """Batched forward and vp / Q gradients == one shot at a time (Q differs).""" + sources, receivers = _geometry() + models = _models() + d_b, g_b = _run(impl, models, sources, receivers) + for b in range(B): + d_1, g_1 = _run(impl, [m[b] for m in models], sources[b:b + 1], receivers[b:b + 1]) + ed = _rel(d_b[b:b + 1], d_1) + eg = [_rel(gb[b], g1) for gb, g1 in zip(g_b, g_1)] + print(f"[{impl}] shot {b}: rel(data)={ed:.2e} rel(grad vp, Q)={eg[0]:.2e}, {eg[1]:.2e}") + # The coupling bug sat at ~1e-3; what is left is float noise (cuFFT + # batch plans, atomicAdd order in the compiled gradient). + assert ed < 1e-5, f"[{impl}] shot {b} data rel={ed:.2e}" + assert max(eg) < 1e-4, f"[{impl}] shot {b} grad rel={eg}" + + +def test_visco_per_shot_exponent_is_per_element(): + """prepare_models: one exponent per batch element for per-shot models, + each the same reduction a single-shot run takes; one scalar otherwise.""" + eq = ViscoAcoustic(spatial_order=SO, device="cpu", backend="torch") + vp, Q, om = [m.cpu() for m in _models()] + gamma = torch.arctan(1.0 / Q) / np.pi + + eq.prepare_models([vp[:, None], Q[:, None], om[:, None]]) # eager per-shot layout + assert eq._gbar.shape == (B, 1, 1, 1) + for b in range(B): + eq_1 = ViscoAcoustic(spatial_order=SO, device="cpu", backend="torch") + eq_1.prepare_models([vp[b], Q[b], om[b]]) + assert torch.equal(eq._gbar[b].reshape(()), eq_1._gbar) + + eq.prepare_models([vp[0], Q[0], om[0]]) # shared (nz, nx) + assert eq._gbar.ndim == 0 and torch.equal(eq._gbar, torch.mean(gamma[0])) + + rep = lambda m: m[0][None, None].repeat(B, 1, 1, 1) # shared, as impl='c' lays it out + eq.prepare_models([rep(vp), rep(Q), rep(om)]) + assert eq._gbar.ndim == 0 + assert torch.equal(eq._gbar, torch.mean(rep(gamma))) + + +@needs_cuda +def test_visco_shared_model_c_unchanged(): + """impl='c': a shared model keeps (nz, nx) filter grids and equals an + identical per-shot batch bit-for-bit; per-shot Q gives (B, nz, nx) grids.""" + sources, receivers = _geometry() + vp, Q, om = _models() + prop = _prop("c") + d_shared = prop(_wavelet(), sources, receivers, models=[vp[0], Q[0], om[0]]).cpu() + assert all(g.ndim == 2 for g in prop.equation.c_eq_aux(prop)) + rep = lambda m: m[0][None].repeat(B, 1, 1) + d_rep = _prop("c")(_wavelet(), sources, receivers, models=[rep(vp), rep(Q), rep(om)]).cpu() + assert torch.equal(d_shared, d_rep) + + prop = _prop("c") + prop(_wavelet(), sources, receivers, models=[vp, Q, om]) + shapes = [tuple(g.shape) for g in prop.equation.c_eq_aux(prop)] + assert shapes == [(B, *shapes[1]), shapes[1], (B, *shapes[1])], shapes