Add void density (#154)

* Add load void data

* Add void densities support

csiborgtools/flow/flow_model.py

@@ -221,11 +221,11 @@ class BaseFlowValidationModel(ABC):
         self.z_xrange = jnp.asarray(z_xrange)
         self.mu_xrange = jnp.asarray(mu_xrange)
 
-    def _set_void_data(self, RA, dec, kind, h, order):
+    def _set_void_data(self, RA, dec, profile, kind, h, order):
         """Create the void interpolator."""
-        # h is the MOND model value of local H0 to convert the radial grid to
+        # h is the MOND model value of local H0 to convert the radial grid
-        # Mpc / h
+        # to Mpc / h
-        rLG_grid, void_grid = load_void_data(kind)
+        rLG_grid, void_grid = load_void_data(profile, kind)
         void_grid = jnp.asarray(void_grid, dtype=jnp.float32)
         rLG_grid = jnp.asarray(rLG_grid, dtype=jnp.float32)
 
@@ -244,9 +244,17 @@ class BaseFlowValidationModel(ABC):
                                  model_axis.ra.rad, model_axis.dec.rad)
         phi = jnp.asarray(phi * 180 / np.pi, dtype=jnp.float32)
 
-        self.void_interpolator = lambda rLG: interpolate_void(
+        if kind == "density":
+            void_grid = jnp.log(void_grid)
+            self.void_log_rho_interpolator = lambda rLG: interpolate_void(
                 rLG, self.r_xrange, phi, void_grid, rgrid_min, rgrid_max,
                 rLG_min, rLG_max, order)
+        elif kind == "vrad":
+            self.void_vrad_interpolator = lambda rLG: interpolate_void(
+                rLG, self.r_xrange, phi, void_grid, rgrid_min, rgrid_max,
+                rLG_min, rLG_max, order)
+        else:
+            raise ValueError(f"Unknown kind: `{kind}`.")
 
     @property
     def ndata(self):
@@ -264,21 +272,22 @@ class BaseFlowValidationModel(ABC):
     def los_density(self, **kwargs):
         if self.is_void_data:
             # Currently we have no densities for the void.
-            return jnp.ones((1, self.ndata, len(self.r_xrange)))
+            # return jnp.ones((1, self.ndata, len(self.r_xrange)))
+            raise NotImplementedError("Only log-density for the void.")
 
         return self._los_density
 
     def log_los_density(self, **kwargs):
         if self.is_void_data:
-            # Currently we have no densities for the void.
+            # We want the shape to be `(1, n_objects, n_radial_steps)``.
-            return jnp.zeros((1, self.ndata, len(self.r_xrange)))
+            return self.void_log_rho_interpolator(kwargs["rLG"])[None, ...]
 
         return self._log_los_density
 
     def los_velocity(self, **kwargs):
         if self.is_void_data:
             # We want the shape to be `(1, n_objects, n_radial_steps)``.
-            return self.void_interpolator(kwargs["rLG"])[None, ...]
+            return self.void_vrad_interpolator(kwargs["rLG"])[None, ...]
 
         return self._los_velocity
 
@@ -428,7 +437,7 @@ def sample_calibration(Vext_min, Vext_max, Vmono_min, Vmono_max, beta_min,
         Vext = jnp.zeros(3)
 
         if sample_Vmag_vax:
-            Vext_mag = sample("Vext_axis_mag", Uniform(0.0, Vext_max))
+            Vext_mag = sample("Vext_axis_mag", Uniform(Vext_min, Vext_max))
             # In the direction if (l, b) = (117, 4)
             Vext = Vext_mag * jnp.asarray([0.4035093, -0.01363162, 0.91487396])
 
@@ -524,7 +533,8 @@ class PV_LogLikelihood(BaseFlowValidationModel):
 
         # This must be done before we convert to radians.
         if void_kwargs is not None:
-            self._set_void_data(RA=RA, dec=dec, **void_kwargs)
+            self._set_void_data(RA=RA, dec=dec,  kind="density", **void_kwargs)
+            self._set_void_data(RA=RA, dec=dec,  kind="vrad", **void_kwargs)
 
         # Convert RA/dec to radians.
         RA, dec = np.deg2rad(RA), np.deg2rad(dec)

csiborgtools/flow/void_model.py

@@ -29,35 +29,46 @@ from tqdm import tqdm
 ###############################################################################
 
 
-def load_void_data(kind):
+def load_void_data(profile, kind):
     """
     Load the void velocities from Sergij & Indranil's files for a given kind
     of void profile per observer.
 
     Parameters
     ----------
+    profile : str
+        Void profile to load. One of "exp", "gauss", "mb".
     kind : str
-        The kind of void profile to load. One of "exp", "gauss", "mb".
+        Data kind, either "density" or "vrad".
 
     Returns
     -------
     velocities : 3-dimensional array of shape (nLG, nrad, nphi)
     """
-    if kind not in ["exp", "gauss", "mb"]:
+    if profile not in ["exp", "gauss", "mb"]:
-        raise ValueError("kind must be one of 'exp', 'gauss', 'mb'")
+        raise ValueError("profile must be one of 'exp', 'gauss', 'mb'")
+
+    if kind not in ["density", "vrad"]:
+        raise ValueError("kind must be one of 'density', 'vrad'")
 
     fdir = "/mnt/extraspace/rstiskalek/catalogs/IndranilVoid"
 
-    kind = kind.upper()
+    if kind == "density":
-    fdir = join(fdir, f"{kind}profile")
+        fdir = join(fdir, "rho_data")
+        tag = "rho"
+    else:
+        tag = "v_pec"
+
+    profile = profile.upper()
+    fdir = join(fdir, f"{profile}profile")
 
     files = glob(join(fdir, "*.dat"))
-    rLG = [int(search(rf'v_pec_{kind}profile_rLG_(\d+)', f).group(1))
+    rLG = [int(search(rf'{tag}_{profile}profile_rLG_(\d+)', f).group(1))
            for f in files]
     rLG = np.sort(rLG)
 
-    for i, ri in enumerate(tqdm(rLG, desc="Loading void observer data")):
+    for i, ri in enumerate(tqdm(rLG, desc=f"Loading void `{kind}`observer data")):  # noqa
-        f = join(fdir, f"v_pec_{kind}profile_rLG_{ri}.dat")
+        f = join(fdir, f"{tag}_{profile}profile_rLG_{ri}.dat")
         data_i = np.genfromtxt(f).T
 
         if i == 0:

scripts/flow/flow_validation.py

@@ -232,7 +232,7 @@ def run_model(model, nsteps, nburn,  model_kwargs, out_folder,
 ###############################################################################
 
 def get_distmod_hyperparams(catalogue, sample_alpha, sample_mag_dipole):
-    alpha_min = -1.0
+    alpha_min = -10 if "IndraniVoid" in ARGS.simname else -1.0
     alpha_max = 10.0
 
     if catalogue in ["LOSS", "Foundation", "Pantheon+", "Pantheon+_groups", "Pantheon+_zSN"]:  # noqa
@@ -307,7 +307,7 @@ if __name__ == "__main__":
     num_epochs = 50
     inference_method = "mike"
     mag_selection = None
-    sample_alpha = False if "IndranilVoid_" in ARGS.simname or ARGS.simname == "no_field" else True  # noqa
+    sample_alpha = False if ARGS.simname == "no_field" else True
     sample_beta = None
     no_Vext = None
     sample_Vmag_vax = False
@@ -357,10 +357,10 @@ if __name__ == "__main__":
             raise ValueError(
                 "`IndranilVoid` does not have multiple realisations.")
 
-        kind = ARGS.simname.split("_")[-1]
+        profile = ARGS.simname.split("_")[-1]
-        h = select_void_h(kind)
+        h = select_void_h(profile)
         rdist = np.arange(0, 165, 0.5)
-        void_kwargs = {"kind": kind, "h": h, "order": 1, "rdist": rdist}
+        void_kwargs = {"profile": profile, "h": h, "order": 1, "rdist": rdist}
     else:
         void_kwargs = None
         h = 1.