Aquila-Consortium/csiborgtools
0
0
Fork 0
mirror of https://github.com/Richard-Sti/csiborgtools.git synced 2024-12-22 01:28:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add Manticore (#130)

Browse source 
* Add CB2X paths

* Update imports

* Add CB2 params

* Add CB2x fields

* Add CB2X field

* Add parallelized

* Add **kwargs

* Update nb

* Minor updates
This commit is contained in:
Richard Stiskalek 2024-06-21 19:21:06 +01:00 committed by GitHub
parent d3b4bfd29c
commit ffaf92cd4b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 311 additions and 206 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

20
csiborgtools/flow/flow_model.py
View file

@ -29,9 +29,9 @@ import numpyro
import numpyro.distributions as dist
from astropy.cosmology import FlatLambdaCDM
from h5py import File
from jax import jit
from jax import devices, jit
from jax import numpy as jnp
from jax import vmap
from jax import pmap, vmap
from jax.lax import cond, scan
from jax.random import PRNGKey
from numdifftools import Hessian
@ -44,7 +44,7 @@ from tqdm import trange
from ..params import SPEED_OF_LIGHT, simname2Omega_m
from ..utils import fprint, radec_to_galactic
H0 = 100 # km / s / Mpc
H0 = 100 # km / s / Mpc
###############################################################################
@ -1113,7 +1113,7 @@ class TF_PV_validation_model(BaseFlowValidationModel):
mu = self.mu(a, b)
squared_e_mu = self.squared_e_mu(b, e_mu_intrinsic)
def scan_body(ll, i):
def scan_body(i):
# Calculate p(r) and multiply it by the galaxy bias
ptilde = self._f_ptilde_wo_bias(mu[i], squared_e_mu[i])
ptilde *= self._los_density[i]**alpha
@ -1126,12 +1126,16 @@ class TF_PV_validation_model(BaseFlowValidationModel):
ptilde *= calculate_likelihood_zobs(
self._z_obs[i], zobs_pred, sigma_v)
return ll + jnp.log(self._f_simps(ptilde) / pnorm), None
# return ll + jnp.log(self._f_simps(ptilde) / pnorm), None
return jnp.log(self._f_simps(ptilde) / pnorm)
ll = 0.
ll, __ = scan(scan_body, ll, jnp.arange(self.ndata))
def pmap_body(indxs):
return jnp.sum(vmap(scan_body)(indxs))
numpyro.factor("ll", ll)
# NOTE: move this elsewhere?
indxs = jnp.arange(self.ndata).reshape(len(devices()), -1)
ll = pmap(pmap_body)(indxs)
numpyro.factor("ll", jnp.sum(ll))
###############################################################################

2
csiborgtools/params.py
View file

@ -70,6 +70,7 @@ def simname2boxsize(simname):
"csiborg2_main": 676.6,
"csiborg2_varysmall": 676.6,
"csiborg2_random": 676.6,
"csiborg2X": 681.1,
"borg1": 677.7,
"borg2": 676.6,
"borg2_all": 676.6,
@ -91,6 +92,7 @@ def simname2Omega_m(simname):
"csiborg2_main": 0.3111,
"csiborg2_random": 0.3111,
"csiborg2_varysmall": 0.3111,
"csiborg2X": 0.306,
"borg1": 0.307,
"borg2": 0.3111,
"borg2_all": 0.3111,

4
csiborgtools/read/__init__.py
View file

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

35
csiborgtools/read/paths.py
View file

@ -16,7 +16,7 @@
CSiBORG paths manager.
"""
from glob import glob
from os import makedirs
from os import makedirs, listdir
from os.path import isdir, join
from warnings import warn
from re import search
@ -103,6 +103,12 @@ class Paths:
files = glob(join(self.csiborg2_varysmall_srcdir, "chain_*"))
files = [int(search(r'chain_16417_(\d+)', f).group(1))
for f in files]
elif simname == "csiborg2X":
# NOTE this too is preliminary
fname = "/mnt/extraspace/rstiskalek/MANTICORE/resimulations/fields/2MPP_N128_DES_PROD/R512" # noqa
fdirs = listdir(fname)
files = [int(search(r'\d+', fdir).group())
for fdir in fdirs if search(r'\d+', fdir)]
elif simname == "quijote":
files = glob(join(self.quijote_dir, "fiducial_processed",
"chain_*"))
@ -140,6 +146,8 @@ class Paths:
snaps = [int(search(r'snapshot_16417_(\d+)', f).group(1))
for f in snaps]
snaps = sorted(snaps)
elif simname == "csiborg2X":
raise ValueError("Snapshots not available for CSiBORG2X.")
elif simname == "quijote":
snaps = glob(join(self.quijote_dir, "fiducial_processed",
f"chain_{nsim}", "snapshot_*"))
@ -150,7 +158,7 @@ class Paths:
if has_ics:
snaps.insert(0, "ICs")
else:
raise ValueError(f"Unknown simulation name `{simname}`.")
raise ValueError(f"Unsupported simulation name `{simname}`.")
return snaps
def snapshot(self, nsnap, nsim, simname):
@ -168,12 +176,14 @@ class Paths:
fpath = join(self.csiborg2_varysmall_srcdir,
f"chain_16417_{str(nsim).zfill(3)}", "output",
f"snapshot_{str(nsnap).zfill(3)}.hdf5")
elif simname == "csiborg2X":
raise ValueError("Snapshots not available for CSiBORG2X.")
elif simname == "quijote":
fpath = join(self.quijote_dir, "fiducial_processed",
f"chain_{nsim}",
f"snapshot_{str(nsnap).zfill(3)}.hdf5")
else:
raise ValueError(f"Unknown simulation name `{simname}`.")
raise ValueError(f"Unsupported simulation name `{simname}`.")
return fpath
@ -192,12 +202,14 @@ class Paths:
return join(self.csiborg2_varysmall_srcdir,
f"chain_16417_{str(nsim).zfill(3)}", "output",
f"fof_subhalo_tab_{str(nsnap).zfill(3)}.hdf5")
elif simname == "csiborg2X":
raise ValueError("Snapshots not available for CSiBORG2X.")
elif simname == "quijote":
return join(self.quijote_dir, "fiducial_processed",
f"chain_{nsim}",
f"fof_{str(nsnap).zfill(3)}.hdf5")
else:
raise ValueError(f"Unknown simulation name `{simname}`.")
raise ValueError(f"Unsupported simulation name `{simname}`.")
def external_halo_catalogue(self, name):
"""Path to an external halo catalogue."""
@ -227,7 +239,7 @@ class Paths:
return join(self.quijote_dir, "fiducial_processed",
f"chain_{nsim}", "initial_lagpatch.npy")
else:
raise ValueError(f"Unknown simulation name `{simname}`.")
raise ValueError(f"Unsupported simulation name `{simname}`.")
def trees(self, nsim, simname):
"""Path to the halo trees of a simulation snapshot."""
@ -246,7 +258,7 @@ class Paths:
elif simname == "quijote":
raise ValueError("Trees not available for Quijote.")
else:
raise ValueError(f"Unknown simulation name `{simname}`.")
raise ValueError(f"Unsupported simulation name `{simname}`.")
def overlap(self, simname, nsim0, nsimx, min_logmass, smoothed):
"""
@ -344,9 +356,18 @@ class Paths:
return join(self.borg2_dir, f"mcmc_{nsim}.h5")
if simname == "borg1":
#
# NOTE Replace eventually with non-local paths
return f"/mnt/zfsusers/hdesmond/BORG_final/mcmc_{nsim}.h5"
if simname == "csiborg2X":
basedir = "/mnt/extraspace/rstiskalek/MANTICORE/resimulations/fields/2MPP_N128_DES_PROD/R512" # noqa
if kind == "overdensity":
return join(basedir, f"mcmc_{nsim}", "delta_field.hdf5")
elif kind == "velocity":
return join(basedir, f"mcmc_{nsim}", "vel_field.hdf5")
else:
raise ValueError(f"Unsupported CSiBORG2X field: `{kind}`.")
if MAS == "SPH" and kind in ["density", "velocity"]:
if simname == "csiborg1":
return join(self.csiborg1_srcdir, "field",

59
csiborgtools/read/snapshot.py
View file

@ -23,7 +23,7 @@ from os.path import join
import numpy
from h5py import File
from ..params import paths_glamdring, simname2boxsize
from ..params import paths_glamdring, simname2boxsize, simname2Omega_m
from .paths import Paths
from .util import find_boxed
@ -708,13 +708,7 @@ class CSiBORG2Field(BaseField):
@property
def kind(self):
"""
CSiBORG2 run kind.
Returns
-------
str
"""
"""CSiBORG2 run kind."""
return self._kind
@kind.setter
@ -772,6 +766,51 @@ class CSiBORG2Field(BaseField):
return numpy.load(fpath)
class CSiBORG2XField(BaseField):
"""
CSiBORG2X `z = 0` field class.
Parameters
----------
nsim : int
Simulation index.
paths : Paths, optional
Paths object. By default, the paths are set to the `glamdring` paths.
"""
def __init__(self, nsim, paths=None):
super().__init__(nsim, paths, False)
def overdensity_field(self, **kwargs):
fpath = self.paths.field(
"overdensity", None, None, self.nsim, "csiborg2X")
with File(fpath, "r") as f:
field = f["delta_cic"][...].astype(numpy.float32)
return field
def density_field(self, **kwargs):
field = self.overdensity_field()
omega0 = simname2Omega_m("csiborg2X")
rho_mean = omega0 * 277.53662724583074 # Msun / kpc^3
field += 1
field *= rho_mean
return field
def velocity_field(self, **kwargs):
fpath = self.paths.field(
"velocity", None, None, self.nsim, "csiborg2X")
with File(fpath, "r") as f:
v0 = f["v_0"][...]
v1 = f["v_1"][...]
v2 = f["v_2"][...]
field = numpy.array([v0, v1, v2])
return field
def radial_velocity_field(self, **kwargs):
raise RuntimeError("The radial velocity field is not available.")
###############################################################################
# BORG1 field class #
###############################################################################
@ -800,7 +839,7 @@ class BORG1Field(BaseField):
def density_field(self):
field = self.overdensity_field()
omega0 = 0.307
omega0 = simname2Omega_m("borg1")
rho_mean = omega0 * 277.53662724583074 # Msun / kpc^3
field += 1
field *= rho_mean
@ -841,7 +880,7 @@ class BORG2Field(BaseField):
def density_field(self):
field = self.overdensity_field()
omega0 = 0.3111
omega0 = simname2Omega_m("borg2")
rho_mean = omega0 * 277.53662724583074 # h^2 Msun / kpc^3
field += 1
field *= rho_mean

387
notebooks/flow/flow_calibration.ipynb
View file

File diff suppressed because one or more lines are too long

8
scripts/field_los.py
View file

@ -126,9 +126,11 @@ def get_field(simname, nsim, kind, MAS, grid):
# Open the field reader.
if simname == "csiborg1":
field_reader = csiborgtools.read.CSiBORG1Field(nsim)
elif "csiborg2" in simname:
elif "csiborg2_" in simname:
simkind = simname.split("_")[-1]
field_reader = csiborgtools.read.CSiBORG2Field(nsim, simkind)
elif simname == "csiborg2X":
field_reader = csiborgtools.read.CSiBORG2XField(nsim)
elif simname == "Carrick2015":
folder = "/mnt/extraspace/rstiskalek/catalogs"
warn(f"Using local paths from `{folder}`.", RuntimeWarning)
@ -158,9 +160,9 @@ def get_field(simname, nsim, kind, MAS, grid):
# Read in the field.
if kind == "density":
field = field_reader.density_field(MAS, grid)
field = field_reader.density_field(MAS=MAS, grid=grid)
elif kind == "velocity":
field = field_reader.velocity_field(MAS, grid)
field = field_reader.velocity_field(MAS=MAS, grid=grid)
else:
raise ValueError(f"Unknown field kind: `{kind}`.")

2
scripts/field_los.sh
View file

@ -1,5 +1,5 @@
nthreads=1
memory=32
memory=7
on_login=0
queue="berg"
env="/mnt/users/rstiskalek/csiborgtools/venv_csiborg/bin/python"