Improvements to unit tests & docstrings

pyproject.toml

@@ -72,3 +72,9 @@ docstring-code-format = true  # Formats code in docstrings
 # If you plan a src/ layout later, this keeps installs sane even without code yet
 [tool.setuptools]
 package-dir = {"" = "src"}
+
+[tool.pytest.ini_options]
+filterwarnings = [
+    "ignore:.*Force converting.*:UserWarning",
+]
+

src/astrix/functs.py

@@ -99,9 +99,12 @@ def ecef2geodet(ecef: Array) -> np.ndarray:
 
     warn_if_not_numpy(ecef)
     ecef_np = np.array(ecef)
+    x, y, z = ecef_np[:, 0], ecef_np[:, 1], ecef_np[:, 2]
+    if len(x) == 1:
+        x, y, z = x.tolist(), y.tolist(), z.tolist()
     geodet = np.array(
         t_ecef2geodet.transform(
-            ecef_np[:, 0], ecef_np[:, 1], ecef_np[:, 2], radians=False
+            x, y, z, radians=False
         )
     ).T
     return geodet
@@ -122,15 +125,18 @@ def geodet2ecef(geodet: Array) -> np.ndarray:
 
     warn_if_not_numpy(geodet)
     geodet_np = np.array(geodet)
+    x, y, z = geodet_np[:, 0], geodet_np[:, 1], geodet_np[:, 2]
+    if len(x) == 1:
+        x, y, z = x.tolist(), y.tolist(), z.tolist()
     ecef = np.array(
         t_geodet2ecef.transform(
-            geodet_np[:, 0], geodet_np[:, 1], geodet_np[:, 2], radians=False
+            x, y, z, radians=False
         )
     ).T
-
     return ecef
 
 
+
 @backend_jit(["inverse", "xp"])
 def apply_rot(r: Rotation, v: Array, inverse: bool = False, xp: ArrayNS = np) -> Array:
     """Apply a scipy Rotation to a set of vectors.
@@ -455,14 +461,17 @@ def interp_haversine(
 
 
 def ned_rotation(geodet: Array, xp: Backend = None) -> Rotation:
-    """Get the rotation from ECEF base to the North-East-Down frame at given
-    geodetic locations.
+    """Get the orientation of the local North-East-Down (NED) frame relative 
+    to the ECEF frame at given geodetic locations.
+
+    This rotation represents the NED-to-ECEF orientation. Applying this rotation
+    to a vector in NED coordinates transforms it to ECEF coordinates.
 
     Args:
-        pos_geodet (Array): Nx3 array of lat, long, alt [deg, deg, m]
+        geodet (Array): Nx3 array of lat, long, alt [deg, deg, m]
 
     Returns:
-        Rotation: scipy Rotation object representing the NED frame rotation
+        Rotation: scipy Rotation object representing the NED frame orientation
     """
 
     xp = coerce_ns(xp)
@@ -613,11 +622,16 @@ def refraction_correction_bennett(
 
     Refraction is scaled by an exponential atmosphere model with scale height of 7.5 km.
     Inputs/outputs are degrees.
+
+    Note: Clips refraction correction from 0.01 to 90 degrees elevation.
     """
 
     xp = coerce_ns(backend)
     alpha = 1 - xp.exp(-alt / 7500)  # Scale height of atmosphere ~7.5 km
 
+    # Clip elevation to avoid singularities at or below horizon
+    el = xp.clip(el, 0.001, 90.0)
+
     Rm_ = 1 / xp.tan(xp.deg2rad(el + (7.31 / (el + 4.4))))
     Rm = Rm_ - 0.06 * xp.sin(xp.deg2rad(14.7 * Rm_ + 13))
     return (Rm / 60) * alpha  # in degrees
@@ -636,14 +650,22 @@ def project_velocity_to_az_el(
         r = right
         d = down (positive down)
 
+    Args
+    ----
+    pos_frd : (N, 3) array
+        Position vectors in FRD frame (m).
+    vel_frd : (N, 3) array
+        Velocity vectors in FRD frame (m/s).
+    backend : Backend, optional
+        Array backend to use. Defaults to None.
+
     Returns
     -------
-    az_rate_deg : (...)
-        Azimuth angular rate in degrees/s.
-    el_rate_deg : (...)
-        Elevation angular rate in degrees/s (positive upward).
-    range_rate  : (...)
-        Radial velocity (m/s), positive away from observer.
+    rates : (N, 3) array
+        Angular and radial rates:
+            - Column 0: Azimuth angular rate (deg/s).
+            - Column 1: Elevation angular rate (deg/s, positive upward).
+            - Column 2: Radial velocity (m/s, positive away from observer).
     """
 
     xp = coerce_ns(backend)
@@ -655,8 +677,12 @@ def project_velocity_to_az_el(
     r = pos[..., 1]
     d = pos[..., 2]
 
+    horiz_sq = f**2 + r**2
+    horiz = xp.sqrt(horiz_sq)
+    R_sq = horiz_sq + d**2
+    R = xp.sqrt(R_sq)
+
     az = xp.atan2(r, f)
-    horiz = xp.sqrt(f**2 + r**2)
     el = xp.atan2(-d, horiz)
 
     e_az = xp.stack(
@@ -681,16 +707,18 @@ def project_velocity_to_az_el(
     v_az = xp.sum(vel * e_az, axis=-1)
     v_el = xp.sum(vel * e_el, axis=-1)
 
-    # Range
-    R = xp.sqrt(f**2 + r**2 + d**2)
+    # Range rate (m/s)
+    range_rate = xp.sum(vel * (pos / R[..., xp.newaxis]), axis=-1)
 
     # Angular rates (rad/s)
-    az_rate = v_az / R
+    # az_rate = v_az / (R * cos(el)) = v_az / horiz
+    az_rate = v_az / xp.where(horiz < 1e-6, 1e-6, horiz)
     el_rate = v_el / R
 
     # Convert to deg/s
     deg = 180.0 / xp.asarray(np.pi)
     az_rate_deg = az_rate * deg
     el_rate_deg = el_rate * deg
 
-    return xp.array([az_rate_deg, el_rate_deg])
+    return xp.stack([az_rate_deg, el_rate_deg, range_rate], axis=-1)
+

src/astrix/plots/plot3d.py

@@ -1008,7 +1008,7 @@ def add_legend(self, labels: list[tuple[str, str]], pos_x: float = 0.7):
             t = self.p.add_text(
                 "- " + label.split("-")[-1],
                 position=(x0, y0 - i * line_h),
-                font_size=10,
+                font_size=12,
                 color=rgb,
                 font="courier",
             )

src/astrix/spatial/location.py

@@ -96,7 +96,7 @@ class Point(Location):
 
     Examples
     --------
-    Single static point:
+    Single time-invariant (constant) point:
 
     >>> from astrix import Point
     >>> p1 = Point([-5047162.4, 2568329.79, -2924521.17])  # Brisbane ECEF (m)
@@ -106,7 +106,7 @@ class Point(Location):
     >>> p2.ecef
     array([[-5047162.4, 2568329.79, -2924521.17]])
 
-    Multiple static points:
+    Multiple points:
 
     >>> pts = Point([
     ...     [-5047162.4, 2568329.79, -2924521.17],  # Brisbane
@@ -116,7 +116,7 @@ class Point(Location):
     >>> len(pts)
     2
 
-    Dynamic point with time:
+    Multiple points each with associated time:
 
     >>> from datetime import datetime, timezone
     >>> from astrix import Time

src/astrix/spatial/ray.py

@@ -509,22 +509,48 @@ def convert_to(self, backend: BackendArg) -> Ray:
             unit_converted, origin_converted, time_converted, frame_converted, xp
         )
 
-    def correct_refraction(self, alt: float = 100e3) -> Ray:
+    def correct_refraction(
+        self, alt: float = 100e3, mode: str = "observed_to_true"
+    ) -> Ray:
         """Apply atmospheric refraction correction to the Ray object using Bennett's formula.
-        Altitude sets the exponential scale height (metres) for the correction.
+        The ``alt`` parameter is the altitude in metres supplied to the model's
+        exponential attenuation term; it does not set the scale height itself.
+
+        Args:
+            alt (float, optional): Altitude in metres at which to evaluate the
+                correction scaling. A value of 0 applies the sea-level scaling,
+                and larger values reduce the correction according to the underlying
+                Bennett model. Defaults to 100e3 (100 km).
+            mode (str, optional): Direction of correction. Either 'observed_to_true'
+                (subtract correction from apparent elevation) or 'true_to_observed'
+                (add correction to true elevation). Defaults to 'observed_to_true'.
 
         Returns:
             Ray: New Ray object with refraction-corrected direction vectors.
 
         Notes:
             - Assumes standard atmospheric conditions.
         """
-
+
+        if mode not in ["observed_to_true", "true_to_observed"]:
+            raise ValueError(
+                "mode must be either 'observed_to_true' or 'true_to_observed'"
+            )
+
         rays_ned = self.to_ned()
         head_el = rays_ned.az_el
-        el_corrected = head_el[:,1] + refraction_correction_bennett(head_el[:, 1], alt, backend=self._xp)
+
+        correction = refraction_correction_bennett(
+            head_el[:, 1], alt, backend=self._xp
+        )
+
+        if mode == "observed_to_true":
+            el_corrected = head_el[:, 1] - correction
+        else:
+            el_corrected = head_el[:, 1] + correction
+
         dir_corrected = vec_from_az_el(
-            self._xp.c_[head_el[:, 0], el_corrected], backend=self._xp
+            self._xp.stack((head_el[:, 0], el_corrected), axis=1), backend=self._xp
         )
         rays_ned_corrected = Ray._constructor(
             dir_corrected,

tests/core/test_functs.py

@@ -2,6 +2,7 @@
 import warnings
 import numpy as np
 from scipy.spatial.transform import Rotation as R
+from astrix._backend_utils import _convert_rot_backend
 
 from astrix.functs import (
     great_circle_distance,
@@ -190,27 +191,46 @@ def test_angles_and_vectors_roundtrip(xp):
 
 
 def test_apply_rot_and_refraction(xp):
-    if xp.__name__.startswith("jax"):
-        pytest.skip("Rotation.as_matrix returns NumPy; skip inverse check on JAX backend")
-
     rot = R.from_euler("z", 90, degrees=True)
-    vec = np.asarray([[1.0, 0.0, 0.0]])
-    out = apply_rot(rot, vec, xp=np)
-    assert np.allclose(out, np.asarray([[0.0, 1.0, 0.0]]))
-    inv = apply_rot(rot, out, inverse=True, xp=np)
-    assert np.allclose(inv, vec)
+    rot = _convert_rot_backend(rot, xp)
+    vec = xp.asarray([[1.0, 0.0, 0.0]])
+    out = apply_rot(rot, vec, xp=xp)
+    assert xp.allclose(out, xp.asarray([[0.0, 1.0, 0.0]]))
+    inv = apply_rot(rot, out, inverse=True, xp=xp)
+    assert xp.allclose(inv, vec)
+
+    correction = refraction_correction_bennett(xp.asarray([10.0]), alt=0.0, backend=xp)
+    assert xp.allclose(correction, xp.asarray([0.0]))
+    high_alt = refraction_correction_bennett(xp.asarray([10.0]), alt=20000.0, backend=xp)
+    assert xp.all(high_alt > correction)
 
-    correction = refraction_correction_bennett(np.asarray([10.0]), alt=0.0, backend=np)
-    assert np.allclose(correction, 0.0)
-    high_alt = refraction_correction_bennett(np.asarray([10.0]), alt=20000.0, backend=np)
-    assert (high_alt > correction).all()
 
 
 def test_project_velocity_to_az_el(xp):
+    # Test case 1: Pure horizontal motion at el=0
     pos = xp.asarray([[1.0, 0.0, 0.0]])
     vel = xp.asarray([[0.0, 1.0, 0.0]])  # yawing right
 
     rates = project_velocity_to_az_el(pos, vel, backend=xp)
-    assert rates.shape == (2, 1)
-    assert xp.allclose(rates[1], 0.0, atol=1e-8)  # no elevation rate
-    assert xp.allclose(rates[0], xp.asarray([57.2957795]), atol=1e-6)
+    assert rates.shape == (1, 3)
+    assert xp.allclose(rates[0, 1], 0.0, atol=1e-8)  # no elevation rate
+    assert xp.allclose(rates[0, 2], 0.0, atol=1e-8)  # no range rate
+    assert xp.allclose(rates[0, 0], xp.asarray([57.2957795]), atol=1e-6)
+
+    # Test case 2: Motion at el=-45 degrees (looking down)
+    # R=sqrt(2), horiz=1, el=-45 deg
+    pos2 = xp.asarray([[1.0, 0.0, 1.0]])
+    vel2 = xp.asarray([[0.0, 1.0, 0.0]])  # motion in 'right' direction
+    # v_az = 1, horiz = 1 => az_rate = 1 rad/s
+    rates2 = project_velocity_to_az_el(pos2, vel2, backend=xp)
+    assert xp.allclose(rates2[0, 0], xp.asarray([57.2957795]), atol=1e-6)
+    assert xp.allclose(rates2[0, 1], 0.0, atol=1e-8)
+    assert xp.allclose(rates2[0, 2], 0.0, atol=1e-8)
+
+    # Test case 3: Radial motion
+    pos3 = xp.asarray([[1.0, 0.0, 0.0]])
+    vel3 = xp.asarray([[2.0, 0.0, 0.0]])  # moving away
+    rates3 = project_velocity_to_az_el(pos3, vel3, backend=xp)
+    assert xp.allclose(rates3[0, 0], 0.0, atol=1e-8)
+    assert xp.allclose(rates3[0, 1], 0.0, atol=1e-8)
+    assert xp.allclose(rates3[0, 2], 2.0, atol=1e-8)

tests/core/test_ray.py

@@ -4,6 +4,7 @@
 import numpy as np
 from .helpers import to_text
 from astrix import Ray, Time, Point, Path, Frame, RotationSequence
+from astrix.spatial.frame import ned_frame
 from scipy.spatial.transform import Rotation as R
 
 @pytest.mark.filterwarnings("ignore:.*deprecated.*:DeprecationWarning")
@@ -124,3 +125,26 @@ def test_ray_moving_frame(xp):
     )
     assert np.allclose(ray_ecef.unit_rel, xp.array([[0.0, 1.0, 0.0],
                                                     [0, -0.70710678, 0.70710678]]))
+
+
+def test_ray_refraction(xp):
+    # Create a ray at some location
+    origin = Point.from_geodet([0.0, 0.0, 0.0], backend=xp)
+    # Pointing at elevation 10 degrees in NED frame
+    ray = Ray.from_az_el(xp.asarray([[0.0, 10.0]]), frame=ned_frame(origin), backend=xp)
+
+    # Observed to true: apparent is 10, true should be < 10
+    ray_true = ray.correct_refraction(mode="observed_to_true")
+    assert ray_true.to_ned().az_el[0, 1] < 10.0
+
+    # True to observed: true is 10, apparent should be > 10
+    ray_obs = ray.correct_refraction(mode="true_to_observed")
+    assert ray_obs.to_ned().az_el[0, 1] > 10.0
+
+    # Check default mode is observed_to_true
+    ray_default = ray.correct_refraction()
+    assert xp.allclose(ray_default.unit_rel, ray_true.unit_rel)
+
+    # Check error on invalid mode
+    with pytest.raises(ValueError):
+        ray.correct_refraction(mode="invalid")