Change check of shape and test for different lengths of vectors

cgalelli · Sep 13, 2024 · 027b38d · 027b38d
1 parent b1c5b45
commit 027b38d
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 11 deletions.
diff --git a/gammapy/stats/tests/test_variability.py b/gammapy/stats/tests/test_variability.py
@@ -278,6 +278,7 @@ def test_dcf():
             [1.7e-11, 2e-12],
             [0.4e-11, 1.8e-12],
             [1e-11, 1.5e-12],
+            [1.1e-11, 0.7e-12],
         ]
     )
 
@@ -297,6 +298,7 @@ def test_dcf():
             [0.07e-11, 0.14e-12],
             [0.12e-11, 0.18e-12],
             [0.1e-11, 0.13e-12],
+            [0.4e-11, 0.2e-12],
         ]
     )
 
@@ -308,7 +310,14 @@ def test_dcf():
     )
     time2 = (
         np.array(
-            [6.31157019e08, 6.31160619e08, 6.31164219e08, 6.31171419e08, 6.31178419e08]
+            [
+                6.31156919e08,
+                6.31160519e08,
+                6.31163819e08,
+                6.31171419e08,
+                6.31178219e08,
+                6.31178819e08,
+            ]
         )
         * u.s
     )
@@ -320,21 +329,21 @@ def test_dcf():
     assert_allclose(
         dcf,
         [
-            [-0.391848, -0.883385],
-            [-0.161262, -0.111955],
-            [0.207653, 0.270868],
-            [-0.312846, -0.429436],
+            [-0.332361, -1.009638],
+            [-0.065639, 0.313054],
+            [0.215329, 0.133132],
+            [-0.373906, -0.56788],
         ],
         rtol=1e-5,
     )
 
     assert_allclose(
         dcf_err,
         [
-            [0.86587, 0.67408],
-            [0.318534, 0.248418],
-            [0.360007, 0.360725],
-            [0.754145, 1.38026],
+            [0.34881, 0.553113],
+            [0.236548, 0.174538],
+            [0.401313, 0.361757],
+            [0.820525, 1.204655],
         ],
         rtol=1e-5,
     )

diff --git a/gammapy/stats/variability.py b/gammapy/stats/variability.py
@@ -323,8 +323,10 @@ def discrete_correlation(flux1, flux_err1, flux2, flux_err2, time1, time2, tau,
     flux1 = np.rollaxis(flux1, axis, 0)
     flux2 = np.rollaxis(flux2, axis, 0)
 
-    if (np.squeeze(flux1.shape) != np.squeeze(flux2.shape)).all():
-        raise ValueError("flux1 and flux2 must have the same squeezed shape.")
+    if np.squeeze(flux1).shape[1:] != np.squeeze(flux2).shape[1:]:
+        raise ValueError(
+            "flux1 and flux2 must have the same squeezed shape, apart from the chosen axis."
+        )
 
     tau = tau.to(time1.unit)
     time2 = time2.to(time1.unit)