Tests on mock PV catalogues (#123)

* Rename nb

* Add nb

* Update scrtip

* Update script

* Add reading of CB2 mocks

* CB2 mocks support

* Add mock generator

* Add mock support

* Add CB2 mock support

* Update nb for mocks

* Update script

* Update nb

* Update nbs

csiborgtools/flow/flow_model.py

@@ -263,6 +263,13 @@ class DataLoader:
                 arr = np.empty(len(grp["RA"]), dtype=dtype)
                 for key in grp.keys():
                     arr[key] = grp[key][:]
+        elif "CB2_" in catalogue:
+            with File(catalogue_fpath, 'r') as f:
+
+                dtype = [(key, np.float32) for key in f.keys()]
+                arr = np.empty(len(f["RA"]), dtype=dtype)
+                for key in f.keys():
+                    arr[key] = f[key][:]
         else:
             raise ValueError(f"Unknown catalogue: `{catalogue}`.")
 
@@ -854,7 +861,7 @@ class SD_PV_validation_model(BaseFlowValidationModel):
     def predict_zcosmo_from_calibration(self, **kwargs):
         raise NotImplementedError("This method is not implemented yet.")
 
-    def __call__(self, sample_alpha=True, sample_beta=True):
+    def __call__(self, sample_alpha=True, sample_beta=True, sample_h=True):
         """
         The simple distance NumPyro PV validation model.
 
@@ -866,6 +873,9 @@ class SD_PV_validation_model(BaseFlowValidationModel):
         sample_beta : bool, optional
             Whether to sample the velocity bias parameter `beta`, otherwise
             it is fixed to 1.
+        sample_h : bool, optional
+            Whether to sample the location bias parameter `h`, otherwise
+            it is fixed to 1.
         """
         Vx = numpyro.sample("Vext_x", self._Vext)
         Vy = numpyro.sample("Vext_y", self._Vext)
@@ -875,7 +885,7 @@ class SD_PV_validation_model(BaseFlowValidationModel):
         beta = numpyro.sample("beta", self._beta) if sample_beta else 1.0
         sigma_v = numpyro.sample("sigma_v", self._sigma_v)
 
-        h = numpyro.sample("h", self._h)
+        h = numpyro.sample("h", self._h) if sample_h else 1.0
 
         Vext_rad = project_Vext(Vx, Vy, Vz, self._RA, self._dec)
 
@@ -1417,9 +1427,21 @@ def get_model(loader, zcmb_max=None, verbose=True):
             zCMB[mask], mag[mask], eta[mask], e_mag[mask], e_eta[mask],
             loader.rdist, loader._Omega_m)
     elif kind == "SFI_groups":
-        keys = ["RA", "DEC", "z_CMB", "r_hMpc", "e_r_hMpc"]
+        keys = ["RA", "DEC", "zCMB", "r_hMpc", "e_r_hMpc"]
         RA, dec, zCMB, r_hMpc, e_r_hMpc = (loader.cat[k] for k in keys)
 
+        mask = (zCMB < zcmb_max)
+        model = SD_PV_validation_model(
+            los_overdensity[mask], los_velocity[mask], RA[mask], dec[mask],
+            zCMB[mask], r_hMpc[mask], e_r_hMpc[mask], loader.rdist,
+            loader._Omega_m)
+    elif "CB2_" in kind:
+        keys = ["RA", "DEC", "zobs", "r_hMpc"]
+        RA, dec, zCMB, r_hMpc = (loader.cat[k] for k in keys)
+
+        # Set fiducially to be 10% of the distance
+        e_r_hMpc = 0.1 * r_hMpc
+
         mask = (zCMB < zcmb_max)
         model = SD_PV_validation_model(
             los_overdensity[mask], los_velocity[mask], RA[mask], dec[mask],

notebooks/flow_calibration.py

@@ -28,6 +28,11 @@ def read_samples(catalogue, simname, ksmooth, include_calibration=False,
     print(f"\nReading {catalogue} fitted to {simname} with ksmooth = {ksmooth}.", flush=True)  # noqa
     paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
     nsims = paths.get_ics(simname)
+
+    # The last simulation was used to draw the mocks.
+    if catalogue in ["CB2_small", "CB2_large"]:
+        nsims = nsims[:-5]
+
     FDIR_LG = "/mnt/extraspace/rstiskalek/csiborg_postprocessing/peculiar_velocity/observer"  # noqa
 
     Vx, Vy, Vz, beta, sigma_v, alpha = [], [], [], [], [], []
@@ -39,6 +44,8 @@ def read_samples(catalogue, simname, ksmooth, include_calibration=False,
         a, b, e_mu_intrinsic = [], [], []
     elif catalogue == "SFI_groups":
         h = []
+    elif catalogue in ["CB2_small", "CB2_large"]:
+        pass
     else:
         raise ValueError(f"Catalogue {catalogue} not recognized.")
 
@@ -90,6 +97,8 @@ def read_samples(catalogue, simname, ksmooth, include_calibration=False,
                 e_mu_intrinsic.append(f[f"sim_{nsim}/e_mu_intrinsic"][:])
             elif catalogue == "SFI_groups":
                 h.append(f[f"sim_{nsim}/h"][:])
+            elif catalogue in ["CB2_small", "CB2_large"]:
+                pass
             else:
                 raise ValueError(f"Catalogue {catalogue} not recognized.")
 
@@ -103,6 +112,8 @@ def read_samples(catalogue, simname, ksmooth, include_calibration=False,
         a, b, e_mu_intrinsic = np.hstack(a), np.hstack(b), np.hstack(e_mu_intrinsic)  # noqa
     elif catalogue == "SFI_groups":
         h = np.hstack(h)
+    elif catalogue in ["CB2_small", "CB2_large"]:
+        pass
     else:
         raise ValueError(f"Catalogue {catalogue} not recognized.")
 

scripts/field_los.py

@@ -71,6 +71,12 @@ def get_los(catalogue_name, simname, comm):
             with File(fpath, 'r') as f:
                 RA = f["RA"][:]
                 dec = f["DEC"][:]
+        elif "CB2_" in catalogue_name:
+            kind = catalogue_name.split("_")[-1]
+            fname = f"/mnt/extraspace/rstiskalek/catalogs/PV_mock_CB2_17417_{kind}.hdf5"  # noqa
+            with File(fname, 'r') as f:
+                RA = f["RA"][:]
+                dec = f["DEC"][:]
         else:
             raise ValueError(f"Unknown field name: `{catalogue_name}`.")
 

scripts/field_los.sh

@@ -1,5 +1,5 @@
-nthreads=5
-memory=40
+nthreads=6
+memory=42
 on_login=0
 queue="berg"
 env="/mnt/users/rstiskalek/csiborgtools/venv_csiborg/bin/python"

scripts/flow_validation.py

@@ -54,6 +54,9 @@ def get_model(args, nsim_iterator, get_model_kwargs):
     elif args.catalogue in ["LOSS", "Foundation", "Pantheon+", "SFI_gals",
                             "2MTF", "SFI_groups", "SFI_gals_masked"]:
         fpath = join(folder, "PV_compilation_Supranta2019.hdf5")
+    elif "CB2_" in args.catalogue:
+        kind = args.catalogue.split("_")[-1]
+        fpath = join(folder, f"PV_mock_CB2_17417_{kind}.hdf5")
     else:
         raise ValueError(f"Unknown catalogue: `{args.catalogue}`.")
 
@@ -199,7 +202,9 @@ if __name__ == "__main__":
     nsims = paths.get_ics(args.simname)
 
     get_model_kwargs = {"zcmb_max": 0.06}
-    model_kwargs = {"sample_alpha": True}
+    model_kwargs = {"sample_alpha": True, "sample_beta": True}
+    if "CB2_" in args.catalogue:
+        model_kwargs["sample_h"] = False
 
     # Create the dumping folder.
     if comm.Get_rank() == 0:

scripts/flow_validation.sh

@@ -7,7 +7,7 @@ queue="berg"
 env="/mnt/users/rstiskalek/csiborgtools/venv_csiborg/bin/python"
 file="flow_validation.py"
 
-catalogue="SFI_groups"
+catalogue="CB2_small"
 simname="csiborg2_main"