diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml new file mode 100644 index 0000000..f49e72e --- /dev/null +++ b/.github/workflows/tests.yml @@ -0,0 +1,27 @@ +name: Tests + +on: + push: + branches: [main] + pull_request: + +jobs: + test: + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v4 + + - name: Install uv + uses: astral-sh/setup-uv@v5 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.12" + + - name: Install dependencies + run: uv sync --group dev + + - name: Run tests + run: uv run pytest tests/ diff --git a/.gitignore b/.gitignore index 1ca5add..e33c601 100644 --- a/.gitignore +++ b/.gitignore @@ -3,4 +3,5 @@ __pycache__/ .DS_Store docs/_build/ dist/ -.vscode/ \ No newline at end of file +.vscode/ +uv.lock \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 5a6cded..72c9be9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -74,3 +74,9 @@ exclude_lines = [ "if __name__ == .__main__.:", "if TYPE_CHECKING:", ] + +[dependency-groups] +dev = [ + "pytest>=8.3.5", + "pytest-cov>=5.0.0", +] diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..4c25277 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,55 @@ +""" +Shared test fixtures for pygenray tests. +""" +import matplotlib +matplotlib.use('Agg') + + +def pytest_addoption(parser): + """Register --regenerate-physics CLI flag for physics regression tests.""" + parser.addoption( + '--regenerate-physics', action='store_true', default=False, + help='Regenerate physics regression fixture and skip comparison.', + ) + +import numpy as np +import pytest +import xarray as xr + +from pygenray.ray_objects import Ray, RayFan + + +def _make_ray(launch_angle: float, source_depth: float, n_bottom: int = 0, + n_surface: int = 0, N: int = 10, R: float = 10000.0) -> Ray: + """Helper: build a synthetic Ray without running the ODE solver. + + The y array uses the positive-z convention expected by Ray.__init__: + y[0,:] = travel times + y[1,:] = depth (positive = deeper) + y[2,:] = ray parameter sin(θ)/c (positive for downward ray in ODE) + """ + r = np.linspace(0.0, R, N) + t = r / 1500.0 + # Depth increases linearly (simulating a shallow downward ray) + z_ode = np.linspace(source_depth, source_depth + R * 0.01, N) + p_ode = np.ones(N) * np.sin(np.radians(abs(launch_angle) + 1e-3)) / 1500.0 + y = np.vstack([t, z_ode, p_ode]) # shape (3, N) + return Ray(r=r, y=y, n_bottom=n_bottom, n_surface=n_surface, + launch_angle=launch_angle, source_depth=source_depth) + + +@pytest.fixture +def simple_ray(): + """Single synthetic Ray with 10 range steps.""" + return _make_ray(launch_angle=-10.0, source_depth=100.0) + + +@pytest.fixture +def simple_rayfan(): + """Small RayFan of 3 synthetic Rays — no ray tracing needed.""" + rays = [ + _make_ray(launch_angle=-5.0, source_depth=100.0, n_bottom=0), + _make_ray(launch_angle=5.0, source_depth=150.0, n_bottom=1), + _make_ray(launch_angle=-10.0, source_depth=200.0, n_bottom=0), + ] + return RayFan(rays) diff --git a/tests/fixtures/munk_regression.npz b/tests/fixtures/munk_regression.npz new file mode 100644 index 0000000..58b6fdb Binary files /dev/null and b/tests/fixtures/munk_regression.npz differ diff --git a/tests/test_environment.py b/tests/test_environment.py new file mode 100644 index 0000000..a639236 --- /dev/null +++ b/tests/test_environment.py @@ -0,0 +1,180 @@ +""" +Tests for pygenray.environment: munk_ssp, OceanEnvironment2D, eflat, eflatinv. +""" +import numpy as np +import pytest +import xarray as xr +from matplotlib import pyplot as plt + +from pygenray.environment import ( + OceanEnvironment2D, + eflat, + eflatinv, + munk_ssp, +) + + +# --------------------------------------------------------------------------- +# munk_ssp +# --------------------------------------------------------------------------- + +class TestMunkSSP: + def test_output_shape_matches_input(self): + z = np.arange(0, 5000, 10) + c = munk_ssp(z) + assert c.shape == z.shape + + def test_minimum_at_sofar_depth(self): + sofar = 1300.0 + z = np.arange(0, 6000, 1) + c = munk_ssp(z, sofar_depth=sofar) + # Minimum should be at the SOFAR depth + assert z[np.argmin(c)] == pytest.approx(sofar, abs=2.0) + + def test_default_params_near_1500_at_sofar(self): + sofar = 1300.0 + c_sofar = munk_ssp(np.array([sofar])) + assert c_sofar[0] == pytest.approx(1500.0, abs=5.0) + + def test_scalar_input(self): + c = munk_ssp(np.array([0.0])) + assert c.shape == (1,) + + +# --------------------------------------------------------------------------- +# OceanEnvironment2D construction +# --------------------------------------------------------------------------- + +class TestOceanEnvironment2DConstruction: + def test_default_init_attributes_exist(self): + env = OceanEnvironment2D() + for attr in ('sound_speed', 'bathymetry', 'dcdz', 'bottom_angle', + 'bottom_angle_interp'): + assert hasattr(env, attr), f"Missing attribute: {attr}" + + def test_default_sound_speed_is_2d(self): + env = OceanEnvironment2D() + assert env.sound_speed.ndim == 2 + assert set(env.sound_speed.dims) == {'range', 'depth'} + + def test_default_flat_earth_attributes_exist(self): + env = OceanEnvironment2D(flat_earth_transform=True) + assert hasattr(env, 'sound_speed_fe') + assert hasattr(env, 'bathymetry_fe') + + def test_flat_earth_false_no_fe_attributes(self): + env = OceanEnvironment2D(flat_earth_transform=False) + assert not hasattr(env, 'sound_speed_fe') + assert not hasattr(env, 'bathymetry_fe') + + def test_custom_1d_sound_speed(self): + z = np.arange(0.0, 3000.0, 10.0) + c_vals = munk_ssp(z) + ssp = xr.DataArray(c_vals, dims=['depth'], coords={'depth': z}) + bathy = xr.DataArray( + np.ones(20) * 4000.0, dims=['range'], + coords={'range': np.linspace(0, 50e3, 20)} + ) + env = OceanEnvironment2D(sound_speed=ssp, bathymetry=bathy, + flat_earth_transform=False) + assert env.sound_speed.ndim == 1 + assert 'depth' in env.sound_speed.dims + + def test_custom_2d_sound_speed(self): + z = np.arange(0.0, 3000.0, 50.0) + r = np.linspace(0.0, 50e3, 20) + c_2d = np.outer(np.ones(len(r)), munk_ssp(z)) + ssp = xr.DataArray(c_2d, dims=['range', 'depth'], + coords={'range': r, 'depth': z}) + env = OceanEnvironment2D(sound_speed=ssp, flat_earth_transform=False) + assert env.sound_speed.ndim == 2 + + def test_custom_bathymetry_stored(self): + bathy_vals = np.ones(20) * 3500.0 + r = np.linspace(0.0, 50e3, 20) + bathy = xr.DataArray(bathy_vals, dims=['range'], coords={'range': r}) + env = OceanEnvironment2D(bathymetry=bathy, flat_earth_transform=False) + np.testing.assert_array_equal(env.bathymetry.values, bathy_vals) + + # --- invalid inputs --- + + def test_sound_speed_not_dataarray_raises_type_error(self): + with pytest.raises(TypeError): + OceanEnvironment2D(sound_speed=np.ones(100)) + + def test_sound_speed_3d_raises_value_error(self): + da = xr.DataArray( + np.ones((5, 10, 20)), + dims=['range', 'depth', 'extra'], + coords={'range': np.arange(5), 'depth': np.arange(10), + 'extra': np.arange(20)} + ) + with pytest.raises(ValueError): + OceanEnvironment2D(sound_speed=da) + + def test_sound_speed_missing_depth_dim_raises_value_error(self): + da = xr.DataArray(np.ones(50), dims=['range'], + coords={'range': np.arange(50)}) + with pytest.raises(ValueError): + OceanEnvironment2D(sound_speed=da) + + def test_2d_sound_speed_missing_range_dim_raises_value_error(self): + da = xr.DataArray( + np.ones((10, 20)), + dims=['depth', 'extra'], + coords={'depth': np.arange(10), 'extra': np.arange(20)} + ) + with pytest.raises(ValueError): + OceanEnvironment2D(sound_speed=da) + + def test_bathymetry_not_dataarray_raises_type_error(self): + with pytest.raises(TypeError): + OceanEnvironment2D(bathymetry=np.ones(50)) + + def test_bathymetry_missing_range_dim_raises_value_error(self): + da = xr.DataArray(np.ones(50), dims=['depth'], + coords={'depth': np.arange(50)}) + with pytest.raises(ValueError): + OceanEnvironment2D(bathymetry=da) + + +# --------------------------------------------------------------------------- +# eflat / eflatinv round-trip +# --------------------------------------------------------------------------- + +class TestEflat: + LAT = 35.0 + + def test_depth_roundtrip(self): + dep = np.array([100.0, 500.0, 1000.0, 2000.0, 4000.0]) + depf, _ = eflat(dep, self.LAT) + dep_rec, _ = eflatinv(depf, np.array([self.LAT])) + np.testing.assert_allclose(dep_rec, dep, atol=1.0, + err_msg="Depth round-trip outside 1 m tolerance") + + def test_sound_speed_roundtrip(self): + dep = np.array([100.0, 500.0, 1000.0, 2000.0]) + cs = np.array([1500.0, 1490.0, 1480.0, 1510.0]) + depf, csf = eflat(dep, self.LAT, cs) + _, cs_rec = eflatinv(depf, np.array([self.LAT]), csf) + np.testing.assert_allclose(cs_rec, cs, rtol=1e-4, + err_msg="Sound speed round-trip outside 0.01% tolerance") + + def test_eflat_increases_depth(self): + """Flat-earth transformation should increase effective depths.""" + dep = np.array([100.0, 1000.0, 3000.0]) + depf, _ = eflat(dep, self.LAT) + assert np.all(depf > dep) + + +# --------------------------------------------------------------------------- +# OceanEnvironment2D.plot smoke test +# --------------------------------------------------------------------------- + +class TestOceanEnvironment2DPlot: + def test_plot_runs_without_error(self): + env = OceanEnvironment2D() + fig, ax = plt.subplots() + plt.sca(ax) + env.plot() + plt.close('all') diff --git a/tests/test_physics.py b/tests/test_physics.py new file mode 100644 index 0000000..62cae0d --- /dev/null +++ b/tests/test_physics.py @@ -0,0 +1,324 @@ +""" +Physics correctness tests for pygenray ray tracing. + +All tests use flatearth=False to avoid flat-earth correction complicating +the comparison with analytical solutions. +""" +import os +import pathlib + +import numpy as np +import pytest +import xarray as xr + +from pygenray.environment import OceanEnvironment2D, munk_ssp +from pygenray.launch_rays import shoot_ray, shoot_rays + +FIXTURE_DIR = pathlib.Path(__file__).parent / 'fixtures' + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _const_c_env(c0=1500.0, z_max=5000.0, r_max=100e3, + bathy_depth=4500.0, nz=200, nr=20): + """Build a range-independent constant sound-speed environment.""" + z = np.linspace(0.0, z_max, nz) + r = np.linspace(0.0, r_max, nr) + c_2d = np.full((nr, nz), c0) + ssp = xr.DataArray(c_2d, dims=['range', 'depth'], + coords={'range': r, 'depth': z}) + bathy = xr.DataArray(np.full(nr, bathy_depth), dims=['range'], + coords={'range': r}) + return OceanEnvironment2D(sound_speed=ssp, bathymetry=bathy, + flat_earth_transform=False) + + +def _linear_gradient_env(c0=1500.0, g=0.05, z_max=5000.0, r_max=100e3, + bathy_depth=4500.0, nz=500, nr=50): + """Build a range-independent linear-gradient environment: c(z) = c0 + g*z.""" + z = np.linspace(0.0, z_max, nz) + r = np.linspace(0.0, r_max, nr) + c_1d = c0 + g * z + c_2d = np.outer(np.ones(nr), c_1d) + ssp = xr.DataArray(c_2d, dims=['range', 'depth'], + coords={'range': r, 'depth': z}) + bathy = xr.DataArray(np.full(nr, bathy_depth), dims=['range'], + coords={'range': r}) + return OceanEnvironment2D(sound_speed=ssp, bathymetry=bathy, + flat_earth_transform=False) + + +def _munk_env(r_max=100e3, nr=50, nz=600, bathy_depth=5000.0): + """Build a range-independent Munk-profile environment.""" + z = np.linspace(0.0, 6000.0, nz) + r = np.linspace(0.0, r_max, nr) + c_1d = munk_ssp(z) + c_2d = np.outer(np.ones(nr), c_1d) + ssp = xr.DataArray(c_2d, dims=['range', 'depth'], + coords={'range': r, 'depth': z}) + bathy = xr.DataArray(np.full(nr, bathy_depth), dims=['range'], + coords={'range': r}) + return OceanEnvironment2D(sound_speed=ssp, bathymetry=bathy, + flat_earth_transform=False) + + +# --------------------------------------------------------------------------- +# A. Snell's invariant in range-independent constant-c medium +# --------------------------------------------------------------------------- + +class TestSnellInvariant: + """ + In a constant sound-speed medium (dc/dz = 0), the ray-parameter + p = sin(θ)/c satisfies dp/dx = 0 exactly. The stored value ray.p = -p_ode + should therefore be constant to within ODE integration error. + """ + + @pytest.mark.parametrize("user_angle", [-5.0, -10.0, -15.0]) + def test_p_constant_along_ray(self, user_angle): + c0 = 1500.0 + env = _const_c_env(c0=c0) + ray = shoot_ray(200.0, 0.0, user_angle, 30e3, 60, + env, rtol=1e-9, flatearth=False, debug=False) + assert ray is not None, "shoot_ray returned None unexpectedly" + # |p| should be constant; sign may flip at boundary reflections + abs_p = np.abs(ray.p) + p_std = np.std(abs_p) + p_mean = np.mean(abs_p) + assert p_std / p_mean < 1e-5, ( + f"|p| not constant in constant-c medium: std/mean = {p_std/p_mean:.2e}" + ) + + +# --------------------------------------------------------------------------- +# B. Constant sound speed — straight-line rays +# --------------------------------------------------------------------------- + +class TestConstantSSPStraightLine: + """ + In a homogeneous medium the ray paths are straight lines. + Analytical travel time: t = R / (c · cos θ) + Analytical final depth: z = z0 + R · tan θ (positive downward, ODE convention) + Stored convention: ray.z = -(ODE z) + """ + + def test_travel_time_analytical(self): + c0 = 1500.0 + user_angle = -10.0 # downward + theta_ode = abs(user_angle) # internal ODE angle (degrees) + z0 = 200.0 # source depth [m] + R = 20e3 # receiver range [m] + env = _const_c_env(c0=c0, r_max=R + 1e3) + + ray = shoot_ray(z0, 0.0, user_angle, R, 50, + env, rtol=1e-9, flatearth=False, debug=False) + assert ray is not None + + t_analytical = R / (c0 * np.cos(np.radians(theta_ode))) + assert abs(ray.t[-1] - t_analytical) / t_analytical < 1e-3 + + def test_final_depth_analytical(self): + c0 = 1500.0 + user_angle = -10.0 + theta_ode = abs(user_angle) + z0 = 200.0 + R = 20e3 + env = _const_c_env(c0=c0, r_max=R + 1e3) + + ray = shoot_ray(z0, 0.0, user_angle, R, 50, + env, rtol=1e-9, flatearth=False, debug=False) + assert ray is not None + + z_ode_end = z0 + R * np.tan(np.radians(theta_ode)) + z_expected = -z_ode_end # stored sign convention + assert abs(ray.z[-1] - z_expected) / abs(z_expected) < 1e-3 + + def test_p_constant_in_const_c(self): + c0 = 1500.0 + user_angle = -10.0 + theta_ode = abs(user_angle) + z0 = 200.0 + R = 20e3 + env = _const_c_env(c0=c0, r_max=R + 1e3) + + ray = shoot_ray(z0, 0.0, user_angle, R, 50, + env, rtol=1e-9, flatearth=False, debug=False) + assert ray is not None + + p_expected = -np.sin(np.radians(theta_ode)) / c0 # stored sign + np.testing.assert_allclose(ray.p, p_expected, + rtol=1e-5, atol=0, + err_msg="p not constant in homogeneous medium") + + +# --------------------------------------------------------------------------- +# C. Linear gradient — turning depth +# --------------------------------------------------------------------------- + +class TestLinearGradientTurningDepth: + """ + For c(z) = c0 + g·z the Hamiltonian is conserved: + H = sqrt(1/c(z)² - p_ode²) = sqrt(1/c_source² - p0²) + At the turning point (ray horizontal, dz/dx = 0) we have p_ode = 0, so: + c(z_turn) = 1 / sqrt(1/c_source² - p0²) + where p0 = sin(θ_ode) / c_source. + + Substituting: c_source = c0 + g·z_source, p0 = sin(θ)/c_source + c(z_turn) = c_source / cos(θ) + z_turn = (c_source/cos(θ) − c0) / g + """ + + C0 = 1500.0 + G = 0.05 # m/s/m + Z_SRC = 200.0 # source depth [m] + THETA = 20.0 # degrees downward (user passes -20) + + def _z_turn_analytical(self): + c_source = self.C0 + self.G * self.Z_SRC + return (c_source / np.cos(np.radians(self.THETA)) - self.C0) / self.G + + def test_turning_depth_approx(self): + z_turn = self._z_turn_analytical() + env = _linear_gradient_env(c0=self.C0, g=self.G) + + ray = shoot_ray(self.Z_SRC, 0.0, -self.THETA, 80e3, 400, + env, rtol=1e-9, flatearth=False, debug=False) + assert ray is not None + + # ray.z uses sign convention: negative = deep + z_turn_numerical = -np.min(ray.z) # max depth reached + assert abs(z_turn_numerical - z_turn) < 50.0, ( + f"Turning depth: expected {z_turn:.1f} m, got {z_turn_numerical:.1f} m" + ) + + def test_hamiltonian_conserved_linear_gradient(self): + """The Hamiltonian H = sqrt(1/c(z)² − p_ode²) is conserved.""" + env = _linear_gradient_env(c0=self.C0, g=self.G) + + ray = shoot_ray(self.Z_SRC, 0.0, -self.THETA, 80e3, 400, + env, rtol=1e-9, flatearth=False, debug=False) + assert ray is not None + + z_ode = -ray.z # positive depth (ODE convention) + p_ode = -ray.p # positive for downward ray + c_along = self.C0 + self.G * z_ode + + H = np.sqrt(1.0 / c_along**2 - p_ode**2) + H_std = np.std(H) + H_mean = np.mean(H) + assert H_std / H_mean < 1e-4, ( + f"Hamiltonian not conserved: std/mean = {H_std/H_mean:.2e}" + ) + + +# --------------------------------------------------------------------------- +# D. Hamiltonian conservation in range-independent Munk profile +# --------------------------------------------------------------------------- + +class TestMunkHamiltonianConservation: + """ + In any range-independent environment the Hamiltonian + H = sqrt(1/c(z)² − p_ode²) + is an exact invariant of the ray equations. Numerical integration + should preserve it to within ODE tolerance. + """ + + @pytest.mark.parametrize("user_angle", [-5.0, -10.0, -15.0]) + def test_hamiltonian_conserved_munk(self, user_angle): + env = _munk_env(r_max=100e3) + z_src = 1000.0 # below SOFAR axis — ray stays in the water column + + ray = shoot_ray(z_src, 0.0, user_angle, 100e3, 200, + env, rtol=1e-9, flatearth=False, debug=False) + assert ray is not None, "shoot_ray returned None; ray may have exited domain" + + z_ode = -ray.z + p_ode = -ray.p + c_along = munk_ssp(z_ode) + + arg = 1.0 / c_along**2 - p_ode**2 + # Clamp small negatives from floating-point near turning point + arg = np.clip(arg, 0.0, None) + H = np.sqrt(arg) + + # Exclude points at/near the turning point where H → 0 + mask = H > 1e-6 / 1500.0 + if mask.sum() < 5: + pytest.skip("Too few valid points away from turning point") + + H_valid = H[mask] + H_std = np.std(H_valid) + H_mean = np.mean(H_valid) + assert H_std / H_mean < 1e-3, ( + f"Hamiltonian not conserved in Munk profile: std/mean = {H_std/H_mean:.2e}" + ) + + +# --------------------------------------------------------------------------- +# E. Regression / golden-file tests +# --------------------------------------------------------------------------- + +def _regenerate_fixture(): + """Run shoot_rays and save regression fixture. Call manually to regenerate.""" + env = _munk_env(r_max=50e3, nr=30, nz=400) + angles = [-8.0, -4.0, 0.0, 4.0, 8.0] + rf = shoot_rays(1300.0, 0.0, angles, 50e3, 50, env, + n_processes=1, debug=False, flatearth=False) + FIXTURE_DIR.mkdir(parents=True, exist_ok=True) + np.savez( + FIXTURE_DIR / 'munk_regression.npz', + ts=rf.ts, zs=rf.zs, ps=rf.ps, + n_botts=rf.n_botts, n_surfs=rf.n_surfs, + thetas=rf.thetas, + ) + return rf + + +class TestMunkRegression: + """ + Golden-file regression: compare shoot_rays output against a saved fixture. + + To regenerate the fixture (e.g., after intentional physics changes): + python -c "from tests.test_physics import _regenerate_fixture; _regenerate_fixture()" + or run pytest with --regenerate-physics flag. + """ + + FIXTURE = FIXTURE_DIR / 'munk_regression.npz' + ANGLES = [-8.0, -4.0, 0.0, 4.0, 8.0] + + def _run_rays(self): + env = _munk_env(r_max=50e3, nr=30, nz=400) + return shoot_rays(1300.0, 0.0, self.ANGLES, 50e3, 50, env, + n_processes=1, debug=False, flatearth=False) + + def test_regression(self, request): + regenerate = request.config.getoption('--regenerate-physics', + default=False) + if regenerate or not self.FIXTURE.exists(): + rf = self._run_rays() + self.FIXTURE.parent.mkdir(parents=True, exist_ok=True) + np.savez( + self.FIXTURE, + ts=rf.ts, zs=rf.zs, ps=rf.ps, + n_botts=rf.n_botts, n_surfs=rf.n_surfs, + thetas=rf.thetas, + ) + if regenerate: + pytest.skip("Fixture regenerated; skipping comparison") + # First-run: fixture just created, nothing to compare + return + + ref = np.load(self.FIXTURE) + rf = self._run_rays() + + np.testing.assert_allclose(rf.ts, ref['ts'], atol=1e-6, + err_msg="Travel times changed") + np.testing.assert_allclose(rf.zs, ref['zs'], atol=0.1, + err_msg="Depths changed") + np.testing.assert_allclose(rf.ps, ref['ps'], atol=0.1, + err_msg="Ray parameters changed") + np.testing.assert_array_equal(rf.n_botts, ref['n_botts']) + np.testing.assert_array_equal(rf.n_surfs, ref['n_surfs']) + + diff --git a/tests/test_ray_objects.py b/tests/test_ray_objects.py new file mode 100644 index 0000000..509dd68 --- /dev/null +++ b/tests/test_ray_objects.py @@ -0,0 +1,278 @@ +""" +Tests for pygenray.ray_objects: Ray, RayFan. +""" +import numpy as np +import pytest +import scipy.io +from matplotlib import pyplot as plt + +from pygenray.ray_objects import Ray, RayFan + + +# --------------------------------------------------------------------------- +# Ray +# --------------------------------------------------------------------------- + +class TestRay: + N = 10 + R = 10000.0 + + def _make_ray(self, launch_angle=-10.0, source_depth=100.0, + n_bottom=0, n_surface=0): + r = np.linspace(0.0, self.R, self.N) + t = r / 1500.0 + z_ode = np.linspace(source_depth, source_depth + self.R * 0.01, self.N) + p_ode = np.ones(self.N) * np.sin(np.radians(abs(launch_angle) + 1e-3)) / 1500.0 + y = np.vstack([t, z_ode, p_ode]) + return Ray(r=r, y=y, n_bottom=n_bottom, n_surface=n_surface, + launch_angle=launch_angle, source_depth=source_depth), y + + def test_attribute_shapes(self): + ray, _ = self._make_ray() + assert ray.r.shape == (self.N,) + assert ray.t.shape == (self.N,) + assert ray.z.shape == (self.N,) + assert ray.p.shape == (self.N,) + + def test_z_sign_convention(self): + """ray.z should equal -y[1,:].""" + ray, y = self._make_ray() + np.testing.assert_array_equal(ray.z, -y[1, :]) + + def test_p_sign_convention(self): + """ray.p should equal -y[2,:].""" + ray, y = self._make_ray() + np.testing.assert_array_equal(ray.p, -y[2, :]) + + def test_launch_angle_stored(self): + ray, _ = self._make_ray(launch_angle=-15.0) + assert ray.launch_angle == pytest.approx(-15.0) + + def test_source_depth_stored(self): + ray, _ = self._make_ray(source_depth=250.0) + assert ray.source_depth == pytest.approx(250.0) + + def test_optional_launch_angle_not_set(self): + r = np.linspace(0.0, self.R, self.N) + t = r / 1500.0 + y = np.vstack([t, np.ones(self.N) * 100.0, np.ones(self.N) * 0.1]) + ray = Ray(r=r, y=y, n_bottom=0, n_surface=0) + assert not hasattr(ray, 'launch_angle') + + def test_optional_source_depth_not_set(self): + r = np.linspace(0.0, self.R, self.N) + t = r / 1500.0 + y = np.vstack([t, np.ones(self.N) * 100.0, np.ones(self.N) * 0.1]) + ray = Ray(r=r, y=y, n_bottom=0, n_surface=0) + assert not hasattr(ray, 'source_depth') + + def test_n_bottom_n_surface_stored(self): + ray, _ = self._make_ray(n_bottom=3, n_surface=1) + assert ray.n_bottom == 3 + assert ray.n_surface == 1 + + def test_plot_smoke(self): + ray, _ = self._make_ray() + plt.figure() + ray.plot() + plt.close('all') + + +# --------------------------------------------------------------------------- +# RayFan +# --------------------------------------------------------------------------- + +class TestRayFan: + M = 3 + N = 10 + R = 10000.0 + + def _make_rays(self, M=None, N=None, R=None): + M = M or self.M + N = N or self.N + R = R or self.R + rays = [] + for i in range(M): + r = np.linspace(0.0, R, N) + theta = float(-5 + i * 5) + t = r / 1500.0 + z_ode = np.linspace(100.0 + i * 50, 200.0 + i * 50, N) + p_ode = np.ones(N) * np.sin(np.radians(abs(theta) + 1e-3)) / 1500.0 + y = np.vstack([t, z_ode, p_ode]) + rays.append( + Ray(r=r, y=y, n_bottom=i % 2, n_surface=0, + launch_angle=theta, source_depth=100.0 + i * 50) + ) + return rays + + # --- Construction --- + + def test_shapes(self, simple_rayfan): + rf = simple_rayfan + assert rf.thetas.shape == (self.M,) + assert rf.rs.shape == (self.M, self.N) + assert rf.ts.shape == (self.M, self.N) + assert rf.zs.shape == (self.M, self.N) + assert rf.ps.shape == (self.M, self.N) + assert rf.n_botts.shape == (self.M,) + assert rf.n_surfs.shape == (self.M,) + assert rf.source_depths.shape == (self.M,) + + def test_ray_ids_set_on_construction(self, simple_rayfan): + assert hasattr(simple_rayfan, 'ray_ids') + assert len(simple_rayfan.ray_ids) == self.M + + # --- compute_rayids --- + + def test_compute_rayids_returns_strings(self, simple_rayfan): + simple_rayfan.compute_rayids() + assert all(isinstance(rid, str) for rid in simple_rayfan.ray_ids) + + def test_compute_rayids_length(self, simple_rayfan): + simple_rayfan.compute_rayids() + assert len(simple_rayfan.ray_ids) == len(simple_rayfan.thetas) + + # --- __len__ --- + + def test_len(self, simple_rayfan): + assert len(simple_rayfan) == self.M + + # --- __getitem__ int --- + + def test_getitem_int_returns_ray(self, simple_rayfan): + ray = simple_rayfan[0] + assert isinstance(ray, Ray) + + def test_getitem_int_correct_index(self, simple_rayfan): + ray = simple_rayfan[1] + np.testing.assert_array_equal(ray.r, simple_rayfan.rs[1]) + + def test_getitem_negative_int(self, simple_rayfan): + ray = simple_rayfan[-1] + assert isinstance(ray, Ray) + np.testing.assert_array_equal(ray.r, simple_rayfan.rs[-1]) + + def test_getitem_out_of_bounds_raises_index_error(self, simple_rayfan): + with pytest.raises(IndexError): + _ = simple_rayfan[100] + + # --- __getitem__ slice --- + + def test_getitem_slice_returns_rayfan(self, simple_rayfan): + result = simple_rayfan[0:2] + assert isinstance(result, RayFan) + assert len(result) == 2 + + def test_getitem_slice_correct_thetas(self, simple_rayfan): + result = simple_rayfan[1:] + np.testing.assert_array_equal(result.thetas, simple_rayfan.thetas[1:]) + + # --- __getitem__ bool mask --- + + def test_getitem_bool_mask_returns_rayfan(self, simple_rayfan): + mask = np.array([True, False, True]) + result = simple_rayfan[mask] + assert isinstance(result, RayFan) + assert len(result) == 2 + + def test_getitem_bool_mask_correct_subset(self, simple_rayfan): + mask = np.array([False, True, False]) + result = simple_rayfan[mask] + np.testing.assert_array_equal(result.thetas, simple_rayfan.thetas[1:2]) + + # --- __getitem__ int array --- + + def test_getitem_int_array_returns_rayfan(self, simple_rayfan): + idx = np.array([0, 2]) + result = simple_rayfan[idx] + assert isinstance(result, RayFan) + assert len(result) == 2 + np.testing.assert_array_equal(result.thetas, + simple_rayfan.thetas[np.array([0, 2])]) + + # --- __add__ --- + + def test_add_correct_length(self): + rays_a = self._make_rays(M=2) + rays_b = self._make_rays(M=3) + rf_a = RayFan(rays_a) + rf_b = RayFan(rays_b) + result = rf_a + rf_b + assert len(result) == 5 + + def test_add_rs_preserved(self): + rays_a = self._make_rays(M=2) + rays_b = self._make_rays(M=1) + rf_a = RayFan(rays_a) + rf_b = RayFan(rays_b) + result = rf_a + rf_b + # All rays should have the same range array + for i in range(len(result)): + np.testing.assert_array_equal(result.rs[i], rf_a.rs[0]) + + def test_add_incompatible_ranges_raises_value_error(self): + rays_a = self._make_rays(M=1, R=10000.0) + rays_b = self._make_rays(M=1, R=20000.0) + rf_a = RayFan(rays_a) + rf_b = RayFan(rays_b) + with pytest.raises(ValueError): + _ = rf_a + rf_b + + def test_add_non_rayfan_raises_type_error(self, simple_rayfan): + with pytest.raises(TypeError): + _ = simple_rayfan + 42 + + # --- save_mat --- + + def test_save_mat_creates_file(self, simple_rayfan, tmp_path): + path = str(tmp_path / 'test_rayfan.mat') + simple_rayfan.save_mat(path) + assert (tmp_path / 'test_rayfan.mat').exists() + + def test_save_mat_loadable(self, simple_rayfan, tmp_path): + path = str(tmp_path / 'test_rayfan.mat') + simple_rayfan.save_mat(path) + data = scipy.io.loadmat(path) + assert 'rayfan' in data + + def test_save_mat_contains_required_keys(self, simple_rayfan, tmp_path): + path = str(tmp_path / 'test_rayfan.mat') + simple_rayfan.save_mat(path) + data = scipy.io.loadmat(path) + rayfan = data['rayfan'] + # MATLAB struct stored as structured array; check dtype field names + expected_keys = {'thetas', 'xs', 'ts', 'zs', 'ps', + 'n_botts', 'n_surfs', 'source_depths'} + actual_keys = set(rayfan.dtype.names) + assert expected_keys <= actual_keys + + def test_save_mat_values_match(self, simple_rayfan, tmp_path): + path = str(tmp_path / 'test_rayfan.mat') + simple_rayfan.save_mat(path) + data = scipy.io.loadmat(path) + rayfan = data['rayfan'] + saved_thetas = rayfan['thetas'][0, 0].flatten() + np.testing.assert_allclose(saved_thetas, simple_rayfan.thetas, + atol=1e-10) + + # --- Plotting smoke tests --- + + def test_plot_ray_fan_smoke(self, simple_rayfan): + plt.figure() + simple_rayfan.plot_ray_fan() + plt.close('all') + + def test_plot_time_front_smoke(self, simple_rayfan): + plt.figure() + simple_rayfan.plot_time_front() + plt.close('all') + + def test_plot_time_front_include_lines_smoke(self, simple_rayfan): + plt.figure() + simple_rayfan.plot_time_front(include_lines=True) + plt.close('all') + + def test_plot_depth_v_angle_smoke(self, simple_rayfan): + plt.figure() + simple_rayfan.plot_depth_v_angle() + plt.close('all')