Add wrap_longitude helper for post-processing trajectory longitudes

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",

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

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)