Add manticore snapshot

csiborgtools/read/__init__.py

@@ -16,11 +16,11 @@ from .catalogue import (CSiBORG1Catalogue, CSiBORG2Catalogue,
                         CSiBORG2SUBFINDCatalogue,                               # noqa
                         CSiBORG2MergerTreeReader, QuijoteCatalogue,             # noqa
                         MDPL2Catalogue, fiducial_observers)                     # noqa
-from .snapshot import (CSiBORG1Snapshot, CSiBORG2Snapshot, QuijoteSnapshot,     # noqa
-                       CSiBORG1Field, CSiBORG2Field, CSiBORG2XField,            # noqa
-                       QuijoteField, BORG2Field, BORG1Field, TNG300_1Field,     # noqa
-                       Carrick2015Field, Lilow2024Field, CLONESField,           # noqa
-                       CF4Field)                                                # noqa
+from .snapshot import (CSiBORG1Snapshot, CSiBORG2Snapshot, CSiBORG2XSnapshot,   # noqa
+                       QuijoteSnapshot, CSiBORG1Field, CSiBORG2Field,           # noqa
+                       CSiBORG2XField, QuijoteField, BORG2Field, BORG1Field,    # noqa
+                       TNG300_1Field, Carrick2015Field, Lilow2024Field,         # noqa
+                       CLONESField, CF4Field)                                   # noqa
 from .obs import (SDSS, MCXCClusters, PlanckClusters, TwoMPPGalaxies,           # noqa
                   TwoMPPGroups, ObservedCluster, match_array_to_no_masking,     # noqa
                   cols_to_structured, read_pantheonplus_data)                   # noqa

csiborgtools/read/paths.py

@@ -120,6 +120,8 @@ class Paths:
                 files = [int(search(r'mcmc_(\d+)', f).group(1)) for f in files]
                 # Downsample at the moment to only 20, instead of 40
                 files = files[::2]
+            else:
+                raise ValueError(f"Unknown MANTICORE simulation `{simname}`.")
         elif simname == "csiborg2X":
             # NOTE this too is preliminary
             fname = "/mnt/extraspace/rstiskalek/MANTICORE/resimulations/fields/2MPP_N128_DES_PROD/R512"  # noqa
@@ -216,9 +218,21 @@ class Paths:
                          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.")
+            raise ValueError("Snapshots not available for CSiBORG2X based on "
+                             "Stephen's ICs.")
         elif "manticore" in simname:
-            raise ValueError("Snapshots not available for MANTICORE.")
+            fdir = self.manticore_dir
+            if simname == "manticore_2MPP_N128_DES_V1":
+                if nsnap != 1:
+                    raise ValueError("Only snapshot 1 is available for "
+                                     "MANTICORE 2MPP_N128_DES_V1.")
+
+                fdir = join(fdir, "2MPP_N128_DES_V1", "resimulations", "R512")
+                nsnap_str = str(nsnap).zfill(4)
+                fpath = join(fdir, f"mcmc_{nsim}", "swift_monofonic",
+                             f"snap_{nsnap_str}", f"snap_{nsnap_str}.hdf5")
+            else:
+                raise ValueError(f"Unknown MANTICORE simulation `{simname}`.")
         elif simname == "quijote":
             fpath = join(self.quijote_dir, "fiducial_processed",
                          f"chain_{nsim}",

csiborgtools/read/snapshot.py

@@ -517,6 +517,91 @@ class CSiBORG2Snapshot(BaseSnapshot):
                             }
 
 
+###############################################################################
+#                         CSiBORG2x snapshot class                            #
+###############################################################################
+
+
+class CSiBORG2XSnapshot(BaseSnapshot):
+    """
+    CSiBORG2X snapshot from the SWIFT N-body simulations provided by Stuart
+    based on the Manticore ICs. The snapshots are at `z = 0` and correspond to
+    `manticore_2MPP_N128_DES_V1`.
+
+    Parameters
+    ----------
+    nsim : int
+        Simulation index.
+    nsnap : int
+        Snapshot index.
+    kind : str
+        CSiBORG2 run kind. One of `main`, `random`, or `varysmall`.
+    paths : Paths, optional
+        Paths object.
+    keep_snapshot_open : bool, optional
+        Whether to keep the snapshot file open when reading halo particles.
+        This is useful for repeated access to the snapshot.
+    flip_xz : bool, optional
+        Whether to flip the x- and z-axes to undo the MUSIC bug so that the
+        coordinates are consistent with observations.
+    """
+    def __init__(self, nsim, paths=None, keep_snapshot_open=False):
+        nsnap = 1
+        flip_xz = False
+        super().__init__(nsim, nsnap, paths, keep_snapshot_open, flip_xz)
+
+        fpath = self.paths.snapshot(self.nsnap, self.nsim,
+                                    "manticore_2MPP_N128_DES_V1")
+
+        self._snapshot_path = fpath
+        self._simname = "manticore_2MPP_N128_DES_V1"
+
+    def _get_particles(self, kind):
+        with File(self._snapshot_path, "r") as f:
+            h = f["Cosmology"].attrs["h"][0]
+
+            if kind == "Masses":
+                npart = f["Header"].attrs["NumPart_Total"][1]
+                mpart = f["Header"].attrs["InitialMassTable"][1] * 1e10 * h
+                x = np.ones(npart, dtype=np.float32) * mpart
+            else:
+                x = f[f"DMParticles/{kind}"][...]
+
+            # Convert coordinates to Mpc / h
+            if kind == "Coordinates":
+                x *= h
+        return x
+
+    def coordinates(self):
+        return self._get_particles("Coordinates")
+
+    def velocities(self):
+        return self._get_particles("Velocities")
+
+    def masses(self):
+        return self._get_particles("Masses")
+
+    def particle_ids(self):
+        raise NotImplementedError("Recovering particle IDs is not implemented "
+                                  "for CSiBORG2X.")
+
+    def _get_halo_particles(self, halo_id, kind, is_group):
+        raise NotImplementedError("Recovering halo particles is not "
+                                  "implemented for CSiBORG2X.")
+
+    def halo_coordinates(self, halo_id, is_group=True):
+        return self._get_halo_particles(halo_id, "Coordinates", is_group)
+
+    def halo_velocities(self, halo_id, is_group=True):
+        return self._get_halo_particles(halo_id, "Velocities", is_group)
+
+    def halo_masses(self, halo_id, is_group=True):
+        return self._get_halo_particles(halo_id, "Masses", is_group) * 1e10
+
+    def _make_hid2offset(self):
+        raise NotImplementedError("Recovering halo offsets is not implemented"
+                                  "for CSiBORG2X.")
+
 ###############################################################################
 #                          CSiBORG2 snapshot class                            #
 ###############################################################################