Commit 49e5fff8 authored by Loic Hausammann's avatar Loic Hausammann

ZoomIn: remove plot script

parent e0084a85
#!/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")
def addMetals(f):
def _metallicity(field, data):
if len(data["PartType0", "Metals"].shape) == 1:
return data["PartType0", "Metals"]
else:
# in_units("") turned out to be crucial!;
# otherwise code_metallicity will be used
return data["PartType0", "Metals"][:, -1].in_units("")
# We are creating ("Gas", "Metallicity") here, different from
# ("Gas", "metallicity") which is auto-generated by yt
# but doesn't work properly
f.add_field(("PartType0", "Metallicity"), function=_metallicity,
display_name="Metal mass fraction", particle_type=True,
take_log=True, units="")
def _density_squared(field, data):
return data[("PartType0", "density")]**2
f.add_field(("PartType0", "density_squared"), function=_density_squared,
particle_type=True, units="g**2/cm**6")
#!/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, "a")
# 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"]
# 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),
("InitialMass", "BirthMasses", h),
("StarInitMass", "BirthMasses", h),
("Velocities", "Velocities", 1. / a**0.5),
("Density", "Densities", 1. / h**2),
("Entropies", "Entropies", 1.),
("InternalEnergy", "InternalEnergies", 1. / a**2),
("SmoothingLength", "SmoothingLengths", h),
("Metals", "SmoothedElementAbundances", 1.)
]
# 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_m"]
head["HubbleParam"] = cosmo["h"][0]
head["Time"] = float(a)
f.close()
#!/usr/bin/env python3
import yt
import projection_plot
import phase_plot
import add_fields
import star_plot
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import AxesGrid
from yt.units import kpc
swift = "swift_final.hdf5"
gear = "h050_final.hdf5"
width = 50 * kpc
do_dmo = True
do_hydro = True
do_stars = True
do_feedback = True
do_plot = {
# hydro
"projection_density": do_hydro,
"projection_temperature": do_hydro,
"phase_2d": do_hydro,
# stars
"SFR": do_stars,
# feedback
"abundances": do_feedback,
}
# Generate the figures
figsize = (100, 100)
figures = {
"projection_density": plt.figure(figsize=figsize),
"projection_temperature": plt.figure(figsize=figsize),
"phase_2d": 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),
"phase_2d": AxesGrid(
figures["phase_2d"], 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),
}
# Data
data = {
"SFR": ([], []),
"abundances": ([], []),
}
names = []
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["phase_2d"]:
phase_plot.save2DPlot(figures["phase_2d"])
if do_plot["SFR"]:
data["SFR"][1].extend(names)
star_plot.saveSFRPlot(data["SFR"])
if do_plot["abundances"]:
data["abundances"][1].extend(names)
star_plot.saveAbundancesPlot(data["abundances"])
def doPlot(filename, i, name, center):
names.append(name)
f = yt.load(filename)
if center is None:
center = f.find_max("Density")[1]
f.center = center
f.width = width * f.scale_factor
if (do_plot["projection_temperature"] or do_plot["phase_2d"]):
add_fields.addTemperature(f)
add_fields.addMassDeposit(f)
if (do_plot["projection_metals"] or do_plot["metals"]):
add_fields.addMetals(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"])
# 2D Phase plot
if do_plot["phase_2d"]:
phase_plot.do2DPlot(f, name, i, figures["phase_2d"],
axes["phase_2d"])
if do_plot["SFR"]:
p = star_plot.doSFRPlot(f, name, i)
data["SFR"][0].append(p)
if do_plot["abundances"]:
p = star_plot.doAbundancesPlot(f, name, i)
data["abundances"][0].append(p)
return center
center = None
# center = np.array([1724.33547783, 1802.56263082, 1785.09893269])
center = doPlot(gear, 1, "GEAR", center=center)
center = None
doPlot(swift, 0, "SWIFT", center=center)
savePlot()
#!/usr/bin/env python3
import yt
from yt.units import Msun, amu, cm, K, kpc
from matplotlib.offsetbox import AnchoredText
import matplotlib.pyplot as plt
import matplotlib.lines as ln
import numpy as np
limits_temperature = (1e1 * K, 1e5 * K)
limits_density = (1e-7 * amu / cm**3, 1e7 * amu / cm**3)
limits_mass = (1e3 * Msun, 1e7 * Msun)
def save2DPlot(fig):
fig.savefig("phase_2d.png", bbox_inches="tight",
pad_inches=0.03, dpi=300)
def do2DPlot(f, name, i, fig, axes):
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)
# Make a grid
x_ticks = plot.axes.get_xticks()
y_ticks = plot.axes.get_yticks()
for x in x_ticks:
N = len(y_ticks)
line = ln.Line2D([x]*N, y_ticks, linestyle=":", linewidth=0.6,
color="k", alpha=0.7)
plot.axes.add_line(line)
for y in y_ticks:
N = len(x_ticks)
line = ln.Line2D(x_ticks, [y]*N, linestyle=":", linewidth=0.6,
color="k", alpha=0.7)
plot.axes.add_line(line)
#!/usr/bin/env python3
import yt
from yt.units import kpc, g, cm, K
from matplotlib.offsetbox import AnchoredText
import numpy as np
cmap_density = "algae"
cmap_temperature = "magma"
cmap_mass = "inferno"
cmap_metals = "coolwarm"
limits_density = (1e-10 * g / cm**2, 1e2 * g / cm**2)
limits_temperature = (10 * K, 1e5 * K)
limits_metals = None
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=f.center,
width=f.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")
a = f.scale_factor
# compute the projection
p = yt.ProjectionPlot(f, direction, field, center=f.center,
weight_field="density",
width=f.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)
#!/usr/bin/env python3
import matplotlib.pyplot as plt
import numpy as np
from yt.units import Msun
from yt.utilities.cosmology import Cosmology
from multiprocessing import Pool
from functools import partial
from h5py import File
fe_sol_abund = 1.76603e-3
mg_sol_abund = 9.24316e-4
def saveSFRPlot(profiles):
plt.figure(figsize=(8, 8))
plt.tight_layout(pad=0.4, w_pad=0.5, h_pad=0.5)
markers = ["s", "o"]
sfr_fig = plt.figure()
mstar_fig = plt.figure()
# get the boundary
t_max = None
for p in profiles[0]:
if t_max is None:
t_max = p[1].max()
else:
tmp = p[1].max()
if t_max < tmp:
t_max = tmp
t_max = float(t_max)
for i, p in enumerate(profiles[0]):
masses = p[0]
form_time = p[1]
time_range = [0, t_max]
n_bins = 100
hist, bins = np.histogram(form_time, bins=n_bins, range=time_range)
inds = np.digitize(form_time, bins=bins)
time = (bins[:-1] + bins[1:])/2
sfr = np.array([masses[inds == j+1].sum()/(bins[j+1]-bins[j])
for j in range(len(time))])
sfr[sfr == 0] = np.nan
sfr /= 1e9 # convert M / Gyr -> M / yr
plt.figure(sfr_fig.number)
plt.plot(time, sfr, linestyle="-", marker=markers[i],
markeredgecolor="none", linewidth=1.2, alpha=0.8)
mstars = np.array([masses[inds == j+1].sum()
for j in range(len(time))])
mstars = np.cumsum(mstars) / 1e6
plt.figure(mstar_fig.number)
plt.plot(time, mstars, linestyle="-",
markeredgecolor="none", linewidth=1.2, alpha=0.8)
plt.figure(sfr_fig.number)
plt.xlabel('Time [Gyr]')
plt.ylabel('SFR [M$_\odot$ yr$^{-1}$]')
plt.legend(profiles[1])
plt.savefig("sfr.png")
plt.figure(mstar_fig.number)
plt.xlabel('Time [Gyr]')
plt.ylabel('Stellar mass [$10^6$ M$_\odot$]')
plt.legend(profiles[1])
plt.savefig("mstars.png")
def doSFRPlot(f, name, i):
sp = f.sphere(f.center, f.width)
if name == "GEAR":
masses = (sp["PartType1", "StarInitMass"] * 1e10 * Msun).in_units("Msun")
a = sp["PartType1", "StarFormationTime"]
else:
masses = (sp["PartType4", "StarInitMass"] * 1e10 * Msun).in_units("Msun")
a = sp["PartType4", "BirthScaleFactors"]
cosmo = Cosmology(hubble_constant=f.hubble_constant,
omega_matter=f.omega_matter,
omega_lambda=f.omega_lambda)
z = np.array(1. / a - 1.)
with Pool() as p:
formation_time = p.map(
partial(cosmo.lookback_time, z_f=1e6), z)
formation_time = f.arr(list(formation_time), "s").in_units("Gyr")
return masses, formation_time
def saveAbundancesPlot(profiles):
plt.figure(figsize=(8, 8))
plt.tight_layout(pad=0.4, w_pad=0.5, h_pad=0.5)
markers = ["s", "o"]
for i, p in enumerate(profiles[0]):
mg = np.log10(p[0] / mg_sol_abund)
fe = np.log10(p[1] / fe_sol_abund)
plt.plot(fe, mg - fe, linestyle="", marker=markers[i],
markeredgecolor="none", alpha=0.8, markersize=3)
plt.xlabel('[Fe / H]')
plt.ylabel('[Mg / Fe]')
plt.legend(profiles[1])
plt.savefig("abundances.png")
plt.xlim([-4, 0.5])
plt.ylim([-1, 1.5])
plt.savefig("abundances_zoom.png")
def getElement(f, sp, name, el):
if name == "GEAR":
metals = sp["PartType1", "StarMetals"]
h = File(f.filename_template, "r")
names = h["Header"].attrs["ChimieLabels"]
names = names.split(b",")
# remove garbage
names = names[:-2]
# bytes to unicode
names = [name.decode() for name in names]
h.close()
else:
metals = sp["PartType4", "ElementAbundances"]
h = File(f.filename_template, "r")
sub = h["SubgridScheme"].attrs
names = []
for i in range(metals.shape[1]):
name = sub["Element %i" % i].decode()
names.append(name)
ind = names.index(el)
return metals[:, ind]
def doAbundancesPlot(f, name, i):
sp = f.sphere(f.center, f.width)
mg = getElement(f, sp, name, "Mg")
fe = getElement(f, sp, name, "Fe")