From 1cd3f10848060471889384996f489de3fb54d18b Mon Sep 17 00:00:00 2001
From: rstiskalek <richard.stiskalek@protonmail.com>
Date: Sun, 6 Nov 2022 09:05:51 +0000
Subject: [PATCH] add num density

---
 csiborgtools/match/__init__.py    |   1 +
 csiborgtools/match/num_density.py | 107 ++++++++++++++++++++++++++++++
 2 files changed, 108 insertions(+)
 create mode 100644 csiborgtools/match/num_density.py

diff --git a/csiborgtools/match/__init__.py b/csiborgtools/match/__init__.py
index a2914f0..e09466d 100644
--- a/csiborgtools/match/__init__.py
+++ b/csiborgtools/match/__init__.py
@@ -14,4 +14,5 @@
 # 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
 
 from .match import brute_spatial_separation  # noqa
+from .num_density import (binned_counts, number_density)  # noqa
 # from .correlation import (get_randoms_sphere, sphere_angular_tpcf) # noqa
diff --git a/csiborgtools/match/num_density.py b/csiborgtools/match/num_density.py
new file mode 100644
index 0000000..c515fbb
--- /dev/null
+++ b/csiborgtools/match/num_density.py
@@ -0,0 +1,107 @@
+# 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.
+"""
+Calculation of number density functions.
+"""
+
+import numpy
+
+
+def binned_counts(x, bins):
+    """
+    Calculate number of samples in bins.
+
+    Parameters
+    ----------
+    x : 1-dimensional array
+        Samples' values.
+    bins : 1-dimensional array
+        Bin edges of shape `(n_edges, )`.
+
+    Returns
+    -------
+    centres : 1-dimensional array
+        Bin centres of shape `(n_edges - 1, )`.
+    counts : 1-dimensional array
+        Bin counts of shape `(n_edges - 1, )`.
+    """
+    if not isinstance(bins, numpy.ndarray) and bins.ndim == 1:
+        raise TypeError("`bins` must a 1-dimensional array.")
+
+    n_bins = bins.size
+    # Bin centres
+    centres = numpy.asarray(
+        [0.5 * (bins[i + 1] + bins[i]) for i in range(n_bins - 1)])
+    # Bin counts
+    out = numpy.full(n_bins - 1, numpy.nan, dtype=int)
+    for i in range(n_bins - 1):
+        out[i] = numpy.sum((x >= bins[i]) & (x < bins[i + 1]))
+    return centres, out
+
+
+def number_density(data, feat, bins, max_dist, to_log10):
+    """
+    Calculate volume-limited number density of a feature `feat` from array
+    `data`, normalised also by the bin width.
+
+    Parameters
+    ----------
+    data : structured array
+        Input array of halos.
+    feat : str
+        Parameter whose number density to calculate.
+    bins : 1-dimensional array
+        Bin edges. Note that if `to_log10` then the edges must be specified
+        in the logarithmic space, not linear.
+    max_dist : float
+        Maximum radial distance of the volume limited sample.
+    to_log10 : bool
+        Whether to take a logarithm of base 10 of the feature. If so, then the
+        bins must also be logarithmic.
+
+    Returns
+    -------
+    centres : 1-dimensional array
+        Bin centres of shape `(n_edges - 1, )`. If `to_log10` then converts the
+        bin centres back to linear space.
+    nd : 1-dimensional array
+        Number density of shape `(n_edges - 1, )`.
+    nd_err : 1-dimensional array
+        Poissonian uncertainty of `nd` of shape `(n_edges - 1, )`.
+    """
+    # Extract the param and optionally convert to log10
+    x = data[feat]
+    x = numpy.log10(x) if to_log10 else x
+    # Get only things within distance from the origin
+    rdist = (data["peak_x"]**2 + data["peak_y"]**2 + data["peak_z"]**2)**0.5
+    x = x[rdist < max_dist]
+
+    # Make sure bins equally spaced
+    dbins = numpy.diff(bins)
+    dbin = dbins[0]
+    if not numpy.alltrue(dbins == dbin):
+        raise ValueError("Bins must be equally spaced. Currently `{}`."
+                         .format(bins))
+
+    # Encompassed volume around the origin
+    volume = 4 * numpy.pi / 3 * max_dist**3
+    # Poissonian statistics
+    bin_centres, counts = binned_counts(x, bins)
+    nd = counts / volume / dbin
+    nd_err = counts**0.5 / volume / dbin
+    # Convert bins to linear space if log10
+    if to_log10:
+        bin_centres = 10**bin_centres
+    return bin_centres, nd, nd_err