diff --git a/src/pastax/__init__.py b/src/pastax/__init__.py index bc2a075..01f41a0 100644 --- a/src/pastax/__init__.py +++ b/src/pastax/__init__.py @@ -7,6 +7,7 @@ degrees_to_meters, haversine, meters_to_degrees, + wrap_longitude, ) from .grid import Grid from .interpolation import bilinear_interp_2d, linear_interp_1d, spatiotemporal_interp @@ -41,6 +42,7 @@ "haversine", "meters_to_degrees", "degrees_to_meters", + "wrap_longitude", # interpolation "linear_interp_1d", "bilinear_interp_2d", diff --git a/src/pastax/geo.py b/src/pastax/geo.py index 975be14..697ab41 100644 --- a/src/pastax/geo.py +++ b/src/pastax/geo.py @@ -10,6 +10,7 @@ "haversine", "meters_to_degrees", "degrees_to_meters", + "wrap_longitude", ] EARTH_RADIUS: float = 6_371_008.8 @@ -110,3 +111,43 @@ def degrees_to_meters( meters = rad * EARTH_RADIUS lon_scale = jnp.cos(jnp.radians(lat_deg)) return meters.at[..., 0].multiply(lon_scale) + + +def wrap_longitude( + lon: Float[Array, "..."], + period: float = 360.0, + lower: float = -180.0, +) -> Float[Array, "..."]: + r"""Wrap longitude(s) into the half-open window ``[lower, lower + period)``. + + The solver integrates position without normalising it, so a particle that + drifts across the antimeridian accumulates an unbounded longitude + (``181°``, ``200°``, …) — which is the correct *continuous* representation + for a trajectory. This is a **post-processing / display** helper that folds + such longitudes back into a canonical window. The defaults give the + ``[-180, 180)`` convention; pass ``lower=0.0`` for ``[0, 360)``. + + Element-wise and shape-preserving; pure arithmetic, so it is safe under + ``jit`` / ``vmap`` / ``grad``. For a ``[lon, lat]`` trajectory of shape + ``(..., 2)``, wrap only the longitude column, e.g.:: + + traj = traj.at[..., 0].set(wrap_longitude(traj[..., 0])) + + .. warning:: + + Do **not** feed wrapped longitudes into quantities computed by finite + differences along a trajectory (velocity, separation, dispersion): the + ``±period`` discontinuity at the window edge corrupts them. Use the raw + (unwrapped) longitude for those, and wrap only for presentation. + :func:`haversine` is unaffected — it is periodic in longitude. + + Args: + lon: Longitude value(s) in degrees, any shape. + period: Longitude period in degrees (default ``360.0``). + lower: Lower edge of the target window (default ``-180.0``); the window + is ``[lower, lower + period)``. + + Returns: + Longitudes folded into ``[lower, lower + period)``, same shape as ``lon``. + """ + return jnp.mod(lon - lower, period) + lower diff --git a/tests/test_geo.py b/tests/test_geo.py index 6db1829..ffdded0 100644 --- a/tests/test_geo.py +++ b/tests/test_geo.py @@ -9,6 +9,7 @@ degrees_to_meters, haversine, meters_to_degrees, + wrap_longitude, ) @@ -99,3 +100,42 @@ def test_meters_to_degrees_unchanged_at_normal_lat(self): deg = jnp.degrees(disp_m / EARTH_RADIUS) expected = deg.at[0].divide(jnp.cos(jnp.radians(lat))) assert jnp.allclose(got, expected, rtol=1e-6) + + +class TestWrapLongitude: + def test_default_window_minus180_to_180(self): + lon = jnp.array([45.0, 181.0, -181.0, 200.0, 540.0]) + got = wrap_longitude(lon) + # 45->45, 181->-179, -181->179, 200->-160, 540->180->-180 + assert jnp.allclose(got, jnp.array([45.0, -179.0, 179.0, -160.0, -180.0])) + + def test_in_window_values_unchanged(self): + lon = jnp.array([-180.0, -90.0, 0.0, 90.0, 179.999]) + assert jnp.allclose(wrap_longitude(lon), lon) + + def test_lower_zero_gives_0_360(self): + lon = jnp.array([-10.0, 10.0, 370.0, 360.0]) + got = wrap_longitude(lon, lower=0.0) + # -10->350, 10->10, 370->10, 360->0 + assert jnp.allclose(got, jnp.array([350.0, 10.0, 10.0, 0.0])) + + def test_idempotent(self): + lon = jnp.array([181.0, 540.0, -400.0, 12.3]) + once = wrap_longitude(lon) + assert jnp.allclose(wrap_longitude(once), once) + + def test_shape_preserved_and_lonlat_column(self): + # a [lon, lat] trajectory: wrap only the longitude column + traj = jnp.array([[178.0, 10.0], [181.0, 11.0], [184.0, 12.0]]) + wrapped = traj.at[..., 0].set(wrap_longitude(traj[..., 0])) + assert wrapped.shape == traj.shape + assert jnp.allclose(wrapped[:, 0], jnp.array([178.0, -179.0, -176.0])) + assert jnp.allclose(wrapped[:, 1], traj[:, 1]) # latitude untouched + + def test_jit_vmap_grad_safe(self): + assert float(jax.jit(wrap_longitude)(jnp.array(181.0))) == pytest.approx(-179.0) + out = jax.vmap(wrap_longitude)(jnp.array([181.0, -181.0, 200.0])) + assert jnp.allclose(out, jnp.array([-179.0, 179.0, -160.0])) + # derivative is 1 away from the wrap seam, and finite + g = jax.grad(lambda x: wrap_longitude(x))(jnp.array(200.0)) + assert jnp.isfinite(g) and float(g) == pytest.approx(1.0)