Fix catalogues and add CSiBORG 2 catalogue support (#98)

* Add box selection to snapshot

* Edit imports

* Rename catalogue

* Add beautified catalogues

* Fix paths issue

* Create params to avoid circular import

* Add CSiBORG2 catalogue support

* Remove print statement

* Edit docs

csiborgtools/__init__.py

@@ -18,51 +18,7 @@ from .utils import (center_of_mass, delta2ncells, number_counts,
                     periodic_distance, periodic_distance_two_points,            # noqa
                     binned_statistic, cosine_similarity, fprint,                # noqa
                     hms_to_degrees, dms_to_degrees, great_circle_distance)      # noqa
-
-# Arguments to csiborgtools.read.Paths.
-paths_glamdring = {
-    "csiborg1_srcdir": "/mnt/extraspace/rstiskalek/csiborg1",
-    "csiborg2_main_srcdir": "/mnt/extraspace/rstiskalek/csiborg2_main",
-    "csiborg2_varysmall_srcdir": "/mnt/extraspace/rstiskalek/csiborg2_varysmall",   # noqa
-    "csiborg2_random_srcdir": "/mnt/extraspace/rstiskalek/csiborg2_random",         # noqa
-    "postdir": "/mnt/extraspace/rstiskalek/csiborg_postprocessing/",
-    "quijote_dir": "/mnt/extraspace/rstiskalek/quijote",
-    }
-
-
-neighbour_kwargs = {"rmax_radial": 155 / 0.705,
-                    "nbins_radial": 50,
-                    "rmax_neighbour": 100.,
-                    "nbins_neighbour": 150,
-                    "paths_kind": paths_glamdring}
-
-
-def simname2boxsize(simname):
-    """
-    Return boxsize in `Mpc/h` for a given simname.
-
-    Parameters
-    ----------
-    simname : str
-        Simulation name.
-
-    Returns
-    -------
-    boxsize : float
-    """
-    d = {"csiborg1": 677.7,
-         "csiborg2_main": 676.6,
-         "csiborg2_varysmall": 676.6,
-         "csiborg2_random": 676.6,
-         "quijote": 1000.
-         }
-
-    boxsize = d.get(simname, None)
-
-    if boxsize is None:
-        raise ValueError("Unknown simname: {}".format(simname))
-
-    return boxsize
+from .params import paths_glamdring, simname2boxsize                            # noqa
 
 
 ###############################################################################

csiborgtools/params.py

@@ -0,0 +1,62 @@
+# Copyright (C) 2023 Richard Stiskalek
+# This program is free software; you can redistribute it and/or modify it
+# under the terms of the GNU General Public License as published by the
+# Free Software Foundation; either version 3 of the License, or (at your
+# option) any later version.
+#
+# This program is distributed in the hope that it will be useful, but
+# WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General
+# Public License for more details.
+#
+# You should have received a copy of the GNU General Public License along
+# with this program; if not, write to the Free Software Foundation, Inc.,
+# 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
+"""
+Various user parameters for csiborgtools.
+"""
+
+
+def simname2boxsize(simname):
+    """
+    Return boxsize in `Mpc/h` for a given simname.
+
+    Parameters
+    ----------
+    simname : str
+        Simulation name.
+
+    Returns
+    -------
+    boxsize : float
+    """
+    d = {"csiborg1": 677.7,
+         "csiborg2_main": 676.6,
+         "csiborg2_varysmall": 676.6,
+         "csiborg2_random": 676.6,
+         "quijote": 1000.
+         }
+
+    boxsize = d.get(simname, None)
+
+    if boxsize is None:
+        raise ValueError("Unknown simname: {}".format(simname))
+
+    return boxsize
+
+
+paths_glamdring = {
+    "csiborg1_srcdir": "/mnt/extraspace/rstiskalek/csiborg1",
+    "csiborg2_main_srcdir": "/mnt/extraspace/rstiskalek/csiborg2_main",
+    "csiborg2_varysmall_srcdir": "/mnt/extraspace/rstiskalek/csiborg2_varysmall",   # noqa
+    "csiborg2_random_srcdir": "/mnt/extraspace/rstiskalek/csiborg2_random",         # noqa
+    "postdir": "/mnt/extraspace/rstiskalek/csiborg_postprocessing/",
+    "quijote_dir": "/mnt/extraspace/rstiskalek/quijote",
+    }
+
+
+# neighbour_kwargs = {"rmax_radial": 155 / 0.705,
+#                     "nbins_radial": 50,
+#                     "rmax_neighbour": 100.,
+#                     "nbins_neighbour": 150,
+#                     "paths_kind": paths_glamdring}

csiborgtools/read/__init__.py

@@ -12,8 +12,7 @@
 # You should have received a copy of the GNU General Public License along
 # with this program; if not, write to the Free Software Foundation, Inc.,
 # 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
-from .catalogue import (CSiBORGCatalogue, QuijoteCatalogue,                     # noqa
-                        fiducial_observers)                                     # noqa
+from .catalogue import CSiBORG1Catalogue, CSiBORG2Catalogue, QuijoteCatalogue   # noqa
 from .snapshot import (CSIBORG1Snapshot, CSIBORG2Snapshot, QuijoteSnapshot,     # noqa
                        CSiBORG1Field, CSiBORG2Field, QuijoteField)              # noqa
 from .obs import (SDSS, MCXCClusters, PlanckClusters, TwoMPPGalaxies,           # noqa

csiborgtools/read/snapshot.py

@@ -18,10 +18,14 @@ should be implemented things such as flipping x- and z-axes, to make sure that
 observed RA-dec can be mapped into the simulation box.
 """
 from abc import ABC, abstractmethod, abstractproperty
-import numpy
 
+import numpy
 from h5py import File
 
+from ..params import paths_glamdring, simname2boxsize
+from .paths import Paths
+from .util import find_boxed
+
 ###############################################################################
 #                          Base snapshot class                                #
 ###############################################################################
@@ -65,6 +69,30 @@ class BaseSnapshot(ABC):
         """
         return self._nsnap
 
+    @property
+    def simname(self):
+        """
+        Simulation name.
+
+        Returns
+        -------
+        str
+        """
+        if self._simname is None:
+            raise ValueError("Simulation name not set.")
+        return self._simname
+
+    @property
+    def boxsize(self):
+        """
+        Simulation boxsize in `cMpc/h`.
+
+        Returns
+        -------
+        float
+        """
+        return simname2boxsize(self.simname)
+
     @property
     def paths(self):
         """
@@ -74,6 +102,8 @@ class BaseSnapshot(ABC):
         -------
         Paths
         """
+        if self._paths is None:
+            self._paths = Paths(**paths_glamdring)
         return self._paths
 
     @abstractproperty
@@ -191,6 +221,27 @@ class BaseSnapshot(ABC):
         """
         pass
 
+    def select_box(self, center, boxwidth):
+        """
+        Find particle coordinates of particles within a box of size `boxwidth`
+        centered on `center`.
+
+        Parameters
+        ----------
+        center : 1-dimensional array
+            Center of the box.
+        boxwidth : float
+            Width of the box.
+
+        Returns
+        -------
+        pos : 2-dimensional array
+        """
+        pos = self.coordinates()
+        mask = find_boxed(pos, center, boxwidth, self.boxsize)
+
+        return pos[mask]
+
 
 ###############################################################################
 #                          CSiBORG1 snapshot class                            #
@@ -208,13 +259,14 @@ class CSIBORG1Snapshot(BaseSnapshot):
         Simulation index.
     nsnap : int
         Snapshot index.
-    paths : Paths
+    paths : Paths, optional
         Paths object.
     """
-    def __init__(self, nsim, nsnap, paths):
+    def __init__(self, nsim, nsnap, paths=None):
         super().__init__(nsim, nsnap, paths)
         self._snapshot_path = self.paths.snapshot(
             self.nsnap, self.nsim, "csiborg1")
+        self._simname = "csiborg1"
 
     def _get_particles(self, kind):
         with File(self._snapshot_path, "r") as f:
@@ -285,17 +337,18 @@ class CSIBORG2Snapshot(BaseSnapshot):
         Simulation index.
     nsnap : int
         Snapshot index.
-    paths : Paths
-        Paths object.
     kind : str
         CSiBORG2 run kind. One of `main`, `random`, or `varysmall`.
+    paths : Paths, optional
+        Paths object.
     """
-    def __init__(self, nsim, nsnap, paths, kind):
+    def __init__(self, nsim, nsnap, kind, paths=None):
         super().__init__(nsim, nsnap, paths)
         self.kind = kind
 
         self._snapshot_path = self.paths.snapshot(
             self.nsnap, self.nsim, f"csiborg2_{self.kind}")
+        self._simname = f"csiborg2_{self.kind}"
 
     @property
     def kind(self):
@@ -426,13 +479,14 @@ class QuijoteSnapshot(CSIBORG1Snapshot):
         Simulation index.
     nsnap : int
         Snapshot index.
-    paths : Paths
+    paths : Paths, optional
         Paths object.
     """
-    def __init__(self, nsim, nsnap, paths):
+    def __init__(self, nsim, nsnap, paths=None):
         super().__init__(nsim, nsnap, paths)
         self._snapshot_path = self.paths.snapshot(self.nsnap, self.nsim,
                                                   "quijote")
+        self._simname = "quijote"
 
     def _make_hid2offset(self):
         catalogue_path = self.paths.snapshot_catalogue(

csiborgtools/read/util.py

@@ -0,0 +1,81 @@
+# Copyright (C) 2022 Richard Stiskalek
+# This program is free software; you can redistribute it and/or modify it
+# under the terms of the GNU General Public License as published by the
+# Free Software Foundation; either version 3 of the License, or (at your
+# option) any later version.
+#
+# This program is distributed in the hope that it will be useful, but
+# WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General
+# Public License for more details.
+#
+# You should have received a copy of the GNU General Public License along
+# with this program; if not, write to the Free Software Foundation, Inc.,
+# 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
+import numpy
+from numba import jit
+
+###############################################################################
+#                Select particles in a box utility functions                  #
+###############################################################################
+
+
+@jit(nopython=True, fastmath=True, boundscheck=False)
+def pbc_distance(x1, x2, boxsize):
+    """Calculate periodic distance between two points."""
+    delta = abs(x1 - x2)
+    return min(delta, boxsize - delta)
+
+
+@jit(nopython=True, fastmath=True, boundscheck=False)
+def find_next_particle(start_index, end_index, pos, x0, y0, z0,
+                       half_width, boxsize):
+    """
+    Find the next particle in a box of size `half_width` centered on `x0`,
+    `y0`, `z0`, where the periodic simulation box size is `boxsize`.
+    """
+    for i in range(start_index, end_index):
+        x, y, z = pos[i]
+        if ((pbc_distance(x, x0, boxsize) < half_width) and (pbc_distance(y, y0, boxsize) < half_width) and (pbc_distance(z, z0, boxsize) < half_width)):  # noqa
+            return i
+
+    return None
+
+
+def find_boxed(pos, center, subbox_size, boxsize):
+    """
+    Find indices of `pos` in a box of size `subbox_size` centered on
+    `center`, where the simulation box size is `boxsize`.
+
+    Parameters
+    ----------
+    pos : 2-dimensional array of shape (nsamples, 3)
+        Positions of all particles in the simulation.
+    center : 1-dimensional array
+        Center of the sub-box.
+    subbox_size : float
+        Size of the sub-box.
+    boxsize : float
+        Size of the simulation box.
+
+    Returns
+    -------
+    indxs : 1-dimensional array of shape
+    """
+    if isinstance(center, list):
+        center = numpy.asanyarray(center)
+
+    half_width = subbox_size / 2.
+
+    indxs, start_index, end_index = [], 0, len(pos)
+    while True:
+        i = find_next_particle(start_index, end_index, pos,
+                               *center, half_width, boxsize)
+
+        if i is None:
+            break
+
+        indxs.append(i)
+        start_index = i + 1
+
+    return indxs