Implemented analytical counts with rp cut

Author: misharash

pycorr/correlation_function.py

@@ -364,7 +364,7 @@ def is_same(array1, array2):
                 size2 = counts['D1D2'].size2
             if boxsize is None:
                 raise ValueError('boxsize must be provided for analytic two-point counts {}.'.format(label12))
-            counts[label12] = AnalyticTwoPointCounter(mode, edges, boxsize, size1=size1, size2=size2)
+            counts[label12] = AnalyticTwoPointCounter(mode, edges, boxsize, size1=size1, size2=size2, selection_attrs=selection_attrs)
             continue
         if log: logger.info('Computing two-point counts {}.'.format(label12))
         twopoint_kwargs = {'twopoint_weights': twopoint_weights[label12], 'weight_type': weight_type[label12], 'selection_attrs': selection_attrs[label12]}

pycorr/twopoint_counter.py

@@ -1429,7 +1429,7 @@ class AnalyticTwoPointCounter(BaseTwoPointCounter):
     """
     name = 'analytic'
 
-    def __init__(self, mode, edges, boxsize, size1=10, size2=None, los='z'):
+    def __init__(self, mode, edges, boxsize, size1=10, size2=None, los='z', selection_attrs=None):
         r"""
         Initialize :class:`AnalyticTwoPointCounter`, and set :attr:`wcounts` and :attr:`sep`.
 
@@ -1474,17 +1474,42 @@ def __init__(self, mode, edges, boxsize, size1=10, size2=None, los='z'):
         self._set_default_seps()
         self._set_reversible()
         self.wnorm = self.normalization()
+        self._set_selection(selection_attrs)
         self.run()
 
     def run(self):
         """Set analytical two-point counts."""
+        if any(name != "rp" for name in self.selection_attrs.keys()):
+            raise TwoPointCounterError("Analytic random counts not implemented for selections other than rp")
+        rp_selection = self.selection_attrs.get("rp", None)
+        if rp_selection is not None and self.mode not in ['s', 'smu']:
+            raise TwoPointCounterError("Analytic random counts not implemented for rp selection with mode {}".format(self.mode))
         if self.mode == 's':
             v = 4. / 3. * np.pi * self.edges[0]**3
             dv = np.diff(v, axis=0)
+            if rp_selection:
+                v_rpcut = [4. / 3. * np.pi * (self.edges[0]**3 - (self.edges[0]**2 - np.fmin(rp_cut, self.edges[0])**1.5)) for rp_cut in rp_selection]
+                v_rpcut = v_rpcut[1] - v_rpcut[0]
+                dv_rpcut = np.diff(v_rpcut, axis=0) # volume in each bin removed by rp selection
+                dv_total = dv - dv_rpcut
+                dv = np.where(dv_total >= 1e-8 * dv, dv_total, 0) # ensure that the volume is not negative and further remove small positive values that may arise due to rounding errors; assumes that dv is accurate
         elif self.mode == 'smu':
             # we bin in mu
             v = 2. / 3. * np.pi * self.edges[0][:, None]**3 * self.edges[1]
             dv = np.diff(np.diff(v, axis=0), axis=-1)
+            if rp_selection:
+                s, mu = self.edges[0][:, None], self.edges[1][None, :]
+                sin_theta = np.sqrt(1 - mu**2) * np.ones_like(s)
+                v_rpcut = []
+                for rp_cut in rp_selection:
+                    ss = s * np.ones_like(mu); c = ss * sin_theta > rp_cut; ss[c] = rp_cut / sin_theta[c] # this prevents division by zero
+                    r = ss * sin_theta
+                    h = ss * mu
+                    v_rpcut.append(2. / 3. * np.pi * (s**3 - (s**2 - r**2)**1.5 + r**2 * h))
+                v_rpcut = v_rpcut[1] - v_rpcut[0]
+                dv_rpcut = np.diff(np.diff(v_rpcut, axis=0), axis=-1) # volume in each bin removed by rp selection
+                dv_total = dv - dv_rpcut
+                dv = np.where(dv_total >= 1e-8 * dv, dv_total, 0) # ensure that the volume is not negative and further remove small positive values that may arise due to rounding errors; assumes that dv is accurate
         elif self.mode == 'rppi':
             v = np.pi * self.edges[0][:, None]**2 * self.edges[1]
             dv = np.diff(np.diff(v, axis=0), axis=1)