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

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,13 +269,19 @@ 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
-        val = match_to_no_selection(val, parser_args)
+        if parser_args.survey != "TNG300-1":
+            val = match_to_no_selection(val, parser_args)
 
     if verbose:
         print(f"Saving to ... `{fout}`.")
@@ -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()
 
-    paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
-    nsims = get_nsims(args, paths)
+    if args.simname == "TNG300-1" and args.survey != "TNG300-1":
+        raise ValueError("TNG300-1 simulation is only supported for TNG300-1 survey.")  # noqa
 
-    pos = open_galaxy_positions(args.survey, MPI.COMM_WORLD)
+    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, 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