Attempt at absolute calibration and more (#146)

* Update nb * Update defaults * Add absolute calibration to fpath * Add basic absolute calibration * Switch to log density * Feed in r directly * Add more complex spacing option * Update script * Comment out abs calibration for now * Update script * Add MB profile * Add script * Update submit for all profiles * Update * Add option * Add iterator over sims * Update defaults on alpha * Update script * Parallelise script * Switch to feasible init * Switch to median init * Add vext option * Remove file if exists * Add optional sample_Vext * Add param * rm print * Add more iterator * Update script * Update nb * Update nb
2025-05-13 22:21:12 +00:00 · 2024-09-16 12:12:32 +02:00 · 2024-09-16 12:12:32 +02:00 · 9c1aa26428
commit 9c1aa26428
parent 32e36afdc3
11 changed files with 568 additions and 160 deletions
--- a/scripts/field_prop/field_los.py
+++ b/scripts/field_prop/field_los.py
@ -23,16 +23,19 @@ from os import makedirs, remove, rmdir
 from os.path import exists, join
 from warnings import warn

-import csiborgtools
 import numpy as np
 from astropy import units as u
 from astropy.coordinates import SkyCoord
+from astropy.cosmology import FlatLambdaCDM
 from astropy.io import fits
 from h5py import File
 from mpi4py import MPI
 from numba import jit
+from scipy.interpolate import interp1d
 from taskmaster import work_delegation  # noqa

+import csiborgtools
+
 sys.path.append("../")
 from utils import get_nsims  # noqa

@ -41,6 +44,35 @@ from utils import get_nsims  # noqa
 ###############################################################################


+def make_spacing(rmax, dr, dense_mu_min, dense_mu_max, dmu, Om0):
+    """
+    Make radial spacing that at low distance is with constant spacing in
+    distance modulus and at higher distances is with constant spacing in
+    comoving distance.
+    """
+    # Create interpolant to go from distance modulus to comoving distance.
+    cosmo = FlatLambdaCDM(H0=100, Om0=Om0)
+
+    z_range = np.linspace(0, 0.1, 1000000)[1:]
+    r_range = cosmo.comoving_distance(z_range).value
+    mu_range = cosmo.distmod(z_range).value
+
+    mu2r = interp1d(mu_range, r_range, kind='cubic')
+
+    # Create the spacing in distance modulus.
+    mu = np.arange(dense_mu_min, dense_mu_max, dmu)
+    rmin_dense = mu2r(np.min(mu))
+    rmax_dense = mu2r(np.max(mu))
+
+    # Create the spacing in comoving distance below and above.
+    rlow = np.arange(0, rmin_dense, dr)
+    rmed = mu2r(mu)
+    rhigh = np.arange(rmax_dense, rmax, dr)[1:]
+
+    # Combine the spacings.
+    return np.hstack([rlow, rmed, rhigh])
+
+
 def get_los(catalogue_name, simname, comm):
    """
    Get the line of sight RA/dec coordinates for the given catalogue.
@ -296,8 +328,8 @@ def replace_nan_with_last_finite(x, y, apply_decay):
    return y, rmax


-def interpolate_field(pos, simname, nsim, MAS, grid, dump_folder, rmax,
-                      dr, smooth_scales, verbose=False):
+def interpolate_field(pos, simname, nsim, MAS, grid, dump_folder, r,
+                      smooth_scales, verbose=False):
    """
    Interpolate the density and velocity fields along the line of sight.

@ -315,10 +347,8 @@ def interpolate_field(pos, simname, nsim, MAS, grid, dump_folder, rmax,
        Grid resolution.
    dump_folder : str
        Folder where the temporary files are stored.
-    rmax : float
-        Maximum distance along the line of sight.
-    dr : float
-        Distance spacing along the line of sight.
+    r : 1-dimensional array
+        Radial spacing.
    smooth_scales : list
        Smoothing scales.

@ -335,7 +365,7 @@ def interpolate_field(pos, simname, nsim, MAS, grid, dump_folder, rmax,
              flush=True)
    density = get_field(simname, nsim, "density", MAS, grid)
    rdist, finterp = csiborgtools.field.evaluate_los(
-        density, sky_pos=pos, boxsize=boxsize, rmax=rmax, dr=dr,
+        density, sky_pos=pos, boxsize=boxsize, rdist=r,
        smooth_scales=smooth_scales, verbose=verbose,
        interpolation_method="linear")

@ -361,9 +391,8 @@ def interpolate_field(pos, simname, nsim, MAS, grid, dump_folder, rmax,
    velocity = get_field(simname, nsim, "velocity", MAS, grid)
    rdist, finterp = csiborgtools.field.evaluate_los(
        velocity[0], velocity[1], velocity[2],
-        sky_pos=pos, boxsize=boxsize, rmax=rmax, dr=dr,
-        smooth_scales=smooth_scales, verbose=verbose,
-        interpolation_method="linear")
+        sky_pos=pos, boxsize=boxsize, rdist=r, smooth_scales=smooth_scales,
+        verbose=verbose, interpolation_method="linear")

    rmax_velocity = np.full((3, len(pos), len(smooth_scales)), np.nan)
    for k in range(3):
@ -398,16 +427,15 @@ if __name__ == "__main__":
    parser.add_argument("--grid", type=int, help="Grid resolution.")
    args = parser.parse_args()

-    rmax = 200
-    if args.catalogue == "CF4_GroupAll":
-        dr = 1
-    else:
-        dr = 0.75
+    Om0 = csiborgtools.simname2Omega_m(args.simname)
+    # r = make_spacing(200, 0.75, 23.25, 34, 0.01, Om0)
+    r = np.arange(0, 200, 0.75)

    # smooth_scales = [0, 2, 4, 6, 8]
    smooth_scales = [0]

-    print(f"Running catalogue {args.catalogue} for simulation {args.simname}.")
+    print(f"Running catalogue {args.catalogue} for simulation {args.simname} "
+          f"with {len(r)} radial points.")

    comm = MPI.COMM_WORLD
    paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
@ -429,7 +457,7 @@ if __name__ == "__main__":

    def main(nsim):
        interpolate_field(pos, args.simname, nsim, args.MAS, args.grid,
-                          dump_folder, rmax, dr, smooth_scales,
+                          dump_folder, r, smooth_scales,
                          verbose=comm.Get_size() == 1)

    work_delegation(main, nsims, comm, master_verbose=True)
--- a/scripts/field_prop/field_los.sh
+++ b/scripts/field_prop/field_los.sh
@ -19,9 +19,10 @@ fi


 # for simname in "csiborg1" "csiborg2_main" "csiborg2X" "Lilow2024" "Carrick2015" "CF4"; do
-for simname in "CLONES"; do
+for simname in "Carrick2015"; do
    # for catalogue in "2MTF" "SFI_gals" "CF4_TFR"; do
-    for catalogue in "Foundation" "2MTF" "SFI_gals" "CF4_TFR"; do
+    # for catalogue in "Foundation" "2MTF" "SFI_gals" "CF4_TFR"; do
+    for catalogue in "CF4_TFR"; do
        pythoncm="$env $file --catalogue $catalogue --nsims $nsims --simname $simname --MAS $MAS --grid $grid"
        if [ $on_login -eq 1 ]; then
            echo $pythoncm
--- a/scripts/field_prop/field_los_indranil_void.py
+++ b/scripts/field_prop/field_los_indranil_void.py
@ -12,29 +12,35 @@
 # with this program; if not, write to the Free Software Foundation, Inc.,
 # 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
 """Script to interpolate the Indranil void profiles for lines of sight."""
-from os.path import join
+from argparse import ArgumentParser
+from os.path import join, exists
+from os import remove

 import csiborgtools
 import numpy as np
 from astropy.coordinates import SkyCoord, angular_separation
+from field_los import get_los
 from h5py import File
 from mpi4py import MPI
 from scipy.interpolate import RegularGridInterpolator
 from tqdm import trange

-from field_los import get_los
-

 def interpolate_indranil_void(kind, nsims, RA, dec, rmax, dr, dump_folder,
                              catalogue):
-    if kind not in ["exp", "gauss"]:
+    if kind not in ["exp", "gauss", "mb"]:
        raise ValueError(f"Unknown void kind: `{kind}`.")

    kind = kind.upper()
    fdir = join("/mnt/extraspace/rstiskalek/catalogs", "IndranilVoid",
                f"{kind}profile")

-    fname_out = join(dump_folder, f"los_{catalogue}_IndranilVoid_{kind}.hdf5")
+    fname_out = join(
+        dump_folder, f"los_{catalogue}_IndranilVoid_{kind.lower()}.hdf5")
+    if exists(fname_out):
+        print(f"Fname `{fname_out}` already exists. Removing.")
+        remove(fname_out)
+
    print(f"Writing to `{fname_out}`.")
    for k in trange(len(nsims), desc="LG observers"):
        nsim = nsims[k]
@ -69,7 +75,7 @@ def interpolate_indranil_void(kind, nsims, RA, dec, rmax, dr, dump_folder,

        # Write the output, homogenous density.
        density = np.ones_like(result)
-        with File(fname_out, 'w') as f_out:
+        with File(fname_out, 'a') as f_out:
            f_out.create_dataset(f"rdist_{k}", data=r_eval * 0.674)
            f_out.create_dataset(f"density_{k}", data=density)
            f_out.create_dataset(f"velocity_{k}", data=result)
@ -81,19 +87,21 @@ def interpolate_indranil_void(kind, nsims, RA, dec, rmax, dr, dump_folder,


 if __name__ == "__main__":
-    kind = "gauss"
+    parser = ArgumentParser()
+    parser.add_argument("--kind", type=str, choices=["exp", "gauss", "mb"],
+                        help="Kind of void profile.")
+    parser.add_argument("--catalogue", type=str, help="Catalogue name.")
+    args = parser.parse_args()
+
    rmax = 165
-    dr = 1
+    dr = 0.75

    comm = MPI.COMM_WORLD
    paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
-    nsims = paths.get_ics(f"IndranilVoid_{kind}")
-
+    nsims = paths.get_ics(f"IndranilVoid_{args.kind}")
    out_folder = "/mnt/extraspace/rstiskalek/csiborg_postprocessing/field_los"

-    for catalogue in ["LOSS", "Foundation", "2MTF", "SFI_gals", "CF4_TFR"]:
-        print(f"Running kind `{kind}` for catalogue `{catalogue}`.")
-
-        RA, dec = get_los(catalogue, "", comm).T
-        interpolate_indranil_void(
-            kind, nsims, RA, dec, rmax, dr, out_folder, catalogue)
+    print(f"Running kind `{args.kind}` for catalogue `{args.catalogue}`.")
+    RA, dec = get_los(args.catalogue, "", comm).T
+    interpolate_indranil_void(
+        args.kind, nsims, RA, dec, rmax, dr, out_folder, args.catalogue)
--- a/scripts/field_prop/field_los_indranil_void.sh
+++ b/scripts/field_prop/field_los_indranil_void.sh
@ -0,0 +1,32 @@
+nthreads=1
+memory=7
+on_login=${1}
+queue="berg"
+env="/mnt/users/rstiskalek/csiborgtools/venv_csiborg/bin/python"
+file="field_los_indranil_void.py"
+
+
+if [ "$on_login" != "1" ] && [ "$on_login" != "0" ]
+then
+    echo "'on_login' (1) must be either 0 or 1."
+    exit 1
+fi
+
+
+for kind in "exp" "gauss" "mb"; do
+    for catalogue in "2MTF" "SFI_gals" "CF4_TFR"; do
+        pythoncm="$env $file --kind $kind --catalogue $catalogue"
+
+        if [ $on_login -eq 1 ]; then
+            echo $pythoncm
+            eval $pythoncm
+        else
+            cm="addqueue -q $queue -n $nthreads -m $memory $pythoncm"
+            echo "Submitting:"
+            echo $cm
+            echo
+            eval $cm
+        fi
+    done
+done
+
--- a/scripts/flow/flow_validation.py
+++ b/scripts/flow/flow_validation.py
@ -22,6 +22,19 @@ from argparse import ArgumentParser, ArgumentTypeError
 def none_or_int(value):
    if value.lower() == "none":
        return None
+
+    if "_" in value:
+        args = value.split("_")
+        if len(args) == 2:
+            k0, kf = args
+            dk = 1
+        elif len(args) == 3:
+            k0, kf, dk = args
+        else:
+            raise ArgumentTypeError(f"Invalid length of arguments: `{value}`.")
+
+        return [int(k) for k in range(int(k0), int(kf), int(dk))]
+
    try:
        return int(value)
    except ValueError:
@ -73,17 +86,17 @@ def print_variables(names, variables):
    print(flush=True)


-def get_models(get_model_kwargs, mag_selection, verbose=True):
+def get_models(ksim, get_model_kwargs, mag_selection, verbose=True):
    """Load the data and create the NumPyro models."""
    paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
    folder = "/mnt/extraspace/rstiskalek/catalogs/"

    nsims = paths.get_ics(ARGS.simname)
-    if ARGS.ksim is None:
+    if ksim is None:
        nsim_iterator = [i for i in range(len(nsims))]
    else:
-        nsim_iterator = [ARGS.ksim]
-        nsims = [nsims[ARGS.ksim]]
+        nsim_iterator = [ksim]
+        nsims = [nsims[ksim]]

    if verbose:
        print(f"{'Simulation:':<20} {ARGS.simname}")
@ -114,7 +127,7 @@ def get_models(get_model_kwargs, mag_selection, verbose=True):
        models[i] = csiborgtools.flow.get_model(
            loader, mag_selection=mag_selection[i], **get_model_kwargs)

-    print(f"\n{'Num. radial steps':<20} {len(loader.rdist)}\n", flush=True)
+    fprint(f"num. radial steps is {len(loader.rdist)}")
    return models


@ -143,7 +156,7 @@ def run_model(model, nsteps, nburn,  model_kwargs, out_folder,
        raise AttributeError("The models must have an attribute `ndata` "
                             "indicating the number of data points.") from e

-    nuts_kernel = NUTS(model, init_strategy=init_to_median(num_samples=1000))
+    nuts_kernel = NUTS(model, init_strategy=init_to_median(num_samples=10000))
    mcmc = MCMC(nuts_kernel, num_warmup=nburn, num_samples=nsteps)
    rng_key = jax.random.PRNGKey(42)

@ -283,11 +296,14 @@ if __name__ == "__main__":
    num_epochs = 50
    inference_method = "mike"
    mag_selection = None
-    sample_alpha = True
+    sample_alpha = False if "IndranilVoid_" in ARGS.simname else True
    sample_beta = None
+    sample_Vext = None
    sample_Vmono = False
    sample_mag_dipole = False
+    absolute_calibration = None
    calculate_harmonic = False if inference_method == "bayes" else True
+    sample_h = True if absolute_calibration is not None else False

    fname_kwargs = {"inference_method": inference_method,
                    "smooth": ARGS.ksmooth,
@ -297,8 +313,10 @@ if __name__ == "__main__":
                    "mag_selection": mag_selection,
                    "sample_alpha": sample_alpha,
                    "sample_beta": sample_beta,
+                    "sample_Vext": sample_Vext,
                    "sample_Vmono": sample_Vmono,
                    "sample_mag_dipole": sample_mag_dipole,
+                    "absolute_calibration": absolute_calibration,
                    }

    main_params = {"nsteps": nsteps, "nburn": nburn,
@ -310,6 +328,8 @@ if __name__ == "__main__":
                   "num_epochs": num_epochs,
                   "inference_method": inference_method,
                   "sample_mag_dipole": sample_mag_dipole,
+                   "absolute_calibration": absolute_calibration,
+                   "sample_h": sample_h,
                   }
    print_variables(main_params.keys(), main_params.values())

@ -335,8 +355,11 @@ if __name__ == "__main__":
                               "Vmono_min": -1000, "Vmono_max": 1000,
                               "beta_min": -10.0, "beta_max": 10.0,
                               "sigma_v_min": 1.0, "sigma_v_max": 750.,
+                               "h_min": 0.01, "h_max": 5.0,
+                               "sample_Vext": True if sample_Vext is None else sample_Vext,  # noqa
                               "sample_Vmono": sample_Vmono,
                               "sample_beta": sample_beta,
+                               "sample_h": sample_h,
                               }
    print_variables(
        calibration_hyperparams.keys(), calibration_hyperparams.values())
@ -353,17 +376,34 @@ if __name__ == "__main__":

    ###########################################################################

-    get_model_kwargs = {"zcmb_min": zcmb_min, "zcmb_max": zcmb_max}
-    models = get_models(get_model_kwargs, mag_selection)
-    model_kwargs = {
-        "models": models,
-        "field_calibration_hyperparams": calibration_hyperparams,
-        "distmod_hyperparams_per_model": distmod_hyperparams_per_catalogue,
-        "inference_method": inference_method,
+    get_model_kwargs = {
+        "zcmb_min": zcmb_min,
+        "zcmb_max": zcmb_max,
+        "absolute_calibration": absolute_calibration,
+        "calibration_fpath": "/mnt/extraspace/rstiskalek/catalogs/PV/CF4/CF4_TF_calibration.hdf5",  # noqa
        }

-    model = csiborgtools.flow.PV_validation_model
+    # In case we want to run multiple simulations independently.
+    if not isinstance(ARGS.ksim, list):
+        ksim_iterator = [ARGS.ksim]
+    else:
+        ksim_iterator = ARGS.ksim

-    run_model(model, nsteps, nburn, model_kwargs, out_folder,
-              calculate_harmonic, nchains_harmonic, num_epochs, kwargs_print,
-              fname_kwargs)
+    for i, ksim in enumerate(ksim_iterator):
+        if len(ksim_iterator) > 1:
+            print(f"{'Current simulation:':<20} {i + 1} ({ksim}) out of {len(ksim_iterator)}.")  # noqa
+
+        fname_kwargs["nsim"] = ksim
+        models = get_models(ksim, get_model_kwargs, mag_selection)
+        model_kwargs = {
+            "models": models,
+            "field_calibration_hyperparams": calibration_hyperparams,
+            "distmod_hyperparams_per_model": distmod_hyperparams_per_catalogue,
+            "inference_method": inference_method,
+            }
+
+        model = csiborgtools.flow.PV_validation_model
+
+        run_model(model, nsteps, nburn, model_kwargs, out_folder,
+                  calculate_harmonic, nchains_harmonic, num_epochs,
+                  kwargs_print, fname_kwargs)
--- a/scripts/flow/flow_validation.sh
+++ b/scripts/flow/flow_validation.sh
@ -37,42 +37,36 @@ else
 fi


-# for simname in "CLONES"; do
-# for simname in "csiborg2_main" "csiborg2X" ; do
-# for simname in "carrick2015" "lilow2024" "csiborg2_main" "csiborg2x" "cf4" "clones"; do
-for simname in "Carrick2015" "csiborg2_main"; do
-# for simname in "Carrick2015" "csiborg2X" "csiborg2_main"; do
-# for simname in "Carrick2015"; do
-    # for catalogue in "LOSS" "Foundation" "2MTF" "SFI_gals" "CF4_TFR_i" "CF4_TFR_w1"; do
-    for catalogue in "SFI_gals" "2MTF" "CF4_TFR_i"; do
+for simname in "IndranilVoid_exp" "IndranilVoid_gauss" "IndranilVoid_mb"; do
+    # for catalogue in "2MTF" "SFI_gals" "CF4_TFR_i" "CF4_TFR_w1"; do
    # for catalogue in "CF4_TFR_i" "CF4_TFR_w1"; do
-    # for catalogue in "2MTF" "SFI" "CF4_TFR_not2MTForSFI_i"; do
-    # for catalogue in "2MTF" "SFI_gals" "CF4_TFR_i"; do
-    # for catalogue in "CF4_TFR_w1"; do
-    # for catalogue in "CF4_GroupAll"; do
-    for ksmooth in 1 2 3 4; do
-        for ksim in "none"; do
+    for catalogue in "2MTF" "SFI_gals" "CF4_TFR_i" "CF4_TFR_w1"; do
+        # for ksim in "none"; do
+        # for ksim in 0; do
+        # for ksim in $(seq 0 5 500); do
+        for ksim in "0_100_5" "100_200_5" "200_300_5" "300_400_5" "400_500_5"; do
        # for ksim in {0..500}; do
-            pythoncm="$env $file --catalogue $catalogue --simname $simname --ksim $ksim --ksmooth $ksmooth --ndevice $ndevice --device $device"
+            for ksmooth in 0; do
+                pythoncm="$env $file --catalogue $catalogue --simname $simname --ksim $ksim --ksmooth $ksmooth --ndevice $ndevice --device $device"

-            if [ "$on_login" == "1" ]; then
-                echo $pythoncm
-                eval $pythoncm
-            else
-                if [ "$device" == "gpu" ]; then
-                    cm="addqueue -q $queue -s -m $memory --gpus 1 --gputype $gputype $pythoncm"
+                if [ "$on_login" == "1" ]; then
+                    echo $pythoncm
+                    eval $pythoncm
                else
-                    cm="addqueue -s -q $queue -n 1 -m $memory $pythoncm"
+                    if [ "$device" == "gpu" ]; then
+                        cm="addqueue -q $queue -s -m $memory --gpus 1 --gputype $gputype $pythoncm"
+                    else
+                        cm="addqueue -s -q $queue -n 1 -m $memory $pythoncm"
+                    fi
+                    echo "Submitting:"
+                    echo $cm
+                    eval $cm
                fi
-                echo "Submitting:"
-                echo $cm
-                eval $cm
-            fi

-            echo
-            sleep 0.01
+                echo
+                sleep 0.001

+            done
        done
    done
-done
 done