import json import pickle from dataclasses import dataclass from enum import Enum from pathlib import Path from pprint import pprint from subprocess import run from sys import argv from typing import List, Tuple import h5py import numpy as np import pandas as pd import pynbody from matplotlib import pyplot as plt from matplotlib.axes import Axes from matplotlib.colors import LogNorm from matplotlib.figure import Figure from matplotlib.image import AxesImage from mpl_toolkits.axes_grid1.anchored_artists import AnchoredSizeBar from numpy import log10 from pynbody.array import SimArray from pynbody.snapshot import FamilySubSnap from pynbody.snapshot.ramses import RamsesSnap from scipy import constants from cache import HDFCache from cic import cic_from_radius, cic_range from find_center import find_center from halo_mass_profile import halo_mass_profile, property_profile from nfw import fit_nfw from paths import auriga_dir, richings_dir, auriga_dir_new, richings_dir_new from ramses import load_ramses_data, get_slice_argument, load_slice_data from read_vr_files import read_velo_halos from readfiles import read_file, read_halo_file, ParticlesMeta from slices import create_2d_slice, filter_3d from utils import read_swift_config, figsize_from_page_fraction class Mode(Enum): richings = "richings" auriga6 = "auriga" mode = Mode(argv[1]) cache = HDFCache(Path("auriga_cache.hdf5")) def dir_name_to_parameter(dir_name: str): return map( int, dir_name.lstrip("auriga6_halo") .lstrip("auriga6_music20_") .lstrip("richings21_") .lstrip("bary_") .lstrip("ramses_") .split("_"), ) def levelmax_to_softening_length(levelmax: int) -> float: box_size = 100 return box_size / 30 / 2 ** levelmax def main(): fig1: Figure = plt.figure(figsize=figsize_from_page_fraction()) ax1: Axes = fig1.gca() fig2: Figure = plt.figure(figsize=figsize_from_page_fraction()) ax2: Axes = fig2.gca() axs_baryon: List[List[Axes]] fig4: Figure fig4, axs_baryon = plt.subplots( nrows=2, ncols=4, sharex="all", sharey="all", figsize=figsize_from_page_fraction(columns=2, height_to_width=0.55), layout='constrained' ) fig5: Figure = plt.figure(figsize=figsize_from_page_fraction()) ax5: Axes = fig5.gca() fig6: Figure = plt.figure(figsize=figsize_from_page_fraction()) ax6: Axes = fig6.gca() baryon_plot_counter = 0 vminmax = {} extents = [] for ax in [ax1, ax2]: ax.set_xlabel(r"R [Mpc]") ax1.set_ylabel(r"M [$10^{10} \mathrm{M}_\odot$]") ax2.set_ylabel("density [$\\frac{10^{10} \\mathrm{M}_\\odot}{Mpc^3}$]") part_numbers = [] reference_file = Path(f"auriga_reference_{mode.value}.pickle") centers = {} halo_props = { "R_200crit": {}, "Mass_200crit": {}, "cNFW_200crit": {}, "FOF_Mass": {}, "FOF_Size": {}, } halo_props["Mass_200crit"]["Au6_paper"] = 104.39 halo_props["R_200crit"]["Au6_paper"] = 213.83 / 1000 @dataclass class Result: title: str rho: np.ndarray levels: Tuple[int, int, int] images = [] vmin = np.Inf vmax = -np.Inf root_dir = auriga_dir if mode == Mode.auriga6 else richings_dir if True: root_dir = auriga_dir_new if mode == Mode.auriga6 else richings_dir_new i = 0 mapping = {} for dir in sorted(root_dir.glob("*")): dir = dir.resolve() print("----------------------------") if not dir.is_dir() or "bak" in dir.name: continue is_ramses = "ramses" in dir.name has_baryons = "bary" in dir.name or is_ramses is_by_adrian = "arj" in dir.name is_resim = dir.name[0].isdigit() or is_by_adrian print(dir.name) if not is_by_adrian: levelmin, levelmin_TF, levelmax = dir_name_to_parameter(dir.name) print(levelmin, levelmin_TF, levelmax) # if not is_resim: # continue # if levelmax < 12 and not is_by_adrian: # continue if mode == Mode.auriga6: print("sdfds") if (levelmin, levelmin_TF, levelmax) == (7, 9, 9): print("Aaaa") continue elif mode == Mode.richings: if not has_baryons: continue # if levelmax != 11: # continue # if not is_ramses: # continue if dir.name in ["bary_ramses_7_10_11"]: continue input_file = dir / "output_0007.hdf5" print(input_file) if mode == Mode.richings and not is_resim: input_file = dir / "output_0004.hdf5" if is_by_adrian or is_ramses: input_file = dir / "output_0000.hdf5" softening_length = None else: try: swift_conf = read_swift_config(dir) # print_wall_time(dir) except FileNotFoundError: print(f"file not found: {dir}") continue gravity_conf = swift_conf["Gravity"] 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, ideal_softening_length): if softening_length < ideal_softening_length: print( f"softening length smaller than calculated: {softening_length}<{ideal_softening_length} ({levelmax=})") else: raise ValueError( f"softening length for levelmax {levelmax} should be {ideal_softening_length} " f"but is {softening_length}" ) print(input_file) if mode == Mode.richings and is_by_adrian: h = 0.6777 with h5py.File(dir / "Richings_object_z0.h5") as f: df = pd.DataFrame(f["Coordinates"][:] / h, columns=["X", "Y", "Z"]) particles_meta = ParticlesMeta(particle_mass=1.1503e7 / 1e10) center = np.array([60.7, 29, 64]) / h softening_length = None elif "ramses" in dir.name: h = 0.6777 hr_coordinates, particles_meta, center = load_ramses_data(dir / "output_00009") df = pd.DataFrame(hr_coordinates, columns=["X", "Y", "Z"]) softening_length = None else: df, particles_meta = read_file(input_file) df_halos = read_halo_file(input_file.with_name("fof_" + input_file.name)) halo_id = 1 while True: particles_in_halo = df.loc[df["FOFGroupIDs"] == halo_id] if len(particles_in_halo) > 1: break halo_id += 1 halo = df_halos.loc[halo_id] part_numbers.append(len(df) * particles_meta.particle_mass) # halo = halos.loc[1] center = np.array([halo.X, halo.Y, halo.Z]) halo_props["FOF_Mass"][dir.name] = halo["Masses"] halo_props["FOF_Size"][dir.name] = halo["Sizes"] if mode == Mode.auriga6: vr_halo = read_velo_halos(dir, veloname="velo_out").loc[1] print(vr_halo["R_200crit"]) halo_props["R_200crit"][dir.name] = vr_halo["R_200crit"] halo_props["Mass_200crit"][dir.name] = vr_halo["Mass_200crit"] halo_props["cNFW_200crit"][dir.name] = vr_halo["cNFW_200crit"] center2 = np.array([vr_halo.X, vr_halo.Y, vr_halo.Z]) print("center-diff", center2 - center) print("center-diff", center) print("center-diff", center2) center = find_center(df, center) log_radial_bins, bin_masses, bin_densities, center = halo_mass_profile( df[["X", "Y", "Z"]].to_numpy(), center, particles_meta, plot=False, num_bins=100, rmin=0.002, rmax=6.5 ) i_min_border = np.argmax( 0.01 < log_radial_bins ) # first bin outside of specific radius i_max_border = np.argmax(1.5 < log_radial_bins) popt = fit_nfw( log_radial_bins[i_min_border:i_max_border], bin_densities[i_min_border:i_max_border], ) # = rho_0, r_s print(popt) # # Plot NFW profile # ax.loglog( # log_radial_bins[i_min_border:i_max_border], # nfw(log_radial_bins[i_min_border:i_max_border], *popt), # linestyle="dotted" # ) centers[dir.name] = center if is_by_adrian: with reference_file.open("wb") as f: pickle.dump([log_radial_bins, bin_masses, bin_densities], f) if is_by_adrian: label = "Reference" else: label = f"({levelmin}, {levelmin_TF}, {levelmax})" ax1.loglog(log_radial_bins[:-1], bin_masses, label=label, c=f"C{i}") ax2.loglog(log_radial_bins[:-1], bin_densities, label=label, c=f"C{i}") mapping[i] = dir.name if reference_file.exists() and not is_by_adrian: with reference_file.open("rb") as f: data: List[np.ndarray] = pickle.load(f) ref_log_radial_bins, ref_bin_masses, ref_bin_densities = data mass_deviation: np.ndarray = np.abs(bin_masses - ref_bin_masses) density_deviation: np.ndarray = np.abs(bin_densities - ref_bin_densities) ax1.loglog(log_radial_bins[:-1], mass_deviation, c=f"C{i}", linestyle="dotted") ax2.loglog( log_radial_bins[:-1], density_deviation, c=f"C{i}", linestyle="dotted" ) accuracy = mass_deviation / ref_bin_masses print(accuracy) print("mean accuracy", accuracy.mean()) if softening_length: for ax in [ax1, ax2]: ax.axvline(4 * softening_length, color=f"C{i}", linestyle="dotted") # for ax in [ax1, ax2]: # ax.axvline(vr_halo.Rvir, color=f"C{i}", linestyle="dashed") coords = df[["X", "Y", "Z"]].to_numpy() # shift: (-6, 0, -12) # if not is_by_adrian: # xshift = Xc - Xc_adrian # yshift = Yc - Yc_adrian # zshift = Zc - Zc_adrian # print("shift", xshift, yshift, zshift) coords_centered = coords - center rho, extent = cic_from_radius(coords_centered[::, 0], coords_centered[::, 2], 500, 0, 0, 1.5, periodic=False) vmin = min(vmin, rho.min()) vmax = max(vmax, rho.max()) res = Result( rho=rho, title=f"levelmin={levelmin}, levelmin_TF={levelmin_TF}, levelmax={levelmax}" if not is_by_adrian else "Reference", levels=(levelmin, levelmin_TF, levelmax) if not is_by_adrian else (100, 100, 100), ) res.title = res.title + "\n" + dir.name images.append(res) i += 1 if has_baryons: interpolation_method = "nearest" # "linear" bary_file = dir / "output_00009" if is_ramses else input_file if is_ramses: s: RamsesSnap = pynbody.load(str(bary_file)) gas_data: FamilySubSnap = s.gas 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(log10(p_array / rho_array ** (5 / 3))), } else: with h5py.File(input_file) as f: pt0 = f["PartType0"] bary_coords = pt0["Coordinates"][:] mass_array = pt0["Masses"][:] bary_properties = { "InternalEnergies": pt0["InternalEnergies"][:], "Densities": pt0["Densities"][:], "Pressures": pt0["Pressures"][:], # "Entropies": log10(pt0["Densities"][:] / pt0["Densities"][:] ** (5 / 3)), "Entropies": pt0["Entropies"][:] } bary_properties["Temperatures"] = bary_properties["InternalEnergies"] radius = 1.9 resolution = 500 # xrange[0], xrange[-1], yrange[0], yrange[-1] extent = [center[0] - radius, center[0] + radius, center[1] - radius, center[1] + radius] extents.append(extent) # extent = [42, 62, 50, 70] ramses_done = False labels = { "cic": "$\\rho_{\\mathrm{DM}}$ [$10^{10}\\mathrm{M}_\\odot \\mathrm{Mpc}^{-3}$]", "Densities": "$\\rho_{\\mathrm{gas}}$ [$10^{10}\\mathrm{M}_\\odot \\mathrm{Mpc}^{-3}$]", "Entropies": "$s$", "Temperatures": "$T$ [K]", } for ii, property in enumerate(["cic", "Densities", "Entropies", "Temperatures"]): print("property:", property) key = f"grid_{resolution}_{property}_{interpolation_method}_{radius}" cached_grid = cache.get(key, str(bary_file)) if cached_grid is not None: grid = cached_grid else: print("grid not yet cached, calculating now") if property == "cic": coords_in_box = filter_3d(coords, extent, zlimit=(center[2] - .1, center[2] + .1)) rho, _ = cic_range(coords_in_box[::, 0], coords_in_box[::, 1], resolution, *extent, periodic=False) grid = 1.1 + rho.T else: if not is_ramses: grid = create_2d_slice(center, coords=bary_coords, resolution=resolution, property_name=property, property_data=bary_properties[property], extent=extent, method=interpolation_method) else: frac_center = center / 100 frac_extent = np.array(extent) / 100 print(frac_extent) print(frac_center) args, imager_dir = get_slice_argument( frac_extent, frac_center, bary_file, interpolation_method, depth=.001 ) print(" ".join(args)) if not ramses_done: run(args, cwd=imager_dir) ramses_done = True property_map = { "Densities": "rhomap", "Entropies": "Smap", "Temperatures": "Tmap" } fname = imager_dir / f"snapshot_{property_map[property]}_zproj_zobs-0p00.bin" grid = load_slice_data(fname).T cache.set(key, grid, str(bary_file), compressed=True) if property == "Densities" and is_ramses: # 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) ax_baryon = axs_baryon[baryon_plot_counter][ii] if property in vminmax: minmax = vminmax[property] else: minmax = np.min(grid), np.max(grid) vminmax[property] = minmax print("minmax", minmax) if is_ramses: np.save(f"/home/lukas/tmp/ramses_grid_{property}", grid) img: AxesImage = ax_baryon.imshow( grid, norm=LogNorm(vmin=minmax[0], vmax=minmax[1]), interpolation="none", origin="lower", extent=extent, ) if baryon_plot_counter == 0: fig4.colorbar(img, ax=ax_baryon, location="top", label=labels[property]) # ax_baryon.set_title(property) # ax_baryon.set_xlabel("X") # ax_baryon.set_ylabel("Y") ax_baryon.set_aspect("equal") # exit() baryon_plot_counter += 1 continue r, prof = property_profile(bary_coords, center, mass_array, bary_properties, num_bins=100, rmin=0.002, rmax=6.5) integrator_name = "Ramses" if is_ramses else "Swift" label = f"{integrator_name} {levelmin}, {levelmin_TF}, {levelmax}" ax5.set_title("Densities") ax6.set_title("Pressures") ax5.loglog(r[1:], prof["Densities"], label=label) ax6.loglog(r[1:], prof["Pressures"], label=label) if mode==Mode.richings: extents = np.asarray(extents) print(extents) global_xlim = extents[:, 0].max(), extents[:, 1].min() global_ylim = extents[:, 2].max(), extents[:, 3].min() for ax in axs_baryon.flatten(): ax.set_xlim(global_xlim) ax.set_ylim(global_ylim) ax.set_xticks([]) ax.set_yticks([]) scalebar = AnchoredSizeBar( ax.transData, 1, "1 Mpc", "lower left", # pad=0.1, color="white", frameon=False, # size_vertical=1 ) ax.add_artist(scalebar) fig3: Figure = plt.figure( # just a bit more than 2/3 so that the two rows don't overlap figsize=figsize_from_page_fraction(columns=2, height_to_width=33 / 48) ) axes: List[Axes] = fig3.subplots(2, 3, sharex="all", sharey="all").flatten() # images.sort(key=lambda r: r.levels, reverse=True) for result, ax in zip(images, axes): data = 1.1 + result.rho vmin_scaled = 1.1 + vmin vmax_scaled = 1.1 + vmax img = ax.imshow( data.T, norm=LogNorm(vmin=vmin_scaled, vmax=vmax_scaled), extent=extent, origin="lower", cmap="Greys", interpolation="none" ) ax.text( 0.5, 0.95, result.title, horizontalalignment="center", verticalalignment="top", transform=ax.transAxes, ) for ax in [ax1, ax2, ax5, ax6]: ax.legend() for fig in [fig1, fig2, fig3, fig5, fig6]: fig.tight_layout() fig.subplots_adjust(wspace=0, hspace=0) axs_baryon[0][0].set_ylabel("Swift") axs_baryon[1][0].set_ylabel("Ramses") fig1.savefig(Path(f"~/tmp/{mode.value}1.pdf").expanduser()) fig2.savefig(Path(f"~/tmp/{mode.value}2.pdf").expanduser()) fig3.savefig(Path(f"~/tmp/{mode.value}3.pdf").expanduser()) fig4.savefig(Path(f"~/tmp/{mode.value}4.pdf").expanduser()) pprint(centers) if mode == Mode.auriga6: with open("halo_props.json","w") as f: json.dump(halo_props,f, indent=4) plt.show() print(part_numbers) print(mapping) if __name__ == '__main__': main()