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

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
-memory=32
-on_login=${1}
+nthreads=5
+memory=40
+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="borg2"
-MAS="SPH"
-grid=1024
+simname=${1}
 
-
-pythoncm="$env $file --field $field --simname $simname --MAS $MAS --grid $grid"
+pythoncm="$env $file --simname $simname"
 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)