From 6d7c3cb4f1d4b2135ea4c4844cc54d7a2d9793e6 Mon Sep 17 00:00:00 2001
From: Joachim Moeyens <moeyensj@gmail.com>
Date: Wed, 22 Nov 2023 12:23:15 -0800
Subject: [PATCH] Assume missing off-diagonal covariance values are zero when
 calculating chi2

---
 adam_core/coordinates/residuals.py            | 19 +++++++
 adam_core/coordinates/tests/test_residuals.py | 51 ++++++++++++++++++-
 2 files changed, 68 insertions(+), 2 deletions(-)

diff --git a/adam_core/coordinates/residuals.py b/adam_core/coordinates/residuals.py
index db17d5ce..d4f36beb 100644
--- a/adam_core/coordinates/residuals.py
+++ b/adam_core/coordinates/residuals.py
@@ -1,3 +1,4 @@
+import warnings
 from typing import List, Tuple, Union
 
 import numpy as np
@@ -293,6 +294,9 @@ def calculate_chi2(
     For residuals with no covariance (all non-diagonal covariance elements are zero) this
     is exactly equivalent to chi2.
 
+    If the off-diagonal covariance elements (the covariate terms) are missing (represented by NaNs),
+    then they will assumed to be 0.0.
+
     Parameters
     ----------
     residuals : `~numpy.ndarray` (N, D)
@@ -310,8 +314,23 @@ def calculate_chi2(
     [1] Carpino, M. et al. (2003). Error statistics of asteroid optical astrometric observations.
         Icarus, 166, 248-270. https://doi.org/10.1016/S0019-1035(03)00051-4
     """
+    # Raise error if any of the diagonal elements are nan
+    if np.any(np.isnan(np.diagonal(covariances, axis1=1, axis2=2))):
+        raise ValueError("Covariance matrix has NaNs on the diagonal.")
+
+    # Warn if any of the non-diagonal elements are nan
+    print(covariances)
+    if np.any(np.isnan(covariances)):
+        warnings.warn(
+            "Covariance matrix has NaNs on the off-diagonal (these will be assumed to be 0.0).",
+            UserWarning,
+        )
+        covariances = np.where(np.isnan(covariances), 0.0, covariances)
+        print(covariances)
+
     W = np.linalg.inv(covariances)
     chi2 = np.einsum("ij,ji->i", np.einsum("ij,ijk->ik", residuals, W), residuals.T)
+    print(chi2)
     return chi2
 
 
diff --git a/adam_core/coordinates/tests/test_residuals.py b/adam_core/coordinates/tests/test_residuals.py
index 7533a899..011b0547 100644
--- a/adam_core/coordinates/tests/test_residuals.py
+++ b/adam_core/coordinates/tests/test_residuals.py
@@ -38,6 +38,49 @@ def test_calculate_chi2():
     np.testing.assert_allclose(calculate_chi2(residuals, covariances), [9])
 
 
+def test_calculate_chi2_missing_diagonal_covariance_values():
+    # Lets make sure we raise an error if the covariance matrix has NaNs on the diagonal
+    residuals = np.array(
+        [
+            [1, 2, 3],
+            [2, 1, 1],
+        ]
+    )
+    covariances = np.array(
+        [
+            [[np.nan, 0, 0], [0, np.nan, 0], [0, 0, 1]],
+            [[np.nan, 0, 0], [0, np.nan, 0], [0, 0, 1]],
+        ]
+    )
+
+    with pytest.raises(
+        ValueError, match=r"Covariance matrix has NaNs on the diagonal."
+    ):
+        calculate_chi2(residuals, covariances)
+
+
+def test_calculate_chi2_missing_off_diagonal_covariance_values():
+    # Lets make sure we raise an error if the covariance matrix has NaNs on the off-diagonal
+    residuals = np.array(
+        [
+            [1, 2, 3],
+            [2, 1, 1],
+        ]
+    )
+    covariances = np.array(
+        [
+            [[1, np.nan, 0], [np.nan, 1, 0], [0, 0, 1]],
+            [[1, np.nan, 0], [np.nan, 1, 0], [0, 0, 1]],
+        ]
+    )
+
+    with pytest.warns(
+        UserWarning,
+        match=r"Covariance matrix has NaNs on the off-diagonal \(these will be assumed to be 0.0\).",
+    ):
+        np.testing.assert_allclose(calculate_chi2(residuals, covariances), [14, 6])
+
+
 def test_calculate_chi2_mahalanobis():
     # Test that the calculate_chi2 is equivalent to the Mahalanobis distance squared
     observed = np.array([[1, 1, 1], [2, 2, 2]])
@@ -357,7 +400,7 @@ def test_Residuals_calculate_missing_covariance_values(
     np.testing.assert_almost_equal(actual_probabilities[3], 1.0)
 
 
-def test_Residuals_calculate_missing_covariate_values(
+def test_Residuals_calculate_missing_off_diagonal_covariance_values(
     observed_array, predicted_array, expected_residuals
 ):
     # Test that Residuals.calculate correctly identifies the number of degrees of freedom,
@@ -423,7 +466,11 @@ def test_Residuals_calculate_missing_covariate_values(
         frame="ecliptic",
     )
 
-    residuals = Residuals.calculate(observed, predicted)
+    with pytest.warns(
+        UserWarning,
+        match=r"Covariance matrix has NaNs on the off-diagonal \(these will be assumed to be 0.0\).",
+    ):
+        residuals = Residuals.calculate(observed, predicted)
 
     assert len(residuals) == 4
     assert residuals.to_array().shape == (4, 6)