Power spectrum stuff (#16)

* add autok1d

* update TODO

* add basic cross pk

* del blank line

* fix small bugs

* fix loop bug

* rm nbs

* edit docstring

* Update TODO

* get rid of verbose flag

* take all ICs

* attempt better memory

* simplify ASCII dump

* rm verbose statement to get rid of bugs!

* close fortran files

* add MAS options

* add boxsize setter

* Move comment up

* paths chagne

* Update TODO

* Update TODO

* Forgotten units in Xpk

* Attempt at fixing the distance bug

* Update TODO

* Update TODO

* Remove comment

* add summary reader

* Update gitignore

.gitignore

@@ -17,3 +17,4 @@ csiborgtools.egg-info/*
 scripts/playground_*
 scripts/*.ipynb
 Pylians3/*
+scripts/plot_correlation.ipynb

README.md

@@ -1,20 +1,27 @@
 # CSiBORG tools
 
+## Questions to answer
+- How well can observed clusters be matched to CSiBORG? Do their masses agree?
+- Is the number of clusters in CSiBORG consistent?
+- Are any observed clusters suspiciously missing in CSiBORG?
+
+
 ## Short-term TODO
-- [ ] Get `powmes` working and compare its power spectrum to `Pylians`.
-- [ ] Check alignment of the angular momentum
-- [x] Make sure I am not taking halos outside of the well-resolved region.
-- [x] Calculate the gravitational potential.
-- [x] Add the gravitational field calculation.
+- [x] Add code to calculate the cross-correlation for resolved region only.
+- [ ] Calculate the spectra for all 101 boxes and visualise them.
+- [ ] See about the $z=70$ particles.
 
 
 ## Long-term TODO
 - [ ] Calculate the cross-correlation in CSiBORG. Should see the scale of the constraints?
 - [ ] Find the distribution of particles in the first snapshot
 - [ ] Implement a custom model for matchin galaxies to halos.
-- [x] Calculate DM environmental properties.
 
 
 ## Open questions
 - What scaling of the search region? No reason for it to be a multiple of $R_{200c}$.
 - Begin extracting the DM environmental properties at galaxy positions?
+
+
+## Notes
+- New sims with part files in initial snapshot: `/mnt/extraspace/hdesmond/ramses_out_7468_new`. Also numbers 7588, 8020, 8452, 8836.

csiborgtools/field/density.py

@@ -12,9 +12,12 @@
 # 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.
-
+"""
+Density field and cross-correlation calculations.
+"""
 import numpy
 import MAS_library as MASL
+import Pk_library as PKL
 import smoothing_library as SL
 from warnings import warn
 from tqdm import trange
@@ -22,15 +25,24 @@ from ..units import (BoxUnits, radec_to_cartesian)
 
 
 class DensityField:
-    """
+    r"""
     Density field calculations. Based primarily on routines of Pylians [1].
 
     Parameters
     ----------
     particles : structured array
-        Particle array. Must contain keys `['x', 'y', 'z', 'M']`.
+        Particle array. Must contain keys `['x', 'y', 'z', 'M']`. Particle
+        coordinates are assumed to be :math:`\in [0, 1]` or in box units
+        otherwise.
+    boxsize : float
+        Box length. Multiplies `particles` positions to fix the power spectum
+        units.
     box : :py:class:`csiborgtools.units.BoxUnits`
         The simulation box information and transformations.
+    MAS : str, optional
+        Mass assignment scheme. Options are Options are: 'NGP' (nearest grid
+        point), 'CIC' (cloud-in-cell), 'TSC' (triangular-shape cloud), 'PCS'
+        (piecewise cubic spline).
 
     References
     ----------
@@ -39,11 +51,13 @@ class DensityField:
     _particles = None
     _boxsize = None
     _box = None
+    _MAS = None
 
-    def __init__(self, particles, box):
+    def __init__(self, particles, boxsize, box, MAS="CIC"):
         self.particles = particles
         self.box = box
-        self._boxsize = 1.
+        self.boxsize = boxsize
+        self.MAS = MAS
 
     @property
     def particles(self):
@@ -64,6 +78,22 @@ class DensityField:
                              "containing `['x', 'y', 'z', 'M']`.")
         self._particles = particles
 
+    @property
+    def boxsize(self):
+        """
+        Box length. Determines the power spectrum units.
+
+        Returns
+        -------
+        boxsize : float
+        """
+        return self._boxsize
+
+    @boxsize.setter
+    def boxsize(self, boxsize):
+        """Set `self.boxsize`."""
+        self._boxsize = boxsize
+
     @property
     def box(self):
         """
@@ -83,15 +113,24 @@ class DensityField:
         self._box = box
 
     @property
-    def boxsize(self):
+    def MAS(self):
         """
-        Boxsize.
+        The mass-assignment scheme.
 
         Returns
         -------
-        boxsize : float
+        MAS : str
         """
-        return self._boxsize
+        return self._MAS
+
+    @MAS.setter
+    def MAS(self, MAS):
+        """Sets `self.MAS`."""
+        opts = ["NGP", "CIC", "TSC", "PCS"]
+        if MAS not in opts:
+            raise ValueError("Invalid MAS `{}`. Options are: `{}`."
+                             .format(MAS, opts))
+        self._MAS = MAS
 
     @staticmethod
     def _force_f32(x, name):
@@ -127,11 +166,10 @@ class DensityField:
         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)
+        MASL.MA(pos, rho, self.boxsize, self.MAS, W=weights, verbose=verbose)
         return rho
 
     def overdensity_field(self, grid, verbose=True):
@@ -176,7 +214,8 @@ class DensityField:
         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")
+        return MASL.potential(
+            delta, self.box._omega_m, self.box._aexp, self.MAS)
 
     def tensor_field(self, grid, verbose=True):
         """
@@ -196,8 +235,8 @@ class DensityField:
             the relevant tensor components.
         """
         delta = self.overdensity_field(grid, verbose)
-        return MASL.tidal_tensor(delta, self.box._omega_m, self.box._aexp,
-                                 "CIC")
+        return MASL.tidal_tensor(
+            delta, self.box._omega_m, self.box._aexp, self.MAS)
 
     def gravitational_field(self, grid, verbose=True):
         """
@@ -219,7 +258,28 @@ class DensityField:
         """
         delta = self.overdensity_field(grid, verbose)
         return MASL.grav_field_tensor(
-            delta, self.box._omega_m, self.box._aexp, "CIC")
+            delta, self.box._omega_m, self.box._aexp, self.MAS)
+
+    def auto_powerspectrum(self, grid, verbose=True):
+        """
+        Calculate the auto 1-dimensional power spectrum.
+
+        Parameters
+        ----------
+        grid : int
+            The grid size.
+        verbose : float, optional
+            A verbosity flag. By default `True`.
+
+        Returns
+        -------
+        pk : py:class`Pk_library.Pk`
+            Power spectrum object.
+        """
+        delta = self.overdensity_field(grid, verbose)
+        return PKL.Pk(
+            delta, self.boxsize, axis=1, MAS=self.MAS, threads=1,
+            verbose=verbose)
 
     def smooth_field(self, field, scale, threads=1):
         """

csiborgtools/read/__init__.py

@@ -18,3 +18,4 @@ from .make_cat import (HaloCatalogue, CombinedHaloCatalogue)  # noqa
 from .readobs import (PlanckClusters, MCXCClusters, TwoMPPGalaxies,  # noqa
                       TwoMPPGroups, SDSS)  # noqa
 from .outsim import (dump_split, combine_splits, make_ascii_powmes)  # noqa
+from .summaries import PKReader  # noqa

csiborgtools/read/outsim.py

@@ -18,11 +18,9 @@ I/O functions for analysing the CSiBORG realisations.
 
 
 import numpy
-from os.path import (join, dirname, basename, isfile)
-import gc
+from os.path import (join, isfile)
 from os import remove
 from tqdm import trange
-from astropy.io import ascii
 from astropy.table import Table
 
 I64 = numpy.int64
@@ -125,10 +123,9 @@ def combine_splits(n_splits, part_reader, cols_add, remove_splits=False,
     return out
 
 
-def make_ascii_powmes(particles, fout, verbose=True):
+def make_ascii_powmes(particles, fout):
     """
     Write an ASCII file with appropriate formatting for POWMES.
-    This is an extremely memory inefficient implementation.
 
     Parameters
     ----------
@@ -136,8 +133,6 @@ def make_ascii_powmes(particles, fout, verbose=True):
         Array of particles.
     fout : str
         File path to store the ASCII file.
-    verbose : bool, optional
-        Verbosity flag. By default `True`.
 
     Returns
     -------
@@ -146,27 +141,11 @@ def make_ascii_powmes(particles, fout, verbose=True):
     out = Table()
     for p in ('x', 'y', 'z', 'M'):
         out[p] = particles[p]
+    Npart = particles.size
+
     # If fout exists, remove
     if isfile(fout):
         remove(fout)
 
-    # Write the temporaty file
-    ftemp = join(dirname(fout), "_" + basename(fout))
-    if verbose:
-        print("Writing temporary file `{}`...".format(ftemp))
-    ascii.write(out, ftemp, overwrite=True, delimiter=",", fast_writer=True)
-
-    del out
-    gc.collect()
-
-    # Write to the first line the number of particles
-    if verbose:
-        print("Writing the full file `{}`...".format(fout))
-    with open(ftemp, 'r') as fread, open(fout, 'w') as fwrite:
-        fwrite.write(str(particles.size) + '\n')
-        for i, line in enumerate(fread):
-            if i == 0:
-                continue
-            fwrite.write(line)
-
-    remove(ftemp)
+    with open(fout, 'wb') as f:
+        numpy.savetxt(f, out, delimiter=',', header=str(Npart), comments='')

csiborgtools/read/readsim.py

@@ -249,6 +249,8 @@ class CSiBORGPaths:
         files = [f.split("/")[-1] for f in files]
         # Remove files with inverted ICs
         files = [f for f in files if "_inv" not in f]
+        # Remove the new files with z = 70 only
+        files = [f for f in files if "_new" not in f]
         # Remove the filename with _old
         files = [f for f in files if "OLD" not in f]
         ids = [int(f.split("_")[-1]) for f in files]
@@ -596,6 +598,9 @@ class ParticleReader:
                     out[fname][i:i + j] = self.read_sp(fdtype, partfiles[cpu])
                 else:
                     dum[i:i + j] = self.read_sp(fdtype, partfiles[cpu])
+        # Close the fortran files
+        for partfile in partfiles:
+            partfile.close()
 
         return out
 
@@ -645,6 +650,8 @@ class ParticleReader:
             j = nparts[cpu]
             ff = self.open_unbinding(cpu)
             clumpid[i:i + j] = ff.read_ints()
+            # Close
+            ff.close()
 
         return clumpid
 

csiborgtools/read/summaries.py

@@ -0,0 +1,168 @@
+# Copyright (C) 2022 Richard Stiskalek, Harry Desmond
+# 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.
+"""
+Tools for summarising various results.
+"""
+import numpy
+import joblib
+from tqdm import tqdm
+
+
+class PKReader:
+    """
+    A shortcut object for reading in the power spectrum files.
+
+    Parameters
+    ----------
+    ic_ids : list of int
+        IC IDs to be read.
+    hw : float
+        Box half-width.
+    fskel : str, optional
+        The skeleton path. By default
+        `/mnt/extraspace/rstiskalek/csiborg/crosspk/out_{}_{}_{}.p`, where
+        the formatting options are `ic0, ic1, hw`.
+    dtype : dtype, optional
+        Output precision. By default `numpy.float32`.
+    """
+    def __init__(self, ic_ids, hw, fskel=None, dtype=numpy.float32):
+        self.ic_ids = ic_ids
+        self.hw = hw
+        if fskel is None:
+            fskel = "/mnt/extraspace/rstiskalek/csiborg/crosspk/out_{}_{}_{}.p"
+        self.fskel = fskel
+        self.dtype = dtype
+
+    @staticmethod
+    def _set_klim(kmin, kmax):
+        """
+        Sets limits on the wavenumber to 0 and infinity if `None`s provided.
+        """
+        if kmin is None:
+            kmin = 0
+        if kmax is None:
+            kmax = numpy.infty
+        return kmin, kmax
+
+    def read_autos(self, kmin=None, kmax=None):
+        """
+        Read in the autocorrelation power spectra.
+
+        Parameters
+        ----------
+        kmin : float, optional
+            The minimum wavenumber. By default `None`, i.e. 0.
+        kmin : float, optional
+            The maximum wavenumber. By default `None`, i.e. infinity.
+
+        Returns
+        -------
+        ks : 1-dimensional array
+            Array of wavenumbers.
+        pks : 2-dimensional array of shape `(len(self.ic_ids), ks.size)`
+            Autocorrelation of each simulation.
+        """
+        kmin, kmax = self._set_klim(kmin, kmax)
+        ks, pks, sel = None, None, None
+        for i, nsim in enumerate(self.ic_ids):
+            pk = joblib.load(self.fskel.format(nsim, nsim, self.hw))
+            # Get cuts and pre-allocate arrays
+            if i == 0:
+                x = pk.k3D
+                sel = (kmin < x) & (x < kmax)
+                ks = x[sel].astype(self.dtype)
+                pks = numpy.full((len(self.ic_ids), numpy.sum(sel)), numpy.nan,
+                                 dtype=self.dtype)
+            pks[i, :] = pk.Pk[sel, 0, 0]
+
+        return ks, pks
+
+    def read_single_cross(self, ic0, ic1, kmin=None, kmax=None):
+        """
+        Read cross-correlation between IC IDs `ic0` and `ic1`.
+
+        Parameters
+        ----------
+        ic0 : int
+            The first IC ID.
+        ic1 : int
+            The second IC ID.
+        kmin : float, optional
+            The minimum wavenumber. By default `None`, i.e. 0.
+        kmin : float, optional
+            The maximum wavenumber. By default `None`, i.e. infinity.
+
+        Returns
+        -------
+        ks : 1-dimensional array
+            Array of wavenumbers.
+        xpk : 1-dimensional array of shape `(ks.size, )`
+            Cross-correlation.
+        """
+        if ic0 == ic1:
+            raise ValueError("Requested cross correlation for the same ICs.")
+        kmin, kmax = self._set_klim(kmin, kmax)
+        # Check their ordering. The latter must be larger.
+        ics = (ic0, ic1)
+        if ic0 > ic1:
+            ics = ics[::-1]
+
+        pk = joblib.load(self.fskel.format(*ics, self.hw))
+        ks = pk.k3D
+        sel = (kmin < ks) & (ks < kmax)
+        ks = ks[sel].astype(self.dtype)
+        xpk = pk.XPk[sel, 0, 0].astype(self.dtype)
+
+        return ks, xpk
+
+    def read_cross(self, kmin=None, kmax=None):
+        """
+        Read cross-correlation between all IC pairs.
+
+        Parameters
+        ----------
+        kmin : float, optional
+            The minimum wavenumber. By default `None`, i.e. 0.
+        kmin : float, optional
+            The maximum wavenumber. By default `None`, i.e. infinity.
+
+        Returns
+        -------
+        ks : 1-dimensional array
+            Array of wavenumbers.
+        xpks : 3-dimensional array of shape (`nics, nics - 1, ks.size`)
+            Cross-correlations. The first column is the the IC and is being
+            cross-correlated with the remaining ICs, in the second column.
+        """
+        nics = len(self.ic_ids)
+
+        ks, xpks = None, None
+        for i, ic0 in enumerate(tqdm(self.ic_ids)):
+            k = 0
+            for ic1 in self.ic_ids:
+                # We don't want cross-correlation
+                if ic0 == ic1:
+                    continue
+                x, y = self.read_single_cross(ic0, ic1, kmin, kmax)
+                # If in the first iteration pre-allocate arrays
+                if ks is None:
+                    ks = x
+                    xpks = numpy.full((nics, nics - 1, ks.size), numpy.nan,
+                                      dtype=self.dtype)
+                xpks[i, k, :] = y
+                # Bump up the iterator
+                k += 1
+
+        return ks, xpks

scripts/run_crosspk.py

@@ -0,0 +1,168 @@
+# 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.
+"""
+MPI script to calculate the matter cross power spectrum between CSiBORG
+IC realisations. Units are Mpc/h.
+"""
+from argparse import ArgumentParser
+import numpy
+import joblib
+from datetime import datetime
+from itertools import combinations
+from os.path import join
+from os import remove
+from gc import collect
+from mpi4py import MPI
+import Pk_library as PKL
+try:
+    import csiborgtools
+except ModuleNotFoundError:
+    import sys
+    sys.path.append("../")
+    import csiborgtools
+import utils
+
+
+parser = ArgumentParser()
+parser.add_argument("--grid", type=int)
+parser.add_argument("--halfwidth", type=float, default=0.5)
+args = parser.parse_args()
+
+# Get MPI things
+comm = MPI.COMM_WORLD
+rank = comm.Get_rank()
+nproc = comm.Get_size()
+MAS = "CIC"  # mass asignment scheme
+
+paths = csiborgtools.read.CSiBORGPaths()
+ics = paths.ic_ids
+n_sims = len(ics)
+
+# File paths
+ftemp = join(utils.dumpdir, "temp_crosspk",
+             "out_{}_{}" + "_{}".format(args.halfwidth))
+fout = join(utils.dumpdir, "crosspk",
+            "out_{}_{}" + "_{}.p".format(args.halfwidth))
+
+
+jobs = csiborgtools.fits.split_jobs(n_sims, nproc)[rank]
+for n in jobs:
+    print("Rank {}@{}: saving {}th delta.".format(rank, datetime.now(), n))
+    # Set the paths
+    n_sim = ics[n]
+    paths.set_info(n_sim, paths.get_maximum_snapshot(n_sim))
+    # Set reader and the box
+    reader = csiborgtools.read.ParticleReader(paths)
+    box = csiborgtools.units.BoxUnits(paths)
+    # Read particles
+    particles = reader.read_particle(["x", "y", "z", "M"], verbose=False)
+    # Halfwidth -- particle selection
+    if args.halfwidth < 0.5:
+        hw = args.halfwidth
+        mask = ((0.5 - hw < particles['x']) & (particles['x'] < 0.5 + hw)
+                & (0.5 - hw < particles['y']) & (particles['y'] < 0.5 + hw)
+                & (0.5 - hw < particles['z']) & (particles['z'] < 0.5 + hw))
+        # Subselect the particles
+        particles = particles[mask]
+        # Rescale to range [0, 1]
+        for p in ('x', 'y', 'z'):
+            particles[p] = (particles[p] - 0.5 + hw) / (2 * hw)
+
+        length = box.box2mpc(2 * hw) * box.h
+    else:
+        mask = None
+        length = box.box2mpc(1) * box.h
+    # Calculate the overdensity field
+    field = csiborgtools.field.DensityField(particles, length, box, MAS)
+    delta = field.overdensity_field(args.grid, verbose=False)
+    aexp = box._aexp
+
+    # Try to clean up memory
+    del field, particles, box, reader, mask
+    collect()
+
+    # Dump the results
+    with open(ftemp.format(n_sim, "delta") + ".npy", "wb") as f:
+        numpy.save(f, delta)
+    joblib.dump([aexp, length], ftemp.format(n_sim, "lengths") + ".p")
+
+    # Try to clean up memory
+    del delta
+    collect()
+
+
+comm.Barrier()
+
+# Get off-diagonal elements and append the diagoal
+combs = [c for c in combinations(range(n_sims), 2)]
+for i in range(n_sims):
+    combs.append((i, i))
+prev_delta = [-1, None, None, None]  # i, delta, aexp, length
+
+jobs = csiborgtools.fits.split_jobs(len(combs), nproc)[rank]
+for n in jobs:
+    i, j = combs[n]
+    print("Rank {}@{}: combination {}.".format(rank, datetime.now(), (i, j)))
+
+    # If i same as last time then don't have to load it
+    if prev_delta[0] == i:
+        delta_i = prev_delta[1]
+        aexp_i = prev_delta[2]
+        length_i = prev_delta[3]
+    else:
+        with open(ftemp.format(ics[i], "delta") + ".npy", "rb") as f:
+            delta_i = numpy.load(f)
+        aexp_i, length_i = joblib.load(ftemp.format(ics[i], "lengths") + ".p")
+        # Store in prev_delta
+        prev_delta[0] = i
+        prev_delta[1] = delta_i
+        prev_delta[2] = aexp_i
+        prev_delta[3] = length_i
+
+    # Get jth delta
+    with open(ftemp.format(ics[j], "delta") + ".npy", "rb") as f:
+        delta_j = numpy.load(f)
+    aexp_j, length_j = joblib.load(ftemp.format(ics[j], "lengths") + ".p")
+
+    # Verify the difference between the scale factors! Say more than 1%
+    daexp = abs((aexp_i - aexp_j) / aexp_i)
+    if daexp > 0.01:
+        raise ValueError(
+            "Boxes {} and {} final snapshot scale factors disagree by "
+            "`{}` percent!".format(ics[i], ics[j], daexp * 100))
+    # Check how well the boxsizes agree
+    dlength = abs((length_i - length_j) / length_i)
+    if dlength > 0.001:
+        raise ValueError("Boxes {} and {} box sizes disagree by `{}` percent!"
+                         .format(ics[i], ics[j], dlength * 100))
+
+    # Calculate the cross power spectrum
+    Pk = PKL.XPk([delta_i, delta_j], length_i, axis=1, MAS=[MAS, MAS],
+                 threads=1)
+    joblib.dump(Pk, fout.format(ics[i], ics[j]))
+
+    del delta_i, delta_j, Pk
+    collect()
+
+
+# Clean up the temp files
+comm.Barrier()
+if rank == 0:
+    print("Cleaning up the temporary files...")
+    for ic in ics:
+        remove(ftemp.format(ic, "delta") + ".npy")
+        remove(ftemp.format(ic, "lengths") + ".p")
+
+    print("All finished!")