Density field tests (#110)

* Add imports

* Remove file

* Add boxsize argument

* Add script

* Update script

* Edit script

* Add nbs

csiborgtools/field/__init__.py

@@ -15,6 +15,8 @@
 from .density import (DensityField, PotentialField, TidalTensorField,           # noqa
                       VelocityField, radial_velocity, power_spectrum,           # noqa
                       overdensity_field)                                        # noqa
+from .enclosed_mass import (particles_enclosed_mass,                            # noqa
+                            particles_enclosed_momentum, field_enclosed_mass)   # noqa
 from .interp import (evaluate_cartesian, evaluate_sky, field2rsp,               # noqa
                      fill_outside, make_sky, observer_peculiar_velocity,        # noqa
                      smoothen_field, field_at_distance)                         # noqa

csiborgtools/field/enclosed_mass.py

@@ -0,0 +1,182 @@
+# 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.
+
+
+import numpy
+from numba import jit
+
+
+###############################################################################
+#                 Enclosed mass at each distance from particles               #
+###############################################################################
+
+
+@jit(nopython=True, boundscheck=False)
+def _enclosed_mass(rdist, mass, rmax, start_index):
+    enclosed_mass = 0.
+
+    for i in range(start_index, len(rdist)):
+        if rdist[i] <= rmax:
+            enclosed_mass += mass[i]
+        else:
+            break
+
+    return enclosed_mass, i
+
+
+def particles_enclosed_mass(rdist, mass, distances):
+    """
+    Calculate the enclosed mass at each distance from a set of particles. Note
+    that the particles must be sorted by distance from the center of the box.
+
+    Parameters
+    ----------
+    rdist : 1-dimensional array
+        Sorted distance of particles from the center of the box.
+    mass : 1-dimensional array
+        Sorted mass of particles.
+    distances : 1-dimensional array
+        Distances at which to calculate the enclosed mass.
+
+    Returns
+    -------
+    enclosed_mass : 1-dimensional array
+        Enclosed mass at each distance.
+    """
+    enclosed_mass = numpy.full_like(distances, 0.)
+    start_index = 0
+    for i, dist in enumerate(distances):
+        if i > 0:
+            enclosed_mass[i] += enclosed_mass[i - 1]
+
+        m, start_index = _enclosed_mass(rdist, mass, dist, start_index)
+        enclosed_mass[i] += m
+
+    return enclosed_mass
+
+
+###############################################################################
+#                     Enclosed mass from a density field                      #
+###############################################################################
+
+
+@jit(nopython=True)
+def _cell_rdist(i, j, k, Ncells, boxsize):
+    """Radial distance of the center of a cell from the center of the box."""
+    xi = boxsize / Ncells * (i + 0.5) - boxsize / 2
+    yi = boxsize / Ncells * (j + 0.5) - boxsize / 2
+    zi = boxsize / Ncells * (k + 0.5) - boxsize / 2
+
+    return (xi**2 + yi**2 + zi**2)**0.5
+
+
+@jit(nopython=True, boundscheck=False)
+def _field_enclosed_mass(field, rmax, boxsize):
+    Ncells = field.shape[0]
+    cell_volume = (1000 * boxsize / Ncells)**3
+
+    mass = 0.
+    volume = 0.
+    for i in range(Ncells):
+        for j in range(Ncells):
+            for k in range(Ncells):
+                if _cell_rdist(i, j, k, Ncells, boxsize) < rmax:
+                    mass += field[i, j, k]
+                    volume += 1.
+
+    return mass * cell_volume, volume * cell_volume
+
+
+def field_enclosed_mass(field, distances, boxsize):
+    """
+    Calculate the approximate enclosed mass within a given radius from a
+    density field, counts the mass in cells and volume of cells whose
+    centers are within the radius.
+
+    Parameters
+    ----------
+    field : 3-dimensional array
+        Density field in units of `h^2 Msun / kpc^3`.
+    rmax : 1-dimensional array
+        Radii to calculate the enclosed mass at in `Mpc / h`.
+    boxsize : float
+        Box size in `Mpc / h`.
+
+    Returns
+    -------
+    enclosed_mass : 1-dimensional array
+        Enclosed mass at each distance.
+    enclosed_volume : 1-dimensional array
+        Enclosed grid-like volume at each distance.
+    """
+    enclosed_mass = numpy.zeros_like(distances)
+    enclosed_volume = numpy.zeros_like(distances)
+    for i, dist in enumerate(distances):
+        enclosed_mass[i], enclosed_volume[i] = _field_enclosed_mass(
+            field, dist, boxsize)
+
+    return enclosed_mass, enclosed_volume
+
+
+###############################################################################
+#              Enclosed momentum at each distance from particles              #
+###############################################################################
+
+
+@jit(nopython=True, boundscheck=False)
+def _enclosed_momentum(rdist, mass, vel, rmax, start_index):
+    bulk_momentum = numpy.zeros(3, dtype=rdist.dtype)
+
+    for i in range(start_index, len(rdist)):
+        if rdist[i] <= rmax:
+            bulk_momentum += mass[i] * vel[i]
+        else:
+            break
+
+    return bulk_momentum, i
+
+
+def particles_enclosed_momentum(rdist, mass, vel, distances):
+    """
+    Calculate the enclosed momentum at each distance. Note that the particles
+    must be sorted by distance from the center of the box.
+
+    Parameters
+    ----------
+    rdist : 1-dimensional array
+        Sorted distance of particles from the center of the box.
+    mass : 1-dimensional array
+        Sorted mass of particles.
+    vel : 2-dimensional array
+        Sorted velocity of particles.
+    distances : 1-dimensional array
+        Distances at which to calculate the enclosed momentum.
+
+    Returns
+    -------
+    bulk_momentum : 2-dimensional array
+        Enclosed momentum at each distance.
+    """
+    bulk_momentum = numpy.zeros((len(distances), 3))
+    start_index = 0
+    for i, dist in enumerate(distances):
+        if i > 0:
+            bulk_momentum[i] += bulk_momentum[i - 1]
+
+        v, start_index = _enclosed_momentum(rdist, mass, vel, dist,
+                                            start_index)
+        bulk_momentum[i] += v
+
+    return bulk_momentum

csiborgtools/utils.py

@@ -229,7 +229,7 @@ def dms_to_degrees(degrees, arcminutes=None, arcseconds=None):
     return degrees + (arcminutes or 0) / 60 + (arcseconds or 0) / 3600
 
 
-def real2redshift(pos, vel, observer_location, observer_velocity, box,
+def real2redshift(pos, vel, observer_location, observer_velocity, boxsize,
                   periodic_wrap=True, make_copy=True):
     r"""
     Convert real-space position to redshift space position.
@@ -244,8 +244,8 @@ def real2redshift(pos, vel, observer_location, observer_velocity, box,
         Observer location in `Mpc / h`.
     observer_velocity: 1-dimensional array `(3,)`
         Observer velocity in `km / s`.
-    box : py:class:`csiborg.read.CSiBORG1Box`
-        Box units.
+    boxsize : float
+        Box size in `Mpc / h`.
     periodic_wrap : bool, optional
         Whether to wrap around the box, particles may be outside the default
         bounds once RSD is applied.
@@ -278,7 +278,6 @@ def real2redshift(pos, vel, observer_location, observer_velocity, box,
         vel += observer_velocity
 
     if periodic_wrap:
-        boxsize = box.box2mpc(1.)
         pos = periodic_wrap_grid(pos, boxsize)
 
     return pos

data/csiborg_descr.txt

@@ -1,37 +0,0 @@
-Various column names
-
-Clump file columns:
-    "index", "lev", "parent", "ncell", "peak_x", "peak_y", "peak_z", "rho-", "rho+", "rho_av", "mass_cl", "relevance"
-
-Mergertree file columns:
-    "clump", "progenitor", "prog outputnr", "desc mass", "desc npart", "desc x", "desc y", "desc z", "desc vx", "desc vy", "desc vz"
-
-
-%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
-
-
-The merger trees are stored in `output_XXXXX/mergertree_XXXXX.txtYYYYY` files. Each file contains 11 columns:
-
-* `clump`:            clump ID of a clump at this output number
-* `progenitor`:       the progenitor clump ID in output number "prog_outputnr"
-* `prog_outputnr`:    the output number of when the progenitor was an alive clump
-* `desc mass`:        mass of the current clump. 
-* `desc npart`:       number of particles of the current clump. 
-* `desc x,y,z`:       x, y, z position of current clump.
-* `desc vx, vy, vz`:  x, y, z velocities of current clump.
-
-desc_mass and desc_npart will be either inclusive or exclusive, depending on how you set the `use_exclusive_mass` parameter. 
-(See below for details)
-
-**How to read the output:**
-
-* A clump > 0 has progenitor > 0: Standard case. A direct progenitor from the adjacent previous snapshot was identified for this clump.
-* A clump > 0 has progenitor = 0: no progenitor could be established and the clump is treated as newly formed.
-* A clump > 0 has progenitor < 0: it means that no direct progenitor could be found in the adjacent previous snapshot, but a progenitor was identified from an earlier, non-adjacent snapshot.
-* A clump < 0 has progenitor > 0: this progenitor merged into this clump, but is not this clump's main progenitor.
-* A clump < 0 has progenitor < 0: this shouldn't happen.
-
-### Visualisation
-
-`ramses/utils/py/mergertreeplot.py` is a python 2 script to plot the merger trees as found by this patch. 
- Details on options and usage are at the start of the script as a comment.

scripts/field_prop.sh

@@ -1,13 +1,13 @@
-nthreads=6
+nthreads=1
 memory=64
 on_login=${1}
 queue="berg"
 env="/mnt/zfsusers/rstiskalek/csiborgtools/venv_csiborg/bin/python"
 file="field_prop.py"
-kind="radvel"
-simname="csiborg2_random"
-nsims="-1"
-MAS="SPH"
+kind="density"
+simname="csiborg1"
+nsims="9844"
+MAS="PCS"
 grid=1024
 
 

scripts/mass_enclosed.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2022 Richard Stiskalek
+# 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
@@ -27,7 +27,6 @@ from gc import collect
 import csiborgtools
 import numpy
 from tqdm import tqdm
-from numba import jit
 
 from datetime import datetime
 
@@ -132,167 +131,6 @@ def get_particles(reader, boxsize, get_velocity=True, verbose=True):
     return dist, mass
 
 
-###############################################################################
-#                Calculate the enclosed mass at each distance                 #
-###############################################################################
-
-
-@jit(nopython=True, boundscheck=False)
-def _enclosed_mass(rdist, mass, rmax, start_index):
-    enclosed_mass = 0.
-
-    for i in range(start_index, len(rdist)):
-        if rdist[i] <= rmax:
-            enclosed_mass += mass[i]
-        else:
-            break
-
-    return enclosed_mass, i
-
-
-def enclosed_mass(rdist, mass, distances):
-    """
-    Calculate the enclosed mass at each distance.
-
-    Parameters
-    ----------
-    rdist : 1-dimensional array
-        Distance of particles from the center of the box.
-    mass : 1-dimensional array
-        Mass of particles.
-    distances : 1-dimensional array
-        Distances at which to calculate the enclosed mass.
-
-    Returns
-    -------
-    enclosed_mass : 1-dimensional array
-        Enclosed mass at each distance.
-    """
-    enclosed_mass = numpy.full_like(distances, 0.)
-    start_index = 0
-    for i, dist in enumerate(distances):
-        if i > 0:
-            enclosed_mass[i] += enclosed_mass[i - 1]
-
-        m, start_index = _enclosed_mass(rdist, mass, dist, start_index)
-        enclosed_mass[i] += m
-
-    return enclosed_mass
-
-
-###############################################################################
-#                Calculate enclosed mass from a density field                 #
-###############################################################################
-
-
-@jit(nopython=True)
-def _cell_rdist(i, j, k, Ncells, boxsize):
-    """Radial distance of the center of a cell from the center of the box."""
-    xi = boxsize / Ncells * (i + 0.5) - boxsize / 2
-    yi = boxsize / Ncells * (j + 0.5) - boxsize / 2
-    zi = boxsize / Ncells * (k + 0.5) - boxsize / 2
-
-    return (xi**2 + yi**2 + zi**2)**0.5
-
-
-@jit(nopython=True, boundscheck=False)
-def _field_enclosed_mass(field, rmax, boxsize):
-    Ncells = field.shape[0]
-    cell_volume = (1000 * boxsize / Ncells)**3
-
-    mass = 0.
-    volume = 0.
-    for i in range(Ncells):
-        for j in range(Ncells):
-            for k in range(Ncells):
-                if _cell_rdist(i, j, k, Ncells, boxsize) < rmax:
-                    mass += field[i, j, k]
-                    volume += 1.
-
-    return mass * cell_volume, volume * cell_volume
-
-
-def field_enclosed_mass(field, distances, boxsize):
-    """
-    Calculate the approximate enclosed mass within a given radius from a
-    density field.
-
-    Parameters
-    ----------
-    field : 3-dimensional array
-        Density field in units of `h^2 Msun / kpc^3`.
-    rmax : 1-dimensional array
-        Radii to calculate the enclosed mass at in `Mpc / h`.
-    boxsize : float
-        Box size in `Mpc / h`.
-
-    Returns
-    -------
-    enclosed_mass : 1-dimensional array
-        Enclosed mass at each distance.
-    enclosed_volume : 1-dimensional array
-        Enclosed grid-like volume at each distance.
-    """
-    enclosed_mass = numpy.zeros_like(distances)
-    enclosed_volume = numpy.zeros_like(distances)
-    for i, dist in enumerate(distances):
-        enclosed_mass[i], enclosed_volume[i] = _field_enclosed_mass(
-            field, dist, boxsize)
-
-    return enclosed_mass, enclosed_volume
-
-
-###############################################################################
-#              Calculate the enclosed momentum at each distance               #
-###############################################################################
-
-
-@jit(nopython=True, boundscheck=False)
-def _enclosed_momentum(rdist, mass, vel, rmax, start_index):
-    bulk_momentum = numpy.zeros(3, dtype=rdist.dtype)
-
-    for i in range(start_index, len(rdist)):
-        if rdist[i] <= rmax:
-            bulk_momentum += mass[i] * vel[i]
-        else:
-            break
-
-    return bulk_momentum, i
-
-
-def enclosed_momentum(rdist, mass, vel, distances):
-    """
-    Calculate the enclosed momentum at each distance.
-
-    Parameters
-    ----------
-    rdist : 1-dimensional array
-        Distance of particles from the center of the box.
-    mass : 1-dimensional array
-        Mass of particles.
-    vel : 2-dimensional array
-        Velocity of particles.
-    distances : 1-dimensional array
-        Distances at which to calculate the enclosed momentum.
-
-    Returns
-    -------
-    bulk_momentum : 2-dimensional array
-        Enclosed momentum at each distance.
-    """
-    bulk_momentum = numpy.zeros((len(distances), 3))
-    start_index = 0
-    for i, dist in enumerate(distances):
-        if i > 0:
-            bulk_momentum[i] += bulk_momentum[i - 1]
-
-        v, start_index = _enclosed_momentum(rdist, mass, vel, dist,
-                                            start_index)
-        bulk_momentum[i] += v
-
-    return bulk_momentum
-
-
 ###############################################################################
 #                       Main & command line interface                         #
 ###############################################################################
@@ -316,7 +154,7 @@ def main_borg(args, folder):
         else:
             raise ValueError(f"Unknown simname: `{args.simname}`.")
 
-        cumulative_mass[i, :], cumulative_volume[i, :] = field_enclosed_mass(
+        cumulative_mass[i, :], cumulative_volume[i, :] = csiborgtools.field.field_enclosed_mass(  # noqa
             field, distances, boxsize)
 
     # Finally save the output
@@ -343,12 +181,14 @@ def main_csiborg(args, folder):
         rdist, mass, vel = get_particles(reader, boxsize,  verbose=False)
 
         # Calculate masses
-        cumulative_mass[i, :] = enclosed_mass(rdist, mass, distances)
-        mass135[i] = enclosed_mass(rdist, mass, [135])[0]
+        cumulative_mass[i, :] = csiborgtools.field.particles_enclosed_mass(
+            rdist, mass, distances)
+        mass135[i] = csiborgtools.field.particles_enclosed_mass(
+            rdist, mass, [135])[0]
         masstot[i] = numpy.sum(mass)
 
         # Calculate velocities
-        cumulative_velocity[i, ...] = enclosed_momentum(
+        cumulative_velocity[i, ...] = csiborgtools.field.particles_enclosed_momentum(  # noqa
             rdist, mass, vel, distances)
         for j in range(3):  # Normalize the momentum to get velocity out of it.
             cumulative_velocity[i, :, j] /= cumulative_mass[i, :]