|
| 1 | +"""Test: replicate ivp_1D_simulation.py via CLI commands. |
| 2 | +
|
| 3 | +Exercises: 1D grid, heterogeneous medium (array .npy files), |
| 4 | +custom p0, sparse sensor mask, custom CFL. |
| 5 | +""" |
| 6 | + |
| 7 | +import json |
| 8 | + |
| 9 | +import numpy as np |
| 10 | +import pytest |
| 11 | +from click.testing import CliRunner |
| 12 | + |
| 13 | +from kwave.cli.main import cli |
| 14 | +from kwave.data import Vector |
| 15 | +from kwave.kgrid import kWaveGrid |
| 16 | +from kwave.kmedium import kWaveMedium |
| 17 | +from kwave.ksensor import kSensor |
| 18 | +from kwave.ksource import kSource |
| 19 | +from kwave.kspaceFirstOrder import kspaceFirstOrder |
| 20 | + |
| 21 | + |
| 22 | +def _invoke(runner, args, session_dir): |
| 23 | + result = runner.invoke(cli, ["--session-dir", str(session_dir)] + args, catch_exceptions=False) |
| 24 | + assert result.exit_code == 0, f"Command failed: {args}\n{result.output}" |
| 25 | + output = result.output.strip() |
| 26 | + try: |
| 27 | + return json.loads(output) |
| 28 | + except json.JSONDecodeError: |
| 29 | + pass |
| 30 | + # For run command: find the last top-level JSON object |
| 31 | + depth = 0 |
| 32 | + last_start = None |
| 33 | + for i, ch in enumerate(output): |
| 34 | + if ch == "{" and depth == 0: |
| 35 | + last_start = i |
| 36 | + if ch == "{": |
| 37 | + depth += 1 |
| 38 | + elif ch == "}": |
| 39 | + depth -= 1 |
| 40 | + if last_start is not None: |
| 41 | + return json.loads(output[last_start:]) |
| 42 | + raise ValueError(f"Could not parse JSON from output: {output[:200]}") |
| 43 | + |
| 44 | + |
| 45 | +# --- Build the 1D IVP arrays (same as ivp_1D_simulation.py) --- |
| 46 | + |
| 47 | +Nx = 512 |
| 48 | +dx = 0.05e-3 |
| 49 | + |
| 50 | + |
| 51 | +def _make_sound_speed(): |
| 52 | + c = 1500 * np.ones(Nx) |
| 53 | + c[: Nx // 3] = 2000 |
| 54 | + return c |
| 55 | + |
| 56 | + |
| 57 | +def _make_density(): |
| 58 | + rho = 1000 * np.ones(Nx) |
| 59 | + rho[4 * Nx // 5 :] = 1500 |
| 60 | + return rho |
| 61 | + |
| 62 | + |
| 63 | +def _make_p0(): |
| 64 | + p0 = np.zeros(Nx) |
| 65 | + x0, width = 280, 100 |
| 66 | + pulse = 0.5 * (np.sin(np.arange(width + 1) * np.pi / width - np.pi / 2) + 1) |
| 67 | + p0[x0 : x0 + width + 1] = pulse |
| 68 | + return p0 |
| 69 | + |
| 70 | + |
| 71 | +def _make_sensor_mask(): |
| 72 | + mask = np.zeros(Nx) |
| 73 | + mask[Nx // 4] = 1 |
| 74 | + mask[3 * Nx // 4] = 1 |
| 75 | + return mask |
| 76 | + |
| 77 | + |
| 78 | +@pytest.fixture |
| 79 | +def session_dir(tmp_path): |
| 80 | + return tmp_path / "kwave_test_session" |
| 81 | + |
| 82 | + |
| 83 | +@pytest.fixture |
| 84 | +def data_dir(tmp_path): |
| 85 | + """Directory for pre-built .npy files (simulating what an agent would prepare).""" |
| 86 | + d = tmp_path / "arrays" |
| 87 | + d.mkdir() |
| 88 | + np.save(d / "sound_speed.npy", _make_sound_speed()) |
| 89 | + np.save(d / "density.npy", _make_density()) |
| 90 | + np.save(d / "p0.npy", _make_p0()) |
| 91 | + np.save(d / "sensor_mask.npy", _make_sensor_mask()) |
| 92 | + return d |
| 93 | + |
| 94 | + |
| 95 | +@pytest.fixture |
| 96 | +def runner(): |
| 97 | + return CliRunner() |
| 98 | + |
| 99 | + |
| 100 | +class TestCLI1DIVP: |
| 101 | + """Replicate ivp_1D_simulation.py end-to-end via CLI.""" |
| 102 | + |
| 103 | + def test_cli_matches_python_api(self, runner, session_dir, data_dir): |
| 104 | + # -- CLI flow -- |
| 105 | + _invoke(runner, ["session", "init"], session_dir) |
| 106 | + |
| 107 | + _invoke( |
| 108 | + runner, |
| 109 | + [ |
| 110 | + "phantom", |
| 111 | + "load", |
| 112 | + "--grid-size", |
| 113 | + "512", |
| 114 | + "--spacing", |
| 115 | + "0.05e-3", |
| 116 | + "--sound-speed", |
| 117 | + str(data_dir / "sound_speed.npy"), |
| 118 | + "--density", |
| 119 | + str(data_dir / "density.npy"), |
| 120 | + "--cfl", |
| 121 | + "0.3", |
| 122 | + ], |
| 123 | + session_dir, |
| 124 | + ) |
| 125 | + |
| 126 | + _invoke( |
| 127 | + runner, |
| 128 | + [ |
| 129 | + "source", |
| 130 | + "define", |
| 131 | + "--type", |
| 132 | + "initial-pressure", |
| 133 | + "--p0-file", |
| 134 | + str(data_dir / "p0.npy"), |
| 135 | + ], |
| 136 | + session_dir, |
| 137 | + ) |
| 138 | + |
| 139 | + _invoke( |
| 140 | + runner, |
| 141 | + [ |
| 142 | + "sensor", |
| 143 | + "define", |
| 144 | + "--mask", |
| 145 | + str(data_dir / "sensor_mask.npy"), |
| 146 | + "--record", |
| 147 | + "p", |
| 148 | + ], |
| 149 | + session_dir, |
| 150 | + ) |
| 151 | + |
| 152 | + plan_resp = _invoke(runner, ["plan"], session_dir) |
| 153 | + assert plan_resp["status"] == "ok" |
| 154 | + assert plan_resp["result"]["grid"]["N"] == [512] |
| 155 | + assert plan_resp["result"]["grid"]["Nt"] > 0 |
| 156 | + |
| 157 | + run_resp = _invoke(runner, ["run"], session_dir) |
| 158 | + assert run_resp["status"] == "ok" |
| 159 | + |
| 160 | + # Load CLI results |
| 161 | + cli_p = np.load(run_resp["result"]["outputs"]["p"]["path"]) |
| 162 | + |
| 163 | + # -- Direct Python API (the example) -- |
| 164 | + sound_speed = _make_sound_speed() |
| 165 | + density = _make_density() |
| 166 | + kgrid = kWaveGrid(Vector([Nx]), Vector([dx])) |
| 167 | + kgrid.makeTime(sound_speed, cfl=0.3) |
| 168 | + medium = kWaveMedium(sound_speed=sound_speed, density=density) |
| 169 | + source = kSource() |
| 170 | + source.p0 = _make_p0() |
| 171 | + sensor = kSensor(mask=_make_sensor_mask()) |
| 172 | + result = kspaceFirstOrder(kgrid, medium, source, sensor, backend="python", quiet=True) |
| 173 | + |
| 174 | + # -- Compare -- |
| 175 | + assert cli_p.shape == result["p"].shape, f"Shape mismatch: {cli_p.shape} vs {result['p'].shape}" |
| 176 | + np.testing.assert_allclose(cli_p, result["p"], rtol=0, atol=0) |
0 commit comments