From 0fe535c2ccbe408c37a6c54629c507c2788485b9 Mon Sep 17 00:00:00 2001
From: Aaron LI <aly@aaronly.me>
Date: Wed, 4 Oct 2017 20:57:56 +0800
Subject: clusters: Use "hmf" to calculate halo mass functions/distributions

* New dependency "hmf" (halo mass functions) module
* Calculate halo mass distributions/functions (dndlnm) with respect to masses
  and redshifts, instead of use the previous data file ("ps_data")
* New section "[extragalactic][psformalism]" in configurations
* New functions to write and read the dndlnm data

TODO:
* update the method to sample (mass, redshift) for clusters from the dndlnm
  data
---
 fg21sim/utils/io.py | 46 ++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 46 insertions(+)

(limited to 'fg21sim/utils')

diff --git a/fg21sim/utils/io.py b/fg21sim/utils/io.py
index 2449ca5..b3666b6 100644
--- a/fg21sim/utils/io.py
+++ b/fg21sim/utils/io.py
@@ -380,3 +380,49 @@ def write_fits_healpix(outfile, hpmap, header=None, float32=False,
     ], header=hdr)
     hdu.writeto(outfile, checksum=checksum)
     logger.info("Wrote HEALPix map to FITS file: %s" % outfile)
+
+
+def write_dndlnm(outfile, dndlnm, z, mass, clobber=False):
+    """
+    Write the halo mass distribution data into file in NumPy's ".npz"
+    format, which packs the ``dndlnm``, ``z``, and ``mass`` arrays.
+
+    Parameters
+    ----------
+    outfile : str
+        The output file to store the dndlnm data, in ".npz" format.
+    dndlnm : 2D float `~numpy.ndarray`
+        Shape: (len(z), len(mass))
+        Differential mass function in terms of natural log of M.
+        Unit: [Mpc^-3] (the little "h" is folded into the values)
+    z : 1D float `~numpy.ndarray`
+        Redshifts where the halo mass distribution is calculated.
+    mass : 1D float `~numpy.ndarray`
+        (Logarithmic-distributed) masses points.
+        Unit: [Msun] (the little "h" is folded into the values)
+    clobber : bool, optional
+        Whether to overwrite the existing output file?
+    """
+    _create_dir(outfile)
+    _check_existence(outfile, clobber=clobber, remove=True)
+    np.savez(outfile, dndlnm=dndlnm, z=z, mass=mass)
+
+
+def read_dndlnm(infile):
+    """
+    Read the halo mass distribution data from the above saved file.
+
+    Parameters
+    ----------
+    infile : str
+        The ".npz" file from which to read the dndlnm data.
+
+    Returns
+    -------
+    (dndlnm, z, mass)
+    """
+    with np.load(infile) as npzfile:
+        dndlnm = npzfile["dndlnm"]
+        z = npzfile["z"]
+        mass = npzfile["mass"]
+    return (dndlnm, z, mass)
-- 
cgit v1.2.2