feat: add 2D analytical solution, num_samples config, n param, 4 new tests

- sim_wave.py: add analytical_solution_2d (2D FFT exact solution for wave/KG)
- sim_wave.py: add n param to WaveSimulator for API compat with other simulators
- wave.yaml: add num_samples (was hardcoded 1000 in gen_wave.py) and n: 1
- gen_wave.py: use config.num_samples instead of hardcoded 1000
- test_sim_wave.py: add seed reproducibility, 2D leapfrog vs analytical,
  2D analytical at t=0, and periodicity test (single mode returns to IC after T=1/c)

Tests: 14/14 pass

Author: gpartin

pdebench/data_gen/configs/wave.yaml

@@ -12,6 +12,8 @@ output_path: 1D_Wave_c1.0
 
 name: 1d_wave
 
+num_samples: 1000
+
 sim:
   c: 1.0
   chi: 0.0       # 0 = wave equation, >0 = Klein-Gordon
@@ -20,6 +22,7 @@ sim:
   xdim: 1024
   ndim: 1
   n_modes: 5
+  n: 1
   seed: "???"
 
 plot:

pdebench/data_gen/gen_wave.py

@@ -164,7 +164,7 @@ def main(config: DictConfig):
     base_name = f"{config.sim.ndim}D_Wave_c{config.sim.c}{chi_str}"
     config.output_path = str((output_dir / base_name).with_suffix(".h5"))
 
-    num_samples = 1000
+    num_samples = config.num_samples
 
     log.info(f"Generating {num_samples} samples -> {config.output_path}")
     log.info(f"PDE: d2u/dt2 = {config.sim.c}^2 * Lap(u) - {config.sim.chi}^2 * u")

pdebench/data_gen/src/sim_wave.py

@@ -37,6 +37,7 @@ class WaveSimulator:
     :param xdim: number of spatial grid points per dimension
     :param ndim: spatial dimensionality (1 or 2)
     :param n_modes: number of Fourier modes in IC generation
+    :param n: number of batches (unused; present for API compatibility)
     :param seed: random seed for IC generation
     """
 
@@ -49,6 +50,7 @@ def __init__(
         xdim: int = 1024,
         ndim: int = 1,
         n_modes: int = 5,
+        n: int = 1,  # noqa: ARG002
         seed: int = 0,
     ):
         self.c = c
@@ -219,3 +221,42 @@ def analytical_solution_1d(
         result[i] = np.fft.ifft(u_hat_t).real
 
     return result
+
+
+def analytical_solution_2d(
+    x: np.ndarray,
+    t: np.ndarray,
+    u0: np.ndarray,
+    c: float,
+    chi: float = 0.0,
+) -> np.ndarray:
+    """
+    Compute analytical solution for 2D wave/KG equation via 2D FFT.
+
+    Each Fourier mode (kx, ky) oscillates at frequency
+    omega = sqrt(c^2 * (2*pi)^2 * (kx^2 + ky^2) + chi^2).
+
+    :param x: 1D spatial grid (Nx,), same for both dimensions
+    :param t: time points (Nt,)
+    :param u0: initial condition (Nx, Nx)
+    :param c: wave speed
+    :param chi: mass parameter
+    :return: solution array (Nt, Nx, Nx)
+    """
+    Nx = u0.shape[0]
+    u0_hat = np.fft.fft2(u0)
+    kx = np.fft.fftfreq(Nx, d=1.0 / Nx)
+    ky = np.fft.fftfreq(Nx, d=1.0 / Nx)
+    KX, KY = np.meshgrid(kx, ky, indexing="ij")
+
+    omega = np.sqrt(
+        c**2 * (2 * np.pi) ** 2 * (KX**2 + KY**2) + chi**2 + 0j
+    ).real
+
+    result = np.zeros((len(t), Nx, Nx), dtype=np.float64)
+    for i, ti in enumerate(t):
+        # du/dt(0) = 0 => only cosine term
+        u_hat_t = u0_hat * np.cos(omega * ti)
+        result[i] = np.fft.ifft2(u_hat_t).real
+
+    return result

tests/test_sim_wave.py

@@ -5,7 +5,7 @@
 import numpy as np
 import pytest
 
-from pdebench.data_gen.src.sim_wave import WaveSimulator, analytical_solution_1d
+from pdebench.data_gen.src.sim_wave import WaveSimulator, analytical_solution_1d, analytical_solution_2d
 
 
 # ---------------------------------------------------------------------------
@@ -141,3 +141,85 @@ def test_analytical_solution_kg_at_t0():
     result = analytical_solution_1d(x, t, u0, c=1.0, chi=5.0)
 
     np.testing.assert_allclose(result[0], u0, atol=1e-12)
+
+
+# ---------------------------------------------------------------------------
+# Seed reproducibility
+# ---------------------------------------------------------------------------
+
+
+def test_seed_reproducibility():
+    """Same seed must produce bit-identical output across two calls."""
+    sim_a = WaveSimulator(c=1.0, chi=0.0, xdim=32, tdim=11, t=0.5, seed=42)
+    sim_b = WaveSimulator(c=1.0, chi=0.0, xdim=32, tdim=11, t=0.5, seed=42)
+    np.testing.assert_array_equal(sim_a.generate_sample(), sim_b.generate_sample())
+
+
+# ---------------------------------------------------------------------------
+# 2D analytical solution
+# ---------------------------------------------------------------------------
+
+
+def test_wave_2d_matches_analytical():
+    """
+    2D leapfrog should match the 2D analytical solution to within 1% nRMSE.
+
+    IC: u0(x, y) = cos(2*pi*x) * cos(2*pi*y), du/dt=0.
+    Exact: u(x, y, t) = cos(2*pi*x) * cos(2*pi*y) * cos(2*pi*sqrt(2)*c*t).
+    """
+    Nx = 64
+    sim = WaveSimulator(c=1.0, chi=0.0, xdim=Nx, tdim=11, t=0.3, ndim=2, seed=0)
+
+    x1d = np.linspace(0, 1, Nx, endpoint=False)
+    X, Y = np.meshgrid(x1d, x1d, indexing="ij")
+    u0_2d = np.cos(2 * np.pi * X) * np.cos(2 * np.pi * Y)
+
+    sim._random_fourier_ic_2d = lambda rng: u0_2d.copy()  # noqa: ARG005
+
+    numerical = sim.generate_sample()  # (11, Nx, Nx) float32
+    exact = analytical_solution_2d(x1d, sim.t_save, u0_2d, c=1.0, chi=0.0)
+
+    u_range = exact.max() - exact.min()
+    nrmse = np.sqrt(np.mean((numerical.astype(np.float64) - exact) ** 2)) / u_range
+    assert nrmse < 0.01, f"2D nRMSE={nrmse:.4f} exceeds 1% tolerance"
+
+
+def test_analytical_solution_2d_at_t0():
+    """2D analytical solution at t=0 must equal u0 exactly."""
+    Nx = 32
+    x = np.linspace(0, 1, Nx, endpoint=False)
+    X, Y = np.meshgrid(x, x, indexing="ij")
+    u0 = np.sin(2 * np.pi * X) + 0.3 * np.cos(4 * np.pi * Y)
+    t = np.array([0.0, 0.5, 1.0])
+
+    result = analytical_solution_2d(x, t, u0, c=1.0, chi=0.0)
+
+    np.testing.assert_allclose(result[0], u0, atol=1e-12)
+
+
+# ---------------------------------------------------------------------------
+# Periodicity: single mode returns to IC after one period
+# ---------------------------------------------------------------------------
+
+
+def test_wave_periodicity_single_mode():
+    """
+    After one full period T = 1/c, a single k=1 cosine IC returns to itself.
+
+    u(x, t) = cos(2*pi*x) * cos(2*pi*c*t) has period T = 1/c.
+    The leapfrog nRMSE at t=T should be < 0.1% (O(dt^2) per step).
+    """
+    c = 1.0
+    Nx = 256
+    # tdim=3 saves t=0, t=T/2, t=T
+    sim = WaveSimulator(c=c, chi=0.0, xdim=Nx, tdim=3, t=1.0 / c, seed=0)
+    u0 = np.cos(2 * np.pi * sim.x)
+
+    sim._random_fourier_ic_1d = lambda rng: u0.copy()  # noqa: ARG005
+
+    result = sim.generate_sample()  # (3, Nx): t=0, t=T/2, t=T
+    u_final = result[-1].astype(np.float64)
+
+    u_range = u0.max() - u0.min()
+    nrmse = np.sqrt(np.mean((u_final - u0) ** 2)) / u_range
+    assert nrmse < 0.001, f"Periodicity nRMSE={nrmse:.6f} exceeds 0.1%"