fix: add nans parameter to ballet_centroid for handling NaN values (#6)

fix: add nans parameter to ballet_centroid for handling NaN values (#6)

Author: lgrcia

noxfile.py

@@ -7,7 +7,7 @@
 
 @nox.session(python=ALL_PYTHON_VS)
 def test(session):
-    session.install(".[test]")
+    session.install(".[test, jax]")
     session.run("pytest", "-n", "auto", *session.posargs)
 
 

pyproject.toml

@@ -14,7 +14,7 @@ classifiers = [
     "Programming Language :: Python",
     "Programming Language :: Python :: 3",
 ]
-version = "0.0.10b"
+version = "0.0.11b"
 dependencies = [
     "numpy",
     "scikit-image",

src/eloy/centroid.py

@@ -57,7 +57,7 @@ def photutils_centroid(data, coords, cutout=21, centroid_fun=None):
     return centroid_coords
 
 
-def ballet_centroid(data, coords, cnn):
+def ballet_centroid(data, coords, cnn, nans=False):
     """
     Compute centroids for sources using a CNN-based model.
 
@@ -69,11 +69,19 @@ def ballet_centroid(data, coords, cnn):
         Array of (x, y) coordinates for sources.
     cnn : object
         CNN model with a `centroid` method that accepts cutouts.
+    nans : bool, optional
+        If True, NaN values in the output will be replaced with the original coordinates.
 
     Returns
     -------
     np.ndarray
         Array of refined centroid coordinates.
     """
     cutouts = utils.cutout(data, coords, (15, 15), fill_value=np.median(data))
-    return coords - 15 / 2 + cnn.centroid(cutouts)
+    centroids = np.array(coords - 15 / 2 + cnn.centroid(cutouts))
+    if not nans:
+        is_nan = np.isnan(centroids).any(axis=1)
+        centroids[is_nan] = coords[is_nan]
+        return centroids
+    else:
+        return centroids

src/eloy/utils.py

@@ -8,6 +8,7 @@
 from astropy.nddata import Cutout2D
 import numpy as np
 from astropy.nddata import Cutout2D
+from astropy.nddata.utils import NoOverlapError
 
 
 def cutout(data, coords, shape, wcs=None, fill_value=np.nan):
@@ -31,10 +32,14 @@ def cutout(data, coords, shape, wcs=None, fill_value=np.nan):
         Array of cutout images.
     """
     values = []
+    dummy = np.zeros(shape, dtype=data.dtype)
     for coords in coords:
-        cutout = Cutout2D(
-            data, coords, shape, wcs=wcs, fill_value=fill_value, mode="partial"
-        )
+        try:
+            cutout = Cutout2D(
+                data, coords, shape, wcs=wcs, fill_value=fill_value, mode="partial"
+            )
+        except NoOverlapError:
+            cutout = dummy
         values.append(cutout.data)
     return np.array(values)
 

tests/test_centroids.py

@@ -1,6 +1,7 @@
 from eloy import utils, centroid, alignment
 from photutils.centroids import centroid_com
 import numpy as np
+import pytest
 
 
 def test_centroid():
@@ -19,3 +20,17 @@ def test_centroid_out():
     data = np.random.rand(50, 50)
     coords = np.array([[-1, 1], [20, 20]])
     centroid.photutils_centroid(data, coords)
+
+
+def test_ballet_nans():
+    """Test ballet centroid with nans"""
+    pytest.importorskip("jax")
+    from eloy.centroid import ballet_centroid, Ballet
+
+    cnn = Ballet()
+
+    im = np.random.rand(20, 20)
+    coords = np.array([[10.0, 10.0], [-20.0, -20.0]])
+
+    assert not np.any(np.isnan(ballet_centroid(im, coords, cnn, nans=False)))
+    assert np.any(np.isnan(ballet_centroid(im, coords, cnn, nans=True)))