Add dumping to ASCII of halos (#89)

* Rename file

* Add argument

* Add ASCII positions paths

* Add halo positions dumping

* Update halo catalogues

csiborgtools/read/halo_cat.py

@@ -256,6 +256,10 @@ class BaseCatalogue(ABC):
 
     @observer_velocity.setter
     def observer_velocity(self, obs_vel):
+        if obs_vel is None:
+            self._observer_velocity = None
+            return
+
         assert isinstance(obs_vel, (list, tuple, numpy.ndarray))
         obs_vel = numpy.asanyarray(obs_vel)
         assert obs_vel.shape == (3,)

csiborgtools/read/paths.py

@@ -449,6 +449,26 @@ class Paths:
         fname = f"parts_{str(nsim).zfill(5)}.h5"
         return join(fdir, fname)
 
+    def ascii_positions(self, nsim, kind):
+        """
+        Path to ASCII files containing the positions of particles or halos.
+
+        Parameters
+        ----------
+        nsim : int
+            IC realisation index.
+        kind : str
+            Kind of data to extract. Must be one of `particles`,
+            `particles_rsp`, `halos`, `halos_rsp`.
+        """
+        assert kind in ["particles", "particles_rsp", "halos", "halos_rsp"]
+
+        fdir = join(self.postdir, "ascii_positions")
+        try_create_directory(fdir)
+        fname = f"pos_{kind}_{str(nsim).zfill(5)}.txt"
+
+        return join(fdir, fname)
+
     def structfit(self, nsnap, nsim, simname):
         """
         Path to the halo catalogue from `fit_halos.py`.

scripts/dump_to_ascii.py

@@ -0,0 +1,116 @@
+# 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
+# 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.
+"""Convert the HDF5 CSiBORG particle file to an ASCII file."""
+from argparse import ArgumentParser
+
+import h5py
+import numpy
+from mpi4py import MPI
+from taskmaster import work_delegation
+from tqdm import trange
+
+import csiborgtools
+from utils import get_nsims
+
+
+def positions_to_ascii(positions, output_filename, boxsize=None,
+                       chunk_size=50_000, verbose=True):
+    """
+    Convert array of positions to an ASCII file. If `boxsize` is given,
+    multiples the positions by it.
+    """
+    total_size = len(positions)
+
+    if verbose:
+        print(f"Number of rows to write: {total_size}")
+
+    with open(output_filename, 'w') as out_file:
+        # Write the header
+        out_file.write("#px py pz\n")
+
+        # Loop through data in chunks
+        for i in trange(0, total_size, chunk_size,
+                        desc=f"Writing to ... `{output_filename}`",
+                        disable=not verbose):
+
+            end = i + chunk_size
+            if end > total_size:
+                end = total_size
+
+            data_chunk = positions[i:end]
+            # Convert to positions Mpc / h
+            data_chunk = data_chunk[:, :3]
+
+            if boxsize is not None:
+                data_chunk *= boxsize
+
+            chunk_str = "\n".join([f"{x:.4f} {y:.4f} {z:.4f}"
+                                   for x, y, z in data_chunk])
+            out_file.write(chunk_str + "\n")
+
+
+def extract_positions(nsim, paths, kind):
+    """
+    Extract either the particle or halo positions.
+    """
+    if kind == "particles":
+        fname = paths.particles(nsim, args.simname)
+        return h5py.File(fname, 'r')["particles"]
+
+    if kind == "particles_rsp":
+        raise NotImplementedError("RSP of particles is not implemented yet.")
+
+    fpath = paths.observer_peculiar_velocity("PCS", 512, nsim)
+    vpec_observer = numpy.load(fpath)["observer_vp"][0, :]
+    cat = csiborgtools.read.CSiBORGHaloCatalogue(
+        nsim, paths, load_fitted=True, load_initial=False,
+        observer_velocity=vpec_observer)
+
+    if kind == "halos":
+        return cat.position()
+
+    if kind == "halos_rsp":
+        return cat.redshift_space_position()
+
+    raise ValueError(f"Unknown kind `{kind}`. Allowed values are: "
+                     "`particles`, `particles_rsp`, `halos`, `halos_rsp`.")
+
+
+def main(nsim, paths, kind):
+    boxsize = 677.7 if "particles" in kind else None
+    pos = extract_positions(nsim, paths, kind)
+    output_filename = paths.ascii_positions(nsim, kind)
+    positions_to_ascii(pos, output_filename, boxsize=boxsize)
+
+
+if __name__ == "__main__":
+    parser = ArgumentParser()
+    parser.add_argument("--kind", type=str, required=True,
+                        choices=["particles", "particles_rsp", "halos", "halos_rsp"],  # noqa
+                        help="Kind of data to extract.")
+    parser.add_argument("--nsims", type=int, nargs="+", default=None,
+                        help="IC realisations. If `-1` processes all.")
+    parser.add_argument("--simname", type=str, default="csiborg",
+                        choices=["csiborg"],
+                        help="Simulation name")
+    args = parser.parse_args()
+
+    paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
+    nsims = get_nsims(args, paths)
+
+    def _main(nsim):
+        main(nsim, paths, args.kind)
+
+    work_delegation(main, nsims, MPI.COMM_WORLD)

scripts/field_sample.py

@@ -180,6 +180,8 @@ if __name__ == "__main__":
                         help="Field in RSP?")
     parser.add_argument("--nrand", type=int, required=True,
                         help="Number of rand. positions to evaluate the field")
+    parser.add_argument("--simname", type=str, default="csiborg",
+                        choices=["csiborg"], help="Simulation name")
     args = parser.parse_args()
 
     paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)

scripts/particles_to_ascii.py

@@ -1,82 +0,0 @@
-# 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
-# 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.
-"""Convert the HDF5 CSiBORG particle file to an ASCII file."""
-from argparse import ArgumentParser
-
-import csiborgtools
-import h5py
-
-from mpi4py import MPI
-
-from utils import get_nsims
-from tqdm import trange
-
-from taskmaster import work_delegation
-
-
-def h5_to_ascii(nsim, paths, chunk_size=50_000, verbose=True):
-    """
-    Convert the HDF5 CSiBORG particle file to an ASCII file. Outputs only
-    particle positions in Mpc / h. Ignores the unequal particle masses.
-    """
-    fname = paths.particles(nsim, args.simname)
-    boxsize = 677.7
-
-    fname_out = fname.replace(".h5", ".txt")
-
-    with h5py.File(fname, 'r') as f:
-        dataset = f["particles"]
-        total_size = dataset.shape[0]
-
-        if verbose:
-            print(f"Number of rows to write: {total_size}")
-
-        with open(fname_out, 'w') as out_file:
-            # Write the header
-            out_file.write("#px py pz\n")
-
-            # Loop through data in chunks
-            for i in trange(0, total_size, chunk_size,
-                            desc=f"Writing to ... `{fname_out}`",
-                            disable=not verbose):
-                end = i + chunk_size
-                if end > total_size:
-                    end = total_size
-
-                data_chunk = dataset[i:end]
-                # Convert to positions Mpc / h
-                data_chunk = data_chunk[:, :3] * boxsize
-
-                chunk_str = "\n".join([f"{x:.4f} {y:.4f} {z:.4f}"
-                                       for x, y, z in data_chunk])
-                out_file.write(chunk_str + "\n")
-
-
-if __name__ == "__main__":
-    parser = ArgumentParser()
-    parser.add_argument("--nsims", type=int, nargs="+", default=None,
-                        help="IC realisations. If `-1` processes all.")
-    parser.add_argument("--simname", type=str, default="csiborg",
-                        choices=["csiborg"],
-                        help="Simulation name")
-    args = parser.parse_args()
-
-    paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
-    nsims = get_nsims(args, paths)
-
-    def main(nsim):
-        h5_to_ascii(nsim, paths, verbose=MPI.COMM_WORLD.Get_size() == 1)
-
-    work_delegation(main, nsims, MPI.COMM_WORLD)