GEAR: add zoomin scripts

examples/GEAR/ZoomIn/add_fields.py

+#!/usr/bin/env python3
+
+import numpy as np
+import yt.utilities.physical_constants as constants
+from yt.mods import YTArray
+
+
+def addMassDeposit(f):
+    f.add_deposited_particle_field(("all", "Masses"), "sum")
+
+
+def addTemperature(f):
+    """
+    Add the temperature field.
+
+    Parameters
+    ----------
+
+    f: yt dataset
+        The dataset that needs the temperature
+    """
+
+    def calc_mu_table_local(temperature):
+        """
+        From grackle/src/python/utilities/convenience.py: Calculate a tabulated
+        approximation to mean molecular weight (valid for data that used
+        Grackle 2.0 or below)
+        """
+        tt = np.array([1.0e+01, 1.0e+02, 1.0e+03, 1.0e+04, 1.3e+04, 2.1e+04,
+                       3.4e+04, 6.3e+04, 1.0e+05, 1.0e+09])
+        mt = np.array([1.18701555, 1.15484424, 1.09603514, 0.9981496,
+                       0.96346395, 0.65175895, 0.6142901, 0.6056833, 0.5897776,
+                       0.58822635])
+        logttt = np.log(temperature)
+        # linear interpolation in log-log space
+        logmu = np.interp(logttt, np.log(tt), np.log(mt))
+        return np.exp(logmu)
+
+    temperature_values = []
+    mu_values = []
+    T_over_mu_values = []
+    current_temperature = 1e1
+    final_temperature = 1e7
+    dlogT = 0.1
+
+    while current_temperature < final_temperature:
+        temperature_values.append(current_temperature)
+        current_mu = calc_mu_table_local(current_temperature)
+        mu_values.append(current_mu)
+        T_over_mu_values.append(current_temperature/current_mu)
+        current_temperature = np.exp(np.log(current_temperature)+dlogT)
+
+    def convert_T_over_mu_to_T(T_over_mu):
+        logT_over_mu = np.log(T_over_mu)
+        # linear interpolation in log-log space
+        logT = np.interp(logT_over_mu, np.log(T_over_mu_values),
+                         np.log(temperature_values))
+        return np.exp(logT)
+
+    def _Temperature_3(field, data):
+        gamma = 5.0/3.0
+        T_over_mu = (data["PartType0", "InternalEnergy"] * (gamma-1)
+                     * constants.mass_hydrogen_cgs /
+                     constants.boltzmann_constant_cgs).in_units('K').d  # T/mu
+        return YTArray(convert_T_over_mu_to_T(T_over_mu), 'K')  # now T
+    f.add_field(("PartType0", "Temperature"), function=_Temperature_3,
+                particle_type=True, force_override=True, units="K")

examples/GEAR/ZoomIn/cleanupSwift.py

+#!/usr/bin/env python3
+
+# ./translate_particles.py filename output_name
+from h5py import File
+import sys
+from shutil import copyfile
+
+NPartType = 6
+filename = sys.argv[-2]
+out = sys.argv[-1]
+
+copyfile(filename, out)
+
+f = File(out)
+
+# read cosmological parameters
+a = f["Cosmology"].attrs["Scale-factor"][0]
+h = f["Cosmology"].attrs["h"][0]
+
+# add little h in the boxsize
+f["Header"].attrs["BoxSize"] = f["Header"].attrs["BoxSize"][0] * h
+
+# avoid the snapshot to be taken for SWIFT
+del f["Header"].attrs["Code"]
+
+# Delete problematic fields
+for i in range(NPartType):
+    name = "PartType{}/ElementAbundances".format(i)
+    if name in f:
+        del f[name]
+
+# Update the fields (name + cosmo)
+for i in range(NPartType):
+    name = "PartType%i" % i
+    if name not in f:
+        continue
+
+    if "SmoothingLengths" in f[name]:
+        grp = f[name + "/SmoothingLengths"]
+        grp[:] *= 1.823
+
+    # fix issue due to the name of densities
+    fields = [
+        ("Coordinates", "Coordinates", h),
+        ("Masses", "Masses", h),
+        ("Velocities", "Velocities", a**0.5),
+        ("Density", "Densities", 1. / h**2),
+        ("Entropies", "Entropies", 1.),
+        ("InternalEnergy", "InternalEnergies", 1.),
+        ("SmoothingLength", "SmoothingLengths", h)
+    ]
+
+    # create links
+    for field in fields:
+        if field[1] in f[name] and field[0] not in f[name]:
+            f[name + "/" + field[0]] = f[name + "/" + field[1]]
+
+    # apply unit transform
+    for field in fields:
+        field_name = name + "/" + field[0]
+        if field[0] in f[name]:
+            f[field_name][:] *= field[2]
+
+
+# update cosmological parameters in order to be consistent with GEAR
+cosmo = f["Cosmology"].attrs
+head = f["Header"].attrs
+head["Redshift"] = float(cosmo["Redshift"])
+head["OmegaLambda"] = cosmo["Omega_lambda"]
+head["Omega0"] = cosmo["Omega_b"]
+head["HubbleParam"] = cosmo["h"][0]
+head["Time"] = float(a)
+
+f.close()

examples/GEAR/ZoomIn/make_image.py

+#!/usr/bin/env python3
+
+import yt
+from yt.units import Msun, kpc, s, km
+import unyt
+import sys
+import matplotlib
+import projection_plot
+import phase_plot
+import add_fields
+import matplotlib.pyplot as plt
+from mpl_toolkits.axes_grid1 import AxesGrid
+from matplotlib.offsetbox import AnchoredText
+matplotlib.use('Agg')
+
+# Parameters
+
+snap = int(sys.argv[-1])
+
+swift = "swift/snapshot_%04i.hdf5" % snap
+gear = "gear/snapshot_%04i.hdf5" % snap
+
+
+do_plot = {
+    "projection_density": False,
+    "projection_temperature": False,
+    "projection_mass": True,
+    "phase": False
+}
+
+# Generate the figures
+
+figsize = (100, 100)
+figures = {
+    "projection_density": plt.figure(figsize=figsize),
+    "projection_temperature": plt.figure(figsize=figsize),
+    "projection_mass": plt.figure(figsize=figsize),
+    "phase": plt.figure(figsize=figsize)
+}
+
+# define global variables
+scale_factor = None
+
+# Generate the axes
+
+fig_size = (0.01, 0.01, 0.99, 0.99)
+subgrid = (1, 2)
+axes = {
+    "projection_density": AxesGrid(
+        figures["projection_density"], fig_size, subgrid,
+        add_all=True, share_all=True, cbar_mode="single",
+        cbar_size="2%", cbar_pad=0.02),
+    "projection_temperature": AxesGrid(
+        figures["projection_temperature"], fig_size, subgrid,
+        add_all=True, share_all=True, cbar_mode="single",
+        cbar_size="2%", cbar_pad=0.02),
+    "projection_mass": AxesGrid(
+        figures["projection_mass"], fig_size, subgrid,
+        add_all=True, share_all=True, cbar_mode="single",
+        cbar_size="2%", cbar_pad=0.02),
+    "phase": AxesGrid(
+        figures["phase"], fig_size, subgrid, axes_pad=0.05,
+        add_all=True, share_all=True, cbar_mode="single",
+        cbar_size="2%", cbar_pad=0.05, aspect=False)
+}
+
+
+def savePlot():
+    if do_plot["projection_density"]:
+        projection_plot.savePlot(figures["projection_density"],
+                                 "density")
+    if do_plot["projection_temperature"]:
+        projection_plot.savePlot(figures["projection_temperature"],
+                                 "temperature")
+
+    if do_plot["projection_mass"]:
+        projection_plot.savePlot(figures["projection_mass"],
+                                 "mass")
+
+    if do_plot["phase"]:
+        phase_plot.savePlot(figures["phase"])
+
+
+def doPlot(filename, i, name):
+    f = yt.load(filename)
+    if (do_plot["projection_temperature"] or do_plot["phase"]):
+        add_fields.addTemperature(f)
+
+    global scale_factor
+    if scale_factor is None:
+        scale_factor = f.scale_factor
+    else:
+        if abs(scale_factor - f.scale_factor) > scale_factor * 1e-2:
+            raise Exception("Error: the snapshots are not at the same time")
+
+    # Do density projection plot
+    if do_plot["projection_density"]:
+        projection_plot.doDensityPlot(
+            f, name, i, figures["projection_density"],
+            axes["projection_density"])
+
+    # Do temperature projection plot
+    if do_plot["projection_temperature"]:
+        projection_plot.doTemperaturePlot(
+            f, name, i, figures["projection_temperature"],
+            axes["projection_temperature"])
+
+    # Do mass projection plot
+    if do_plot["projection_mass"]:
+        add_fields.addMassDeposit(f)
+        projection_plot.doMassPlot(
+            f, name, i, figures["projection_mass"],
+            axes["projection_mass"])
+
+    # Phase plot
+    if do_plot["phase"]:
+        phase_plot.doPlot(f, name, i, figures["phase"],
+                          axes["phase"])
+
+
+doPlot(swift, 0, "SWIFT")
+doPlot(gear, 1, "GEAR")
+savePlot()

examples/GEAR/ZoomIn/phase_plot.py

+#!/usr/bin/env python3
+
+import yt
+from yt.units import Msun, amu, cm, K
+from matplotlib.offsetbox import AnchoredText
+import matplotlib.pyplot as plt
+
+
+limits_temperature = (1e1 * K, 1e7 * K)
+limits_density = (1e-7 * amu / cm**3, 1e7 * amu / cm**3)
+limits_mass = (1e3 * Msun, 1e7 * Msun)
+
+def savePlot(fig):
+    fig.savefig("phase.png", bbox_inches="tight",
+                pad_inches=0.03, dpi=300)
+
+
+def doPlot(f, name, i, fig, axes):
+    # limits in g / cm3
+    # u_rho = g / cm**3
+    # limits = [1e-33 * u_rho, 1e-24 * u_rho]
+
+    p = yt.PhasePlot(
+        f, ("PartType0", "density"), ("PartType0", "Temperature"),
+        ("PartType0", "mass"),
+        weight_field=None, x_bins=300, y_bins=300)
+
+    plt.title("Scale Factor %g" % f.scale_factor)
+
+    # Set the unit system
+    p.set_log("density", True)
+    p.set_unit("density", "amu / cm**3")
+    p.set_xlim(limits_density[0],
+               limits_density[1])
+
+    p.set_log("Temperature", True)
+    p.set_unit("Temperature", "K")
+    p.set_ylim(limits_temperature[0],
+               limits_temperature[1])
+
+    p.set_log("mass", True)
+    p.set_unit("mass", "Msun")
+    p.set_zlim("mass", limits_mass[0],
+               limits_mass[1])
+
+    plot = p.plots[("PartType0", "mass")]
+
+    plot.figure = fig
+    plot.axes = axes[i].axes
+
+    # plot color bar
+    if i != 0:
+        plot.cax = axes.cbar_axes[0]
+        # p.hide_axes()
+    p._setup_plots()
+
+
+    at = AnchoredText(name, loc=2, prop=dict(size=6), frameon=True)
+    plot.axes.add_artist(at)

examples/GEAR/ZoomIn/projection_plot.py

+#!/usr/bin/env python3
+
+import yt
+from yt.units import kpc, g, cm, K
+from matplotlib.offsetbox import AnchoredText
+
+cmap_density = "algae"
+cmap_temperature = "magma"
+cmap_mass = "inferno"
+center = (1441.9954833984375000,
+          1666.1169433593750000,
+          1891.6248779296875000)
+small_width = 200 * kpc
+large_width = 2500 * kpc
+limits_density = (1e-10 * g / cm**2, 1e-2 * g / cm**2)
+limits_temperature = (10 * K, 1e5 * K)
+limits_mass = None
+
+
+def savePlot(fig, field):
+    """
+    Save the plot into a file.
+
+    Parameters
+    ----------
+
+    fig: matplotlib figure
+        The figure to save
+
+    field: str
+        The name of the field in the plot
+    """
+    fig.savefig("projection_%s.png" % field, bbox_inches="tight",
+                pad_inches=0.03, dpi=300)
+
+
+def doDensityPlot(f, name, i, fig, axes):
+    """
+    Generate a density projection plot.
+
+    Parameters
+    ----------
+
+    f: yt dataset
+        The data set to use in the projection
+
+    name: str
+        The name to print on the plot
+
+    i: int
+        The index of the plot
+
+    fig: matplotlib figure
+        The figure to use
+
+    axes: list
+        The list of axes to use for the plot
+    """
+    direction = "x"
+    field = ("PartType0", "density")
+
+    # compute the projection
+    p = yt.ProjectionPlot(f, direction, field, center=center,
+                          width=small_width, buff_size=(800, 800))
+
+    # Compute the limits
+    p.set_unit("density", "g / cm**2")
+    data = p.to_fits_data()["density"].data
+    if limits_density[0] > data.min():
+        print("WARNING: data below min", data.min())
+    if limits_density[1] < data.max():
+        print("WARNING: data above max", data.max())
+
+    # Adjust the plot
+    p.set_cmap(field=field, cmap=cmap_density)
+    p.set_log(field, True)
+    p.set_zlim(field, limits_density[0], limits_density[1])
+
+    # Draw it into the correct figure
+    plot = p.plots[field]
+
+    plot.figure = fig
+    plot.axes = axes[i].axes
+
+    # Add code name
+    at = AnchoredText(name, loc=2, prop=dict(size=6), frameon=True)
+    plot.axes.add_artist(at)
+
+    # plot color bar
+    if i != 0:
+        plot.cax = axes.cbar_axes[0]
+        p.hide_axes()
+    p._setup_plots()
+
+    if i == 0:
+        z = 1. / f.scale_factor - 1.
+        text = "Redshift = %.2g" % z
+        prop = {
+            "color": "w"
+        }
+        at = AnchoredText(text, loc=3, prop=prop, frameon=False)
+        plot.axes.add_artist(at)
+
+    # Add code name
+    at = AnchoredText(name, loc=2, prop=dict(size=6), frameon=True)
+    plot.axes.add_artist(at)
+
+
+def doTemperaturePlot(f, name, i, fig, axes):
+    """
+    Generate a temperature projection plot.
+
+    Parameters
+    ----------
+
+    f: yt dataset
+        The data set to use in the projection
+
+    name: str
+        The name to print on the plot
+
+    i: int
+        The index of the plot
+
+    fig: matplotlib figure
+        The figure to use
+
+    axes: list
+        The list of axes to use for the plot
+    """
+    direction = "x"
+    field = ("PartType0", "Temperature")
+
+    # compute the projection
+    p = yt.ProjectionPlot(f, direction, field, center=center,
+                          weight_field="density",
+                          width=small_width, buff_size=(800, 800))
+
+    # Compute the limits
+    p.set_unit("Temperature", "K")
+    data = p.to_fits_data()["Temperature"].data
+    if limits_temperature[0] > data.min():
+        print("WARNING: data below min", data.min())
+    if limits_temperature[1] < data.max():
+        print("WARNING: data above max", data.max())
+
+    # Adjust the plot
+    p.set_cmap(field=field, cmap=cmap_temperature)
+    p.set_log(field, True)
+    p.set_zlim(field, limits_temperature[0], limits_temperature[1])
+
+    # Draw it into the correct figure
+    plot = p.plots[field]
+
+    plot.figure = fig
+    plot.axes = axes[i].axes
+
+    # Add code name
+    at = AnchoredText(name, loc=2, prop=dict(size=6), frameon=True)
+    plot.axes.add_artist(at)
+
+    # plot color bar
+    if i != 0:
+        plot.cax = axes.cbar_axes[0]
+        p.hide_axes()
+    p._setup_plots()
+
+    if i == 0:
+        z = 1. / f.scale_factor - 1.
+        text = "Redshift = %.2g" % z
+        prop = {
+            "color": "w"
+        }
+        at = AnchoredText(text, loc=3, prop=prop, frameon=False)
+        plot.axes.add_artist(at)
+
+    # Add code name
+    at = AnchoredText(name, loc=2, prop=dict(size=6), frameon=True)
+    plot.axes.add_artist(at)
+
+
+def doMassPlot(f, name, i, fig, axes):
+    """
+    Generate a mass projection (including dark matter) plot.
+
+    Parameters
+    ----------
+
+    f: yt dataset
+        The data set to use in the projection
+
+    name: str
+        The name to print on the plot
+
+    i: int
+        The index of the plot
+
+    fig: matplotlib figure
+        The figure to use
+
+    axes: list
+        The list of axes to use for the plot
+    """
+    direction = "y"
+    field = ("all", "Masses")
+
+    # compute the projection
+    p = yt.ParticleProjectionPlot(f, direction, field, center="center",
+                                  width=large_width)
+
+    # # Compute the limits
+    # p.set_unit("masses", "Msun * kpc")
+    # data = p.to_fits_data()["masses"].data
+    # global limits_mass
+    # if limits_mass is None:
+    #     dmin = data.min()
+    #     if data.min() == 0:
+    #         dmin = 1e-6 * data.max()
+    #     else:
+    #         dmin *= 0.5
+    #     dmax = data.max() * 2
+    #     limits_mass = (dmin, dmax)
+    # else:
+    #     if limits_mass[0] > data.min():
+    #         print("WARNING: data below min", data.min())
+    #     if limits_mass[1] < data.max():
+    #         print("WARNING: data above max", data.max())
+
+    # Adjust the plot
+    p.set_cmap(field=field, cmap=cmap_mass)
+    # p.set_log(field, True)
+    # p.set_zlim(field, limits_mass[0], limits_mass[1])
+
+    # Draw it into the correct figure
+    plot = p.plots[field]
+
+    plot.figure = fig
+    plot.axes = axes[i].axes
+
+    # Add code name
+    at = AnchoredText(name, loc=2, prop=dict(size=6), frameon=True)
+    plot.axes.add_artist(at)
+
+    # plot color bar
+    if i != 0:
+        plot.cax = axes.cbar_axes[0]
+        p.hide_axes()
+    p._setup_plots()
+
+    # Add code name
+    at = AnchoredText(name, loc=2, prop=dict(size=6), frameon=True)
+    plot.axes.add_artist(at)