separate ramses plotting script

ramses_plot.py

@@ -0,0 +1,158 @@
+from pathlib import Path
+from subprocess import run
+from typing import List
+
+import numpy as np
+import pynbody
+from matplotlib import pyplot as plt
+from matplotlib.colors import LogNorm
+from pynbody.array import SimArray
+from pynbody.snapshot import FamilySubSnap
+from pynbody.snapshot.ramses import RamsesSnap
+from scipy import constants
+
+ramses_imager = Path("~/cosmoca/RamsesImager/build/RamsesImager").expanduser()
+dir = Path("/media/ssd/cosmos_data/coxeter/richings21_ics/richings21_bary_ramses_7_10_11/")
+
+
+def get_slice_argument(
+        extent: List[float], zcenter: float,
+        ramses_dir: Path, interpolation_method: str,
+        depth: float):
+    xmin, xmax, ymin, ymax = extent
+    # interpolate = interpolation_method == "linear"
+    interpolate = False
+    arguments = {
+        "x": (xmin + xmax) / 2,
+        "y": (ymin + ymax) / 2,
+        "z": zcenter,
+        "w": xmax - xmin,
+        "h": ymax - ymin,
+        "d": depth,
+        "l": 14
+    }
+    args = [str(ramses_imager)]
+    for k, v in arguments.items():
+        args.append(f"-{k} {v}")
+
+    if interpolate:
+        args.append("-i")
+
+    args.append(str(ramses_dir / "info_00009.txt"))
+    return args, ramses_imager.parent
+
+
+def load_slice_data(file: Path):
+    with file.open("rb") as infile:
+        np.fromfile(file=infile, dtype=np.int32, count=1)
+        [nx, ny] = np.fromfile(file=infile, dtype=np.int32, count=2)
+        np.fromfile(file=infile, dtype=np.int32, count=1)
+
+        np.fromfile(file=infile, dtype=np.int32, count=1)
+        data: np.ndarray = np.fromfile(file=infile, dtype=np.float32, count=nx * ny)
+        np.fromfile(file=infile, dtype=np.int32, count=1)
+        print("NEGATIVE", (data < 0).sum())
+        # np.fromfile(file=infile, dtype=np.int32, count=1)
+        # cm_per_px = np.fromfile(file=infile, dtype=np.float64, count=1)[0]
+        # np.fromfile(file=infile, dtype=np.int32, count=1)
+    return data.reshape((nx, ny))
+
+
+def main():
+    file = dir / "output_00009"
+    s: RamsesSnap = pynbody.load(str(file))
+    gas_data: FamilySubSnap = s.gas
+    h = s.properties["h"]
+    a = s.properties["a"]
+    print("RAMSES a", a)
+
+    # dm
+    mass_array: SimArray = s.dm["mass"]
+    coord_array: SimArray = s.dm["pos"]
+    masses = np.asarray(mass_array.in_units("1e10 Msol"))
+    high_res_mass = np.amin(np.unique(masses))  # get lowest mass of particles
+
+    is_high_res_particle = masses == high_res_mass
+
+    coordinates = np.asarray(coord_array.in_units("Mpc"))
+    hr_coordinates = coordinates[is_high_res_particle] / a
+
+    np.save(f"/home/lukas/tmp/ramses_grid_hr_coordinates", hr_coordinates)
+    print(high_res_mass)
+
+    # temperature_array: SimArray = gas_data["temp"]
+    # p_array: SimArray = gas_data["p"].in_units("1e10 Msol Mpc^-3 km^2 s^-2")
+    # rho_array: SimArray = gas_data["rho"].in_units("1e10 Msol Mpc^-3")
+    # coord_array: SimArray = gas_data["pos"].in_units("Mpc")
+    # mass_array = np.asarray(gas_data["mass"].in_units("1e10 Msol"))
+    # bary_coords = np.asarray(coord_array)
+    # bary_properties = {
+    #     "Temperatures": np.asarray(temperature_array.in_units("K")),
+    #     "Pressures": np.asarray(p_array),
+    #     "Densities": np.asarray(rho_array),
+    #     "Entropies": np.asarray(np.log10(p_array / rho_array ** (5 / 3))),
+    # }
+
+    michael_extent = np.asarray([[32.07045664, 34.13128024],
+                                 [37.562513, 39.6233366],
+                                 [34.54953566, 36.61035927]])
+
+    frac_extent = michael_extent[:2].flatten() / 100 / h
+    frac_center = michael_extent[2].mean() / 100 / h
+    # old extents:
+    # frac_extent=[0.46938643, 0.50738643, 0.55012166, 0.58812166]
+    # frac_center=0.52467408
+
+    print("frac")
+
+    print(frac_extent)
+    print(frac_center)
+    args, imager_dir = get_slice_argument(
+        frac_extent, frac_center,
+        file, interpolation_method="nearest",
+        depth=0.00000000001
+        # depth=depth
+    )
+    print(" ".join(args))
+
+    run(args, cwd=imager_dir)
+    property_map = {
+        "Densities": "rhomap",
+        "Entropies": "Smap",
+        "Temperatures": "Tmap"
+    }
+
+    for property in ["cic", "Densities", "Entropies", "Temperatures"]:
+        if property == "cic":
+            ...
+        else:
+
+            fname = imager_dir / f"snapshot_{property_map[property]}_zproj_zobs-0p00.bin"
+            grid = load_slice_data(fname).T
+            if grid.sum() == 0:
+                raise ValueError("all 0")
+            if property == "Densities":
+                # convert g/cm^3 to 1e10 Msol Mpc^-3
+                solar_mass_in_gram = 1.988e33
+                mpc_in_cm = constants.parsec * constants.mega * 100
+                grid = np.asarray(grid)
+                grid *= (mpc_in_cm ** 3 / solar_mass_in_gram / 1e10)
+            print(grid.shape, "shape")
+            print(grid)
+            plt.figure()
+            plt.imshow(
+                grid,
+                norm=LogNorm(),
+                interpolation="none",
+                origin="lower",
+                extent=frac_extent * 100,
+            )
+            plt.colorbar()
+            plt.savefig(f"/home/lukas/tmp/ramses_grid_{property}.pdf")
+
+            np.save(f"/home/lukas/tmp/ramses_grid_{property}", grid)
+    plt.show()
+
+
+if __name__ == '__main__':
+    main()