Add Manticore (#130)

* Add CB2X paths

* Update imports

* Add CB2 params

* Add CB2x fields

* Add CB2X field

* Add parallelized

* Add **kwargs

* Update nb

* Minor updates

csiborgtools/flow/flow_model.py

@@ -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.
+        def pmap_body(indxs):
-        ll, __ = scan(scan_body, ll, jnp.arange(self.ndata))
+            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))
 
 
 ###############################################################################

csiborgtools/params.py

@@ -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,

csiborgtools/read/__init__.py

@@ -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
+                       CSiBORG1Field, CSiBORG2Field, CSiBORG2XField,            # noqa
-                       BORG1Field, TNG300_1Field)                               # 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

csiborgtools/read/paths.py

@@ -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",

csiborgtools/read/snapshot.py

@@ -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."""
-        CSiBORG2 run kind.
-
-        Returns
-        -------
-        str
-        """
         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

scripts/field_los.py

@@ -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}`.")
 

scripts/field_los.sh

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