Add TNG field (#108)

* Add TNG catalogue script

* Add imoprt

* Add TNG box

* Add import

* Fix little bug

* Add TNG300-1

* Add shell scripts

* Organize .gitignore

* Add scatter to file name

* Add argument scatter

* Add threshold on mass

* Update nb

* Add scatter argument

* Add scattering of radial positions

* Update nb

* Update script

.gitignore

@@ -1,32 +1,38 @@
+# Python virtual environments
+venv/
 venv_csiborg/
-share/
-bin/
-.bashrc
+
+# Compiled Python files
 *.pyc
+
+# Jupyter Notebook checkpoints
 */.ipynb_checkpoints/
-plots/*
-.vscode/settings.json
-csiborgtools/fits/_halo_profile.py
-csiborgtools/fits/_filenames.py
-csiborgtools/fits/analyse_voids_25.py
-scripts/*.out
-build/*
+
+# Egg-info directories
 .eggs/*
 csiborgtools.egg-info/*
-Pylians3/*
-scripts/plot_correlation.ipynb
-scripts/*.sh
-venv/
-.trunk/*
-scripts_test/
-scripts_plots/python.sh
-scripts_plots/submit.sh
-scripts_plots/*.out
-scripts_plots/*.sh
-notebooks/test.ipynb
-scripts/mgtree.py
-scripts/makemerger.py
 
-*.out
-*/python.sh
+# Build directories
+build/*
+bin/*
+share/*
+
+# Scripts and their outputs
+scripts/*.out
+scripts_plots/*.out
 scripts_independent/clear.sh
+
+# Specific script files
+*python.sh
+*python3.sh
+
+# IDE settings
+.vscode/settings.json
+
+# Miscellaneous
+.bashrc
+*.out
+
+# Generated plots
+plots/*
+

csiborgtools/__init__.py

@@ -18,7 +18,7 @@ from .utils import (center_of_mass, delta2ncells, number_counts,
                     periodic_distance, periodic_distance_two_points,            # noqa
                     binned_statistic, cosine_similarity, fprint,                # noqa
                     hms_to_degrees, dms_to_degrees, great_circle_distance,      # noqa
-                    radec_to_cartesian)                                         # noqa
+                    radec_to_cartesian, cartesian_to_radec)                     # noqa
 from .params import paths_glamdring, simname2boxsize                            # noqa
 
 

csiborgtools/params.py

@@ -36,7 +36,8 @@ def simname2boxsize(simname):
          "csiborg2_random": 676.6,
          "borg1": 677.7,
          "borg2": 676.6,
-         "quijote": 1000.
+         "quijote": 1000.,
+         "TNG300-1": 205.
          }
 
     boxsize = d.get(simname, None)
@@ -54,6 +55,7 @@ paths_glamdring = {
     "csiborg2_random_srcdir": "/mnt/extraspace/rstiskalek/csiborg2_random",         # noqa
     "postdir": "/mnt/extraspace/rstiskalek/csiborg_postprocessing/",
     "quijote_dir": "/mnt/extraspace/rstiskalek/quijote",
+    "borg1_dir": "/mnt/users/hdesmond/BORG_final",
     "borg2_dir": "/mnt/extraspace/rstiskalek/BORG_STOPYRA_2023",
     "tng300_1_dir": "/mnt/extraspace/rstiskalek/TNG300-1/",
     }

csiborgtools/read/__init__.py

@@ -16,7 +16,7 @@ from .catalogue import (CSiBORG1Catalogue, CSiBORG2Catalogue,
                         CSiBORG2MergerTreeReader, QuijoteCatalogue)             # noqa
 from .snapshot import (CSiBORG1Snapshot, CSiBORG2Snapshot, QuijoteSnapshot,     # noqa
                        CSiBORG1Field, CSiBORG2Field, QuijoteField, BORG2Field,  # noqa
-                       BORG1Field)                                              # noqa
+                       BORG1Field, TNG300_1Field)                               # noqa
 from .obs import (SDSS, MCXCClusters, PlanckClusters, TwoMPPGalaxies,           # noqa
                   TwoMPPGroups, ObservedCluster, match_array_to_no_masking,     # noqa
                   cols_to_structured)                                           # noqa

csiborgtools/read/paths.py

@@ -467,9 +467,10 @@ class Paths:
         fname = f"observer_peculiar_velocity_{simname}_{MAS}_{str(nsim).zfill(5)}_{grid}.npz"  # noqa
         return join(fdir, fname)
 
-    def field_interpolated(self, survey, simname, nsim, kind, MAS, grid):
+    def field_interpolated(self, survey, simname, nsim, kind, MAS, grid,
+                           radial_scatter=None):
         """
-        Path to the files containing the CSiBORG interpolated field for a given
+        Path to the files containing the interpolated field for a given
         survey.
 
         Parameters
@@ -486,13 +487,16 @@ class Paths:
            Mass-assignment scheme.
         grid : int
             Grid size.
+        radial_scatter : float, optional
+            Radial scatter added to the galaxy positions, only supported for
+            TNG300-1.
 
         Returns
         -------
         str
         """
-        if "csiborg" not in simname:
-            raise ValueError("Interpolated field only available for CSiBORG.")
+
+        # # In case the galaxy positions of TNG300-1 were scattered..
 
         if kind not in ["density", "potential", "radvel"]:
             raise ValueError("Unsupported field type.")
@@ -501,7 +505,12 @@ class Paths:
         try_create_directory(fdir)
 
         nsim = str(nsim).zfill(5)
-        return join(fdir, f"{survey}_{simname}_{kind}_{MAS}_{nsim}_{grid}.npz")
+        fname = join(fdir, f"{survey}_{simname}_{kind}_{MAS}_{nsim}_{grid}.npz")  # noqa
+
+        if simname == "TNG300-1" and radial_scatter is not None:
+            fname = fname.replace(".npz", f"_scatter{radial_scatter}.npz")
+
+        return fname
 
     def cross_nearest(self, simname, run, kind, nsim=None, nobs=None):
         """

csiborgtools/read/snapshot.py

@@ -1004,7 +1004,7 @@ class TNG300_1Field(BaseField):
         return density
 
     def density_field(self, MAS, grid):
-        fpath = join(self.paths.tng300_1, "postprocessing", "density_field",
+        fpath = join(self.paths.tng300_1(), "postprocessing", "density_field",
                      f"rho_dm_099_{grid}_{MAS}.npy")
         return numpy.load(fpath)
 

scripts/clear.sh

@@ -0,0 +1,2 @@
+cm="rm *.out"
+$cm

scripts/field_prop.sh

@@ -0,0 +1,24 @@
+nthreads=6
+memory=64
+on_login=${1}
+queue="berg"
+env="/mnt/zfsusers/rstiskalek/csiborgtools/venv_csiborg/bin/python"
+file="field_prop.py"
+kind="radvel"
+simname="csiborg2_random"
+nsims="-1"
+MAS="SPH"
+grid=1024
+
+
+pythoncm="$env $file --nsims $nsims --simname $simname --kind $kind --MAS $MAS --grid $grid"
+if [ $on_login -eq 1 ]; then
+    echo $pythoncm
+    $pythoncm
+else
+    cm="addqueue -q $queue -n $nthreads -m $memory $pythoncm"
+    echo "Submitting:"
+    echo $cm
+    echo
+    eval $cm
+fi

scripts/field_sample.py

@@ -26,11 +26,58 @@ from h5py import File
 from mpi4py import MPI
 from taskmaster import work_delegation
 from tqdm import tqdm
+from numba import jit
 
 from utils import get_nsims
 
 
-def open_galaxy_positions(survey_name, comm):
+@jit(nopython=True, fastmath=True, boundscheck=False)
+def scatter_along_radial_direction(pos, scatter, boxsize):
+    """
+    Scatter galaxy positions along the radial direction. Enforces that the
+    radial position is always on the same side of the box and that the galaxy
+    is still inside the box.
+
+    Parameters
+    ----------
+    pos : 2-dimensional array
+        Galaxy positions in the form of (distance, RA, DEC).
+    scatter : float
+        Scatter to add to the radial positions of galaxies in same units as
+        `distance` (Mpc / h).
+    boxsize : float
+        Box size in `Mpc / h`.
+    """
+    pos_new = numpy.copy(pos)
+
+    for i in range(len(pos)):
+        r0, ra, dec = pos[i]
+        # Convert to radians
+        ra *= numpy.pi / 180
+        dec *= numpy.pi / 180
+
+        # Convert to normalized Cartesian coordinates
+        xnorm = numpy.cos(dec) * numpy.cos(ra)
+        ynorm = numpy.cos(dec) * numpy.sin(ra)
+        znorm = numpy.sin(dec)
+
+        while True:
+            rnew = numpy.random.normal(r0, scatter)
+            if rnew < 0:
+                continue
+
+            xnew = rnew * xnorm + boxsize / 2
+            ynew = rnew * ynorm + boxsize / 2
+            znew = rnew * znorm + boxsize / 2
+
+            if 0 <= xnew < boxsize and 0 <= ynew < boxsize and 0 <= znew < boxsize:  # noqa
+                pos_new[i, 0] = rnew
+                break
+
+    return pos_new
+
+
+def open_galaxy_positions(survey_name, comm, scatter=None):
     """
     Load the survey's galaxy positions , broadcasting them to all ranks.
 
@@ -40,6 +87,9 @@ def open_galaxy_positions(survey_name, comm):
         Name of the survey.
     comm : mpi4py.MPI.Comm
         MPI communicator.
+    scatter : float
+        Scatter to add to the radial positions of galaxies, supportted only in
+        TNG300-1.
 
     Returns
     -------
@@ -71,6 +121,24 @@ def open_galaxy_positions(survey_name, comm):
                 samples["ra"][:] * 180 / numpy.pi,
                 samples["dec"][:] * 180 / numpy.pi],
                                ).T
+        elif survey_name == "TNG300-1":
+            with File("/mnt/extraspace/rstiskalek/TNG300-1/postprocessing/subhalo_catalogue_099.hdf5", 'r') as f:  # noqa
+                pos = numpy.vstack([f["SubhaloPos"][:, 0],
+                                    f["SubhaloPos"][:, 1],
+                                    f["SubhaloPos"][:, 2]],
+                                   ).T
+                boxsize = csiborgtools.simname2boxsize("TNG300-1")
+                pos -= boxsize / 2
+                pos = csiborgtools.cartesian_to_radec(pos)
+                if scatter is not None:
+                    if scatter < 0:
+                        raise ValueError("Scatter must be positive.")
+                    if scatter > 0:
+                        print(f"Adding scatter of {scatter} Mpc / h.",
+                              flush=True)
+                        pos = scatter_along_radial_direction(pos, scatter,
+                                                             boxsize)
+
         else:
             raise NotImplementedError(f"Survey `{survey_name}` not "
                                       "implemented.")
@@ -118,7 +186,7 @@ def evaluate_field(field, pos, boxsize, smooth_scales, verbose=True):
                 field, scale * mpc2box, boxsize=1, make_copy=True)
         else:
             field_smoothed = numpy.copy(field)
-
+        print("Going to evaluate the field....")
         val[:, i] = csiborgtools.field.evaluate_sky(
             field_smoothed, pos=pos, mpc2box=mpc2box)
 
@@ -182,6 +250,8 @@ def main(nsim, parser_args, pos, verbose):
     elif "csiborg2" in parser_args.simname:
         kind = parser_args.simname.split("_")[-1]
         freader = csiborgtools.read.CSiBORG2Field(nsim, kind)
+    elif parser_args.simname == "TNG300-1":
+        freader = csiborgtools.read.TNG300_1Field()
     else:
         raise NotImplementedError(f"Simulation `{parser_args.simname}` is not supported.")  # noqa
 
@@ -199,12 +269,18 @@ def main(nsim, parser_args, pos, verbose):
             "/mnt/extraspace/rstiskalek/GWLSS/",
             f"{parser_args.kind}_{parser_args.MAS}_{parser_args.grid}_{nsim}_H1L1V1-EXTRACT_POSTERIOR_GW170817-1187008600-400.npz")  # noqa
     else:
+        if parser_args.simname == "TNG300-1":
+            scatter = parser_args.scatter
+        else:
+            scatter = None
+
         fout = paths.field_interpolated(
             parser_args.survey, parser_args.simname, nsim, parser_args.kind,
-            parser_args.MAS, parser_args.grid)
+            parser_args.MAS, parser_args.grid, scatter)
 
         # The survey above had some cuts, however for compatibility we want
         # the same shape as the `uncut` survey
+        if parser_args.survey != "TNG300-1":
             val = match_to_no_selection(val, parser_args)
 
     if verbose:
@@ -218,10 +294,10 @@ if __name__ == "__main__":
     parser.add_argument("--nsims", type=int, nargs="+", default=None,
                         help="IC realisations. If `-1` processes all.")
     parser.add_argument("--simname", type=str, default="csiborg1",
-                        choices=["csiborg1", "csiborg2_main", "csiborg2_random", "csiborg2_varysmall"],  # noqa
+                        choices=["csiborg1", "csiborg2_main", "csiborg2_random", "csiborg2_varysmall", "TNG300-1"],  # noqa
                         help="Simulation name")
     parser.add_argument("--survey", type=str, required=True,
-                        choices=["SDSS", "SDSSxALFALFA", "GW170817"],
+                        choices=["SDSS", "SDSSxALFALFA", "GW170817", "TNG300-1"],  # noqa
                         help="Galaxy survey")
     parser.add_argument("--smooth_scales", type=float, nargs="+", default=None,
                         help="Smoothing scales in Mpc / h.")
@@ -233,12 +309,20 @@ if __name__ == "__main__":
                         choices=["NGP", "CIC", "TSC", "PCS", "SPH"],
                         help="Mass assignment scheme.")
     parser.add_argument("--grid", type=int, help="Grid resolution.")
+    parser.add_argument("--scatter", type=float, default=None,
+                        help="Scatter to add to the radial positions of galaxies, supportted only in TNG300-1.")  # noqa
     args = parser.parse_args()
 
+    if args.simname == "TNG300-1" and args.survey != "TNG300-1":
+        raise ValueError("TNG300-1 simulation is only supported for TNG300-1 survey.")  # noqa
+
     paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
+    if args.simname == "TNG300-1":
+        nsims = [0]
+    else:
         nsims = get_nsims(args, paths)
 
-    pos = open_galaxy_positions(args.survey, MPI.COMM_WORLD)
+    pos = open_galaxy_positions(args.survey, MPI.COMM_WORLD, args.scatter)
 
     def _main(nsim):
         main(nsim, args, pos, verbose=MPI.COMM_WORLD.Get_size() == 1)

scripts/field_sample.sh

@@ -0,0 +1,29 @@
+nthreads=1
+memory=32
+on_login=${1}
+queue="berg"
+env="/mnt/zfsusers/rstiskalek/csiborgtools/venv_csiborg/bin/python"
+file="field_sample.py"
+
+
+nsims="-1"
+simname="TNG300-1"
+survey="TNG300-1"
+smooth_scales="0 2 4 8 16"
+kind="density"
+MAS="PCS"
+grid=1024
+scatter=0
+
+
+pythoncm="$env $file --nsims $nsims --simname $simname --survey $survey --smooth_scales $smooth_scales --kind $kind --MAS $MAS --grid $grid --scatter $scatter"
+if [ $on_login -eq 1 ]; then
+    echo $pythoncm
+    $pythoncm
+else
+    cm="addqueue -q $queue -n $nthreads -m $memory $pythoncm"
+    echo "Submitting:"
+    echo $cm
+    echo
+    eval $cm
+fi

scripts/field_shells.sh

@@ -0,0 +1,24 @@
+nthreads=1
+memory=32
+on_login=${1}
+queue="berg"
+env="/mnt/zfsusers/rstiskalek/csiborgtools/venv_csiborg/bin/python"
+file="field_shells.py"
+
+field="overdensity"
+simname="borg2"
+MAS="SPH"
+grid=1024
+
+
+pythoncm="$env $file --field $field --simname $simname --MAS $MAS --grid $grid"
+if [ $on_login -eq 1 ]; then
+    echo $pythoncm
+    $pythoncm
+else
+    cm="addqueue -q $queue -n $nthreads -m $memory $pythoncm"
+    echo "Submitting:"
+    echo $cm
+    echo
+    eval $cm
+fi

scripts/fit_init.sh

@@ -0,0 +1,23 @@
+nthreads=6
+memory=16
+on_login=${1}
+queue="berg"
+env="/mnt/zfsusers/rstiskalek/csiborgtools/venv_csiborg/bin/python"
+file="fit_init.py"
+
+
+simname="csiborg2_varysmall"
+nsims="-1"
+
+
+pythoncm="$env $file --nsims $nsims --simname $simname"
+if [ $on_login -eq 1 ]; then
+    echo $pythoncm
+    $pythoncm
+else
+    cm="addqueue -q $queue -n $nthreads -m $memory $pythoncm"
+    echo "Submitting:"
+    echo $cm
+    echo
+    eval $cm
+fi

scripts/mass_enclosed.sh

@@ -0,0 +1,21 @@
+nthreads=1
+memory=32
+on_login=${1}
+queue="berg"
+env="/mnt/zfsusers/rstiskalek/csiborgtools/venv_csiborg/bin/python"
+file="mass_enclosed.py"
+
+simname="borg2"
+
+
+pythoncm="$env $file --simname $simname"
+if [ $on_login -eq 1 ]; then
+    echo $pythoncm
+    $pythoncm
+else
+    cm="addqueue -q $queue -n $nthreads -m $memory $pythoncm"
+    echo "Submitting:"
+    echo $cm
+    echo
+    eval $cm
+fi

scripts/match_knn.sh

@@ -0,0 +1,29 @@
+#!/bin/bash
+nthreads=50
+memory=7
+queue="cmb"
+env="/mnt/zfsusers/rstiskalek/csiborgtools/venv_csiborg/bin/python"
+file="match_finsnap.py"
+verbose="true"
+nsims="-1"
+onlogin=false
+
+# for run in "mass001" "mass003" "mass005" "mass007" "mass009"
+for run in "mass002" "mass004" "mass006" "mass008"
+do
+for simname in "csiborg"
+do
+pythoncm="$env $file --simname $simname --run $run --nsims $nsims --verbose $verbose"
+
+if $onlogin
+then
+    $pythoncm
+else
+    cm="addqueue -q $queue -n $nthreads -m $memory $pythoncm"
+    echo "Submitting:"
+    echo $cm
+    echo
+    $cm
+fi
+done
+done

scripts/match_overlap_all.sh

@@ -0,0 +1,24 @@
+#!/bin/bash
+nthreads=11
+memory=4
+queue="cmb"
+env="/mnt/zfsusers/rstiskalek/csiborgtools/venv_csiborg/bin/python"
+file="match_all.py"
+
+simname="quijote"
+min_logmass=13.25
+nsim0=0
+kind="max"
+mult=10
+sigma=1
+verbose="false"
+
+
+pythoncm="$env $file --kind $kind --simname $simname --nsim0 $nsim0 --min_logmass $min_logmass --mult $mult --sigma $sigma --verbose $verbose"
+# $pythoncm
+
+cm="addqueue -q $queue -n $nthreads -m $memory $pythoncm"
+echo "Submitting:"
+echo $cm
+echo
+$cm

scripts/match_overlap_single.sh

@@ -0,0 +1,50 @@
+#!/bin/bash
+nthreads=1
+memory=7
+queue="berg"
+env="/mnt/zfsusers/rstiskalek/csiborgtools/venv_csiborg/bin/python"
+verbose="true"
+file="match_overlap_single.py"
+
+simname="csiborg2_main"
+kind="overlap"
+min_logmass=13.25
+mult=5
+sigma=1
+
+# sims=(7444 7468)
+sims=(16417 16517)
+# sims=(0 1)
+# sims=(7468 7588)
+nsims=${#sims[@]}
+
+for i in $(seq 0 $((nsims-1)))
+do
+    for j in $(seq 0 $((nsims-1)))
+    do
+        if [ $i -eq $j ]
+        then
+            continue
+        elif [ $i -gt $j ]
+        then
+            continue
+        else
+            :
+        fi
+
+        nsim0=${sims[$i]}
+        nsimx=${sims[$j]}
+
+        pythoncm="$env $file --kind $kind --nsim0 $nsim0 --nsimx $nsimx  --simname $simname --min_logmass $min_logmass --sigma $sigma --mult $mult --verbose $verbose"
+
+        # $pythoncm
+
+        cm="addqueue -q $queue -n 1x1 -m $memory $pythoncm"
+        echo "Submitting:"
+        echo $cm
+        echo
+        $cm
+        sleep 0.05
+
+    done
+done

scripts_independent/catalogue_tng.py

@@ -0,0 +1,64 @@
+# 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.
+"""
+Script to iteratively load particles of a TNG simulation and construct the DM
+density field.
+"""
+from os.path import join
+import numpy as np
+from h5py import File
+import illustris_python as il
+
+
+if __name__ == "__main__":
+    fdir = "/mnt/extraspace/rstiskalek/TNG300-1/"
+    basepath = join(fdir, "output")
+    out_fname = join(fdir, "postprocessing/subhalo_catalogue_099.hdf5")
+
+    # SUBFIND catalogue
+    fields = ["SubhaloFlag", "SubhaloPos", "SubhaloMassType",
+              "SubhaloGasMetallicity", "SubhaloStarMetallicity",
+              "SubhaloSFR", "SubhaloSpin", "SubhaloStellarPhotometrics"]
+
+    print("Loading the data.....")
+    data = il.groupcat.loadSubhalos(basepath, 99, fields=fields)
+    data["SubhaloPos"] /= 1000.  # Convert to Mpc/h
+    print("Finished loading!")
+
+    # Take only galaxies with stellar mass more than 10^9 Msun / h
+    mask = (data["SubhaloFlag"] == 1) & (data["SubhaloMassType"][:, 4] > 0.1)
+
+    print(f"Writing the subfind dataset to '{out_fname}'")
+    with File(out_fname, 'w') as f:
+        for key in fields:
+            if key == "SubhaloFlag":
+                continue
+
+            f.create_dataset(key, data=data[key][mask])
+
+    # HIH2 supplemetary catalogue
+    print("Loading the HI & H2 supplementary catalogue.")
+    fname = join(fdir, "postprocessing/hih2/hih2_galaxy_099.hdf5")
+    with File(fname, "r") as f:
+        _m_neutral_H = f["m_neutral_H"][:]
+        _id_subhalo = np.array(f["id_subhalo"][:], dtype=int)
+
+    m_neutral_H = np.full(data["count"], np.nan, dtype=float)
+    for i, j in enumerate(_id_subhalo):
+        m_neutral_H[j] = _m_neutral_H[i]
+
+    print("Adding the HI & H2 supplementary catalogue.")
+    with File(out_fname, 'r+') as f:
+        f.create_dataset("m_neutral_H", data=m_neutral_H[mask])

scripts_plots/submit.sh

@@ -0,0 +1,20 @@
+#!/bin/bash
+
+if [ "$#" -ne 1 ]; then
+    echo "Usage: ./script.sh <path_to_file>"
+    exit 1
+fi
+
+env="/mnt/zfsusers/rstiskalek/csiborgtools/venv_csiborg/bin/python"
+queue="berg"
+nthreads=1
+memory=7
+
+file="$1"
+
+cm="addqueue -q $queue -n $nthreads -m $memory $env $file"
+
+echo "Submitting:"
+echo $cm
+echo
+$cm