Multipoles analysis

vega/analysis.py

@@ -183,9 +183,13 @@ def create_global_monte_carlo(self, fiducial_model, seed=None, scale=None, forec
 
         # TODO The corr items need to have an imposed order
         full_data_mask = []
+        full_model_mask = []
         for name in self._corr_items:
             full_data_mask.append(self._data[name].data_mask)
+            full_model_mask.append(self._data[name].model_mask)
+
         full_data_mask = np.concatenate(full_data_mask)
+        full_model_mask = np.concatenate(full_model_mask)
 
         if self._cholesky_global_cov is None:
             masked_cov = self._global_cov[:, full_data_mask]
@@ -195,23 +199,28 @@ def create_global_monte_carlo(self, fiducial_model, seed=None, scale=None, forec
             self._cholesky_global_cov = np.linalg.cholesky(scale * masked_cov)
 
         # TODO The corr items need to have an imposed order
-        masked_fiducial = []
-        for name, data in self._data.items():
-            mask = data.dist_model_coordinates.get_mask_to_other(data.data_coordinates)
-            if data.data_mask.size == fiducial_model[name].size:
-                masked_fiducial.append(fiducial_model[name])
-            elif mask.size == fiducial_model[name].size:
-                masked_fiducial.append(fiducial_model[name][mask])
-            else:
-                raise ValueError('Input fiducial has unknown size. '
-                                 'It must match the data or the model.')
-        masked_fiducial = np.concatenate(masked_fiducial)
+        full_fiducial_model = np.concatenate([fiducial_model[name] for name in self._data])
+        masked_fiducial = full_fiducial_model[full_model_mask]
+        # Naim comment:  unsure about the purpose of the below code
+        # masked_fiducial = []
+        # for name, data in self._data.items():
+        #     mask = data.dist_model_coordinates.get_mask_to_other(data.data_coordinates)
+        #     if data.data_mask.size == fiducial_model[name].size:
+        #         masked_fiducial.append(fiducial_model[name])
+        #     elif mask.size == fiducial_model[name].size:
+        #         masked_fiducial.append(fiducial_model[name][mask])
+        #     else:
+        #         raise ValueError('Input fiducial has unknown size. '
+        #                          'It must match the data or the model.')
+        # masked_fiducial = np.concatenate(masked_fiducial)
 
         if forecast:
-            self.current_mc_mock = masked_fiducial[full_data_mask]
+            self.current_mc_mock = masked_fiducial
+            # self.current_mc_mock = masked_fiducial[full_data_mask]
         else:
             ran_vec = np.random.randn(full_data_mask.sum())
-            self.current_mc_mock = masked_fiducial[full_data_mask] + self._cholesky_global_cov.dot(
+            # self.current_mc_mock = masked_fiducial[full_data_mask] + self._cholesky_global_cov.dot(
+            self.current_mc_mock = masked_fiducial + self._cholesky_global_cov.dot(
                 ran_vec)
 
         return self.current_mc_mock

vega/coordinates.py

@@ -2,6 +2,84 @@
 
 
 class Coordinates:
+    def get_mask_to_other(self, other):
+        raise NotImplementedError
+
+    def get_mask_scale_cuts(self, cuts_config, small_scale_mask=False):
+        """Build mask to apply scale cuts
+
+        Parameters
+        ----------
+        cuts_config : ConfigParser
+            Cuts section from config
+
+        Returns
+        -------
+        Array
+            Mask
+        """
+        # Read the cuts
+        rp_min_cut = cuts_config.getfloat('rp-min', 0.)
+        rp_max_cut = cuts_config.getfloat('rp-max', 300.)
+
+        rt_min_cut = cuts_config.getfloat('rt-min', 0.)
+        rt_max_cut = cuts_config.getfloat('rt-max', 300.)
+
+        r_min_cut = cuts_config.getfloat('r-min', 10.)
+        r_max_cut = cuts_config.getfloat('r-max', 180.)
+
+        mu_min_cut = cuts_config.getfloat('mu-min', -1.)
+        mu_max_cut = cuts_config.getfloat('mu-max', +1.)
+
+        mask = (self.rp_regular_grid > rp_min_cut) & (self.rt_regular_grid > rt_min_cut)
+        mask &= (self.r_regular_grid > r_min_cut)
+
+        if small_scale_mask:
+            return mask
+
+        mask &= (self.rp_regular_grid < rp_max_cut) & (self.rt_regular_grid < rt_max_cut)
+        mask &= (self.r_regular_grid < r_max_cut)
+        mask &= (self.mu_regular_grid > mu_min_cut) & (self.mu_regular_grid < mu_max_cut)
+
+        return mask
+
+    def get_mask_marginalization_scales(self, cuts_config, marginalization_cuts):
+        """Build mask to for bins that are marginalized
+
+        Parameters
+        ----------
+        marginalization_cuts : dict
+            Dictionary with the small-scale marginalization cuts
+
+        Returns
+        -------
+        Array
+            Mask
+        """
+        mask = np.ones_like(self.rp_regular_grid, dtype=bool)
+
+        if 'rtmax' in marginalization_cuts:
+            rtmax = marginalization_cuts['rtmax']
+            mask &= self.rt_regular_grid < rtmax
+        if 'rtmin' in marginalization_cuts:
+            rtmin = marginalization_cuts['rtmin']
+            mask &= self.rt_regular_grid > rtmin
+        if 'rpmax' in marginalization_cuts:
+            rpmax = marginalization_cuts['rpmax']
+            mask &= np.abs(self.rp_regular_grid) < rpmax
+        if 'rpmin' in marginalization_cuts:
+            rpmin = marginalization_cuts['rpmin']
+            mask &= np.abs(self.rp_regular_grid) > rpmin
+
+        if 'all-rmin' in marginalization_cuts:
+            mask = ~self.get_mask_scale_cuts(
+                cuts_config, small_scale_mask=True
+            )
+
+        return mask
+
+
+class RtRpCoordinates(Coordinates):
     """Class to handle Vega coordinate grids
     """
 
@@ -143,75 +221,140 @@ def get_mask_to_other(self, other):
         mask &= (self.rt_grid <= other.rt_max)
         return mask
 
-    def get_mask_scale_cuts(self, cuts_config, small_scale_mask=False):
-        """Build mask to apply scale cuts
+
+class RMuCoordinates(Coordinates):
+    """Class to handle Vega coordinate grids
+    """
+
+    def __init__(
+        self, mu_min, mu_max, r_max, mu_nbins, r_nbins,
+        mu_grid=None, r_grid=None, z_grid=None, z_eff=None,
+        rp_grid=None, rt_grid=None
+    ):
+        """Initialize the coordinate grids.
 
         Parameters
         ----------
-        cuts_config : ConfigParser
-            Cuts section from config
-
-        Returns
-        -------
-        Array
-            Mask
+        mu_min : float
+            Minimum mu
+        mu_max : float
+            Maximum mu
+        r_max : float
+            Maximum r
+        mu_nbins : int
+            Number of mu bins
+        r_nbins : int
+            Number of r bins
+        rp_grid : Array , optional
+            rp grid, by default None
+        rt_grid : Array, optional
+            rt grid, by default None
+        z_grid : Array, optional
+            z grid, by default None
         """
-        # Read the cuts
-        rp_min_cut = cuts_config.getfloat('rp-min', 0.)
-        rp_max_cut = cuts_config.getfloat('rp-max', 300.)
+        self.mu_min = mu_min
+        self.mu_max = mu_max
+        self.r_max = r_max
+        self.mu_nbins = mu_nbins
+        self.r_nbins = r_nbins
 
-        rt_min_cut = cuts_config.getfloat('rt-min', 0.)
-        rt_max_cut = cuts_config.getfloat('rt-max', 300.)
+        self.mu_binsize = (mu_max - mu_min) / mu_nbins
+        self.r_binsize = r_max / r_nbins
 
-        r_min_cut = cuts_config.getfloat('r-min', 10.)
-        r_max_cut = cuts_config.getfloat('r-max', 180.)
+        # Keep it for other parts of the code
+        self.rp_min = r_max * mu_min
+        self.rp_max = r_max * mu_max
+        self.rt_max = r_max
+        self.rp_nbins = r_nbins
+        self.rt_nbins = r_nbins
 
-        mu_min_cut = cuts_config.getfloat('mu-min', -1.)
-        mu_max_cut = cuts_config.getfloat('mu-max', +1.)
+        self.rp_binsize = (self.rp_max - self.rp_min) / self.rp_nbins
+        self.rt_binsize = self.rt_max / self.rt_nbins
 
-        mask = (self.rp_regular_grid > rp_min_cut) & (self.rt_regular_grid > rt_min_cut)
-        mask &= (self.r_regular_grid > r_min_cut)
+        mu_regular_grid = (0.5 + np.arange(mu_nbins)) * self.mu_binsize + self.mu_min
+        r_regular_grid = (0.5 + np.arange(r_nbins)) * self.r_binsize
 
-        if small_scale_mask:
-            return mask
+        r_regular_grid, mu_regular_grid = np.meshgrid(r_regular_grid, mu_regular_grid)
+        self.mu_regular_grid = mu_regular_grid.flatten()
+        self.r_regular_grid = r_regular_grid.flatten()
 
-        mask &= (self.rp_regular_grid < rp_max_cut) & (self.rt_regular_grid < rt_max_cut)
-        mask &= (self.r_regular_grid < r_max_cut)
-        mask &= (self.mu_regular_grid > mu_min_cut) & (self.mu_regular_grid < mu_max_cut)
+        self.mu_grid = self.mu_regular_grid if mu_grid is None else mu_grid
+        self.r_grid = self.r_regular_grid if r_grid is None else r_grid
 
-        return mask
+        if rp_grid is None:
+            self.rp_grid = self.r_grid * self.mu_grid
+        else:
+            self.rp_grid = rp_grid
 
-    def get_mask_marginalization_scales(self, cuts_config, marginalization_cuts):
-        """Build mask to for bins that are marginalized
+        if rt_grid is None:
+            self.rt_grid = self.r_grid * np.sqrt(1.0 - self.mu_grid**2)
+        else:
+            self.rt_grid = rt_grid
+
+        if z_grid is None and z_eff is None:
+            self.z_grid = None
+        else:
+            self.z_grid = z_eff if z_grid is None else z_grid
+
+        self.rp_regular_grid = self.r_regular_grid * self.mu_regular_grid
+        self.rt_regular_grid = self.r_regular_grid * np.sqrt(
+            1.0 - self.mu_regular_grid**2)
+
+    @classmethod
+    def init_from_grids(cls, other, rp_grid, rt_grid, z_grid):
+        """Initialize from other coordinates and new grids
 
         Parameters
         ----------
-        marginalization_cuts : dict
-            Dictionary with the small-scale marginalization cuts
+        other : Coordinates
+            Other coordinates
+        rp_grid : Array
+            rp grid
+        rt_grid : Array
+            rt grid
+        z_grid : Array
+            z grid
 
         Returns
         -------
-        Array
-            Mask
+        Coordinates
+            New coordinates
         """
-        mask = np.ones_like(self.rp_regular_grid, dtype=bool)
+        raise NotImplementedError
 
-        if 'rtmax' in marginalization_cuts:
-            rtmax = marginalization_cuts['rtmax']
-            mask &= self.rt_regular_grid < rtmax
-        if 'rtmin' in marginalization_cuts:
-            rtmin = marginalization_cuts['rtmin']
-            mask &= self.rt_regular_grid > rtmin
-        if 'rpmax' in marginalization_cuts:
-            rpmax = marginalization_cuts['rpmax']
-            mask &= np.abs(self.rp_regular_grid) < rpmax
-        if 'rpmin' in marginalization_cuts:
-            rpmin = marginalization_cuts['rpmin']
-            mask &= np.abs(self.rp_regular_grid) > rpmin
+    @classmethod
+    def init_from_r_mu_grids(cls, r_grid, mu_grid, z_eff=None):
+        """Initialize from r and mu grids
 
-        if 'all-rmin' in marginalization_cuts:
-            mask = ~self.get_mask_scale_cuts(
-                cuts_config, small_scale_mask=True
-            )
+        Parameters
+        ----------
+        r_grid : Array
+            r grid
+        mu_grid : Array
+            mu grid
+
+        Returns
+        -------
+        Coordinates
+            New coordinates
+        """
+        raise NotImplementedError
 
+    def get_mask_to_other(self, other):
+        """Build mask from the current coordinates to the other coordinates.
+
+        Parameters
+        ----------
+        other : Coordinates
+            Other coordinates
+
+        Returns
+        -------
+        Array
+            Mask
+        """
+        assert self.mu_binsize == other.mu_binsize
+        assert self.r_binsize == other.r_binsize
+        mask = (self.mu_grid >= other.mu_min) & (self.mu_grid <= other.mu_max)
+        mask &= (self.r_grid <= other.r_max)
         return mask

vega/correlation_item.py

@@ -47,6 +47,12 @@ def __init__(self, config, model_pk=False):
             if self.tracer2['weights-path'] is None:
                 self.tracer2['weights-path'] = self.tracer1['weights-path']
 
+        self.use_multipoles = config['model'].getboolean('use_multipoles', False)
+        if self.use_multipoles:
+            ells_to_model = config['model'].get('model_multipoles', "0,2")
+            ells_to_model = ells_to_model.split(',')
+            self.ells_to_model = [int(_) for _ in ells_to_model]
+
         self.test_flag = config['data'].getboolean('test', False)
 
         marg_rs = [