Environmental properties (#20)

* rm get_positions

* Add comment

* add halfwidth func

* Update docs

* Add imprt

* Evaluate multiple fields simulatenously

* add halfwidth selection

* Change order of grav field and tensor field

* Add gravitational field norm

* Add eigenvalue calculation

* Sorted eigenvalues

* add init script

* add progress

* Add surveys

* Add more survey flexibility

* Minor changes

* add survey names

* rm name

* Fix list bug

* Fig bugs when running the script

* add phi to dtype

* fix dump bug

* Add comment

* Add smoothing options

* Add further comment

* Update TODO

scripts/run_crosspk.py

@@ -70,27 +70,18 @@ for n in jobs:
     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
+        particles = csiborgtools.read.halfwidth_select(
+            args.halfwidth, particles)
+        length = box.box2mpc(2 * args.halfwidth) * box.h  # Mpc/h
     else:
-        mask = None
-        length = box.box2mpc(1) * box.h
+        length = box.box2mpc(1) * box.h  # Mpc/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
+    del field, particles, box, reader
     collect()
 
     # Dump the results

scripts/run_fieldprop.py

@@ -0,0 +1,120 @@
+# 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 evaluate field properties at the galaxy positions.
+
+NOTE:
+- Calculate for the entire box or just for a smaller region?
+- Add argparser for different options.
+- In the argparser add options to smoothen the field.
+"""
+import numpy
+from datetime import datetime
+from mpi4py import MPI
+from os.path import join
+from os import remove
+try:
+    import csiborgtools
+except ModuleNotFoundError:
+    import sys
+    sys.path.append("../")
+    import csiborgtools
+import utils
+
+halfwidth = 0.5
+MAS = "CIC"
+grid = 256
+
+# Get MPI things
+comm = MPI.COMM_WORLD
+rank = comm.Get_rank()
+nproc = comm.Get_size()
+
+# Galaxy positions
+survey = "SDSS"
+survey = utils.surveys[survey]()()
+pos = numpy.vstack([survey[p] for p in ("DIST", "RA", "DEC")]).T
+pos = pos.astype(numpy.float32)
+
+# File paths
+ftemp = join(utils.dumpdir, "temp_fields", "out_" + survey.name + "_{}.npy")
+fperm = join(utils.dumpdir, "fields", "out_{}.npy".format(survey.name))
+
+# Edit depending on what is calculated
+dtype = {"names": ["delta", "phi"], "formats": [numpy.float32] * 2}
+
+# CSiBORG simulation paths
+paths = csiborgtools.read.CSiBORGPaths()
+ics = paths.ic_ids[:10]
+n_sims = len(ics)
+
+for n in csiborgtools.fits.split_jobs(n_sims, nproc)[rank]:
+    print("Rank {}@{}: working on {}th IC.".format(rank, datetime.now(), n),
+          flush=True)
+    # 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 and select a subset of them
+    particles = reader.read_particle(["x", "y", "z", "M"], verbose=False)
+    if halfwidth < 0.5:
+        particles = csiborgtools.read.halfwidth_select(halfwidth, particles)
+        length = box.box2mpc(2 * halfwidth) * box.h  # Mpc/h
+    else:
+        length = box.box2mpc(1) * box.h  # Mpc/h
+
+    # Initialise the field object and output array
+    field = csiborgtools.field.DensityField(particles, length, box, MAS)
+    out = numpy.full(pos.shape[0], numpy.nan, dtype=dtype)
+
+    # Calculate the overdensity field and interpolate at galaxy positions
+    feval = field.overdensity_field(grid, verbose=False)
+    out["delta"] = field.evaluate_sky(feval, pos=pos, isdeg=True)[0]
+
+    # Potential
+    feval = field.potential_field(grid, verbose=False)
+    out["phi"] = field.evaluate_sky(feval, pos=pos, isdeg=True)[0]
+
+    # Calculate the remaining fields
+    # ...
+    # ...
+
+    # Dump the results
+    with open(ftemp.format(n_sim), "wb") as f:
+        numpy.save(f, out)
+
+# Wait for all ranks to finish
+comm.Barrier()
+if rank == 0:
+    print("Collecting files...", flush=True)
+
+    out = numpy.full((n_sims, pos.shape[0]), numpy.nan, dtype=dtype)
+
+    for n in range(n_sims):
+        n_sim = ics[n]
+        with open(ftemp.format(n_sim), "rb") as f:
+            fin = numpy.load(f, allow_pickle=True)
+            for name in dtype["names"]:
+                out[name][n, ...] = fin[name]
+        # Remove the temporary file
+        remove(ftemp.format(n_sim))
+
+    print("Saving results to `{}`.".format(fperm), flush=True)
+    with open(fperm, "wb") as f:
+        numpy.save(f, out)

scripts/utils.py

@@ -18,11 +18,11 @@ Notebook utility functions.
 
 # from os.path import join
 
-# try:
-#     import csiborgtools
-# except ModuleNotFoundError:
-#     import sys
-#     sys.path.append("../")
+try:
+    import csiborgtools
+except ModuleNotFoundError:
+    import sys
+    sys.path.append("../")
 
 
 Nsplits = 200
@@ -39,3 +39,23 @@ _virgo = {"RA": (12 + 27 / 60) * 15,
           "COMDIST": 16.5}
 
 specific_clusters = {"Coma": _coma, "Virgo": _virgo}
+
+
+###############################################################################
+#                                 Surveys                                     #
+###############################################################################
+
+
+class SDSS:
+    @staticmethod
+    def steps(cls):
+        return [(lambda x: cls[x], ("IN_DR7_LSS",)),
+                (lambda x: cls[x] < 17.6, ("ELPETRO_APPMAG_r", )),
+                (lambda x: cls[x] < 155, ("DIST", ))
+                ]
+
+    def __call__(self):
+        return csiborgtools.read.SDSS(h=1, sel_steps=self.steps)
+
+
+surveys = {"SDSS": SDSS}