Update field calculations (#51)

* Update density field routines

* Add utils

* Add possibility to return 2d array

* Styling

* Fixed density field

* Fix bugs

* add mpc2box

* Fix evaluating sky

* Fix sky map making

* Rename file

* Add paths if only positions

* Add option to dump particles only

* Add comments

csiborgtools/field/density.py

@@ -15,6 +15,7 @@
 """
 Density field and cross-correlation calculations.
 """
+from abc import ABC
 from warnings import warn
 
 import MAS_library as MASL
@@ -23,73 +24,26 @@ import Pk_library as PKL
 import smoothing_library as SL
 from tqdm import trange
 
+from .utils import force_single_precision
+from ..read.utils import radec_to_cartesian
 
-class DensityField:
-    r"""
-    Density field calculations. Based primarily on routines of Pylians [1].
 
-    Parameters
+class BaseField(ABC):
-    ----------
+    """Base class for density field calculations."""
-    particles : structured array
-        Particle array. Must contain keys `['x', 'y', 'z', 'M']`. Particle
-        coordinates are assumed to be :math:`\in [0, 1]` or in box units
-        otherwise.
-    boxsize : float
-        Box length. Multiplies `particles` positions to fix the power spectum
-        units.
-    box : :py:class:`csiborgtools.units.BoxUnits`
-        The simulation box information and transformations.
-    MAS : str, optional
-        Mass assignment scheme. Options are Options are: 'NGP' (nearest grid
-        point), 'CIC' (cloud-in-cell), 'TSC' (triangular-shape cloud), 'PCS'
-        (piecewise cubic spline).
-
-    References
-    ----------
-    [1] https://pylians3.readthedocs.io/
-    """
-    _particles = None
-    _boxsize = None
     _box = None
     _MAS = None
 
-    def __init__(self, particles, boxsize, box, MAS="CIC"):
-        self.particles = particles
-        assert boxsize > 0
-        self.boxsize = boxsize
-        self.box = box
-        assert MAS in ["NGP", "CIC", "TSC", "PCS"]
-        self._MAS = MAS
-
-    @property
-    def particles(self):
-        """
-        Particles structured array.
-
-        Returns
-        -------
-        particles : structured array
-        """
-        return self._particles
-
-    @particles.setter
-    def particles(self, particles):
-        """Set `particles`, checking it has the right columns."""
-        if any(p not in particles.dtype.names for p in ('x', 'y', 'z', 'M')):
-            raise ValueError("`particles` must be a structured array "
-                             "containing `['x', 'y', 'z', 'M']`.")
-        self._particles = particles
-
     @property
     def boxsize(self):
         """
-        Box length. Determines the power spectrum units.
+        Box size. Particle positions are always assumed to be in box units,
+        therefore this is 1.
 
         Returns
         -------
         boxsize : float
         """
-        return self._boxsize
+        return 1.
 
     @property
     def box(self):
@@ -119,323 +73,91 @@ class DensityField:
         -------
         MAS : str
         """
+        if self._MAS is None:
+            raise ValueError("`mas` is not set.")
         return self._MAS
 
-    @staticmethod
+    @MAS.setter
-    def _force_f32(x, name):
+    def MAS(self, MAS):
-        if x.dtype != numpy.float32:
+        assert MAS in ["NGP", "CIC", "TSC", "PCS"]
-            warn("Converting `{}` to float32.".format(name), stacklevel=1)
+        self._MAS = MAS
-            x = x.astype(numpy.float32)
-        return x
 
-    def density_field(self, grid, smooth_scale=None, verbose=True):
+    def evaluate_cartesian(self, *fields, pos):
         """
-        Calculate the density field using a Pylians routine [1, 2]. Enforces
+        Evaluate a scalar field at Cartesian coordinates using CIC
-        float32 precision.
-
-        Parameters
-        ----------
-        grid : int
-            Grid size.
-        smooth_scale : float, optional
-            Scale to smoothen the density field, in units matching
-            `self.boxsize`. By default no smoothing is applied.
-        verbose : bool
-            Verbosity flag.
-
-        Returns
-        -------
-        rho : 3-dimensional array of shape `(grid, grid, grid)`.
-
-        References
-        ----------
-        [1] https://pylians3.readthedocs.io/
-        [2] https://github.com/franciscovillaescusa/Pylians3/blob/master
-            /library/MAS_library/MAS_library.pyx
-        """
-        pos = numpy.vstack([self.particles[p] for p in ('x', 'y', 'z')]).T
-        pos *= self.boxsize
-        pos = self._force_f32(pos, "pos")
-        weights = self._force_f32(self.particles['M'], 'M')
-
-        # Pre-allocate and do calculations
-        rho = numpy.zeros((grid, grid, grid), dtype=numpy.float32)
-        MASL.MA(pos, rho, self.boxsize, self.MAS, W=weights, verbose=verbose)
-        if smooth_scale is not None:
-            rho = self.smooth_field(rho, smooth_scale)
-        return rho
-
-    def overdensity_field(self, grid, smooth_scale=None, verbose=True):
-        r"""
-        Calculate the overdensity field using Pylians routines.
-        Defined as :math:`\rho/ <\rho> - 1`.
-
-        Parameters
-        ----------
-        grid : int
-            Grid size.
-        smooth_scale : float, optional
-            Scale to smoothen the density field, in units matching
-            `self.boxsize`. By default no smoothing is applied.
-        verbose : bool
-            Verbosity flag.
-
-        Returns
-        -------
-        overdensity : 3-dimensional array of shape `(grid, grid, grid)`.
-        """
-        # Get the overdensity
-        delta = self.density_field(grid, smooth_scale, verbose)
-        delta /= delta.mean()
-        delta -= 1
-        return delta
-
-    def potential_field(self, grid, smooth_scale=None, verbose=True):
-        """
-        Calculate the potential field using Pylians routines.
-
-        Parameters
-        ----------
-        grid : int
-            Grid size.
-        smooth_scale : float, optional
-            Scale to smoothen the density field, in units matching
-            `self.boxsize`. By default no smoothing is applied.
-        verbose : bool
-            Verbosity flag.
-
-        Returns
-        -------
-        potential : 3-dimensional array of shape `(grid, grid, grid)`.
-        """
-        delta = self.overdensity_field(grid, smooth_scale, verbose)
-        if verbose:
-            print("Calculating potential from the overdensity..")
-        return MASL.potential(
-            delta, self.box._omega_m, self.box._aexp, self.MAS)
-
-    def gravitational_field(self, grid, smooth_scale=None, verbose=True):
-        """
-        Calculate the gravitational vector field. Note that this method is
-        only defined in a fork of `Pylians`.
-
-        Parameters
-        ----------
-        grid : int
-            Grid size.
-        smooth_scale : float, optional
-            Scale to smoothen the density field, in units matching
-            `self.boxsize`. By default no smoothing is applied.
-        verbose : bool
-            Verbosity flag.
-
-        Returns
-        -------
-        grav_field_tensor : :py:class:`MAS_library.grav_field_tensor`
-            Tidal tensor object, whose attributes `grav_field_tensor.gi`
-            contain the relevant tensor components.
-        """
-        delta = self.overdensity_field(grid, smooth_scale, verbose)
-        return MASL.grav_field_tensor(
-            delta, self.box._omega_m, self.box._aexp, self.MAS)
-
-    def tensor_field(self, grid, smooth_scale=None, verbose=True):
-        """
-        Calculate the tidal tensor field.
-
-        Parameters
-        ----------
-        grid : int
-            Grid size.
-        smooth_scale : float, optional
-            Scale to smoothen the density field, in units matching
-            `self.boxsize`. By default no smoothing is applied.
-        verbose : bool, optional
-            A verbosity flag.
-
-        Returns
-        -------
-        tidal_tensor : :py:class:`MAS_library.tidal_tensor`
-            Tidal tensor object, whose attributes `tidal_tensor.Tij` contain
-            the relevant tensor components.
-        """
-        delta = self.overdensity_field(grid, smooth_scale, verbose)
-        return MASL.tidal_tensor(
-            delta, self.box._omega_m, self.box._aexp, self.MAS)
-
-    def auto_powerspectrum(self, grid, smooth_scale, verbose=True):
-        """
-        Calculate the auto 1-dimensional power spectrum.
-
-        Parameters
-        ----------
-        grid : int
-            Grid size.
-        smooth_scale : float, optional
-            Scale to smoothen the density field, in units matching
-            `self.boxsize`. By default no smoothing is applied.
-        verbose : bool, optional
-            Verbosity flag.
-
-        Returns
-        -------
-        pk : py:class`Pk_library.Pk`
-        """
-        delta = self.overdensity_field(grid, smooth_scale, verbose)
-        return PKL.Pk(
-            delta, self.boxsize, axis=1, MAS=self.MAS, threads=1,
-            verbose=verbose)
-
-    def smooth_field(self, field, smooth_scale, threads=1):
-        """
-        Smooth a field with a Gaussian filter.
-
-        Parameters
-        ----------
-        field : 3-dimensional array of shape `(grid, grid, grid)`
-            The field to be smoothed.
-        smooth_scale : float, optional
-            Scale to smoothen the density field, in units matching
-            `self.boxsize`. By default no smoothing is applied.
-        threads : int, optional
-            Number of threads. By default 1.
-
-        Returns
-        -------
-        smoothed_field : 3-dimensional array of shape `(grid, grid, grid)`
-        """
-        Filter = "Gaussian"
-        grid = field.shape[0]
-        # FFT of the filter
-        W_k = SL.FT_filter(self.boxsize, smooth_scale, grid, Filter, threads)
-        return SL.field_smoothing(field, W_k, threads)
-
-    def evaluate_field(self, *field, pos):
-        """
-        Evaluate the field at Cartesian coordinates using CIC interpolation.
-
-        Parameters
-        ----------
-        field : (list of) 3-dimensional array of shape `(grid, grid, grid)`
-            Density field that is to be interpolated.
-        pos : 2-dimensional array of shape `(n_samples, 3)`
-            Positions to evaluate the density field. The coordinates span range
-            of [0, boxsize].
-
-        Returns
-        -------
-        interp_field : (list of) 1-dimensional array of shape `(n_samples,).
-        """
-        self._force_f32(pos, "pos")
-
-        interp_field = [numpy.zeros(pos.shape[0], dtype=numpy.float32)
-                        for __ in range(len(field))]
-        for i, f in enumerate(field):
-            MASL.CIC_interp(f, self.boxsize, pos, interp_field[i])
-        return interp_field
-
-    def evaluate_sky(self, *field, pos, isdeg=True):
-        """
-        Evaluate the field at given distance, right ascension and declination.
-        Assumes that the observed is in the centre of the box and uses CIC
         interpolation.
 
         Parameters
         ----------
         field : (list of) 3-dimensional array of shape `(grid, grid, grid)`
-            Density field that is to be interpolated. Assumed to be defined
+            Fields to be interpolated.
-            on a Cartesian grid.
         pos : 2-dimensional array of shape `(n_samples, 3)`
-            Spherical coordinates to evaluate the field. Should be distance,
+            Positions to evaluate the density field. Assumed to be in box
+            units.
+
+        Returns
+        -------
+        interp_fields : (list of) 1-dimensional array of shape `(n_samples,).
+        """
+        pos = force_single_precision(pos, "pos")
+
+        nsamples = pos.shape[0]
+        interp_fields = [numpy.full(nsamples, numpy.nan, dtype=numpy.float32)
+                         for __ in range(len(fields))]
+        for i, field in enumerate(fields):
+            MASL.CIC_interp(field, self.boxsize, pos, interp_fields[i])
+
+        if len(fields) == 1:
+            return interp_fields[0]
+        return interp_fields
+
+    def evaluate_sky(self, *fields, pos, isdeg=True):
+        """
+        Evaluate the scalar fields at given distance, right ascension and
+        declination. Assumes an observed in the centre of the box, with
+        distance being in :math:`Mpc`. Uses CIC interpolation.
+
+        Parameters
+        ----------
+        fields : (list of) 3-dimensional array of shape `(grid, grid, grid)`
+            Field to be interpolated.
+        pos : 2-dimensional array of shape `(n_samples, 3)`
+            Spherical coordinates to evaluate the field. Columns are distance,
             right ascension, declination, respectively.
         isdeg : bool, optional
             Whether `ra` and `dec` are in degres. By default `True`.
 
         Returns
         -------
-        interp_field : (list of) 1-dimensional array of shape `(n_samples,).
+        interp_fields : (list of) 1-dimensional array of shape `(n_samples,).
         """
-        # TODO: implement this
+        pos = force_single_precision(pos, "pos")
-        raise NotImplementedError("This method is not yet implemented.")
+        # We first calculate convert the distance to box coordinates and then
-#         self._force_f32(pos, "pos")
+        # convert to Cartesian coordinates.
-#         X = numpy.vstack(
+        X = numpy.copy(pos)
-#             radec_to_cartesian(*(pos[:, i] for i in range(3)), isdeg)).T
+        X[:, 0] = self.box.mpc2box(X[:, 0])
-#         X = X.astype(numpy.float32)
+        X = radec_to_cartesian(pos, isdeg)
-#         # Place the observer at the center of the box
+        # Then we move the origin to match the box coordinates
-#         X += 0.5 * self.boxsize
+        X -= 0.5
-#         return self.evaluate_field(*field, pos=X)
+        return self.evaluate_field(*fields, pos=X)
 
-    @staticmethod
+    def make_sky(self, field, angpos, dist, verbose=True):
-    def gravitational_field_norm(gx, gy, gz):
+        r"""
-        """
+        Make a sky map of a scalar field. The observer is in the centre of the
-        Calculate the norm (magnitude) of a gravitational field.
+        box the field is evaluated along directions `angpos`. Along each
+        direction, the field is evaluated distances `dist_marg` and summed.
+        Uses CIC interpolation.
 
         Parameters
         ----------
-        gx, gy, gz : 1-dimensional arrays of shape `(n_samples,)`
-            Gravitational field Cartesian components.
-
-        Returns
-        -------
-        g : 1-dimensional array of shape `(n_samples,)`
-        """
-        return numpy.sqrt(gx * gx + gy * gy + gz * gz)
-
-    @staticmethod
-    def tensor_field_eigvals(T00, T01, T02, T11, T12, T22):
-        """
-        Calculate the eigenvalues of a symmetric tensor field. Eigenvalues are
-        sorted in increasing order.
-
-        Parameters
-        ----------
-        T00, T01, T02, T11, T12, T22 : 1-dim arrays of shape `(n_samples,)`
-            Tensor field upper components evaluated for each sample.
-
-        Returns
-        -------
-        eigvals : 2-dimensional array of shape `(n_samples, 3)`
-        """
-        n_samples = T00.size
-        # Fill array of shape `(n_samples, 3, 3)` to calculate eigvals
-        Teval = numpy.full((n_samples, 3, 3), numpy.nan, dtype=numpy.float32)
-        Teval[:, 0, 0] = T00
-        Teval[:, 0, 1] = T01
-        Teval[:, 0, 2] = T02
-        Teval[:, 1, 1] = T11
-        Teval[:, 1, 2] = T12
-        Teval[:, 2, 2] = T22
-
-        # Calculate the eigenvalues
-        eigvals = numpy.full((n_samples, 3), numpy.nan, dtype=numpy.float32)
-        for i in range(n_samples):
-            eigvals[i, :] = numpy.linalg.eigvalsh(Teval[i, ...], 'U')
-            eigvals[i, :] = numpy.sort(eigvals[i, :])
-
-        return eigvals
-
-    def make_sky_map(self, ra, dec, field, dist_marg, isdeg=True,
-                     verbose=True):
-        """
-        Make a sky map of a density field. Places the observed in the center of
-        the box and evaluates the field in directions `ra`, `dec`. At each such
-        position evaluates the field at distances `dist_marg` and sums these
-        interpolated values of the field.
-
-        NOTE: Supports only scalar fields.
-
-        Parameters
-        ----------
-        ra, dec : 1-dimensional arrays of shape `(n_pos, )`
-            Directions to evaluate the field. Assumes `dec` is in [-90, 90]
-            degrees (or equivalently in radians).
         field : 3-dimensional array of shape `(grid, grid, grid)`
-            Density field that is to be interpolated. Assumed to be defined
+            Field to be interpolated
-            on a Cartesian grid `[0, self.boxsize]^3`.
+        angpos : 2-dimensional arrays of shape `(ndir, 2)`
-        dist_marg : 1-dimensional array
+            Directions to evaluate the field. Assumed to be RA
+            :math:`\in [0, 360]` and dec :math:`\in [-90, 90]` degrees,
+            respectively.
+        dist : 1-dimensional array
             Radial distances to evaluate the field.
-        isdeg : bool, optional
-            Whether `ra` and `dec` are in degres. By default `True`.
         verbose : bool, optional
             Verbosity flag.
 
@@ -443,23 +165,276 @@ class DensityField:
         -------
         interp_field : 1-dimensional array of shape `(n_pos, )`.
         """
-        # Angular positions at which to evaluate the field
+        assert angpos.ndim == 2 and dist.ndim == 1
-        Nang = ra.size
+        # We loop over the angular directions, at each step evaluating a vector
-        pos = numpy.vstack([ra, dec]).T
+        # of distances. We pre-allocate arrays for speed.
-
+        dir_loop = numpy.full((dist.size, 3), numpy.nan, dtype=numpy.float32)
-        # Now loop over the angular positions, each time evaluating a vector
+        ndir = angpos.shape[0]
-        # of distances. Pre-allocate arrays for speed
+        out = numpy.zeros(ndir, numpy.nan, dtype=numpy.float32)
-        ra_loop = numpy.ones_like(dist_marg)
+        for i in trange(ndir) if verbose else range(ndir):
-        dec_loop = numpy.ones_like(dist_marg)
+            dir_loop[1, :] = angpos[i, 0]
-        pos_loop = numpy.ones((dist_marg.size, 3), dtype=numpy.float32)
+            dir_loop[2, :] = angpos[i, 1]
-
+            out[i] = numpy.sum(self.evaluate_sky(field, dir_loop, isdeg=True))
-        out = numpy.zeros(Nang, dtype=numpy.float32)
-        for i in trange(Nang) if verbose else range(Nang):
-            # Get the position vector for this choice of theta, phi
-            ra_loop[:] = pos[i, 0]
-            dec_loop[:] = pos[i, 1]
-            pos_loop[:] = numpy.vstack([dist_marg, ra_loop, dec_loop]).T
-            # Evaluate and sum it up
-            out[i] = numpy.sum(self.evaluate_sky(field, pos_loop, isdeg)[0, :])
-
         return out
+
+
+###############################################################################
+#                         Density field calculation                           #
+###############################################################################
+
+
+class DensityField(BaseField):
+    r"""
+    Density field calculations. Based primarily on routines of Pylians [1].
+
+    Parameters
+    ----------
+    pos : 2-dimensional array of shape `(N, 3)`
+        Particle position array. Columns must be ordered as `['x', 'y', 'z']`.
+        The positions are assumed to be in box units, i.e. :math:`\in [0, 1 ]`.
+    mass : 1-dimensional array of shape `(N,)`
+        Particle mass array. Assumed to be in box units.
+    box : :py:class:`csiborgtools.read.BoxUnits`
+        The simulation box information and transformations.
+    MAS : str
+        Mass assignment scheme. Options are Options are: 'NGP' (nearest grid
+        point), 'CIC' (cloud-in-cell), 'TSC' (triangular-shape cloud), 'PCS'
+        (piecewise cubic spline).
+
+    References
+    ----------
+    [1] https://pylians3.readthedocs.io/
+    """
+    _pos = None
+    _mass = None
+
+    def __init__(self, pos, mass, box, MAS):
+        self.pos = pos
+        self.mass = mass
+        self.box = box
+        self.MAS = MAS
+
+    @property
+    def pos(self):
+        """
+        Particle position array.
+
+        Returns
+        -------
+        particles : 2-dimensional array
+        """
+        return self._particles
+
+    @pos.setter
+    def pos(self, pos):
+        assert pos.ndim == 2
+        warn("Flipping the `x` and `z` coordinates of the particle positions.",
+             UserWarning, stacklevel=1)
+        pos[:, [0, 2]] = pos[:, [2, 0]]
+        pos = force_single_precision(pos, "particle_position")
+        self._pos = pos
+
+    @property
+    def mass(self):
+        """
+        Particle mass array.
+
+        Returns
+        -------
+        mass : 1-dimensional array
+        """
+        return self._mass
+
+    @mass.setter
+    def mass(self, mass):
+        assert mass.ndim == 1
+        mass = force_single_precision(mass, "particle_mass")
+        self._mass = mass
+
+    def smoothen(self, field, smooth_scale, threads=1):
+        """
+        Smooth a field with a Gaussian filter.
+
+        Parameters
+        ----------
+        field : 3-dimensional array of shape `(grid, grid, grid)`
+            Field to be smoothed.
+        smooth_scale : float, optional
+            Gaussian kernal scale to smoothen the density field, in box units.
+        threads : int, optional
+            Number of threads. By default 1.
+
+        Returns
+        -------
+        smoothed_field : 3-dimensional array of shape `(grid, grid, grid)`
+        """
+        filter_kind = "Gaussian"
+        grid = field.shape[0]
+        # FFT of the filter
+        W_k = SL.FT_filter(self.boxsize, smooth_scale, grid, filter_kind,
+                           threads)
+        return SL.field_smoothing(field, W_k, threads)
+
+    def overdensity_field(self, delta):
+        r"""
+        Calculate the overdensity field from the density field.
+        Defined as :math:`\rho/ <\rho> - 1`. Overwrites the input array.
+
+
+        Parameters
+        ----------
+        delta : 3-dimensional array of shape `(grid, grid, grid)`
+            The density field.
+
+        Returns
+        -------
+        overdensity : 3-dimensional array of shape `(grid, grid, grid)`.
+        """
+        delta /= delta.mean()
+        delta -= 1
+        return delta
+
+    def __call__(self, grid, smooth_scale=None, verbose=True):
+        """
+        Calculate the density field using a Pylians routine [1, 2].
+
+        Parameters
+        ----------
+        grid : int
+            Grid size.
+        smooth_scale : float, optional
+            Gaussian kernal scale to smoothen the density field, in box units.
+        verbose : bool
+            Verbosity flag.
+
+        Returns
+        -------
+        rho : 3-dimensional array of shape `(grid, grid, grid)`.
+            Density field.
+
+        References
+        ----------
+        [1] https://pylians3.readthedocs.io/
+        [2] https://github.com/franciscovillaescusa/Pylians3/blob/master
+            /library/MAS_library/MAS_library.pyx
+        """
+        # Pre-allocate and do calculations
+        rho = numpy.zeros((grid, grid, grid), dtype=numpy.float32)
+        MASL.MA(self.pos, rho, self.boxsize, self.MAS, W=self.mass,
+                verbose=verbose)
+        if smooth_scale is not None:
+            rho = self.smoothen(rho, smooth_scale)
+        return rho
+
+
+###############################################################################
+#                         Potential field calculation                         #
+###############################################################################
+
+
+class PotentialField(BaseField):
+    """
+    Potential field calculation.
+
+    Parameters
+    ----------
+    box : :py:class:`csiborgtools.read.BoxUnits`
+        The simulation box information and transformations.
+    MAS : str
+        Mass assignment scheme. Options are Options are: 'NGP' (nearest grid
+        point), 'CIC' (cloud-in-cell), 'TSC' (triangular-shape cloud), 'PCS'
+        (piecewise cubic spline).
+    """
+    def __init__(self, box, MAS):
+        self.box = box
+        self.MAS = MAS
+
+    def __call__(self, overdensity_field):
+        """
+        Calculate the potential field.
+
+        Parameters
+        ----------
+        overdensity_field : 3-dimensional array of shape `(grid, grid, grid)`
+            The overdensity field.
+
+        Returns
+        -------
+        potential : 3-dimensional array of shape `(grid, grid, grid)`.
+        """
+        return MASL.potential(overdensity_field, self.box._omega_m,
+                              self.box._aexp, self.MAS)
+
+
+###############################################################################
+#                        Tidal tensor field calculation                       #
+###############################################################################
+
+
+class TidalTensorField(BaseField):
+    """
+    Tidal tensor field calculation.
+
+    Parameters
+    ----------
+    box : :py:class:`csiborgtools.read.BoxUnits`
+        The simulation box information and transformations.
+    MAS : str
+        Mass assignment scheme. Options are Options are: 'NGP' (nearest grid
+        point), 'CIC' (cloud-in-cell), 'TSC' (triangular-shape cloud), 'PCS'
+        (piecewise cubic spline).
+    """
+    def __init__(self, box, MAS):
+        self.box = box
+        self.MAS = MAS
+
+    @staticmethod
+    def tensor_field_eigvals(tidal_tensor):
+        """
+        Calculate eigenvalues of the tidal tensor field, sorted in increasing
+        order.
+
+        Parameters
+        ----------
+        tidal_tensor : :py:class:`MAS_library.tidal_tensor`
+            Tidal tensor object, whose attributes `tidal_tensor.Tij` contain
+            the relevant tensor components.
+
+        Returns
+        -------
+        eigvals : 3-dimensional array of shape `(grid, grid, grid)`
+        """
+        n_samples = tidal_tensor.T00.size
+        # We create a array and then calculate the eigenvalues.
+        Teval = numpy.full((n_samples, 3, 3), numpy.nan, dtype=numpy.float32)
+        Teval[:, 0, 0] = tidal_tensor.T00
+        Teval[:, 0, 1] = tidal_tensor.T01
+        Teval[:, 0, 2] = tidal_tensor.T02
+        Teval[:, 1, 1] = tidal_tensor.T11
+        Teval[:, 1, 2] = tidal_tensor.T12
+        Teval[:, 2, 2] = tidal_tensor.T22
+
+        eigvals = numpy.full((n_samples, 3), numpy.nan, dtype=numpy.float32)
+        for i in range(n_samples):
+            eigvals[i, :] = numpy.linalg.eigvalsh(Teval[i, ...], 'U')
+            eigvals[i, :] = numpy.sort(eigvals[i, :])
+
+        return eigvals
+
+    def __call__(self, overdensity_field):
+        """
+        Calculate the tidal tensor field.
+
+        Parameters
+        ----------
+        overdensity_field : 3-dimensional array of shape `(grid, grid, grid)`
+            The overdensity field.
+
+        Returns
+        -------
+        tidal_tensor : :py:class:`MAS_library.tidal_tensor`
+            Tidal tensor object, whose attributes `tidal_tensor.Tij` contain
+            the relevant tensor components.
+        """
+        return MASL.tidal_tensor(overdensity_field, self.box._omega_m,
+                                 self.box._aexp, self.MAS)

csiborgtools/field/utils.py

@@ -0,0 +1,42 @@
+# 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.
+"""
+Utility functions for the field module.
+"""
+from warnings import warn
+
+import numpy
+
+
+def force_single_precision(x, name):
+    """
+    Convert `x` to float32 if it is not already.
+
+    Parameters
+    ----------
+    x : array
+        Array to convert.
+    name : str
+        Name of the array.
+
+    Returns
+    -------
+    x : array
+        Converted array.
+    """
+    if x.dtype != numpy.float32:
+        warn(f"Converting `{name}` to float32.", UserWarning, stacklevel=1)
+        x = x.astype(numpy.float32)
+    return x

csiborgtools/read/box_units.py

@@ -193,6 +193,23 @@ class BoxUnits:
         """
         return length / (self._unit_l / units.kpc.to(units.cm) / self._aexp)
 
+    def mpc2box(self, length):
+        r"""
+        Convert length from :math:`\mathrm{cMpc}` (with :math:`h=0.705`) to
+        box units.
+
+        Parameters
+        ----------
+        length : float
+            Length in :math:`\mathrm{cMpc}`
+
+        Returns
+        -------
+        length : foat
+            Length in box units.
+        """
+        return self.kpc2box(length * 1e3)
+
     def box2mpc(self, length):
         r"""
         Convert length from box units to :math:`\mathrm{cMpc}` (with

csiborgtools/read/paths.py

@@ -326,7 +326,7 @@ class CSiBORGPaths:
         fname = f"radpos_{str(nsim).zfill(5)}_{str(nsnap).zfill(5)}.npz"
         return join(fdir, fname)
 
-    def particle_h5py_path(self, nsim):
+    def particle_h5py_path(self, nsim, with_vel):
         """
         Path to the files containing all particles in a `.hdf5` file. Used for
         the SPH calculation.
@@ -335,6 +335,8 @@ class CSiBORGPaths:
         ----------
         nsim : int
             IC realisation index.
+        with_vel : bool
+            Whether velocities are included.
 
         Returns
         -------
@@ -344,7 +346,10 @@ class CSiBORGPaths:
         if not isdir(fdir):
             makedirs(fdir)
             warn(f"Created directory `{fdir}`.", UserWarning, stacklevel=1)
-        fname = f"particles_{str(nsim).zfill(5)}.h5"
+        if with_vel:
+            fname = f"parts_{str(nsim).zfill(5)}.h5"
+        else:
+            fname = f"parts_pos_{str(nsim).zfill(5)}.h5"
         return join(fdir, fname)
 
     def density_field_path(self, mas, nsim):

csiborgtools/read/readsim.py

@@ -191,7 +191,8 @@ class ParticleReader:
         """
         return numpy.hstack([[0], numpy.cumsum(nparts[:-1])])
 
-    def read_particle(self, nsnap, nsim, pars_extract, verbose=True):
+    def read_particle(self, nsnap, nsim, pars_extract, return_structured=True,
+                      verbose=True):
         """
         Read particle files of a simulation at a given snapshot and return
         values of `pars_extract`.
@@ -204,17 +205,22 @@ class ParticleReader:
             IC realisation index.
         pars_extract : list of str
             Parameters to be extacted.
+        return_structured : bool, optional
+            Whether to return a structured array or a 2-dimensional array. If
+            the latter, then the order of the columns is the same as the order
+            of `pars_extract`. However, enforces single-precision floating
+            point format for all columns.
         verbose : bool, optional
             Verbosity flag while for reading the CPU outputs.
 
         Returns
         -------
-        out : structured array
+        out : array
         """
         # Open the particle files
         nparts, partfiles = self.open_particle(nsnap, nsim, verbose=verbose)
         if verbose:
-            print("Opened {} particle files.".format(nparts.size))
+            print(f"Opened {nparts.size} particle files.")
         ncpu = nparts.size
         # Order in which the particles are written in the FortranFile
         forder = [("x", numpy.float32), ("y", numpy.float32),
@@ -229,27 +235,34 @@ class ParticleReader:
             pars_extract = [pars_extract]
         for p in pars_extract:
             if p not in fnames:
-                raise ValueError(
+                raise ValueError(f"Undefined parameter `{p}`.")
-                    "Undefined parameter `{}`. Must be one of `{}`."
-                    .format(p, fnames))
 
         npart_tot = numpy.sum(nparts)
         # A dummy array is necessary for reading the fortran files.
         dum = numpy.full(npart_tot, numpy.nan, dtype=numpy.float16)
+        # We allocate the output structured/2D array
+        if return_structured:
             # These are the data we read along with types
-        dtype = {"names": pars_extract,
+            formats = [forder[fnames.index(p)][1] for p in pars_extract]
-                 "formats": [forder[fnames.index(p)][1] for p in pars_extract]}
+            dtype = {"names": pars_extract, "formats": formats}
-        # Allocate the output structured array
             out = numpy.full(npart_tot, numpy.nan, dtype)
+        else:
+            par2arrpos = {par: i for i, par in enumerate(pars_extract)}
+            out = numpy.full((npart_tot, len(pars_extract)), numpy.nan,
+                             dtype=numpy.float32)
+
         start_ind = self.nparts_to_start_ind(nparts)
         iters = tqdm(range(ncpu)) if verbose else range(ncpu)
         for cpu in iters:
             i = start_ind[cpu]
             j = nparts[cpu]
-            # trunk-ignore(ruff/B905)
             for (fname, fdtype) in zip(fnames, fdtypes):
                 if fname in pars_extract:
-                    out[fname][i:i + j] = self.read_sp(fdtype, partfiles[cpu])
+                    single_part = self.read_sp(fdtype, partfiles[cpu])
+                    if return_structured:
+                        out[fname][i:i + j] = single_part
+                    else:
+                        out[i:i + j, par2arrpos[fname]] = single_part
                 else:
                     dum[i:i + j] = self.read_sp(fdtype, partfiles[cpu])
         # Close the fortran files
@@ -279,9 +292,8 @@ class ParticleReader:
         """
         nsnap = str(nsnap).zfill(5)
         cpu = str(cpu + 1).zfill(5)
-        fpath = join(self.paths.ic_path(nsim, tonew=False),
+        fpath = join(self.paths.ic_path(nsim, tonew=False), f"output_{nsnap}",
-                     "output_{}".format(nsnap),
+                     f"unbinding_{nsnap}.out{cpu}")
-                     "unbinding_{}.out{}".format(nsnap, cpu))
         return FortranFile(fpath)
 
     def read_clumpid(self, nsnap, nsim, verbose=True):
@@ -313,7 +325,7 @@ class ParticleReader:
             j = nparts[cpu]
             ff = self.open_unbinding(nsnap, nsim, cpu)
             clumpid[i:i + j] = ff.read_ints()
-            # Close
+
             ff.close()
 
         return clumpid

scripts/pre_dumppart.py

@@ -18,9 +18,9 @@ SPH density field calculation.
 
 from datetime import datetime
 from gc import collect
+from distutils.util import strtobool
 
 import h5py
-import numpy
 from mpi4py import MPI
 
 try:
@@ -42,16 +42,22 @@ nproc = comm.Get_size()
 parser = ArgumentParser()
 parser.add_argument("--ics", type=int, nargs="+", default=None,
                     help="IC realisatiosn. If `-1` processes all simulations.")
+parser.add_argument("--with_vel", type=lambda x: bool(strtobool(x)),
+                    help="Whether to include velocities in the particle file.")
 args = parser.parse_args()
 paths = csiborgtools.read.CSiBORGPaths(**csiborgtools.paths_glamdring)
 partreader = csiborgtools.read.ParticleReader(paths)
+if args.with_vel:
     pars_extract = ['x', 'y', 'z', 'vx', 'vy', 'vz', 'M']
+else:
+    pars_extract = ['x', 'y', 'z', 'M']
 if args.ics is None or args.ics == -1:
     ics = paths.get_ics(tonew=False)
 else:
     ics = args.ics
 
-# We MPI loop over individual simulations.
+# MPI loop over individual simulations. We read in the particles from RAMSES
+# files and dump them to a HDF5 file.
 jobs = csiborgtools.fits.split_jobs(len(ics), nproc)[rank]
 for i in jobs:
     nsim = ics[i]
@@ -59,17 +65,11 @@ for i in jobs:
     print(f"{datetime.now()}: Rank {rank} completing simulation {nsim}.",
           flush=True)
 
-    # We read in the particles from RASMSES files, switch from a
+    out = partreader.read_particle(
-    # structured array to 2-dimensional array and dump it.
+        nsnap, nsim, pars_extract, return_structured=False, verbose=nproc == 1)
-    parts = partreader.read_particle(nsnap, nsim, pars_extract,
-                                     verbose=nproc == 1)
-    out = numpy.full((parts.size, len(pars_extract)), numpy.nan,
-                     dtype=numpy.float32)
-    for j, par in enumerate(pars_extract):
-        out[:, j] = parts[par]
 
     with h5py.File(paths.particle_h5py_path(nsim), "w") as f:
         dset = f.create_dataset("particles", data=out)
 
-    del parts, out
+    del out
     collect()

scripts/pre_initmatch.py

@@ -16,13 +16,18 @@
 Script to calculate the particle centre of mass and Lagrangian patch size in
 the initial snapshot. Optinally dumps the particle files, however this requires
 a lot of memory.
+
+TODO:
+    - stop saving the particle IDs. Unnecessary.
+    - Switch to h5py files. This way can save the positions in the particle
+    array only.
 """
 from argparse import ArgumentParser
 from datetime import datetime
 from distutils.util import strtobool
 from gc import collect
 from os import remove
-from os.path import join
+from os.path import isfile, join
 
 import numpy
 from mpi4py import MPI
@@ -129,6 +134,10 @@ for i, nsim in enumerate(paths.get_ics(tonew=True)):
         out = numpy.full(parent_ids.size, numpy.nan, dtype=dtype)
         for i, clid in enumerate(parent_ids):
             fpath = ftemp.format(nsim, clid, "fit")
+            # There is no file if the halo was skipped due to too few
+            # particles.
+            if not isfile(fpath):
+                continue
             with open(fpath, "rb") as f:
                 inp = numpy.load(f)
                 out["index"][i] = clid
@@ -151,8 +160,11 @@ for i, nsim in enumerate(paths.get_ics(tonew=True)):
             out = {}
             for clid in parent_ids:
                 fpath = ftemp.format(nsim, clid, "particles")
+                if not isfile(fpath):
+                    continue
                 with open(fpath, "rb") as f:
                     out.update({str(clid): numpy.load(f)})
+                remove(fpath)
 
             fout = paths.initmatch_path(nsim, "particles")
             print(f"{datetime.now()}: dumping particles to .. `{fout}`.",