Add monopole to VF (#136)

* Add monopole sampling

* Add monopole

* Update script

* Bring back zmax

* Add monopole

* Add condition in case division by 0

* Add monopole to field

* Update nb

* Update script

* Update scripts

* Simplify docs

csiborgtools/field/density.py

@@ -301,8 +301,11 @@ def radial_velocity(rho_vel, observer_velocity):
                 vy = rho_vel[1, i, j, k] - vy0
                 vz = rho_vel[2, i, j, k] - vz0
 
-                radvel[i, j, k] = ((px * vx + py * vy + pz * vz)
+                r = (px**2 + py**2 + pz**2)**0.5
-                                   / numpy.sqrt(px**2 + py**2 + pz**2))
+
+                # There will be at most one cell with r = 0
+                if r > 0:
+                    radvel[i, j, k] = (px * vx + py * vy + pz * vz) / r
     return radvel
 
 

csiborgtools/flow/flow_model.py

@@ -596,12 +596,18 @@ def sample_TFR(e_mu_min, e_mu_max, a_mean, a_std, b_mean, b_std):
 ###############################################################################
 
 
-def sample_calibration(Vext_std, alpha_min, alpha_max, beta_min, beta_max,
+def sample_calibration(Vext_min, Vext_max, Vmono_min, Vmono_max,
-                       sigma_v_min, sigma_v_max, sample_alpha, sample_beta):
+                       alpha_min, alpha_max, beta_min, beta_max, sigma_v_min,
+                       sigma_v_max, sample_Vmono, sample_alpha, sample_beta):
     """Sample the flow calibration."""
-    Vext = sample("Vext", Normal(0, Vext_std).expand([3]))
+    Vext = sample("Vext", Uniform(Vext_min, Vext_max).expand([3]))
     sigma_v = sample("sigma_v", Uniform(sigma_v_min, sigma_v_max))
 
+    if sample_Vmono:
+        Vmono = sample("Vmono", Uniform(Vmono_min, Vmono_max))
+    else:
+        Vmono = 0.0
+
     if sample_alpha:
         alpha = sample("alpha", Uniform(alpha_min, alpha_max))
     else:
@@ -612,7 +618,7 @@ def sample_calibration(Vext_std, alpha_min, alpha_max, beta_min, beta_max,
     else:
         beta = 1.0
 
-    return Vext, sigma_v, alpha, beta
+    return Vext, Vmono, sigma_v, alpha, beta
 
 
 ###############################################################################
@@ -669,7 +675,7 @@ class PV_validation_model(BaseFlowValidationModel):
     def __call__(self, calibration_hyperparams, distmod_hyperparams,
                  store_ll_all=False):
         """NumPyro PV validation model."""
-        Vext, sigma_v, alpha, beta = sample_calibration(**calibration_hyperparams)  # noqa
+        Vext, Vmono, sigma_v, alpha, beta = sample_calibration(**calibration_hyperparams)  # noqa
         cz_err = jnp.sqrt(sigma_v**2 + self.e2_cz_obs)
         Vext_rad = project_Vext(Vext[0], Vext[1], Vext[2], self.RA, self.dec)
 
@@ -694,7 +700,7 @@ class PV_validation_model(BaseFlowValidationModel):
         pnorm = simpson(ptilde, dx=self.dr, axis=-1)
 
         # Calculate z_obs at each distance. Shape is (n_sims, ndata, nxrange)
-        vrad = beta * self.los_velocity + Vext_rad[None, :, None]
+        vrad = beta * self.los_velocity + Vext_rad[None, :, None] + Vmono
         zobs = (1 + self.z_xrange[None, None, :]) * (1 + vrad / SPEED_OF_LIGHT) - 1  # noqa
 
         ptilde *= calculate_likelihood_zobs(self.z_obs, zobs, cz_err)

notebooks/flow/flow_bulk.py

@@ -43,6 +43,24 @@ def read_enclosed_density(simname):
     return r, overdensity
 
 
+def read_enclosed_monopole(simname, rmax=155):
+    fname = join(FDIR, f"enclosed_mass_{simname}.npz")
+
+    if exists(fname):
+        data = np.load(fname)
+    else:
+        raise FileNotFoundError(f"File {fname} not found.")
+
+    r = data["distances"]
+    vmono = data["cumulative_velocity_mono"]
+
+    mask = r < rmax
+    r = r[mask]
+    vmono = vmono[..., mask]
+
+    return r, vmono
+
+
 def read_enclosed_flow(simname):
     fname = join(FDIR, f"enclosed_mass_{simname}.npz")
 

scripts/field_bulk.py

@@ -25,12 +25,16 @@ from datetime import datetime
 from gc import collect
 from os.path import join
 
-import csiborgtools
 import numpy as np
 from astropy import units as u
 from astropy.coordinates import CartesianRepresentation, SkyCoord
 from tqdm import tqdm
 
+import csiborgtools
+from csiborgtools import fprint
+from csiborgtools.field import (field_enclosed_mass, particles_enclosed_mass,
+                                particles_enclosed_momentum)
+
 ###############################################################################
 #                Read in information about the simulation                     #
 ###############################################################################
@@ -41,23 +45,7 @@ def t():
 
 
 def get_reader(simname, paths, nsim):
-    """
+    """Get the appropriate snapshot reader for the simulation."""
-    Get the appropriate snaspshot reader for the simulation.
-
-    Parameters
-    ----------
-    simname : str
-        Name of the simulation.
-    paths : csiborgtools.read.Paths
-        Paths object.
-    nsim : int
-        Simulation index.
-
-    Returns
-    -------
-    reader : instance of csiborgtools.read.BaseSnapshot
-        Snapshot reader.
-    """
     if simname == "csiborg1":
         nsnap = max(paths.get_snapshots(nsim, simname))
         reader = csiborgtools.read.CSiBORG1Snapshot(nsim, nsnap, paths,
@@ -73,49 +61,26 @@ def get_reader(simname, paths, nsim):
 
 
 def get_particles(reader, boxsize, get_velocity=True, verbose=True):
-    """
+    """Get the snapshot particles."""
-    Get the distance of particles from the center of the box and their masses.
+    fprint("reading coordinates and calculating radial distance.", verbose)
-
-    Parameters
-    ----------
-    reader : instance of csiborgtools.read.BaseSnapshot
-        Snapshot reader.
-    boxsize : float
-        Box size in Mpc / h.
-    get_velocity : bool, optional
-        Whether to also return the velocity of particles.
-    verbose : bool
-        Verbosity flag.
-
-    Returns
-    -------
-    dist : 1-dimensional array
-        Distance of particles from the center of the box.
-    mass : 1-dimensional array
-        Mass of particles.
-    vel : 2-dimensional array, optional
-        Velocity of particles.
-    """
-    if verbose:
-        print(f"{t()},: reading coordinates and calculating radial distance.")
     pos = reader.coordinates()
     dtype = pos.dtype
     pos -= boxsize / 2
     dist = np.linalg.norm(pos, axis=1).astype(dtype)
+    collect()
+
+    if get_velocity:
+        fprint("reading velocities.", verbose)
+        vel = reader.velocities().astype(dtype)
+        vrad = np.sum(pos, vel, axis=1) / dist
+
     del pos
     collect()
 
-    if verbose:
+    fprint("reading masses.")
-        print(f"{t()}: reading masses.")
     mass = reader.masses()
 
-    if get_velocity:
+    fprint("sorting arrays.")
-        if verbose:
-            print(f"{t()}: reading velocities.")
-        vel = reader.velocities().astype(dtype)
-
-    if verbose:
-        print(f"{t()}: sorting arrays.")
     indxs = np.argsort(dist)
     dist = dist[indxs]
     mass = mass[indxs]
@@ -126,7 +91,7 @@ def get_particles(reader, boxsize, get_velocity=True, verbose=True):
     collect()
 
     if get_velocity:
-        return dist, mass, vel
+        return dist, mass, vel, vrad
 
     return dist, mass
 
@@ -140,7 +105,7 @@ def main_borg(args, folder):
     paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
     boxsize = csiborgtools.simname2boxsize(args.simname)
     nsims = paths.get_ics(args.simname)
-    distances = np.linspace(0, boxsize / 2, 101)[1:]
+    distances = np.linspace(0, boxsize / 2, 101)
 
     cumulative_mass = np.zeros((len(nsims), len(distances)))
     cumulative_volume = np.zeros((len(nsims), len(distances)))
@@ -154,7 +119,7 @@ def main_borg(args, folder):
         else:
             raise ValueError(f"Unknown simname: `{args.simname}`.")
 
-        cumulative_mass[i, :], cumulative_volume[i, :] = csiborgtools.field.field_enclosed_mass(  # noqa
+        cumulative_mass[i, :], cumulative_volume[i, :] = field_enclosed_mass(
             field, distances, boxsize)
 
     # Finally save the output
@@ -174,35 +139,42 @@ def main_csiborg(args, folder):
     cumulative_mass = np.zeros((len(nsims), len(distances)))
     mass135 = np.zeros(len(nsims))
     masstot = np.zeros(len(nsims))
+    cumulative_vel_mono = np.zeros((len(nsims), len(distances)))
     cumulative_velocity = np.zeros((len(nsims), len(distances), 3))
 
     for i, nsim in enumerate(tqdm(nsims, desc="Simulations")):
         reader = get_reader(args.simname, paths, nsim)
-        rdist, mass, vel = get_particles(reader, boxsize,  verbose=False)
+        rdist, mass, vel, vrad = get_particles(reader, boxsize,  verbose=True)
 
         # Calculate masses
-        cumulative_mass[i, :] = csiborgtools.field.particles_enclosed_mass(
+        cumulative_mass[i, :] = particles_enclosed_mass(rdist, mass, distances)
-            rdist, mass, distances)
+        mass135[i] = particles_enclosed_mass(rdist, mass, [135])[0]
-        mass135[i] = csiborgtools.field.particles_enclosed_mass(
-            rdist, mass, [135])[0]
         masstot[i] = np.sum(mass)
 
+        # Calculate monopole momentum
+        cumulative_vel_mono[i] = particles_enclosed_mass(
+            rdist, vrad * mass, distances)
+
         # Calculate velocities
-        cumulative_velocity[i, ...] = csiborgtools.field.particles_enclosed_momentum(  # noqa
+        cumulative_velocity[i, ...] = particles_enclosed_momentum(
             rdist, mass, vel, distances)
-        for j in range(3):  # Normalize the momentum to get velocity out of it.
+
+        # Normalize the momentum to get velocity out of it.
+        for j in range(3):
             cumulative_velocity[i, :, j] /= cumulative_mass[i, :]
+        cumulative_vel_mono[i, ...] /= cumulative_mass[i, ...]
 
     # Finally save the output
     fname = f"enclosed_mass_{args.simname}.npz"
     fname = join(folder, fname)
     np.savez(fname, enclosed_mass=cumulative_mass, mass135=mass135,
              masstot=masstot, distances=distances,
-             cumulative_velocity=cumulative_velocity)
+             cumulative_velocity=cumulative_velocity,
+             cumulative_velocity_mono=cumulative_vel_mono)
 
 
 def main_from_field(args, folder):
-    """Bulk flow in the Manticore boxes provided by Stuart."""
+    """Bulk flows in 3D fields"""
     paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
     boxsize = csiborgtools.simname2boxsize(args.simname)
     nsims = paths.get_ics(args.simname)
@@ -210,6 +182,7 @@ def main_from_field(args, folder):
 
     cumulative_mass = np.zeros((len(nsims), len(distances)))
     cumulative_volume = np.zeros((len(nsims), len(distances)))
+    cumulative_vel_mono = np.zeros((len(nsims), len(distances)))
     cumulative_vel_x = np.zeros((len(nsims), len(distances)))
     cumulative_vel_y = np.zeros_like(cumulative_vel_x)
     cumulative_vel_z = np.zeros_like(cumulative_vel_x)
@@ -224,18 +197,30 @@ def main_from_field(args, folder):
             raise ValueError(f"Unknown simname: `{args.simname}`.")
 
         density_field = reader.density_field()
-        velocity_field = reader.velocity_field()
+        cumulative_mass[i, :], cumulative_volume[i, :] = field_enclosed_mass(
-
-        cumulative_mass[i, :], cumulative_volume[i, :] = csiborgtools.field.field_enclosed_mass(  # noqa
             density_field, distances, boxsize, verbose=False)
+        del density_field
+        collect()
 
-        cumulative_vel_x[i, :], __ = csiborgtools.field.field_enclosed_mass(
+        velocity_field = reader.velocity_field()
+        radial_velocity_field = csiborgtools.field.radial_velocity(
+            velocity_field, [0., 0., 0.])
+
+        cumulative_vel_mono[i, :], __ = field_enclosed_mass(
+            radial_velocity_field, distances, boxsize, verbose=False)
+        del radial_velocity_field
+        collect()
+
+        cumulative_vel_x[i, :], __ = field_enclosed_mass(
             velocity_field[0], distances, boxsize, verbose=False)
-        cumulative_vel_y[i, :], __ = csiborgtools.field.field_enclosed_mass(
+        cumulative_vel_y[i, :], __ = field_enclosed_mass(
             velocity_field[1], distances, boxsize, verbose=False)
-        cumulative_vel_z[i, :], __ = csiborgtools.field.field_enclosed_mass(
+        cumulative_vel_z[i, :], __ = field_enclosed_mass(
             velocity_field[2], distances, boxsize, verbose=False)
 
+        del velocity_field
+        collect()
+
     if args.simname in ["Carrick2015", "Lilow2024"]:
         # Carrick+2015 and Lilow+2024 box is in galactic coordinates, so we
         # need to convert the bulk flow vector to RA/dec Cartesian
@@ -253,12 +238,14 @@ def main_from_field(args, folder):
     cumulative_vel = np.stack(
         [cumulative_vel_x, cumulative_vel_y, cumulative_vel_z], axis=-1)
     cumulative_vel /= cumulative_volume[..., None]
+    cumulative_vel_mono /= cumulative_volume
 
     # Finally save the output
     fname = f"enclosed_mass_{args.simname}.npz"
     fname = join(folder, fname)
     print(f"Saving to `{fname}`.")
     np.savez(fname, enclosed_mass=cumulative_mass, distances=distances,
+             cumulative_velocity_mono=cumulative_vel_mono,
              cumulative_velocity=cumulative_vel,
              enclosed_volume=cumulative_volume)
 

scripts/field_bulk.sh

@@ -1,6 +1,6 @@
 nthreads=1
-memory=12
+memory=48
-on_login=1
+on_login=0
 queue="berg"
 env="/mnt/zfsusers/rstiskalek/csiborgtools/venv_csiborg/bin/python"
 file="field_bulk.py"

scripts/flow_validation.py

@@ -243,10 +243,12 @@ if __name__ == "__main__":
                    "num_epochs": num_epochs}
     print_variables(main_params.keys(), main_params.values())
 
-    calibration_hyperparams = {"Vext_std": 500,
+    calibration_hyperparams = {"Vext_min": -1000, "Vext_max": 1000,
+                               "Vmono_min": -1000, "Vmono_max": 1000,
                                "alpha_min": -1.0, "alpha_max": 3.0,
                                "beta_min": -1.0, "beta_max": 3.0,
                                "sigma_v_min": 5.0, "sigma_v_max": 750.,
+                               "sample_Vmono": True,
                                "sample_alpha": sample_alpha,
                                "sample_beta": sample_beta,
                                }

scripts/flow_validation.sh

@@ -19,8 +19,7 @@ fi
 
 # Submit a job for each combination of simname, catalogue, ksim
 # for simname in "Lilow2024" "CF4" "CF4gp" "csiborg1" "csiborg2_main" "csiborg2X"; do
-for ksmooth in 0 1 2 3 4; do
+for simname in "csiborg2X"; do
-for simname in "Carrick2015"; do
 # for simname in "csiborg1" "csiborg2_main" "csiborg2X"; do
     for catalogue in "Pantheon+"; do
     # for catalogue in "2MTF"; do
@@ -43,5 +42,3 @@ for simname in "Carrick2015"; do
         done
     done
 done
-
-done