Write particles to ASCII (#117)

* Add option to output high-resolution particles only

* Add ASCII writing of simulatios

csiborgtools/read/snapshot.py

@@ -474,7 +474,7 @@ class CSiBORG2Snapshot(BaseSnapshot):
 
         self._kind = value
 
-    def _get_particles(self, kind):
+    def _get_particles(self, kind, high_resolution_only=False):
         with File(self._snapshot_path, "r") as f:
             if kind == "Masses":
                 npart = f["Header"].attrs["NumPart_Total"][1]
@@ -483,27 +483,28 @@ class CSiBORG2Snapshot(BaseSnapshot):
             else:
                 x = f[f"PartType1/{kind}"][...]
 
-            if x.ndim == 1:
-                x = numpy.hstack([x, f[f"PartType5/{kind}"][...]])
-            else:
-                x = numpy.vstack([x, f[f"PartType5/{kind}"][...]])
+            if not high_resolution_only:
+                if x.ndim == 1:
+                    x = numpy.hstack([x, f[f"PartType5/{kind}"][...]])
+                else:
+                    x = numpy.vstack([x, f[f"PartType5/{kind}"][...]])
 
         if self.flip_xz and kind in ["Coordinates", "Velocities"]:
             x[:, [0, 2]] = x[:, [2, 0]]
 
         return x
 
-    def coordinates(self):
-        return self._get_particles("Coordinates")
+    def coordinates(self, high_resolution_only=False):
+        return self._get_particles("Coordinates", high_resolution_only)
 
-    def velocities(self):
-        return self._get_particles("Velocities")
+    def velocities(self, high_resolution_only=False):
+        return self._get_particles("Velocities", high_resolution_only)
 
-    def masses(self):
-        return self._get_particles("Masses") * 1e10
+    def masses(self, high_resolution_only=False):
+        return self._get_particles("Masses", high_resolution_only) * 1e10
 
-    def particle_ids(self):
-        return self._get_particles("ParticleIDs")
+    def particle_ids(self, high_resolution_only=False):
+        return self._get_particles("ParticleIDs", high_resolution_only)
 
     def _get_halo_particles(self, halo_id, kind, is_group):
         if not is_group:

scripts/output_halocatalogue.py

@@ -19,14 +19,48 @@ galaxies in a survey as an ASCII file.
 from os.path import join
 
 import csiborgtools
-import numpy
+import numpy as np
 from tqdm import tqdm
 
 DIR_OUT = "/mnt/extraspace/rstiskalek/csiborg_postprocessing/ascii_positions"
 
 
-def process_simulation(simname):
-    """Watch out about the distinction between real and redshift space."""
+def write_simulation(simname, high_resolution_only=True):
+    """"
+    Write the positions, velocities and IDs of the particles in a simulation
+    to an ASCII file. The output is `X Y Z VX VY VZ ID`.
+    """
+    if not high_resolution_only:
+        raise RuntimeError("Writing low-res particles is not implemented.")
+
+    paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
+
+    if "csiborg2" in simname:
+        nsims = paths.get_ics(simname)
+        kind = simname.split("_")[-1]
+
+        for nsim in tqdm(nsims, desc="Simulations"):
+            reader = csiborgtools.read.CSiBORG2Snapshot(nsim, 99, kind, paths,
+                                                        flip_xz=False)
+            x = np.hstack(
+                [reader.coordinates(high_resolution_only=True),
+                 reader.velocities(high_resolution_only=True),
+                 reader.particle_ids(high_resolution_only=True).reshape(-1, 1)]
+                 )
+            # Store positions and velocities with 6 decimal places and IDs as
+            # integers.
+            fmt_string = "%1.6f %1.6f %1.6f %1.6f %1.6f %1.6f %d"
+            fname = join(DIR_OUT, f"high_res_particles_{simname}_{nsim}.txt")
+            np.savetxt(fname, x, fmt=fmt_string)
+    else:
+        raise RuntimeError("Simulation not implemented..")
+
+
+def write_halos(simname):
+    """
+    Watch out about the distinction between real and redshift space. The output
+    is `X Y Z MASS`.
+    """
     paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
     if "csiborg2" in simname:
         nsims = paths.get_ics(simname)
@@ -37,16 +71,16 @@ def process_simulation(simname):
             pos = cat["cartesian_pos"]
             mass = cat["totmass"]
             # Stack positions and masses
-            x = numpy.hstack([pos, mass.reshape(-1, 1)])
+            x = np.hstack([pos, mass.reshape(-1, 1)])
 
             # Save to a file
             fname = join(DIR_OUT, f"halos_real_{simname}_{nsim}.txt")
-            numpy.savetxt(fname, x)
+            np.savetxt(fname, x)
     else:
         raise RuntimeError("Simulation not implemented..")
 
 
-def process_survey(survey_name, boxsize):
+def write_survey(survey_name, boxsize):
     """Watch out about the distance definition."""
     if survey_name == "SDSS":
         survey = csiborgtools.SDSS()()
@@ -58,19 +92,21 @@ def process_survey(survey_name, boxsize):
         raise RuntimeError("Survey not implemented..")
 
     # Convert to Cartesian coordinates
-    X = numpy.vstack([dist, ra, dec]).T
+    X = np.vstack([dist, ra, dec]).T
     X = csiborgtools.radec_to_cartesian(X)
 
     # Center the coordinates in the box
     X += boxsize / 2
 
     fname = join(DIR_OUT, f"survey_{survey_name}.txt")
-    numpy.savetxt(fname, X)
+    np.savetxt(fname, X)
 
 
 if __name__ == "__main__":
-    # process_simulation("csiborg2_main")
+    write_simulation("csiborg2_main")
 
-    boxsize = 676.6
-    for survey in ["SDSS", "SDSSxALFALFA"]:
-        process_survey(survey, boxsize)
+    # write_halos("csiborg2_main")
+
+    # boxsize = 676.6
+    # for survey in ["SDSS", "SDSSxALFALFA"]:
+    #     write_survey(survey, boxsize)