Merge pull request #15 from igmhub/optimal-estimator-continuum-margin…

…alization

Optimal estimator continuum marginalization

examples/lyaxlya_cf.ini

@@ -41,6 +41,9 @@ num_bins_rt_model = 50
 rejection_fraction = 0.99
 get-old-distortion = True
 
+# Continuum order for optimal estimator marginalization
+continuum-order = -1
+
 [cosmology]
 # Whether to use the old, less accurate, picca cosmology module
 use-picca-cosmo = False

lya_2pt/distortion.py

@@ -61,6 +61,7 @@ def compute_dmat(healpix_id):
             globals.rp_max, globals.rt_max
             )
 
+        np.random.seed(globals.seed)
         w = np.random.rand(neighbours.size) > globals.rejection_fraction
         num_pairs += neighbours.size
         num_pairs_used += w.sum()

lya_2pt/global_data.py

@@ -22,3 +22,5 @@
 
 get_old_distortion = None
 rejection_fraction = None
+continuum_order = None
+seed = None

lya_2pt/interface.py

@@ -19,7 +19,7 @@
 accepted_options = [
     "input-nside", "output-nside", "num-cpu", "z_min", "z_max", "rp_min", "rp_max", "rt_max",
     "num_bins_rp", "num_bins_rt", "num_bins_rp_model", "num_bins_rt_model",
-    "rejection_fraction", "get-old-distortion"
+    "rejection_fraction", "get-old-distortion", "continuum-order", "seed"
 ]
 
 defaults = {
@@ -36,7 +36,9 @@
     "num_bins_rp_model": 50,
     "num_bins_rt_model": 50,
     "rejection_fraction": 0.99,
-    "get-old-distortion": True
+    "get-old-distortion": True,
+    "continuum-order": -1,
+    "seed": 0,
 }
 
 
@@ -99,6 +101,8 @@ def __init__(self, config):
         globals.num_bins_rt_model = self.settings.getint('num_bins_rt_model')
         globals.rejection_fraction = self.settings.getfloat('rejection_fraction')
         globals.get_old_distortion = self.settings.getboolean('get-old-distortion')
+        globals.continuum_order = self.settings.getint('continuum-order')
+        globals.seed = self.settings.getint('seed')
 
         self.input_nside = self.settings.getint("input-nside")
         self.output_nside = self.settings.getint("output-nside")
@@ -241,6 +245,8 @@ def run(self, healpix_ids=None, mpi_rank=0):
                     raise ValueError(f'HEALPix ID {id} not found. '
                                      f'Currently stored IDs: {self.healpix_ids}')
 
+            healpix_ids = np.array(healpix_ids)
+
         self.xi_output = {}
         if self.config['compute'].getboolean('compute-correlation', False):
             self.reset_global_counter()

lya_2pt/optimal_estimator.py

@@ -209,12 +209,13 @@ def compute_xi_and_fisher(healpix_id):
             globals.rp_max, globals.rt_max
         )
 
-        tracer1.set_inverse_covariance(xi1d_interp)
+        tracer1.set_inverse_covariance(xi1d_interp, globals.continuum_order)
         tracer1.apply_invcov_to_deltas()
 
+        np.random.seed(globals.seed)
         w = np.random.rand(neighbours.size) > globals.rejection_fraction
         for tracer2, angle in zip(neighbours[w], angles[w]):
-            tracer2.set_inverse_covariance(xi1d_interp)
+            tracer2.set_inverse_covariance(xi1d_interp, globals.continuum_order)
             tracer2.apply_invcov_to_deltas()
             compute_xi_and_fisher_pair(tracer1, tracer2, angle, xi_est, fisher_est)
             tracer2.release_inverse_covariance()

lya_2pt/tracer.py

@@ -136,6 +136,10 @@ def __init__(self, healpix_id, los_id, ra, dec, order, deltas,
         self.invcov = None
         self.is_weighted = False
 
+    @property
+    def size(self):
+        return self.z.size
+
     # def add_neighbours(self, neighbours):
     #     """Add neighbours mask
 
@@ -296,7 +300,7 @@ def apply_z_evol_to_weights(self, redshift_evol, reference_z):
                                                    / self.sum_weights)
             self.term3_norm = (self.weights * self.logwave_term**2).sum()
 
-    def set_inverse_covariance(self, xi1d_interp):
+    def set_inverse_covariance(self, xi1d_interp, cont_order=2):
         if self.invcov is not None:
             return
 
@@ -305,11 +309,23 @@ def set_inverse_covariance(self, xi1d_interp):
 
         delta_lambdas = wavelength[:, None] - wavelength[None, :]
         covariance = xi1d_interp((z_ij, delta_lambdas))
-        covariance[np.diag_indices(self.z.size)] += 1 / self.ivar
-        # covariance[np.diag_indices(tracer.z.size)] += 1 / tracer.weights
+        covariance[np.diag_indices(self.size)] += 1 / self.ivar
 
         self.invcov = np.linalg.inv(covariance)
 
+        if cont_order < 0:
+            return
+
+        template_matrix = np.vander(self.log_lambda, cont_order + 1)
+        U, s, _ = np.linalg.svd(template_matrix, full_matrices=False)
+
+        # Remove small singular valued vectors
+        w = s > 1e-6
+        U = U[:, w]  # shape = (self.size, cont_order + 1)
+        Y = self.invcov @ U
+        # Woodbury formula. Note that U and Y are not square matrices.
+        self.invcov -= Y @ np.linalg.inv(U.T @ Y) @ Y.T
+
     def release_inverse_covariance(self):
         self.invcov = None