Potential (#13)

* add basic density field

* Add TODO

* add field smoothing

* update how pos are calculated

* add transforms both ways

* add import

* add sky density

* add make skymap func

* update TODO

* update gitignore

* add potential field calculation

* delete boxsize setter

* add tidal tensor

.gitignore

@@ -16,3 +16,4 @@ build/*
 csiborgtools.egg-info/*
 scripts/playground_*
 scripts/playground.ipynb
+Pylians3/*

csiborgtools/__init__.py

@@ -13,4 +13,4 @@
 # with this program; if not, write to the Free Software Foundation, Inc.,
 # 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
 
-from csiborgtools import (read, match, utils, units, fits)  # noqa
+from csiborgtools import (read, match, utils, units, fits, field)  # noqa

csiborgtools/field/__init__.py

@@ -0,0 +1,16 @@
+# 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.
+
+from .density import DensityField  # noqa

csiborgtools/field/density.py

@@ -0,0 +1,323 @@
+# 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
+import MAS_library as MASL
+import smoothing_library as SL
+from warnings import warn
+from tqdm import trange
+from ..units import (BoxUnits, radec_to_cartesian)
+
+
+class DensityField:
+    """
+    Density field calculations. Based primarily on routines of Pylians [1].
+
+    Parameters
+    ----------
+    particles : structured array
+        Particle array. Must contain keys `['x', 'y', 'z', 'M']`.
+    box : :py:class:`csiborgtools.units.BoxUnits`
+        The simulation box information and transformations.
+
+    References
+    ----------
+    [1] https://pylians3.readthedocs.io/
+    """
+    _particles = None
+    _boxsize = None
+    _box = None
+
+    def __init__(self, particles, box):
+        self.particles = particles
+        self.box = box
+        self._boxsize = 1.
+
+    @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 box(self):
+        """
+        The simulation box information and transformations.
+
+        Returns
+        -------
+        box : :py:class:`csiborgtools.units.BoxUnits`
+        """
+        return self._box
+
+    @box.setter
+    def box(self, box):
+        """Set the simulation box."""
+        if not isinstance(box, BoxUnits):
+            raise TypeError("`box` must be `BoxUnits` instance.")
+        self._box = box
+
+    @property
+    def boxsize(self):
+        """
+        Boxsize.
+
+        Returns
+        -------
+        boxsize : float
+        """
+        return self._boxsize
+
+    @staticmethod
+    def _force_f32(x, name):
+        if x.dtype != numpy.float32:
+            warn("Converting `{}` to float32.".format(name))
+            x = x.astype(numpy.float32)
+        return x
+
+    def density_field(self, grid, verbose=True):
+        """
+        Calculate the density field using a Pylians routine [1, 2]. Enforces
+        float32 precision.
+
+        Parameters
+        ----------
+        grid : int
+            The grid size.
+        verbose : float, optional
+            A verbosity flag. By default `True`.
+
+        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
+        """
+        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')
+        MAS = "CIC"  # Cloud in cell
+
+        # Pre-allocate and do calculations
+        rho = numpy.zeros((grid, grid, grid), dtype=numpy.float32)
+        MASL.MA(pos, rho, self.boxsize, MAS, W=weights, verbose=verbose)
+        return rho
+
+    def overdensity_field(self, grid, verbose=True):
+        r"""
+        Calculate the overdensity field using Pylians routines.
+        Defined as :math:`\rho/ <\rho> - 1`.
+
+        Parameters
+        ----------
+        grid : int
+            The grid size.
+        verbose : float, optional
+            A verbosity flag. By default `True`.
+
+        Returns
+        -------
+        overdensity : 3-dimensional array of shape `(grid, grid, grid)`.
+            Overdensity field.
+        """
+        # Get the overdensity
+        delta = self.density_field(grid, verbose)
+        delta /= delta.mean()
+        delta -= 1
+        return delta
+
+    def potential_field(self, grid, verbose=True):
+        """
+        Calculate the potential field using Pylians routines.
+
+        Parameters
+        ----------
+        grid : int
+            The grid size.
+        verbose : float, optional
+            A verbosity flag. By default `True`.
+
+        Returns
+        -------
+        potential : 3-dimensional array of shape `(grid, grid, grid)`.
+            Potential field.
+        """
+        delta = self.overdensity_field(grid, verbose)
+        if verbose:
+            print("Calculating potential from the overdensity..")
+        return MASL.potential(delta, self.box._omega_m, self.box._aexp, "CIC")
+
+    def tensor_field(self, grid, verbose=True):
+        """
+        Calculate the tidal tensor field.
+
+        Parameters
+        ----------
+        grid : int
+            The grid size.
+        verbose : float, optional
+            A verbosity flag. By default `True`.
+
+        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, verbose)
+        return MASL.tidal_tensor(delta, self.box._omega_m, self.box._aexp,
+                                 "CIC")
+
+    def smooth_field(self, field, 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.
+        scale : float
+            The smoothing scale of the Gaussian filter. Units must match that
+            of `self.boxsize`.
+        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, scale, grid, Filter, threads)
+        return SL.field_smoothing(field, W_k, threads)
+
+    def evaluate_field(self, pos, field):
+        """
+        Evaluate the field at Cartesian coordinates.
+
+        Parameters
+        ----------
+        pos : 2-dimensional array of shape `(n_samples, 3)`
+            Positions to evaluate the density field. The coordinates span range
+            of [0, boxsize].
+        field : 3-dimensional array of shape `(grid, grid, grid)`
+            The density field that is to be interpolated.
+
+        Returns
+        -------
+        interp_field : 1-dimensional array of shape `(n_samples,).
+            Interpolated field at `pos`.
+        """
+        self._force_f32(pos, "pos")
+        density_interpolated = numpy.zeros(pos.shape[0], dtype=numpy.float32)
+        MASL.CIC_interp(field, self.boxsize, pos, density_interpolated)
+        return density_interpolated
+
+    def evaluate_sky(self, pos, field, isdeg=True):
+        """
+        Evaluate the field at given distance, right ascension and declination.
+        Assumes that the observed is in the centre of the box.
+
+        Parameters
+        ----------
+        pos : 2-dimensional array of shape `(n_samples, 3)`
+            Spherical coordinates to evaluate the field. Should be distance,
+            right ascension, declination, respectively.
+        field : 3-dimensional array of shape `(grid, grid, grid)`
+            The density field that is to be interpolated. Assumed to be defined
+            on a Cartesian grid.
+        isdeg : bool, optional
+            Whether `ra` and `dec` are in degres. By default `True`.
+
+        Returns
+        -------
+        interp_field : 1-dimensional array of shape `(n_samples,).
+            Interpolated field at `pos`.
+        """
+        self._force_f32(pos, "pos")
+        X = numpy.vstack(
+            radec_to_cartesian(*(pos[:, i] for i in range(3)), isdeg)).T
+        X = X.astype(numpy.float32)
+        # Place the observer at the center of the box
+        X += 0.5 * self.boxsize
+        return self.evaluate_field(X, field)
+
+    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.
+
+        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)`
+            The density field that is to be interpolated. Assumed to be defined
+            on a Cartesian grid `[0, self.boxsize]^3`.
+        dist_marg : 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. By default `True`.
+
+        Returns
+        -------
+        interp_field : 1-dimensional array of shape `(n_pos, )`.
+        """
+        # Angular positions at which to evaluate the field
+        Nang = ra.size
+        pos = numpy.vstack([ra, dec]).T
+
+        # Now loop over the angular positions, each time evaluating a vector
+        # of distances. Pre-allocate arrays for speed
+        ra_loop = numpy.ones_like(dist_marg)
+        dec_loop = numpy.ones_like(dist_marg)
+        pos_loop = numpy.ones((dist_marg.size, 3), dtype=numpy.float32)
+
+        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(pos_loop, field, isdeg))
+
+        return out

csiborgtools/read/make_cat.py

@@ -163,7 +163,9 @@ class HaloCatalogue:
         data = data[data["m500"] > min_m500]
 
         # Now calculate spherical coordinates
-        d, ra, dec = cartesian_to_radec(data)
+        d, ra, dec = cartesian_to_radec(
+            data["peak_x"], data["peak_y"], data["peak_z"])
+
         data = add_columns(data, [d, ra, dec], ["dist", "ra", "dec"])
 
         # Cut on separation

csiborgtools/units/__init__.py

@@ -13,5 +13,5 @@
 # with this program; if not, write to the Free Software Foundation, Inc.,
 # 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
 
-from .transforms import cartesian_to_radec  # noqa
+from .transforms import cartesian_to_radec, radec_to_cartesian  # noqa
 from .box_units import (BoxUnits)  # noqa

csiborgtools/units/transforms.py

@@ -19,39 +19,49 @@ Various coordinate transformations.
 import numpy
 
 
-def cartesian_to_radec(arr, xpar="peak_x", ypar="peak_y", zpar="peak_z"):
+def cartesian_to_radec(x, y, z):
-    r"""
+    """
-    Extract `x`, `y`, and `z` coordinates from a record array `arr` and
+    Calculate the radial distance, right ascension in [0, 360) degrees and
-    calculate the radial distance :math:`r` in coordinate units, right
+    declination [-90, 90] degrees. Note, the observer should be placed in the
-    ascension :math:`\mathrm{RA} \in [0, 360)` degrees and declination
+    middle of the box.
-    :math:`\delta \in [-90, 90]` degrees.
 
     Parameters
     ----------
-    arr : record array
+    x, y, z : 1-dimensional arrays
-        Record array with the Cartesian coordinates.
+        Cartesian coordinates.
-    xpar : str, optional
-        Name of the x coordinate in the record array.
-    ypar : str, optional
-        Name of the y coordinate in the record array.
-    zpar : str, optional
-        Name of the z coordinate in the record array.
-
     Returns
     -------
-    dist : 1-dimensional array
+    dist, ra, dec : 1-dimensional arrays
-        Radial distance.
+        Radial distance, right ascension and declination.
-    ra : 1-dimensional array
-        Right ascension.
-    dec : 1-dimensional array
-        Declination.
     """
-    x, y, z = arr[xpar], arr[ypar], arr[zpar]
-
     dist = numpy.sqrt(x**2 + y**2 + z**2)
     dec = numpy.rad2deg(numpy.arcsin(z/dist))
     ra = numpy.rad2deg(numpy.arctan2(y, x))
     # Make sure RA in the correct range
     ra[ra < 0] += 360
-
     return dist, ra, dec
+
+
+def radec_to_cartesian(dist, ra, dec, isdeg=True):
+    """
+    Convert distance, right ascension and declination to Cartesian coordinates.
+
+    Parameters
+    ----------
+    dist, ra, dec : 1-dimensional arrays
+        The spherical coordinates.
+    isdeg : bool, optional
+        Whether `ra` and `dec` are in degres. By default `True`.
+
+    Returns
+    -------
+    x, y, z : 1-dimensional arrays
+        Cartesian coordinates.
+    """
+    if isdeg:
+        ra = numpy.deg2rad(ra)
+        dec = numpy.deg2rad(dec)
+    x = dist * numpy.cos(dec) * numpy.cos(ra)
+    y = dist * numpy.cos(dec) * numpy.sin(ra)
+    z = dist * numpy.sin(dec)
+    return x, y, z