gh-910: consistent definition of displacement (#911)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Patrick J. Roddy <patrickjamesroddy@gmail.com>
Co-authored-by: connoraird <c.aird@ucl.ac.uk>

Author: 4 people

glass/points.py

@@ -510,7 +510,7 @@ def displace(
 
     d = xp.atan2(sa * sg, st * ca - ct * sa * cg)
 
-    return lon - d / math.pi * 180, tp / math.pi * 180
+    return lon + d / math.pi * 180, tp / math.pi * 180
 
 
 def displacement(
@@ -528,9 +528,13 @@ def displacement(
 
     Parameters
     ----------
-    from_lon, from_lat
+    from_lon
         Points before displacement.
-    to_lon, to_lat
+    from_lat
+        Points before displacement.
+    to_lon
+        Points after displacement.
+    to_lat
         Points after displacement.
 
     Returns
@@ -550,16 +554,14 @@ def displacement(
         use_compat=False,
     )
 
-    a = (90.0 - to_lat) / 180 * math.pi
-    b = (90.0 - from_lat) / 180 * math.pi
-    g = (from_lon - to_lon) / 180 * math.pi
+    a = uxpx.radians(from_lat)
+    b = uxpx.radians(to_lat)
+    g = uxpx.radians(to_lon - from_lon)
 
     sa, ca = xp.sin(a), xp.cos(a)
     sb, cb = xp.sin(b), xp.cos(b)
     sg, cg = xp.sin(g), xp.cos(g)
 
-    r = xp.atan2(xp.hypot(sa * cb - ca * sb * cg, sb * sg), ca * cb + sa * sb * cg)
-    x = sb * ca - cb * sa * cg
-    y = sa * sg
-    z = xp.hypot(x, y)
-    return r * (x / z + 1j * y / z)
+    r = xp.atan2(xp.hypot(cb * sg, ca * sb - sa * cb * cg), sa * sb + ca * cb * cg)
+    x = xp.atan2(cb * sg, ca * sb - sa * cb * cg)
+    return r * xp.exp(1j * x)

tests/benchmarks/test_points.py

@@ -16,7 +16,6 @@
 
     from glass._types import UnifiedGenerator
     from tests.fixtures.helper_classes import (
-        Compare,
         DataTransformer,
         GeneratorConsumer,
     )
@@ -103,31 +102,22 @@ def function_to_benchmark() -> list[Any]:
 
 
 @pytest.mark.parametrize(
-    ("r_to_alpha", "expected_lon", "expected_lat"),
+    ("r_to_alpha"),
     [
         # Complex
-        (lambda r: r + 0j, 0.0, 5.0),
-        (lambda r: -r + 0j, 0.0, -5.0),
-        (lambda r: 1j * r, -5.0, 0.0),
-        (lambda r: -1j * r, 5.0, 0.0),
+        (lambda r: r + 0j),
         # Real
-        (lambda r: [r, 0], 0.0, 5.0),
-        (lambda r: [-r, 0], 0.0, -5.0),
-        (lambda r: [0, r], -5.0, 0.0),
-        (lambda r: [0, -r], 5.0, 0.0),
+        (lambda r: [r, 0]),
     ],
 )
 @pytest.mark.skipif(
     not hasattr(glass, "displace"),
     reason="test requires glass.displace",
 )
-def test_displace(  # noqa: PLR0913
+def test_displace(
     benchmark: BenchmarkFixture,
-    compare: type[Compare],
     xpb: ModuleType,
     r_to_alpha: Callable[[float], complex | list[float]],
-    expected_lon: float,
-    expected_lat: float,
 ) -> None:
     """Benchmark for glass.displace with complex values."""
     scale_length = 100_000
@@ -146,8 +136,8 @@ def test_displace(  # noqa: PLR0913
         lat0,
         alpha,
     )
-    compare.assert_allclose(lon, expected_lon, atol=1e-15)
-    compare.assert_allclose(lat, expected_lat, atol=1e-15)
+    assert lon.shape == (scale_length,)
+    assert lat.shape == (scale_length,)
 
 
 @pytest.mark.stable

tests/core/test_points.py

@@ -9,6 +9,7 @@
 
 import glass
 import glass.healpix as hp
+from glass._array_api_utils import xp_additions as uxpx
 
 if TYPE_CHECKING:
     from types import ModuleType
@@ -356,11 +357,11 @@ def test_displace_arg_complex(compare: type[Compare], xp: ModuleType) -> None:
 
     # east
     lon, lat = glass.displace(lon0, lat0, xp.asarray(1j * r))
-    compare.assert_allclose([lon, lat], [-d, 0.0], atol=1e-15)
+    compare.assert_allclose([lon, lat], [d, 0.0], atol=1e-15)
 
     # west
     lon, lat = glass.displace(lon0, lat0, xp.asarray(-1j * r))
-    compare.assert_allclose([lon, lat], [d, 0.0], atol=1e-15)
+    compare.assert_allclose([lon, lat], [-d, 0.0], atol=1e-15)
 
 
 def test_displace_arg_real(compare: type[Compare], xp: ModuleType) -> None:
@@ -382,11 +383,11 @@ def test_displace_arg_real(compare: type[Compare], xp: ModuleType) -> None:
 
     # east
     lon, lat = glass.displace(lon0, lat0, xp.asarray([0, r]))
-    compare.assert_allclose([lon, lat], [-d, 0.0], atol=1e-15)
+    compare.assert_allclose([lon, lat], [d, 0.0], atol=1e-15)
 
     # west
     lon, lat = glass.displace(lon0, lat0, xp.asarray([0, -r]))
-    compare.assert_allclose([lon, lat], [d, 0.0], atol=1e-15)
+    compare.assert_allclose([lon, lat], [-d, 0.0], atol=1e-15)
 
 
 def test_displace_abs(
@@ -434,14 +435,14 @@ def test_displacement(
     data = [
         # equator
         (zero, zero, zero, five, deg5 * north),
-        (zero, zero, -five, zero, deg5 * east),
+        (zero, zero, five, zero, deg5 * east),
         (zero, zero, zero, -five, deg5 * south),
-        (zero, zero, five, zero, deg5 * west),
+        (zero, zero, -five, zero, deg5 * west),
         # pole
         (zero, ninety, ninety * 2, ninety - five, deg5 * north),
-        (zero, ninety, -ninety, ninety - five, deg5 * east),
+        (zero, ninety, ninety, ninety - five, deg5 * east),
         (zero, ninety, zero, ninety - five, deg5 * south),
-        (zero, ninety, ninety, ninety - five, deg5 * west),
+        (zero, ninety, -ninety, ninety - five, deg5 * west),
     ]
 
     # test each displacement individually
@@ -457,3 +458,126 @@ def test_displacement(
         urng.uniform(-90.0, 90.0, size=5),
     )
     assert alpha.shape == (20, 5)
+
+
+def test_displacement_zerodist(
+    compare: type[Compare],
+    urng: UnifiedGenerator,
+    xp: ModuleType,
+) -> None:
+    """Check that zero displacement is computed correctly."""
+    lon = urng.uniform(-180.0, 180.0, size=100)
+    lat = urng.uniform(-90.0, 90.0, size=100)
+
+    compare.assert_allclose(
+        glass.displacement(lon, lat, lon, lat),
+        xp.zeros(100),
+    )
+
+
+def test_displacement_consistent(
+    compare: type[Compare],
+    urng: UnifiedGenerator,
+    xp: ModuleType,
+) -> None:
+    """Check displacement is consistent with displace."""
+    n = 1_000
+
+    # magnitude and angle of displacement we want to achieve
+    r = xp.acos(urng.uniform(-1.0, 1.0, size=n))
+    x = urng.uniform(-math.pi, math.pi, size=n)
+
+    # displace at random positions on the sphere
+    from_lon = urng.uniform(-180.0, 180.0, size=n)
+    from_lat = xp.asin(urng.uniform(-1.0, 1.0, size=n)) / math.pi * 180.0
+
+    # compute the intended displacement
+    alpha_in = r * xp.exp(1j * x)
+
+    # displace random points
+    to_lon, to_lat = glass.displace(from_lon, from_lat, alpha_in)
+
+    # measure displacement
+    alpha_out = glass.displacement(from_lon, from_lat, to_lon, to_lat)
+
+    compare.assert_allclose(alpha_out, alpha_in, atol=0.0, rtol=1e-10)
+
+
+def test_displacement_random(
+    compare: type[Compare],
+    urng: UnifiedGenerator,
+    xp: ModuleType,
+) -> None:
+    """Check displacement for random points."""
+    n = 1_000
+
+    # magnitude and angle of displacement we want to achieve
+    r = xp.acos(urng.uniform(-1.0, 1.0, size=n))
+    x = urng.uniform(-math.pi, math.pi, size=n)
+
+    # displacement at random positions on the sphere
+    theta = xp.acos(urng.uniform(-1.0, 1.0, size=n))
+    phi = urng.uniform(-math.pi, math.pi, size=n)
+
+    # rotation matrix that moves (0, 0, 1) to theta and phi
+    zero = xp.zeros(n)
+    one = xp.ones(n)
+    rot_y = xp.stack(
+        [
+            xp.cos(theta), zero, xp.sin(theta),
+            zero, one, zero,
+            -xp.sin(theta), zero, xp.cos(theta),
+        ],
+        axis=1,
+    )  # fmt: skip
+    rot_z = xp.stack(
+        [
+            xp.cos(phi), -xp.sin(phi), zero,
+            xp.sin(phi), xp.cos(phi), zero,
+            zero, zero, one,
+        ],
+        axis=1,
+    )  # fmt: skip
+    rot = xp.reshape(rot_z, (n, 3, 3)) @ xp.reshape(rot_y, (n, 3, 3))
+
+    # meta-check that rotation works by rotating (0, 0, 1) to theta and phi
+    u = xp.stack(
+        [
+            xp.sin(theta) * xp.cos(phi),
+            xp.sin(theta) * xp.sin(phi),
+            xp.cos(theta),
+        ],
+        axis=1,
+    )
+    compare.assert_allclose(rot @ xp.asarray([0.0, 0.0, 1.0]), u)
+
+    # meta-check that recovering theta and phi from vector works
+    compare.assert_allclose(xp.atan2(xp.hypot(u[:, 0], u[:, 1]), u[:, 2]), theta)
+    compare.assert_allclose(xp.atan2(u[:, 1], u[:, 0]), phi)
+
+    # build the displaced points near (0, 0, 1) and rotate near theta and phi
+    v = xp.stack(
+        [
+            xp.sin(r) * xp.cos(math.pi - x),
+            xp.sin(r) * xp.sin(math.pi - x),
+            xp.cos(r),
+        ],
+        axis=1,
+    )
+    v = rot @ xp.reshape(v, (n, 3, 1))
+    v = xp.reshape(v, (n, 3))
+
+    # compute displaced theta and phi
+    theta_d = xp.atan2(xp.hypot(v[:, 0], v[:, 1]), v[:, 2])
+    phi_d = xp.atan2(v[:, 1], v[:, 0])
+
+    # compute longitude and latitude
+    from_lon = uxpx.degrees(phi)
+    from_lat = 90.0 - uxpx.degrees(theta)
+    to_lon = uxpx.degrees(phi_d)
+    to_lat = 90.0 - uxpx.degrees(theta_d)
+
+    # compute displacement and compare to input
+    alpha_in = r * xp.exp(1j * x)
+    alpha_out = glass.displacement(from_lon, from_lat, to_lon, to_lat)
+    compare.assert_allclose(alpha_out, alpha_in, atol=0.0, rtol=1e-10)