Make sure max_n=0 is valid

rmjarvis · Feb 7, 2024 · 93bb668 · 93bb668
1 parent a10290b
commit 93bb668
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 2 deletions.
diff --git a/tests/test_nnn.py b/tests/test_nnn.py
@@ -4677,6 +4677,15 @@ def test_direct_logmultipole_cross12():
     for key in ddd.results:
         np.testing.assert_allclose(ddd2.results[key].weight, ddd.results[key].weight)
 
+    # Make sure max_n=0 works correctly.
+    ddd = treecorr.NNNCorrelation(min_sep=min_sep, max_sep=max_sep, nbins=nbins, max_n=0,
+                                  bin_slop=0, bin_type='LogMultipole')
+    ddd.process(cat1, cat2)
+    assert ddd.ntri.shape[2] == 1
+    np.testing.assert_array_equal(ddd.ntri[:,:,0], true_ntri_122[:,:,max_n])
+    np.testing.assert_allclose(ddd.weight[:,:,0], 75*true_weight_122[:,:,max_n], rtol=1.e-10)
+
+
 
 @timer
 def test_direct_logmultipole_cross():
@@ -4955,6 +4964,14 @@ def test_direct_logmultipole_cross():
     cov = ddd.estimate_cov('sample', func=lambda c: c.weight.ravel())
     cov = ddd.estimate_cov('jackknife', func=lambda c: c.weight.ravel())
 
+    # Make sure max_n=0 works correctly.
+    ddd = treecorr.NNNCorrelation(min_sep=min_sep, max_sep=max_sep, nbins=nbins, max_n=0,
+                                  bin_slop=0, bin_type='LogMultipole')
+    ddd.process(cat1, cat2, cat3)
+    assert ddd.ntri.shape[2] == 1
+    np.testing.assert_array_equal(ddd.ntri[:,:,0], true_ntri_123[:,:,max_n])
+    np.testing.assert_allclose(ddd.weight[:,:,0], true_weight_123[:,:,max_n], rtol=1.e-5)
+
     with assert_raises(ValueError):
         ddd.process(cat1, cat2, cat3, patch_method='global')
 

diff --git a/treecorr/corr3base.py b/treecorr/corr3base.py
@@ -559,7 +559,7 @@ def __init__(self, config=None, *, logger=None, rng=None, **kwargs):
             if self.max_n < 0:
                 raise ValueError("max_n must be non-negative")
             self._ro.min_u = self._ro.max_u = 0
-            self._ro.ubin_size = np.pi / self.max_n  # Use this to compute bu
+            self._ro.ubin_size = 2*np.pi / (2*self.max_n+1)  # Use this to compute bu
             self._ro.min_v = self._ro.max_v = self._ro.nvbins = self._ro.vbin_size = 0
         else:  # pragma: no cover  (Already checked by config layer)
             raise ValueError("Invalid bin_type %s"%self.bin_type)
@@ -1676,7 +1676,7 @@ def _get_minmax_size(self):
                 max_size = 2 * self._max_sep * self.b
             else:
                 # LogMultipole
-                min_size = 2 * self._min_sep * self.b / self.max_n
+                min_size = 2 * self._min_sep * self.b / (2*self.max_n+1)
                 max_size = 2 * self._max_sep * self.b
             return min_size, max_size
         else: