Observer velocity script (#120)

* Rename script

* Delete scripts

* Add script

* Edit script

* Add script

* Update nb

* Update plotting

* Update .gitignore

* Update nb

* Update nb

* Add option to keep beta fixed

.gitignore

@@ -36,3 +36,4 @@ scripts_independent/clear.sh
 # Generated plots
 plots/*
 
+notebooks/test.ipynb

csiborgtools/flow/flow_model.py

@@ -951,7 +951,7 @@ class SN_PV_validation_model(BaseFlowValidationModel):
 
         return zobs_mean, zobs_var
 
-    def __call__(self, sample_alpha=True, fix_calibration=False):
+    def __call__(self, sample_alpha=True, sample_beta=True):
         """
         The supernova NumPyro PV validation model with SALT2 calibration.
 
@@ -960,34 +960,21 @@ class SN_PV_validation_model(BaseFlowValidationModel):
         sample_alpha : bool, optional
             Whether to sample the density bias parameter `alpha`, otherwise
             it is fixed to 1.
-        fix_calibration : str, optional
+        sample_beta : bool, optional
-            Whether to fix the calibration parameters. If not provided, they
+            Whether to sample the velocity bias parameter `beta`, otherwise
-            are sampled. If "Foundation" or "LOSS" is provided, the parameters
+            it is fixed to 1.
-            are fixed to the best inverse parameters for the Foundation or LOSS
+
-            catalogues.
+        Returns
+        -------
+        None
         """
         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._sigma_v)
 
-        if fix_calibration == "Foundation":
-            # Foundation inverse best parameters
-            e_mu_intrinsic = 0.064
-            alpha_cal = 0.135
-            beta_cal = 2.9
-            sigma_v = 149
-            mag_cal = -18.555
-        elif fix_calibration == "LOSS":
-            # LOSS inverse best parameters
-            e_mu_intrinsic = 0.123
-            alpha_cal = 0.123
-            beta_cal = 3.52
-            mag_cal = -18.195
-            sigma_v = 149
-        else:
         e_mu_intrinsic = numpyro.sample("e_mu_intrinsic", self._e_mu)
         mag_cal = numpyro.sample("mag_cal", self._mag_cal)
         alpha_cal = numpyro.sample("alpha_cal", self._alpha_cal)
@@ -1168,7 +1155,7 @@ class TF_PV_validation_model(BaseFlowValidationModel):
 
         return zobs_mean, zobs_var
 
-    def __call__(self, sample_alpha=True):
+    def __call__(self, sample_alpha=True, sample_beta=True):
         """
         The Tully-Fisher NumPyro PV validation model.
 
@@ -1177,12 +1164,19 @@ class TF_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.
+
+        Returns
+        -------
+        None
         """
         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._sigma_v)
 
         e_mu_intrinsic = numpyro.sample("e_mu_intrinsic", self._e_mu)
@@ -1291,7 +1285,7 @@ def get_model(loader, zcmb_max=None, verbose=True):
 ###############################################################################
 
 
-def sample_prior(model, seed, sample_alpha, as_dict=False):
+def sample_prior(model, seed, model_kwargs, as_dict=False):
     """
     Sample a single set of parameters from the prior of the model.
 
@@ -1301,8 +1295,8 @@ def sample_prior(model, seed, sample_alpha, as_dict=False):
         NumPyro model.
     seed : int
         Random seed.
-    sample_alpha : bool
+    model_kwargs : dict
-        Whether to sample the density bias parameter `alpha`.
+        Additional keyword arguments to pass to the model.
     as_dict : bool, optional
         Whether to return the parameters as a dictionary or a list of
         parameters.
@@ -1314,7 +1308,7 @@ def sample_prior(model, seed, sample_alpha, as_dict=False):
         only a dictionary.
     """
     predictive = Predictive(model, num_samples=1)
-    samples = predictive(PRNGKey(seed), sample_alpha=sample_alpha)
+    samples = predictive(PRNGKey(seed), **model_kwargs)
 
     if as_dict:
         return samples
@@ -1327,7 +1321,7 @@ def sample_prior(model, seed, sample_alpha, as_dict=False):
     return x, keys
 
 
-def make_loss(model, keys, sample_alpha=True, to_jit=True):
+def make_loss(model, keys, model_kwargs, to_jit=True):
     """
     Generate a loss function for the NumPyro model, that is the negative
     log-likelihood. Note that this loss function cannot be automatically
@@ -1339,8 +1333,8 @@ def make_loss(model, keys, sample_alpha=True, to_jit=True):
         NumPyro model.
     keys : list
         List of parameter names.
-    sample_alpha : bool, optional
+    model_kwargs : dict
-        Whether to sample the density bias parameter `alpha`.
+        Additional keyword arguments to pass to the model.
     to_jit : bool, optional
         Whether to JIT the loss function.
 
@@ -1353,8 +1347,7 @@ def make_loss(model, keys, sample_alpha=True, to_jit=True):
     def f(x):
         samples = {key: x[i] for i, key in enumerate(keys)}
 
-        loss = -util.log_likelihood(
+        loss = -util.log_likelihood(model, samples, **model_kwargs)["ll"]
-            model, samples, sample_alpha=sample_alpha)["ll"]
 
         loss += cond(samples["sigma_v"] > 0, lambda: 0., lambda: jnp.inf)
         loss += cond(samples["e_mu_intrinsic"] > 0, lambda: 0., lambda: jnp.inf)  # noqa

notebooks/flow_calibration.py

@@ -28,6 +28,7 @@ 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)
+    FDIR_LG = "/mnt/extraspace/rstiskalek/csiborg_postprocessing/peculiar_velocity/observer"  # noqa
 
     Vx, Vy, Vz, beta, sigma_v, alpha = [], [], [], [], [], []
     BIC, AIC, logZ, chi2 = [], [], [], []
@@ -39,17 +40,6 @@ def read_samples(catalogue, simname, ksmooth, include_calibration=False,
     else:
         raise ValueError(f"Catalogue {catalogue} not recognized.")
 
-    if subtract_LG_velocity >= 0:
-        fdir = "/mnt/extraspace/rstiskalek/csiborg_postprocessing/field_shells"
-        fname = join(fdir, f"enclosed_mass_{simname}.npz")
-        if exists(fname):
-            d = np.load(fname)
-            R = d["distances"][subtract_LG_velocity]
-            print(f"Reading off enclosed velocity from R = {R} Mpc / h.")
-            V_LG = d["cumulative_velocity"][:, subtract_LG_velocity, :]
-        else:
-            raise FileNotFoundError(f"File {fname} not found.")
-
     fname = f"/mnt/extraspace/rstiskalek/csiborg_postprocessing/peculiar_velocity/flow_samples_{catalogue}_{simname}_smooth_{ksmooth}.hdf5"  # noqa
     with File(fname, 'r') as f:
         for i, nsim in enumerate(nsims):
@@ -57,15 +47,28 @@ def read_samples(catalogue, simname, ksmooth, include_calibration=False,
             Vy.append(f[f"sim_{nsim}/Vext_y"][:])
             Vz.append(f[f"sim_{nsim}/Vext_z"][:])
 
-            if subtract_LG_velocity >= 0:
-                Vx[-1] += V_LG[i, 0]
-                Vy[-1] += V_LG[i, 1]
-                Vz[-1] += V_LG[i, 2]
-
             alpha.append(f[f"sim_{nsim}/alpha"][:])
             beta.append(f[f"sim_{nsim}/beta"][:])
             sigma_v.append(f[f"sim_{nsim}/sigma_v"][:])
 
+            if subtract_LG_velocity >= 0:
+                fname = join(FDIR_LG, f"{simname}_{nsim}_observer_velocity.npz")  # noqa
+                if not exists(fname):
+                    raise FileNotFoundError(f"File {fname} not found.")
+                d = np.load(fname)
+                R = d["smooth_scales"][subtract_LG_velocity]
+                if i == 0:
+                    print(f"Subtracting LG velocity with kernel {R} Mpc / h.", flush=True)  # noqa
+                Vx_LG, Vy_LG, Vz_LG = d["vobs"][subtract_LG_velocity]
+                if simname == "Carrick2015":
+                    Vx[-1] += beta[-1] * Vx_LG
+                    Vy[-1] += beta[-1] * Vy_LG
+                    Vz[-1] += beta[-1] * Vz_LG
+                else:
+                    Vx[-1] += Vx_LG
+                    Vy[-1] += Vy_LG
+                    Vz[-1] += Vz_LG
+
             BIC.append(f[f"sim_{nsim}/BIC"][...])
             AIC.append(f[f"sim_{nsim}/AIC"][...])
             logZ.append(f[f"sim_{nsim}/logZ"][...])

scripts/field_observer_velocity.py

@@ -0,0 +1,116 @@
+# Copyright (C) 2023 Richard Stiskalek
+# This program is free software; you can redistribute it and/or modify it
+# under the terms of the GNU General Public License as published by the
+# Free Software Foundation; either version 3 of the License, or (at your
+# option) any later version.
+#
+# This program is distributed in the hope that it will be useful, but
+# WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General
+# Public License for more details.
+#
+# You should have received a copy of the GNU General Public License along
+# with this program; if not, write to the Free Software Foundation, Inc.,
+# 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
+"""
+
+"""
+from argparse import ArgumentParser
+from datetime import datetime
+from os.path import join
+from warnings import warn
+
+import csiborgtools
+import numpy as np
+from astropy.coordinates import SkyCoord
+from mpi4py import MPI
+from taskmaster import work_delegation
+
+from utils import get_nsims
+
+FDIR = "/mnt/extraspace/rstiskalek/csiborg_postprocessing/peculiar_velocity/observer"  # noqa
+
+
+def t():
+    return datetime.now()
+
+
+def read_velocity_field(args, nsim):
+    if args.simname == "csiborg1":
+        reader = csiborgtools.read.CSiBORG1Field(nsim)
+        return reader.velocity_field("SPH", 1024)
+    elif "csiborg2" in args.simname:
+        kind = args.simname.split("_")[-1]
+        reader = csiborgtools.read.CSiBORG2Field(nsim, kind)
+        return reader.velocity_field("SPH", 1024)
+    elif args.simname == "Carrick2015":
+        folder = "/mnt/extraspace/rstiskalek/catalogs"
+        warn(f"Using local paths from `{folder}`.", RuntimeWarning)
+        fpath = join(folder, "twompp_velocity_carrick2015.npy")
+        field = np.load(fpath).astype(np.float32)
+
+        # Because the Carrick+2015 data is in the following form:
+        # "The velocities are predicted peculiar velocities in the CMB
+        # frame in Galactic Cartesian coordinates, generated from the
+        # \(\delta_g^*\) field with \(\beta^* = 0.43\) and an external
+        # dipole \(V_\mathrm{ext} = [89,-131,17]\) (Carrick et al Table 3)
+        # has already been added.""
+        field[0] -= 89
+        field[1] -= -131
+        field[2] -= 17
+        field /= 0.43
+        return field
+    else:
+        raise ValueError(f"Unknown simname: `{args.simname}`.")
+
+
+def main(smooth_scales, nsim, args):
+    velocity_field = read_velocity_field(args, nsim)
+    boxsize = csiborgtools.simname2boxsize(args.simname)
+
+    if smooth_scales is None:
+        smooth_scales = [0]
+    smooth_scales = np.asanyarray(smooth_scales) / boxsize
+
+    vobs = csiborgtools.field.observer_peculiar_velocity(
+        velocity_field, smooth_scales=smooth_scales, observer=None,
+        verbose=False)
+
+    # For Carrick+2015 the velocity vector is in the Galactic frame, so we
+    # need to convert it to RA/dec
+    if args.simname == "Carrick2015":
+        coord = SkyCoord(vobs, unit='kpc', frame='galactic',
+                         representation_type='cartesian').transform_to("icrs")
+        vobs = coord.cartesian.xyz.value.T
+
+    fname = join(FDIR, f"{args.simname}_{nsim}_observer_velocity.npz")
+    print(f"Saving to `{fname}`.")
+    np.savez(fname, vobs=vobs, smooth_scales=smooth_scales * boxsize)
+
+
+###############################################################################
+#                       Main & command line interface                         #
+###############################################################################
+
+
+if __name__ == "__main__":
+    parser = ArgumentParser()
+    parser.add_argument("--simname", type=str, help="Simulation name.",
+                        choices=["csiborg1", "csiborg2_main", "csiborg2_varysmall", "csiborg2_random", "Carrick2015"])  # noqa
+    args = parser.parse_args()
+    args.nsims = [-1]
+    comm = MPI.COMM_WORLD
+
+    smooth_scales = [0, 0.5, 1.0, 2.0, 4.0, 40.0]
+    paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
+    nsims = get_nsims(args, paths)
+
+    def main_(nsim):
+        main(smooth_scales, nsim, args)
+
+    work_delegation(main_, nsims, comm, master_verbose=True)
+
+    comm.Barrier()
+
+    if comm.Get_rank() == 0:
+        print("All finished.", flush=True)

scripts/field_observer_velocity.sh

@@ -1,17 +1,13 @@
-nthreads=1
+nthreads=5
-memory=32
+memory=40
-on_login=${1}
+on_login=0
 queue="berg"
 env="/mnt/zfsusers/rstiskalek/csiborgtools/venv_csiborg/bin/python"
-file="field_shells.py"
+file="field_observer_velocity.py"
 
-field="overdensity"
+simname=${1}
-simname="borg2"
-MAS="SPH"
-grid=1024
 
-
+pythoncm="$env $file --simname $simname"
-pythoncm="$env $file --field $field --simname $simname --MAS $MAS --grid $grid"
 if [ $on_login -eq 1 ]; then
     echo $pythoncm
     $pythoncm

scripts/field_shells.py

@@ -1,94 +0,0 @@
-
-# Copyright (C) 2022 Richard Stiskalek
-# This program is free software; you can redistribute it and/or modify it
-# under the terms of the GNU General Public License as published by the
-# Free Software Foundation; either version 3 of the License, or (at your
-# option) any later version.
-#
-# This program is distributed in the hope that it will be useful, but
-# WITHOUT ANY WARRANTY; without even the implied warranty of
-# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General
-# Public License for more details.
-#
-# You should have received a copy of the GNU General Public License along
-# with this program; if not, write to the Free Software Foundation, Inc.,
-# 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
-"""
-NOTE: This script is pretty dodgy.
-
-A script to calculate the mean and standard deviation of a field at different
-distances from the center of the box such that at each distance the field is
-evaluated at uniformly-spaced points on a sphere.
-
-The script is not parallelized in any way but it should not take very long, the
-main bottleneck is reading the data from disk.
-"""
-from argparse import ArgumentParser
-from os.path import join
-
-import csiborgtools
-import numpy
-from tqdm import tqdm
-
-
-def main(args):
-    paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
-    boxsize = csiborgtools.simname2boxsize(args.simname)
-    distances = numpy.linspace(0, boxsize / 2, 101)[1:]
-    nsims = paths.get_ics(args.simname)
-    folder = "/mnt/extraspace/rstiskalek/csiborg_postprocessing/field_shells"
-
-    mus = numpy.zeros((len(nsims), len(distances)))
-    stds = numpy.zeros((len(nsims), len(distances)))
-    for i, nsim in enumerate(tqdm(nsims, desc="Simulations")):
-        # Get the correct field loader
-        if args.simname == "csiborg1":
-            reader = csiborgtools.read.CSiBORG1Field(nsim, paths)
-        elif "csiborg2" in args.simname:
-            kind = args.simname.split("_")[-1]
-            reader = csiborgtools.read.CSiBORG2Field(nsim, kind, paths)
-        elif args.simname == "borg2":
-            reader = csiborgtools.read.BORG2Field(nsim, paths)
-        else:
-            raise ValueError(f"Unknown simname: `{args.simname}`.")
-
-        # Get the field
-        if args.field == "density":
-            field = reader.density_field(args.MAS, args.grid)
-        elif args.field == "overdensity":
-            if args.simname == "borg2":
-                field = reader.overdensity_field()
-            else:
-                field = reader.density_field(args.MAS, args.grid)
-                csiborgtools.field.overdensity_field(field, make_copy=False)
-        elif args.field == "radvel":
-            field = reader.radial_velocity_field(args.MAS, args.grid)
-        else:
-            raise ValueError(f"Unknown field: `{args.field}`.")
-
-        # Evaluate this field at different distances
-        vals = [csiborgtools.field.field_at_distance(field, distance, boxsize)
-                for distance in distances]
-
-        # Calculate the mean and standard deviation
-        mus[i, :] = [numpy.mean(val) for val in vals]
-        stds[i, :] = [numpy.std(val) for val in vals]
-
-    # Finally save the output
-    fname = f"{args.simname}_{args.field}_{args.MAS}_{args.grid}.npz"
-    fname = join(folder, fname)
-    numpy.savez(fname, mean=mus, std=stds, distances=distances)
-
-
-if __name__ == "__main__":
-    parser = ArgumentParser()
-    parser.add_argument("--field", type=str, help="Field type.",
-                        choices=["density", "overdensity", "radvel"])
-    parser.add_argument("--simname", type=str, help="Simulation name.",
-                        choices=["csiborg1", "csiborg2_main", "csiborg2_varysmall", "csiborg2_random", "borg2"])  # noqa
-    parser.add_argument("--MAS", type=str, help="Mass assignment scheme.",
-                        choices=["NGP", "CIC", "TSC", "PCS", "SPH"])
-    parser.add_argument("--grid", type=int, help="Grid size.")
-    args = parser.parse_args()
-
-    main(args)