Add SFI groups (#121)

* Update LOS sampling

* Bring SD back to speed

* Update nb

* Update script

* Update requirements

* Update flow chi2 calc

* Update script

* Add support for SFI groups

* Update plot

* Update script

* Update script

csiborgtools/flow/flow_model.py

@@ -254,7 +254,7 @@ class DataLoader:
                 for key in f.keys():
                     arr[key] = f[key][:]
         elif catalogue in ["LOSS", "Foundation", "SFI_gals", "2MTF",
-                           "Pantheon+"]:
+                           "Pantheon+", "SFI_gals_masked", "SFI_groups"]:
             with File(catalogue_fpath, 'r') as f:
                 grp = f[catalogue]
 
@@ -746,10 +746,13 @@ class SD_PV_validation_model(BaseFlowValidationModel):
         self._r_xrange = r_xrange
 
         # Get the various vmapped functions
+        self._f_ptilde_wo_bias = lambda mu, err: calculate_ptilde_wo_bias(r_xrange, mu, err, r2_xrange, True)   # noqa
+        self._f_simps = lambda y: simps(y, dr)                                                                  # noqa
+        self._f_zobs = lambda beta, Vr, vpec_rad: predict_zobs(r_xrange, beta, Vr, vpec_rad, Omega_m)           # noqa
+
         self._vmap_ptilde_wo_bias = vmap(lambda mu, err: calculate_ptilde_wo_bias(r_xrange, mu, err, r2_xrange, True))                  # noqa
         self._vmap_simps = vmap(lambda y: simps(y, dr))
         self._vmap_zobs = vmap(lambda beta, Vr, vpec_rad: predict_zobs(r_xrange, beta, Vr, vpec_rad, Omega_m), in_axes=(None, 0, 0))    # noqa
-        self._vmap_ll_zobs = vmap(lambda zobs, zobs_pred, sigma_v: calculate_ll_zobs(zobs, zobs_pred, sigma_v), in_axes=(0, 0, None))   # noqa
 
         # Distribution of external velocity components
         self._Vext = dist.Uniform(-500, 500)
@@ -757,12 +760,15 @@ class SD_PV_validation_model(BaseFlowValidationModel):
         self._alpha = dist.LogNormal(*lognorm_mean_std_to_loc_scale(1.0, 0.5))     # noqa
         self._beta = dist.Normal(1., 0.5)
         # Distribution of velocity uncertainty sigma_v
-        self._sv = dist.LogNormal(*lognorm_mean_std_to_loc_scale(150, 100))
+        self._sigma_v = dist.LogNormal(*lognorm_mean_std_to_loc_scale(150, 100))   # noqa
+        self._h = dist.LogNormal(*lognorm_mean_std_to_loc_scale(1., 0.25))
+
+        self._Omega_m = Omega_m
 
     def predict_zobs_single(self, **kwargs):
         raise NotImplementedError("This method is not implemented yet.")
 
-    def __call__(self, sample_alpha=False):
+    def __call__(self, sample_alpha=True, sample_beta=True):
         """
         The simple distance NumPyro PV validation model.
 
@@ -771,29 +777,39 @@ class SD_PV_validation_model(BaseFlowValidationModel):
         sample_alpha : bool, optional
             Whether to sample the density bias parameter `alpha`, otherwise
             it is fixed to 1.
+        sample_beta : bool, optional
+            Whether to sample the velocity bias parameter `beta`, otherwise
+            it is fixed to 1.
         """
         Vx = numpyro.sample("Vext_x", self._Vext)
         Vy = numpyro.sample("Vext_y", self._Vext)
         Vz = numpyro.sample("Vext_z", self._Vext)
 
         alpha = numpyro.sample("alpha", self._alpha) if sample_alpha else 1.0
-        beta = numpyro.sample("beta", self._beta)
+        beta = numpyro.sample("beta", self._beta) if sample_beta else 1.0
-        sigma_v = numpyro.sample("sigma_v", self._sv)
+        sigma_v = numpyro.sample("sigma_v", self._sigma_v)
+
+        h = numpyro.sample("h", self._h)
 
         Vext_rad = project_Vext(Vx, Vy, Vz, self._RA, self._dec)
 
-        # Calculate p(r) and multiply it by the galaxy bias
+        def scan_body(ll, i):
-        ptilde = self._vmap_ptilde_wo_bias(self._r_hMpc, self._e2_rhMpc)
+            # Calculate p(r) and multiply it by the galaxy bias
-        ptilde *= self._los_density**alpha
+            ptilde = self._f_ptilde_wo_bias(
+                h * self._r_hMpc[i], h**2 * self._e2_rhMpc[i])
+            ptilde *= self._los_density[i]**alpha
 
-        # Normalization of p(r)
+            # Normalization of p(r)
-        pnorm = self._vmap_simps(ptilde)
+            pnorm = self._f_simps(ptilde)
 
-        # Calculate p(z_obs) and multiply it by p(r)
+            # Calculate p(z_obs) and multiply it by p(r)
-        zobs_pred = self._vmap_zobs(beta, Vext_rad, self._los_velocity)
+            zobs_pred = self._f_zobs(beta, Vext_rad[i], self._los_velocity[i])
-        ptilde *= self._vmap_ll_zobs(self._z_obs, zobs_pred, sigma_v)
+            ptilde *= calculate_ll_zobs(self._z_obs[i], zobs_pred, sigma_v)
 
-        ll = jnp.sum(jnp.log(self._vmap_simps(ptilde) / pnorm))
+            return ll + jnp.log(self._f_simps(ptilde) / pnorm), None
+
+        ll = 0.
+        ll, __ = scan(scan_body, ll, jnp.arange(self.ndata))
         numpyro.factor("ll", ll)
 
 
@@ -1258,7 +1274,7 @@ def get_model(loader, zcmb_max=None, verbose=True):
             los_overdensity[mask], los_velocity[mask], RA[mask], dec[mask],
             zCMB[mask], mB[mask], x1[mask], c[mask], e_mB[mask], e_x1[mask],
             e_c[mask], loader.rdist, loader._Omega_m)
-    elif kind in ["SFI_gals", "2MTF"]:
+    elif kind in ["SFI_gals", "2MTF", "SFI_gals_masked"]:
         keys = ["RA", "DEC", "z_CMB", "mag", "eta", "e_mag", "e_eta"]
         RA, dec, zCMB, mag, eta, e_mag, e_eta = (loader.cat[k] for k in keys)
 
@@ -1271,6 +1287,15 @@ def get_model(loader, zcmb_max=None, verbose=True):
             los_overdensity[mask], los_velocity[mask], RA[mask], dec[mask],
             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"]
+        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)
     else:
         raise ValueError(f"Catalogue `{kind}` not recognized.")
 

notebooks/flow_calibration.py

@@ -35,8 +35,10 @@ def read_samples(catalogue, simname, ksmooth, include_calibration=False,
 
     if catalogue in ["LOSS", "Foundation", "Pantheon+"]:
         alpha_cal, beta_cal, mag_cal, e_mu_intrinsic = [], [], [], []
-    elif catalogue in ["2MTF", "SFI_gals"]:
+    elif catalogue in ["2MTF", "SFI_gals", "SFI_gals_masked"]:
         a, b, e_mu_intrinsic = [], [], []
+    elif catalogue == "SFI_groups":
+        h = []
     else:
         raise ValueError(f"Catalogue {catalogue} not recognized.")
 
@@ -86,6 +88,8 @@ def read_samples(catalogue, simname, ksmooth, include_calibration=False,
                 a.append(f[f"sim_{nsim}/a"][:])
                 b.append(f[f"sim_{nsim}/b"][:])
                 e_mu_intrinsic.append(f[f"sim_{nsim}/e_mu_intrinsic"][:])
+            elif catalogue == "SFI_groups":
+                h.append(f[f"sim_{nsim}/h"][:])
             else:
                 raise ValueError(f"Catalogue {catalogue} not recognized.")
 
@@ -95,8 +99,10 @@ def read_samples(catalogue, simname, ksmooth, include_calibration=False,
 
     if catalogue in ["LOSS", "Foundation", "Pantheon+"]:
         alpha_cal, beta_cal, mag_cal, e_mu_intrinsic = np.hstack(alpha_cal), np.hstack(beta_cal), np.hstack(mag_cal), np.hstack(e_mu_intrinsic)  # noqa
-    elif catalogue in ["2MTF", "SFI_gals"]:
+    elif catalogue in ["2MTF", "SFI_gals", "SFI_gals_masked"]:
         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)
     else:
         raise ValueError(f"Catalogue {catalogue} not recognized.")
 
@@ -114,9 +120,12 @@ def read_samples(catalogue, simname, ksmooth, include_calibration=False,
         if catalogue in ["LOSS", "Foundation", "Pantheon+"]:
             data += [alpha_cal, beta_cal, mag_cal, e_mu_intrinsic]
             names += ["alpha_cal", "beta_cal", "mag_cal", "e_mu_intrinsic"]
-        elif catalogue in ["2MTF", "SFI_gals"]:
+        elif catalogue in ["2MTF", "SFI_gals", "SFI_gals_masked"]:
             data += [a, b, e_mu_intrinsic]
             names += ["a", "b", "e_mu_intrinsic"]
+        elif catalogue == "SFI_groups":
+            data += [h]
+            names += ["h"]
         else:
             raise ValueError(f"Catalogue {catalogue} not recognized.")
 
@@ -162,6 +171,7 @@ def names_to_latex(names, for_corner=False):
                   "beta": r"$\beta$",
                   "Vmag": r"$V_{\rm ext}$",
                   "sigma_v": r"$\sigma_v$",
+                  "h": r"$h$",
                   }
 
     labels = copy(names)

scripts/field_los.py

@@ -58,7 +58,8 @@ def get_los(catalogue_name, simname, comm):
     if comm.Get_rank() == 0:
         folder = "/mnt/extraspace/rstiskalek/catalogs"
 
-        if catalogue_name in ["LOSS", "Foundation", "SFI_gals", "2MTF",
+        if catalogue_name in ["LOSS", "Foundation", "SFI_gals",
+                              "SFI_gals_masked", "SFI_groups", "2MTF",
                               "Pantheon+"]:
             fpath = join(folder, "PV_compilation_Supranta2019.hdf5")
             with File(fpath, 'r') as f:

scripts/field_los.sh

@@ -1,13 +1,13 @@
-nthreads=3
+nthreads=5
 memory=40
-on_login=${1}
+on_login=0
 queue="berg"
 env="/mnt/users/rstiskalek/csiborgtools/venv_csiborg/bin/python"
 file="field_los.py"
 
+simname=${1}
 catalogue=${2}
 nsims="-1"
-simname="csiborg2_main"
 MAS="SPH"
 grid=1024
 

scripts/flow_validation.py

@@ -52,7 +52,7 @@ def get_model(args, nsim_iterator, get_model_kwargs):
     if args.catalogue == "A2":
         fpath = join(folder, "A2.h5")
     elif args.catalogue in ["LOSS", "Foundation", "Pantheon+", "SFI_gals",
-                            "2MTF"]:
+                            "2MTF", "SFI_groups", "SFI_gals_masked"]:
         fpath = join(folder, "PV_compilation_Supranta2019.hdf5")
     else:
         raise ValueError(f"Unknown catalogue: `{args.catalogue}`.")
@@ -107,10 +107,13 @@ def run_model(model, nsteps, nburn, nchains, nsim, dump_folder,
     # Calculate the chi2
     keys = list(thinned_samples.keys())
     nsamples = len(thinned_samples[keys[0]])
-    zobs_mean, zobs_std = model.predict_zobs(thinned_samples)
+    try:
-    nu = model.ndata - len(keys)
+        zobs_mean, zobs_std = model.predict_zobs(thinned_samples)
-    chi2 = [np.sum((zobs_mean[:, i] - model._z_obs)**2 / zobs_std[:, i]**2) / nu  # noqa
+        nu = model.ndata - len(keys)
-            for i in range(nsamples)]
+        chi2 = [np.sum((zobs_mean[:, i] - model._z_obs)**2 / zobs_std[:, i]**2) / nu  # noqa
+                for i in range(nsamples)]
+    except NotImplementedError:
+        chi2 = [0. for _ in range(nsamples)]
 
     gof = csiborgtools.numpyro_gof(model, mcmc, model_kwargs)
 

scripts/flow_validation.sh

@@ -1,5 +1,5 @@
 memory=4
-on_login=1
+on_login=0
 nthreads=${1}
 ksmooth=${2}
 
@@ -7,8 +7,8 @@ queue="berg"
 env="/mnt/users/rstiskalek/csiborgtools/venv_csiborg/bin/python"
 file="flow_validation.py"
 
-catalogue="LOSS"
+catalogue="SFI_groups"
-simname="Carrick2015"
+simname="csiborg2_main"
 
 
 pythoncm="$env $file --catalogue $catalogue --simname $simname --ksmooth $ksmooth"