Healpy functions & Root solver

ml4gw/__init__.py

@@ -0,0 +1 @@
+from . import constants

ml4gw/constants.py

@@ -34,3 +34,7 @@
 
 # Speed of light in vacuum (:math:`c`), in gigaparsecs per second
 clightGpc = C / 3.0856778570831e22
+
+
+# Maximum allowed nside for HEALPix
+MAX_NSIDE = 1 << 29

ml4gw/optimize.py

@@ -0,0 +1,49 @@
+from typing import Callable
+
+import torch
+
+
+def root_scalar(
+    f: Callable,
+    x0: float,
+    args: tuple = (),
+    fprime: Callable | None = None,
+    maxiter: int = 100,
+    xtol: float = 1e-6,
+):
+    """
+    Find a root of a scalar function.
+
+    Args:
+        f (callable): The function whose root is to be found.
+        x0 (float): Initial guess.
+        args (tuple, optional): Extra arguments passed to the objective
+        function `f` and its derivative(s).
+        fprime (callable, optional): The derivative of the function.
+        xtol (float, optional): The tolerance for the root.
+        maxiter (int, optional): The maximum number of iterations.
+
+    Returns:
+        dict: A dictionary containing the root, and whether the optimization
+        was successful.
+    """
+    if x0 is None:
+        raise ValueError("x0 must be provided")
+    res = {"converged": False, "root": None}
+    for _ in range(maxiter):
+        fx = f(x0, *args)
+        if fprime is not None:
+            fpx = fprime(x0, *args)
+        else:
+            fpx = (f(x0 + xtol, *args) - f(x0 - xtol, *args)) / (2 * xtol)
+        if abs(fpx) < torch.finfo(torch.float).eps:
+            res["root"] = x0
+            res["converged"] = True
+            return res
+        x1 = x0 - fx / fpx
+        if abs(x1 - x0) < xtol:
+            res["root"] = x1
+            res["converged"] = True
+            return res
+        x0 = x1
+    return res

ml4gw/types.py

@@ -1,5 +1,6 @@
 from typing import Union
 
+import numpy.typing as npt
 from jaxtyping import Float
 from torch import Tensor
 
@@ -11,7 +12,6 @@
 NetworkVertices = Float[Tensor, "num_ifos 3"]
 NetworkDetectorTensors = Float[Tensor, "num_ifos 3 3"]
 
-
 TimeSeries1d = Float[Tensor, "time"]
 TimeSeries2d = Float[TimeSeries1d, "channel"]
 TimeSeries3d = Float[TimeSeries2d, "batch"]
@@ -23,3 +23,10 @@
 FrequencySeries1to3d = Union[
     FrequencySeries1d, FrequencySeries2d, FrequencySeries3d
 ]
+
+HealpixIndex = Union[
+    int,
+    float,
+    Tensor,
+    npt.NDArray,
+]

ml4gw/utils/healpix.py

@@ -0,0 +1,36 @@
+from typing import Tuple
+
+import numpy as np
+import torch
+
+from ..constants import PI
+from ..types import HealpixIndex
+
+
+def nest2uniq(
+    nside: HealpixIndex,
+    ipix: HealpixIndex,
+) -> HealpixIndex:
+    return 4 * nside * nside + ipix
+
+
+def nside2npix(nside: HealpixIndex) -> HealpixIndex:
+    return 12 * nside * nside
+
+
+def nside2pixarea(nside: HealpixIndex, degrees: bool = False) -> HealpixIndex:
+    pixarea = 4 * PI / nside2npix(nside)
+
+    if degrees:
+        pixarea = pixarea * (180.0 / PI) ** 2
+
+    return pixarea
+
+
+def lonlat2thetaphi(
+    lon: HealpixIndex, lat: HealpixIndex
+) -> Tuple[HealpixIndex, HealpixIndex]:
+    if isinstance(lon, torch.Tensor) and isinstance(lat, torch.Tensor):
+        return PI / 2.0 - torch.deg2rad(lat), torch.deg2rad(lon)
+    else:
+        return PI / 2.0 - np.deg2rad(lat), np.deg2rad(lon)

pyproject.toml

@@ -4,7 +4,10 @@ version = "0.6.3"
 description = "Tools for training torch models on gravitational wave data"
 readme = "README.md"
 authors = [
-    "Alec Gunny <[email protected]>", "Ethan Marx <[email protected]>", "Will Benoit <[email protected]>", "Deep Chatterjee <[email protected]"
+    "Alec Gunny <[email protected]>",
+    "Ethan Marx <[email protected]>",
+    "Will Benoit <[email protected]>",
+    "Deep Chatterjee <[email protected]",
 ]
 
 [tool.poetry.dependencies]
@@ -18,7 +21,7 @@ torchaudio = "^2.0"
 numpy = "<2.0.0"
 
 [tool.poetry.group.dev.dependencies]
-coverage = {version = "^7.6.10", extras = ["toml"]}
+coverage = { version = "^7.6.10", extras = ["toml"] }
 pre-commit = "^2.16"
 pytest = "^7.0"
 
@@ -28,6 +31,9 @@ bilby = "^2.1"
 jupyter = "^1.0.0"
 gwpy = "^3.0"
 
+# need healpy to test healpy functions
+healpy = "^1.14"
+
 # versions >= 1.9. fix issue with median calculation 
 # https://github.com/scipy/scipy/issues/15601
 scipy = ">=1.9.0,<1.15"

tests/test_optimize.py

@@ -0,0 +1,28 @@
+import numpy as np
+from scipy.optimize import root_scalar as scipy_root_scalar
+
+from ml4gw.optimize import root_scalar
+
+
+def poly(x):
+    return x**5 - 4 * x**4 + 3 * x**3 - 2 * x**2 + x - 1
+
+
+def poly_deriv(x):
+    return 5 * x**4 - 16 * x**3 + 9 * x**2 - 4 * x + 1
+
+
+def test_root_scalar():
+    tol, x0 = 1e-6, 5.0
+
+    result = root_scalar(f=poly, fprime=poly_deriv, xtol=tol, x0=x0)
+    assert result["converged"]
+    assert np.isclose(poly(result["root"]), 0, atol=tol)
+
+    scipy_result = scipy_root_scalar(
+        f=poly, fprime=poly_deriv, xtol=tol, x0=x0
+    )
+    assert scipy_result.converged
+    assert np.isclose(poly(scipy_result.root), 0, atol=tol)
+
+    assert np.isclose(result["root"], scipy_result.root, atol=tol)

tests/utils/test_healpix.py

@@ -0,0 +1,45 @@
+import healpy as hp
+import pytest
+import torch
+from torch.distributions import Uniform
+
+import ml4gw.utils.healpix as mlhp
+
+
+@pytest.fixture(params=[2, 4, 8, 16, 32, 64, 128, 256, 512, 1024])
+def nside(request):
+    return request.param
+
+
+@pytest.fixture()
+def long(request):
+    dist = Uniform(-180, 180)
+    return dist.sample((10,))
+
+
+@pytest.fixture()
+def lat(request):
+    dist = Uniform(-90, 90)
+    return dist.sample((10,))
+
+
+def test_nside2npix(nside):
+    assert mlhp.nside2npix(nside) == hp.nside2npix(nside)
+
+
+def test_nside2pixarea(nside):
+    assert mlhp.nside2pixarea(nside) == hp.nside2pixarea(nside)
+
+
+def test_lonlat2thetaphi(long, lat):
+    theta_ml4gw, phi_ml4gw = mlhp.lonlat2thetaphi(long, lat)
+    theta, phi = hp.pixelfunc.lonlat2thetaphi(long, lat)
+    assert torch.allclose(theta_ml4gw, theta)
+    assert torch.allclose(phi_ml4gw, phi)
+
+
+def test_nest2uniq(nside):
+    ipix = torch.randint(0, mlhp.nside2npix(nside), (10,))
+    assert torch.allclose(
+        mlhp.nest2uniq(nside, ipix), hp.nest2uniq(nside, ipix)
+    )