Density field tests (#110)

* Add imports

* Remove file

* Add boxsize argument

* Add script

* Update script

* Edit script

* Add nbs

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, :]