Merge branch 'main' of github.com:Findus23/halo_comparison

PlotDaubechies.py

@@ -2,6 +2,7 @@
 originally created by Oliver Hahn
 in PlotDaubechies.ipynb
 """
+from math import pi
 from pathlib import Path
 
 import matplotlib.pyplot as plt
@@ -160,8 +161,8 @@ xdb16, phidb16, psidb16 = cascade_algorithm(h_DB16, g_DB16, maxit)
 ###################################
 
 fig: Figure
-fig, ax = plt.subplots(5, 2, figsize=(12, 12))  # ,gridspec_kw = {'wspace':0.025})
+fig, ax = plt.subplots(4, 2, figsize=(8, 12))  # ,gridspec_kw = {'wspace':0.025})
-label = ['Haar', 'DB2', 'DB4', 'DB8', 'DB16']
+labels = ['Haar', 'DB2', 'DB4', 'DB8', 'DB16']
 
 ax[0, 0].set_title('scaling functions $\\varphi$')
 ax[0, 1].set_title('wavelets $\\psi$')
@@ -178,20 +179,22 @@ ax[2, 1].plot(xdb4, psidb4, lw=1)
 ax[3, 0].plot(xdb8, phidb8, lw=1)
 ax[3, 1].plot(xdb8, psidb8, lw=1)
 
-ax[4, 0].plot(xdb16, phidb16, lw=1)
+# ax[4, 0].plot(xdb16, phidb16, lw=1)
-ax[4, 1].plot(xdb16, psidb16, lw=1)
+# ax[4, 1].plot(xdb16, psidb16, lw=1)
-
 for a in ax.flatten():
     a.set_xlabel('t')
 
-for a, i in zip(ax[:, 0], label):
+for a, label in zip(ax[:, 0], labels):
-    a.set_ylabel('$\\varphi_{\\rm ' + i + '}[t]$')
+    a.set_ylabel(r"$\varphi_{\textrm{LABEL}}$".replace("LABEL", label))
+    # a.set_ylabel('$\\varphi_{\\rm ' + i + '}[t]$')
     a.set_ylim([-1.0, 1.5])
 
-for a, i in zip(ax[:, 1], label):
+for a, label in zip(ax[:, 1], labels):
-    a.set_ylabel('$\\psi_{\\rm ' + i + '}[t]$')
+    a.set_ylabel(r"$\psi_{\textrm{LABEL}}$".replace("LABEL", label))
+    # a.set_ylabel('$\\psi_{\\rm ' + i + '}[t]$')
     a.set_ylim([-2, 2])
 
+fig.tight_layout()
 fig.savefig(Path(f"~/tmp/wavelets.pdf").expanduser(), bbox_inches='tight')
 
 # # Spectral Response of Scaling Functions and Wavelets
@@ -231,9 +234,21 @@ kdb8, fphidb8, fpsidb8 = fourier_wavelet(h_DB8, g_DB8, 256)
 ax.plot(kdb8, np.abs(fphidb8) ** 2, label='$\\hat\\varphi_{DB8}$', c="C3")
 ax.plot(kdb8, np.abs(fpsidb8) ** 2, label='$\\hat\\psi_{DB8}$', c="C3", linestyle="dashed")
 
-kdb16, fphidb16, fpsidb16 = fourier_wavelet(h_DB16, g_DB16, 256)
+
-ax.plot(kdb16, np.abs(fphidb16) ** 2, label='$\\hat\\varphi_{DB16}$', c="C4")
+# all k* are np.linspace(0, np.pi, 256), so we can also use them for shannon
-ax.plot(kdb16, np.abs(fpsidb16) ** 2, label='$\\hat\\psi_{DB16}$', c="C4", linestyle="dashed")
+
+def shannon(k):
+    y=np.zeros_like(k)
+    y[k > pi/2] = 1
+    return y
+
+
+ax.plot(kdb8, shannon(kdb8), label='$\\hat\\varphi_{shannon}$', c="C4")
+# ax.plot(kdb8, np.abs(fpsidb8) ** 2, label='$\\hat\\psi_{DB8}$', c="C3", linestyle="dashed")
+
+# kdb16, fphidb16, fpsidb16 = fourier_wavelet(h_DB16, g_DB16, 256)
+# ax.plot(kdb16, np.abs(fphidb16) ** 2, label='$\\hat\\varphi_{DB16}$', c="C4")
+# ax.plot(kdb16, np.abs(fpsidb16) ** 2, label='$\\hat\\psi_{DB16}$', c="C4", linestyle="dashed")
 ax.legend(frameon=False)
 ax.set_xlabel('k')
 ax.set_ylabel('P(k)')

auriga_comparison.py

@@ -17,9 +17,8 @@ from cic import cic_from_radius
 from halo_mass_profile import halo_mass_profile
 from nfw import fit_nfw
 from paths import auriga_dir, richings_dir
-from read_vr_files import read_velo_halos
 from readfiles import read_file, read_halo_file, ParticlesMeta
-from utils import read_swift_config
+from utils import read_swift_config, print_wall_time
 
 
 class Mode(Enum):
@@ -27,11 +26,11 @@ class Mode(Enum):
     auriga6 = 2
 
 
-mode = Mode.auriga6
+mode = Mode.richings
 
 
 def dir_name_to_parameter(dir_name: str):
-    return map(int, dir_name.lstrip("auriga6_halo").lstrip("richings21_").split("_"))
+    return map(int, dir_name.lstrip("auriga6_halo").lstrip("richings21_").lstrip("bary_").split("_"))
 
 
 def levelmax_to_softening_length(levelmax: int) -> float:
@@ -77,7 +76,7 @@ for dir in sorted(root_dir.glob("*")):
     if not is_by_adrian:
         levelmin, levelmin_TF, levelmax = dir_name_to_parameter(dir.name)
         print(levelmin, levelmin_TF, levelmax)
-        # if levelmax != 12:
+        # if levelmax != 9:
         #     continue
 
     input_file = dir / "output_0007.hdf5"
@@ -89,14 +88,20 @@ for dir in sorted(root_dir.glob("*")):
     else:
         try:
             swift_conf = read_swift_config(dir)
+            print_wall_time(dir)
         except FileNotFoundError:
             continue
-        softening_length = swift_conf["Gravity"]["comoving_DM_softening"]
+        gravity_conf = swift_conf["Gravity"]
-        assert softening_length == swift_conf["Gravity"]["max_physical_DM_softening"]
+        softening_length = gravity_conf["comoving_DM_softening"]
+        assert softening_length == gravity_conf["max_physical_DM_softening"]
+        if "max_physical_baryon_softening" in gravity_conf:
+            assert softening_length == gravity_conf["max_physical_baryon_softening"]
+            assert softening_length == gravity_conf["comoving_baryon_softening"]
+
         ideal_softening_length = levelmax_to_softening_length(levelmax)
-        # if not np.isclose(softening_length, levelmax_to_softening_length(levelmax)):
+        if not np.isclose(softening_length, levelmax_to_softening_length(levelmax)):
-        #     raise ValueError(f"softening length for levelmax {levelmax} should be {ideal_softening_length} "
+            raise ValueError(f"softening length for levelmax {levelmax} should be {ideal_softening_length} "
-        #                      f"but is {softening_length}")
+                             f"but is {softening_length}")
     print(input_file)
     if mode == Mode.richings and is_by_adrian:
         h = 0.6777
@@ -108,7 +113,7 @@ for dir in sorted(root_dir.glob("*")):
     else:
         df, particles_meta = read_file(input_file)
         df_halos = read_halo_file(input_file.with_name("fof_" + input_file.name))
-        vr_halo = read_velo_halos(dir, veloname="velo_out").loc[1]
+        # vr_halo = read_velo_halos(dir, veloname="velo_out").loc[1]
         # particles_in_halo = df.loc[df["FOFGroupIDs"] == 3]
 
         halo_id = 1
@@ -162,8 +167,8 @@ for dir in sorted(root_dir.glob("*")):
     if softening_length:
         for ax in [ax1, ax2]:
             ax.axvline(4 * softening_length, color=f"C{i}", linestyle="dotted")
-    for ax in [ax1, ax2]:
+    # for ax in [ax1, ax2]:
-        ax.axvline(vr_halo.Rvir, color=f"C{i}", linestyle="dashed")
+    #     ax.axvline(vr_halo.Rvir, color=f"C{i}", linestyle="dashed")
 
     X, Y, Z = df.X.to_numpy(), df.Y.to_numpy(), df.Z.to_numpy()
 

fix_hdf5_masses.py

@@ -0,0 +1,68 @@
+from math import fabs
+from sys import argv
+
+import h5py
+import numpy as np
+
+gamma = 5 / 3
+YHe = 0.245421
+Tcmb0 = 2.7255
+
+
+def calculate_gas_internal_energy(omegab, hubble_param_, zstart_):
+    astart_ = 1.0 / (1.0 + zstart_)
+    if fabs(1.0 - gamma) > 1e-7:
+        npol = 1.0 / (gamma - 1.)
+    else:
+        npol = 1.0
+    unitv = 1e5
+    adec = 1.0 / (160. * (omegab * hubble_param_ * hubble_param_ / 0.022) ** (2.0 / 5.0))
+    if (astart_ < adec):
+        Tini = Tcmb0 / astart_
+    else:
+        Tini = Tcmb0 / astart_ / astart_ * adec
+    print("Tini", Tini)
+    if Tini > 1.e4:
+        mu = 4.0 / (8. - 5. * YHe)
+    else:
+        mu = 4.0 / (1. + 3. * (1. - YHe))
+    print("mu", mu)
+    ceint_ = 1.3806e-16 / 1.6726e-24 * Tini * npol / mu / unitv / unitv
+    print("ceint", ceint_)
+    return ceint_
+
+
+def calculate_smoothing_length(boxsize, hubble_param_, levelmax):
+    mean_interparte_separation = boxsize / 2 ** levelmax
+    print("smoothing length", mean_interparte_separation)
+    return mean_interparte_separation
+
+
+with h5py.File(argv[1], "r+") as f:
+    omegab = f["Cosmology"].attrs["Omega_b"]
+    h = f["Cosmology"].attrs["h"]
+    zstart = f["Header"].attrs["Redshift"]
+    boxsize = f["Header"].attrs["BoxSize"]
+    levelmax = f["Header"].attrs["Music_levelmax"]
+    internal_energy = calculate_gas_internal_energy(omegab=omegab, hubble_param_=h, zstart_=zstart)
+    smoothing_length = calculate_smoothing_length(boxsize=boxsize, hubble_param_=h, levelmax=levelmax)
+    # exit()
+    bary_mass = f["Header"].attrs["MassTable"][0]
+    bary_count = f["Header"].attrs["NumPart_Total"][0]
+    print("mass table", f["Header"].attrs["MassTable"])
+    pt1 = f["PartType0"]
+    masses_column = pt1.create_dataset(
+        "Masses",
+        data=np.full(bary_count, bary_mass),
+        compression='gzip'
+    )
+    smoothing_length_column = pt1.create_dataset(
+        "SmoothingLength",
+        data=np.full(bary_count, smoothing_length),
+        compression='gzip'
+    )
+    internal_energy_column = pt1.create_dataset(
+        "InternalEnergy",
+        data=np.full(bary_count, internal_energy),
+        compression='gzip'
+    )

halo_plot.py

@@ -8,27 +8,26 @@ from matplotlib.colors import LogNorm
 from matplotlib.figure import Figure
 from matplotlib.patches import Circle
 
-from cic import Extent
+from halo_vis import Coords
 from paths import base_dir, vis_datafile
 from read_vr_files import read_velo_halos
 
 
-def increase_extent_1d(xmin: float, xmax: float, factor: float):
+def coord_to_2d_extent(coords: Coords):
-    xrange = xmax - xmin
+    radius, X, Y, Z = coords
-    xcenter = (xmax + xmin) / 2
+    return X - radius, X + radius, Y - radius, Y + radius
-    return xcenter - xrange / 2 * factor, xcenter + xrange / 2 * factor
 
 
-def increase_extent(extent: Extent, factor: float) -> Extent:
+def in_area(coords: Coords, xobj, yobj, zobj, factor=1.3) -> bool:
-    xmin, xmax, ymin, ymax = extent
+    radius, xcenter, ycenter, zcenter = coords
-    xmin, xmax = increase_extent_1d(xmin, xmax, factor)
+    radius *= factor
-    ymin, ymax = increase_extent_1d(ymin, ymax, factor)
+    return (
-    return xmin, xmax, ymin, ymax
+            (xcenter - radius < xobj < xcenter + radius)
-
+            and
-
+            (ycenter - radius < yobj < ycenter + radius)
-def in_extent(extent: Extent, X, Y, factor=2) -> bool:
+            and
-    xmin, xmax, ymin, ymax = increase_extent(extent, factor)
+            (zcenter - radius < zobj < zcenter + radius)
-    return (xmin < X < xmax) and (ymin < Y < ymax)
+    )
 
 
 def main():
@@ -36,7 +35,7 @@ def main():
     offset = 2
     columns = [128, 256, 512]
     fig: Figure = plt.figure(figsize=(9, 9))
-    axes: List[List[Axes]] = fig.subplots(len(rows), len(columns), sharex=True, sharey=True)
+    axes: List[List[Axes]] = fig.subplots(len(rows), len(columns), sharex="row", sharey="row")
     with h5py.File(vis_datafile) as vis_out:
         vmin, vmax = vis_out["vmin_vmax"]
         print(vmin, vmax)
@@ -46,21 +45,23 @@ def main():
                 halos = read_velo_halos(dir)
                 ax = axes[i][j]
                 rho = np.asarray(vis_out[f"{waveform}_{resolution}_rho"])
-                extent = tuple(vis_out[f"{waveform}_{resolution}_extent"])
+                # radius, X, Y, Z
+                coords: Coords = tuple(vis_out[f"{waveform}_{resolution}_coords"])
                 mass = vis_out[f"{waveform}_{resolution}_mass"][()]  # get scalar value from Dataset
                 main_halo_id = vis_out[f"{waveform}_{resolution}_halo_id"][()]
                 vmin_scaled = (vmin + offset) * mass
                 vmax_scaled = (vmax + offset) * mass
                 rho = (rho + offset) * mass
-
+                extent = coord_to_2d_extent(coords)
                 img = ax.imshow(rho.T, norm=LogNorm(vmin=vmin_scaled, vmax=vmax_scaled), extent=extent,
                                 origin="lower")
+                found_main_halo = False
                 for halo_id, halo in halos.iterrows():
-                    if halo["Vmax"] > 135:
+                    if halo["Vmax"] > 100:
-                        if in_extent(extent, halo.X, halo.Y):
+                        if in_area(coords, halo.X, halo.Y, halo.Z):
                             color = "red" if halo_id == main_halo_id else "white"
                             if halo_id == main_halo_id:
-                                print(halo_id == main_halo_id, halo_id, main_halo_id, halo["Rvir"])
+                                found_main_halo = True
                                 print("plotting main halo")
                             circle = Circle(
                                 (halo.X, halo.Y),
@@ -68,7 +69,7 @@ def main():
                                 linewidth=1, edgecolor=color, fill=None, alpha=.2
                             )
                             ax.add_artist(circle)
-
+                assert found_main_halo
                 print(img)
             #     break
             # break

halo_vis.py

@@ -10,7 +10,7 @@ from paths import base_dir, vis_datafile
 from read_vr_files import read_velo_halo_particles
 from readfiles import read_file
 
-all_in_area = True
+all_in_area = False
 
 Coords = Tuple[float, float, float, float]  # radius, X, Y, Z
 
@@ -21,7 +21,6 @@ def load_halo_data(waveform: str, resolution: int, halo_id: int, coords: Coords)
     df_halo, halo_lookup, unbound = read_velo_halo_particles(dir, recursivly=False)
 
     halo = df_halo.loc[halo_id]
-    if all_in_area:
     if coords:
         radius, X, Y, Z = coords
     else:
@@ -30,6 +29,7 @@ def load_halo_data(waveform: str, resolution: int, halo_id: int, coords: Coords)
         Y = halo["Yc"]
         Z = halo["Zc"]
         coords: Coords = radius, X, Y, Z
+    if all_in_area:
         df = df[df["X"].between(X - radius, X + radius)]
         df = df[df["Y"].between(Y - radius, Y + radius)]
         halo_particles = df[df["Z"].between(Z - radius, Z + radius)]
@@ -67,18 +67,21 @@ def imsave(rho, file_name: str):
 
 
 def main():
-    initial_halo_id = 2
+    waveforms = ["shannon", "DB2", "DB4", "DB8"]
+    initial_halo_id = 1
     first_halo = True
     rhos = {}
     ref_waveform = "shannon"
     ref_resolution = 128
-    coords = None
+    coords = {}
+    for wf in waveforms:
+        coords[wf] = None
     vmin = np.Inf
     vmax = -np.Inf
     if vis_datafile.exists():
         input("confirm to overwrite file")
     with h5py.File(vis_datafile, "w") as vis_out:
-        for waveform in ["shannon", "DB2", "DB4", "DB8"]:
+        for waveform in waveforms:
             for resolution in [128, 256, 512]:
                 if first_halo:
                     assert ref_resolution == resolution
@@ -87,22 +90,22 @@ def main():
                     first_halo = False
                 else:
                     halo_id = map_halo_id(initial_halo_id, ref_waveform, ref_resolution, waveform, resolution)
-                halo, halo_particles, meta, image_coords = load_halo_data(waveform, resolution, halo_id, coords)
+                halo, halo_particles, meta, image_coords = load_halo_data(waveform, resolution, halo_id, coords[waveform])
-                if not coords:
+                if not coords[waveform]:
-                    coords = image_coords
+                    coords[waveform] = image_coords
-                print(coords)
+                print(coords[waveform])
                 print("mass", halo["Mvir"])
                 # print("sleep")
                 # sleep(100)
-                radius, X, Y, Z = coords
+                radius, X, Y, Z = coords[waveform]
-                rho, extent = cic_from_radius(
+                rho, _ = cic_from_radius(
                     halo_particles.X.to_numpy(), halo_particles.Y.to_numpy(),
                     1000, X, Y, radius, periodic=False)
                 rhos[(waveform, resolution)] = rho
                 vmin = min(rho.min(), vmin)
                 vmax = max(rho.max(), vmax)
                 vis_out.create_dataset(f"{waveform}_{resolution}_rho", data=rho, compression='gzip', compression_opts=5)
-                vis_out.create_dataset(f"{waveform}_{resolution}_extent", data=extent)
+                vis_out.create_dataset(f"{waveform}_{resolution}_coords", data=coords[waveform])
                 vis_out.create_dataset(f"{waveform}_{resolution}_mass", data=meta.particle_mass)
                 vis_out.create_dataset(f"{waveform}_{resolution}_halo_id", data=halo_id)
                 imsave(rho, f"out_halo{initial_halo_id}_{waveform}_{resolution}_{halo_id}.png")

hdf5_sample.py

@@ -4,7 +4,7 @@ from sys import argv
 import numpy as np
 from h5py import File
 
-fraction = 0.05
+fraction = 0.01
 num_steps = 60
 
 file = Path(argv[1])

read_vr_files.py

@@ -88,8 +88,12 @@ def read_velo_halos(directory: Path,veloname="vroutput"):
     Returns a dataframe containing all scalar properties of the halos
     (https://velociraptor-stf.readthedocs.io/en/latest/output.html),
     """
-    group_catalog = h5py.File(directory / f"{veloname}.catalog_groups.0")
+    if (directory / f"{veloname}.catalog_groups.0").exists():
-    group_properties = h5py.File(directory / f"{veloname}.properties.0")
+        suffix = ".0"
+    else:
+        suffix = ""
+    group_catalog = h5py.File(directory / f"{veloname}.catalog_groups{suffix}")
+    group_properties = h5py.File(directory / f"{veloname}.properties{suffix}")
     scalar_properties = {}
     for k, v in group_properties.items():
         if not isinstance(v, Dataset):
@@ -124,12 +128,16 @@ def read_velo_halo_particles(
     and returns the halo data from read_velo_halos
     and two dictionaries mapping the halo IDs to sets of particle IDs
     """
+    if (directory / f"vroutput.catalog_particles.0").exists():
+        suffix = ".0"
+    else:
+        suffix = ""
     df = read_velo_halos(directory)
-    particle_catalog = h5py.File(directory / "vroutput.catalog_particles.0")
+    particle_catalog = h5py.File(directory / f"vroutput.catalog_particles{suffix}")
     particle_ids = np.asarray(particle_catalog["Particle_IDs"])
 
     particle_catalog_unbound = h5py.File(
-        directory / "vroutput.catalog_particles.unbound.0"
+        directory / f"vroutput.catalog_particles.unbound{suffix}"
     )
     particle_ids_unbound = particle_catalog_unbound["Particle_IDs"][:]
 

spectra_plot.py

@@ -3,6 +3,7 @@ from pathlib import Path
 from sys import argv
 
 import matplotlib.pyplot as plt
+import numpy as np
 import pandas as pd
 from matplotlib.axes import Axes
 from matplotlib.figure import Figure
@@ -83,6 +84,7 @@ def create_plot(mode):
         len(waveforms), 3, sharex=True, sharey=True,
         constrained_layout=True, figsize=(9, 9),
     )
+    crossings = np.zeros((len(waveforms), len(combination_list)))
     for i, waveform in enumerate(waveforms):
         ax_ics: Axes = axes[i][0]
         ax_z1: Axes = axes[i][1]
@@ -167,7 +169,10 @@ def create_plot(mode):
                 ics_data = spectra_data(waveform, res1, res2, Lbox, 'ics')
                 ics_k = ics_data["k [Mpc]"]
                 ics_pcross = ics_data["Pcross"]
-
+                smaller_res = min(res1, res2)
+                crossing_index = np.searchsorted(ics_k.to_list(), k0 * smaller_res)
+                crossing_value = ics_pcross[crossing_index]
+                crossings[i][j] = crossing_value
                 ax_ics.semilogx(ics_k, ics_pcross, color=colors[j + 3], label=f'{res1} vs {res2}')
 
                 z1_data = spectra_data(waveform, res1, res2, Lbox, 'z=1')
@@ -176,7 +181,6 @@ def create_plot(mode):
 
                 ax_z1.semilogx(z1_k, z1_pcross, color=colors[j + 3], label=f'{res1} vs {res2}')
 
-
                 end_data = spectra_data(waveform, res1, res2, Lbox, 'end')
                 end_k = end_data["k [Mpc]"]
                 end_pcross = end_data["Pcross"]
@@ -191,6 +195,10 @@ def create_plot(mode):
 
         # fig.suptitle(f"Cross Spectra {time}") #Not needed for paper
         # fig.tight_layout()
+    print(crossings)
+    crossings_df = pd.DataFrame(crossings, columns=combination_list, index=waveforms)
+    # print(crossings_df.to_markdown())
+    print(crossings_df.to_latex())
     fig.savefig(Path(f"~/tmp/spectra_{mode}.pdf").expanduser())
 
 

utils.py

@@ -29,3 +29,12 @@ def create_figure() -> Tuple[Figure, Axes]:
 def read_swift_config(dir: Path):
     with (dir / "used_parameters.yml").open() as f:
         return yaml.safe_load(f)
+
+
+def print_wall_time(dir: Path):
+    with(dir / "swift.log").open() as f:
+        last_line = f.readlines()[-1]
+    print(last_line)
+    assert "main: done. Bye." in last_line
+    seconds = float(last_line[1:].split("]")[0])
+    print(f"Runtime: {seconds / 60 / 60:.2f} hours")