better comparisons of gas simulations

auriga_comparison.py

@@ -13,7 +13,7 @@ from matplotlib.axes import Axes
 from matplotlib.colors import LogNorm
 from matplotlib.figure import Figure
 
-from cic import cic_from_radius
+from cic import cic_from_radius, cic_range
 from halo_mass_profile import halo_mass_profile
 from nfw import fit_nfw
 from paths import auriga_dir, richings_dir
@@ -168,9 +168,6 @@ for dir in sorted(root_dir.glob("*")):
     #     linestyle="dotted"
     # )
 
-    if has_baryons:
-        create_2d_slice(input_file, center, property="InternalEnergies")
-
     centers[dir.name] = center
     if is_by_adrian:
         with reference_file.open("wb") as f:
@@ -230,6 +227,35 @@ for dir in sorted(root_dir.glob("*")):
         )
     )
     i += 1
+
+    if has_baryons:
+        fig3, axs_baryon = plt.subplots(nrows=1, ncols=5, sharex="all", sharey="all", figsize=(10, 4))
+        extent = [46, 52, 54, 60]  # xrange[0], xrange[-1], yrange[0], yrange[-1]
+        for ii, property in enumerate(["cic", "Densities", "Entropies", "InternalEnergies", "Temperatures"]):
+            print(property)
+            if property == "cic":
+                grid, _ = cic_range(X + center[0], Y + center[1], 1000, *extent, periodic=False)
+            else:
+                grid = create_2d_slice(input_file, center, property=property, extent=extent)
+            print("minmax", grid.min(), grid.max())
+            assert grid.min() != grid.max()
+            ax_baryon: Axes = axs_baryon[ii]
+            img = ax_baryon.imshow(
+                grid,
+                norm=LogNorm(),
+                interpolation="none",
+                origin="lower",
+                extent=extent,
+            )
+            ax_baryon.set_title(property)
+            # ax_baryon.set_xlabel("X")
+            # ax_baryon.set_ylabel("Y")
+            ax_baryon.set_aspect("equal")
+        fig3.suptitle(input_file.parent.stem)
+        fig3.tight_layout()
+        fig3.savefig(Path("~/tmp/slice.png").expanduser(), dpi=300)
+        plt.show()
+
     # plot_cic(
     #     rho, extent,
     #     title=str(dir.name)

slices.py

@@ -1,5 +1,5 @@
 from pathlib import Path
-from typing import List
+from typing import List, Tuple
 
 import h5py
 import matplotlib.pyplot as plt
@@ -11,56 +11,46 @@ from temperatures import calculate_T
 from utils import create_figure
 
 
+def filter_3d(
+        coords: np.ndarray, data: np.ndarray,
+        extent: List[float]
+) -> Tuple[np.ndarray, np.ndarray]:
+    filter = (
+            (extent[0] < coords[::, 0]) &
+            (coords[::, 0] < extent[1]) &
+
+            (extent[2] < coords[::, 1]) &
+            (coords[::, 1] < extent[3])
+    )
+    print("before", coords.shape)
+    data = data[filter]
+    coords = coords[filter]
+
+    print("after", coords.shape)
+    return coords, data
+
+
 def create_2d_slice(
-        input_file: Path, center: List[float], property: str, axis="Z", thickness=3, method="nearest"
+        input_file: Path, center: List[float], extent,
-):
+        property="InternalEnergies", method="nearest"
-    axis_names = ["X", "Y", "Z"]
+) -> np.ndarray:
-    cut_axis = axis_names.index(axis)
+    cut_axis = 2  # Z
-    limits = {
-        "X": (46, 52),
-        "Y": (54, 60),
-        "Z": (center[cut_axis] - 10, center[cut_axis] + 10)
-    }
     with h5py.File(input_file) as f:
         pt0 = f["PartType0"]
         coords = pt0["Coordinates"][:]
-        energies = pt0["InternalEnergies"][:]
+        data = pt0[property if property != "Temperatures" else "InternalEnergies"][:]
-        entropies = pt0["Entropies"][:]
 
-        print((center[cut_axis] - thickness < coords[::, cut_axis]).shape)
+        coords, data = filter_3d(coords, data, extent)
-        # in_slice = (center[cut_axis] - thickness < coords[::, cut_axis]) & (
+        if property == "Temperatures":
-        #         coords[::, cut_axis] < center[cut_axis] + thickness)
+            print("calculating temperatures")
-        # print("got slice")
+            data = np.array([calculate_T(u) for u in data])
-        # coords_in_slice = coords[in_slice]
-        # data_in_slice = data[in_slice]
-        filter = (
-                (limits["X"][0] < coords[::, 0]) &
-                (coords[::, 0] < limits["X"][1]) &
 
-                (limits["Y"][0] < coords[::, 1]) &
+        xrange = np.linspace(extent[0],extent[1], 1000)
-                (coords[::, 1] < limits["Y"][1]) &
+        yrange = np.linspace(extent[2],extent[3], 1000)
-
-                (limits["Z"][0] < coords[::, 2]) &
-                (coords[::, 2] < limits["Z"][1])
-        )
-        print("before", coords.shape)
-        energies = energies[filter]
-        entropies = entropies[filter]
-        coords = coords[filter]
-
-        print("after", coords.shape)
-        print("calculating temperatures")
-        temperatures = np.array([calculate_T(u) for u in energies])
-
-        other_axis = {"X": ("Y", "Z"), "Y": ("X", "Z"), "Z": ("X", "Y")}
-        x_axis_label, y_axis_label = other_axis[axis]
-        x_axis = axis_names.index(x_axis_label)
-        y_axis = axis_names.index(y_axis_label)
-        xrange = np.linspace(coords[::, x_axis].min(), coords[::, x_axis].max(), 1000)
-        yrange = np.linspace(coords[::, y_axis].min(), coords[::, y_axis].max(), 1000)
         gx, gy, gz = np.meshgrid(xrange, yrange, center[cut_axis])
         print("interpolating")
-        grid = griddata(coords, temperatures, (gx, gy, gz), method=method)[::, ::, 0]
+        grid = griddata(coords, data, (gx, gy, gz), method=method)[::, ::, 0]
+        return grid
         print(grid.shape)
         # stats, x_edge, y_edge, _ = binned_statistic_2d(
         #     coords_in_slice[::, x_axis],
@@ -85,6 +75,4 @@ def create_2d_slice(
         ax.set_aspect("equal")
         fig.colorbar(img, label="Temperatures")
         fig.tight_layout()
-        fig.savefig(Path("~/tmp/slice.png").expanduser(), dpi=300)
         plt.show()
-        exit()

temperatures.py

@@ -16,9 +16,6 @@ const_boltzmann_k_cgs = 1.380649e-16
 
 const_proton_mass = const_proton_mass_cgs / UnitMass_in_cgs
 const_boltzmann_k = const_boltzmann_k_cgs / UnitMass_in_cgs / UnitLength_in_cgs ** 2 * (UnitTime_in_cgs ** 2)
-print(const_proton_mass)
-print(const_boltzmann_k)
-print()
 
 
 @njit
@@ -33,7 +30,11 @@ def calculate_T(u):
     else:
         return T_transition
 
+
 if __name__ == "__main__":
+    print(const_proton_mass)
+    print(const_boltzmann_k)
+    print()
     internal_energies = [6.3726251e+02, 7.7903375e+02, 1.7425287e+04, 6.4113910e+04, 3.8831848e+04,
                          1.1073163e+03, 7.7394878e+03, 7.5230023e+04, 9.1036992e+04, 2.4060946e+00]