diff --git a/examples/GEAR/ZoomIn/add_fields.py b/examples/GEAR/ZoomIn/add_fields.py new file mode 100644 index 0000000000000000000000000000000000000000..43796ba4b48a98036fb8d2001343a94c76449fe1 --- /dev/null +++ b/examples/GEAR/ZoomIn/add_fields.py @@ -0,0 +1,67 @@ +#!/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") diff --git a/examples/GEAR/ZoomIn/cleanupSwift.py b/examples/GEAR/ZoomIn/cleanupSwift.py new file mode 100644 index 0000000000000000000000000000000000000000..5889bca9e321886e148dc5fcaa023421e69ebbce --- /dev/null +++ b/examples/GEAR/ZoomIn/cleanupSwift.py @@ -0,0 +1,74 @@ +#!/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() diff --git a/examples/GEAR/ZoomIn/make_image.py b/examples/GEAR/ZoomIn/make_image.py new file mode 100644 index 0000000000000000000000000000000000000000..e9fbc68a5be9d084a325b779a8d8b7dfc82b54b7 --- /dev/null +++ b/examples/GEAR/ZoomIn/make_image.py @@ -0,0 +1,123 @@ +#!/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() diff --git a/examples/GEAR/ZoomIn/phase_plot.py b/examples/GEAR/ZoomIn/phase_plot.py new file mode 100644 index 0000000000000000000000000000000000000000..7a6917d314a2df31afdd6486a8b54a9d731452ba --- /dev/null +++ b/examples/GEAR/ZoomIn/phase_plot.py @@ -0,0 +1,59 @@ +#!/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) diff --git a/examples/GEAR/ZoomIn/projection_plot.py b/examples/GEAR/ZoomIn/projection_plot.py new file mode 100644 index 0000000000000000000000000000000000000000..e4b35f3ff604f7894aa509b9261f4c69d25f40e8 --- /dev/null +++ b/examples/GEAR/ZoomIn/projection_plot.py @@ -0,0 +1,252 @@ +#!/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)