Commit 07556ab4 authored by Loic Hausammann's avatar Loic Hausammann

ZoomIn: simplify plotting script

parent 776ddf49
#!/usr/bin/env
from yt.extensions.astro_analysis.halo_analysis.api import HaloCatalog
import matplotlib.pyplot as plt
import numpy as np
Nbins = 20
def savePlot(data):
plt.figure()
plt.tight_layout(pad=0.4, w_pad=0.5, h_pad=0.5)
markers = ["s", "o"]
d_min = None
d_max = None
for d in data[0]:
if d_min is None or d.min() < d_min:
d_min = d.min()
if d_max is None or d.max() > d_max:
d_max = d.max()
for i, d in enumerate(data[0]):
plt.hist(d, Nbins, histtype="step", log=True, range=[d_min, d_max])
plt.xlabel("$\mathrm{Mass\ \mathrm{(M_{\odot})}}$", fontsize='large')
plt.ylabel(r"$\mathrm{Number}\ o\mathrm{f}\ \mathrm{Halos}$", fontsize='large')
plt.legend(data[1], loc=4, frameon=True, ncol=2, fancybox=True)
leg = plt.gca().get_legend()
ltext = leg.get_texts()
plt.setp(ltext, fontsize='small')
plt.grid(True)
plt.savefig("halo_distribution.png")
def doPlot(f, name, i):
hc = HaloCatalog(data_ds=f, finder_method="fof")
hc.create()
masses = np.zeros(len(hc.catalog))
for i, halo in enumerate(hc.catalog):
masses[i] = halo["particle_mass"].in_units("Msun")
return masses
#!/usr/bin/env python3 #!/usr/bin/env python3
import yt import yt
import sys
import projection_plot import projection_plot
import velocity_plot
import phase_plot import phase_plot
import halo_distribution
import metal_plot
import add_fields import add_fields
import star_plot import star_plot
import profile_plot
import matplotlib.pyplot as plt import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import AxesGrid from mpl_toolkits.axes_grid1 import AxesGrid
from yt.units import kpc from yt.units import kpc
import numpy as np
# Parameters
swift = "swift_final.hdf5" swift = "swift_final.hdf5"
gear = "h050_final.hdf5" gear = "h050_final.hdf5"
...@@ -27,26 +19,14 @@ do_hydro = True ...@@ -27,26 +19,14 @@ do_hydro = True
do_stars = True do_stars = True
do_feedback = True do_feedback = True
do_plot = { do_plot = {
# DMO
"projection_mass": do_dmo,
"velocity": do_dmo,
"halo_distribution": True, # very slow
"profile": do_dmo,
# hydro # hydro
"projection_density": do_hydro, "projection_density": do_hydro,
"projection_temperature": do_hydro, "projection_temperature": do_hydro,
"phase_1d_density": do_hydro,
"phase_1d_temperature": do_hydro,
"phase_2d": do_hydro, "phase_2d": do_hydro,
"metals": do_hydro,
# stars # stars
"SFR": do_stars, "SFR": do_stars,
"projection_mass_stars": do_stars,
# feedback # feedback
"abundances": do_feedback, "abundances": do_feedback,
"projection_metals": do_feedback,
"iron_distribution": do_feedback,
"mass_distribution": do_feedback
} }
# Generate the figures # Generate the figures
...@@ -55,9 +35,6 @@ figsize = (100, 100) ...@@ -55,9 +35,6 @@ figsize = (100, 100)
figures = { figures = {
"projection_density": plt.figure(figsize=figsize), "projection_density": plt.figure(figsize=figsize),
"projection_temperature": plt.figure(figsize=figsize), "projection_temperature": plt.figure(figsize=figsize),
"projection_metals": plt.figure(figsize=figsize),
"projection_mass": plt.figure(figsize=figsize),
"projection_mass_stars": plt.figure(figsize=figsize),
"phase_2d": plt.figure(figsize=figsize) "phase_2d": plt.figure(figsize=figsize)
} }
...@@ -77,37 +54,17 @@ axes = { ...@@ -77,37 +54,17 @@ axes = {
figures["projection_temperature"], fig_size, subgrid, figures["projection_temperature"], fig_size, subgrid,
add_all=True, share_all=True, cbar_mode="single", add_all=True, share_all=True, cbar_mode="single",
cbar_size="2%", cbar_pad=0.02), 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),
"projection_mass_stars": AxesGrid(
figures["projection_mass_stars"], fig_size, subgrid,
add_all=True, share_all=True, cbar_mode="single",
cbar_size="2%", cbar_pad=0.02),
"phase_2d": AxesGrid( "phase_2d": AxesGrid(
figures["phase_2d"], fig_size, subgrid, axes_pad=0.05, figures["phase_2d"], fig_size, subgrid, axes_pad=0.05,
add_all=True, share_all=True, cbar_mode="single", add_all=True, share_all=True, cbar_mode="single",
cbar_size="2%", cbar_pad=0.05, aspect=False), cbar_size="2%", cbar_pad=0.05, aspect=False),
"projection_metals": AxesGrid(
figures["projection_metals"], fig_size, subgrid,
add_all=True, share_all=True, cbar_mode="single",
cbar_size="2%", cbar_pad=0.02),
} }
# Data # Data
data = { data = {
"phase_1d_density": ([], []),
"phase_1d_temperature": ([], []),
"halo_distribution": ([], []),
"velocity": ([], []),
"metals": ([], []),
"SFR": ([], []), "SFR": ([], []),
"profile": ([], []),
"abundances": ([], []), "abundances": ([], []),
"iron_distribution": ([], []),
"mass_distribution": ([], []),
} }
names = [] names = []
...@@ -121,62 +78,17 @@ def savePlot(): ...@@ -121,62 +78,17 @@ def savePlot():
projection_plot.savePlot(figures["projection_temperature"], projection_plot.savePlot(figures["projection_temperature"],
"temperature") "temperature")
if do_plot["projection_metals"]:
projection_plot.savePlot(figures["projection_metals"],
"metals")
if do_plot["projection_mass"]:
projection_plot.savePlot(figures["projection_mass"],
"mass")
if do_plot["projection_mass_stars"]:
projection_plot.savePlot(figures["projection_mass_stars"],
"mass_stars")
if do_plot["phase_1d_density"]:
data["phase_1d_density"][1].extend(names)
phase_plot.save1DPlotDensity(data["phase_1d_density"])
if do_plot["phase_1d_temperature"]:
data["phase_1d_temperature"][1].extend(names)
phase_plot.save1DPlotTemperature(data["phase_1d_temperature"])
if do_plot["phase_2d"]: if do_plot["phase_2d"]:
phase_plot.save2DPlot(figures["phase_2d"]) phase_plot.save2DPlot(figures["phase_2d"])
# halo distribution
if do_plot["halo_distribution"]:
data["halo_distribution"][1].extend(names)
halo_distribution.savePlot(data["halo_distribution"])
if do_plot["velocity"]:
data["velocity"][1].extend(names)
velocity_plot.save1DPlot(data["velocity"])
if do_plot["metals"]:
data["metals"][1].extend(names)
metal_plot.save1DPlot(data["metals"])
if do_plot["SFR"]: if do_plot["SFR"]:
data["SFR"][1].extend(names) data["SFR"][1].extend(names)
star_plot.saveSFRPlot(data["SFR"]) star_plot.saveSFRPlot(data["SFR"])
if do_plot["profile"]:
data["profile"][1].extend(names)
profile_plot.savePlot(data["profile"])
if do_plot["abundances"]: if do_plot["abundances"]:
data["abundances"][1].extend(names) data["abundances"][1].extend(names)
star_plot.saveAbundancesPlot(data["abundances"]) star_plot.saveAbundancesPlot(data["abundances"])
if do_plot["iron_distribution"]:
data["iron_distribution"][1].extend(names)
star_plot.saveIronDistributionPlot(data["iron_distribution"])
if do_plot["mass_distribution"]:
data["mass_distribution"][1].extend(names)
star_plot.saveMassDistributionPlot(data["mass_distribution"])
def doPlot(filename, i, name, center): def doPlot(filename, i, name, center):
names.append(name) names.append(name)
...@@ -215,72 +127,19 @@ def doPlot(filename, i, name, center): ...@@ -215,72 +127,19 @@ def doPlot(filename, i, name, center):
f, name, i, figures["projection_temperature"], f, name, i, figures["projection_temperature"],
axes["projection_temperature"]) axes["projection_temperature"])
if do_plot["projection_metals"]:
projection_plot.doMetalsPlot(
f, name, i, figures["projection_metals"],
axes["projection_metals"])
# Do mass projection plot
if do_plot["projection_mass"]:
projection_plot.doMassPlot(
f, name, i, figures["projection_mass"],
axes["projection_mass"], "all")
# Do stellar mass projection plot
if do_plot["projection_mass_stars"]:
projection_plot.doMassPlot(
f, name, i, figures["projection_mass_stars"],
axes["projection_mass_stars"], "stars")
# 1D Phase plot
if do_plot["phase_1d_density"]:
p = phase_plot.do1DPlotDensity(f, name, i)
data["phase_1d_density"][0].append(p)
if do_plot["phase_1d_temperature"]:
p = phase_plot.do1DPlotTemperature(f, name, i)
data["phase_1d_temperature"][0].append(p)
# 2D Phase plot # 2D Phase plot
if do_plot["phase_2d"]: if do_plot["phase_2d"]:
phase_plot.do2DPlot(f, name, i, figures["phase_2d"], phase_plot.do2DPlot(f, name, i, figures["phase_2d"],
axes["phase_2d"]) axes["phase_2d"])
# halo distribution
if do_plot["halo_distribution"]:
m = halo_distribution.doPlot(f, name, i)
data["halo_distribution"][0].append(m)
# Velocity plot
if do_plot["velocity"]:
p = velocity_plot.do1DPlot(f, name, i)
data["velocity"][0].append(p)
# metal plot
if do_plot["metals"]:
p = metal_plot.do1DPlot(f, name, i)
data["metals"][0].append(p)
if do_plot["SFR"]: if do_plot["SFR"]:
p = star_plot.doSFRPlot(f, name, i) p = star_plot.doSFRPlot(f, name, i)
data["SFR"][0].append(p) data["SFR"][0].append(p)
if do_plot["profile"]:
p = profile_plot.doPlot(f, name, i)
data["profile"][0].append(p)
if do_plot["abundances"]: if do_plot["abundances"]:
p = star_plot.doAbundancesPlot(f, name, i) p = star_plot.doAbundancesPlot(f, name, i)
data["abundances"][0].append(p) data["abundances"][0].append(p)
if do_plot["iron_distribution"]:
p = star_plot.doIronDistributionPlot(f, name, i)
data["iron_distribution"][0].append(p)
if do_plot["mass_distribution"]:
p = star_plot.doMassDistributionPlot(f, name, i)
data["mass_distribution"][0].append(p)
return center return center
......
#!/usr/bin/env python3
import yt
import matplotlib.pyplot as plt
def save1DPlot(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]):
z = p.x.in_units("")
m = p["Masses"].in_units("Msun").d
plt.plot(z, m, linestyle="-", marker=markers[i],
markeredgecolor='none', linewidth=1.2, alpha=0.8)
plt.semilogx()
plt.semilogy()
plt.xlabel(r"Metal fraction", fontsize='large')
plt.ylabel("$\mathrm{Mass\ (Msun)}$", fontsize='large')
plt.legend(profiles[1], loc=4, frameon=True, ncol=2, fancybox=True)
leg = plt.gca().get_legend()
ltext = leg.get_texts()
plt.setp(ltext, fontsize='small')
plt.grid(True)
plt.savefig("metals.png", bbox_inches='tight', pad_inches=0.03, dpi=300)
def do1DPlot(f, name, i):
sp = f.sphere(f.center, f.width)
# Because ParticleProfilePlot doesn't exist, I will do the following trick.
p = yt.create_profile(sp, ("PartType0", "Metallicity"),
("PartType0", "Masses"),
weight_field=None, n_bins=50,
accumulation=False)
return p
...@@ -17,56 +17,6 @@ def save2DPlot(fig): ...@@ -17,56 +17,6 @@ def save2DPlot(fig):
pad_inches=0.03, dpi=300) pad_inches=0.03, dpi=300)
def save1DPlotDensity(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]):
rho = p.x.in_units("g/cm**3").d
mass = p["Masses"].in_units("Msun").d
mass[mass == 0] = np.nan
plt.plot(rho, mass, linestyle="-", marker=markers[i],
markeredgecolor='none', linewidth=1.2, alpha=0.8)
plt.semilogx()
plt.semilogy()
plt.xlabel("$\mathrm{Density\ (g/cm^3)}$", fontsize='large')
plt.ylabel(r"$\mathrm{Mass}\/\mathrm{(M_{\odot})}$", fontsize='large')
plt.legend(profiles[1], loc=4, frameon=True, ncol=2, fancybox=True)
leg = plt.gca().get_legend()
ltext = leg.get_texts()
plt.setp(ltext, fontsize='small')
plt.grid(True)
plt.savefig("phase_1d_density.png", bbox_inches='tight',
pad_inches=0.03, dpi=300)
def save1DPlotTemperature(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]):
rho = p.x.in_units("K").d
mass = p["Masses"].in_units("Msun").d
mass[mass == 0] = np.nan
plt.plot(rho, mass, linestyle="-", marker=markers[i],
markeredgecolor='none', linewidth=1.2, alpha=0.8)
plt.semilogx()
plt.semilogy()
plt.xlabel("$\mathrm{Temperature\ K}$", fontsize='large')
plt.ylabel(r"$\mathrm{Mass}\/\mathrm{(M_{\odot})}$", fontsize='large')
plt.legend(profiles[1], loc=4, frameon=True, ncol=2, fancybox=True)
leg = plt.gca().get_legend()
ltext = leg.get_texts()
plt.setp(ltext, fontsize='small')
plt.grid(True)
plt.savefig("phase_1d_temperature.png", bbox_inches='tight',
pad_inches=0.03, dpi=300)
def do2DPlot(f, name, i, fig, axes): def do2DPlot(f, name, i, fig, axes):
p = yt.PhasePlot( p = yt.PhasePlot(
f, ("PartType0", "density"), ("PartType0", "Temperature"), f, ("PartType0", "density"), ("PartType0", "Temperature"),
...@@ -118,22 +68,3 @@ def do2DPlot(f, name, i, fig, axes): ...@@ -118,22 +68,3 @@ def do2DPlot(f, name, i, fig, axes):
line = ln.Line2D(x_ticks, [y]*N, linestyle=":", linewidth=0.6, line = ln.Line2D(x_ticks, [y]*N, linestyle=":", linewidth=0.6,
color="k", alpha=0.7) color="k", alpha=0.7)
plot.axes.add_line(line) plot.axes.add_line(line)
def do1DPlotDensity(f, name, i):
sp = f.sphere(f.center, f.width)
# Because ParticleProfilePlot doesn't exist, I will do the following trick.
p = yt.create_profile(sp, ("PartType0", "density"),
("PartType0", "Masses"), weight_field=None,
n_bins=40, accumulation=False)
return p
def do1DPlotTemperature(f, name, i):
sp = f.sphere(f.center, f.width)
# Because ParticleProfilePlot doesn't exist, I will do the following trick.
p = yt.create_profile(sp, ("PartType0", "Temperature"),
("PartType0", "Masses"), weight_field=None,
n_bins=40, accumulation=False)
return p
#!/usr/bin/env python3
import yt
from yt.units import Msun, kpc
import matplotlib.pyplot as plt
limits_mass = (1e3 * Msun, 1e7 * Msun)
def savePlot(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]):
r = p.x.in_units("pc").d
mass = p["Masses"].in_units("Msun").d
plt.plot(r, mass, linestyle="-", marker=markers[i],
markeredgecolor='none', linewidth=1.2, alpha=0.8)
plt.semilogx()
plt.semilogy()
plt.xlabel(r"$\mathrm{Radius}\/\mathrm{(pc)}$", fontsize='large')
plt.ylabel("$\mathrm{Integrated}\ \mathrm{mass}\ (Msun)}$",
fontsize='large')
plt.legend(profiles[1], loc=4, frameon=True, ncol=2, fancybox=True)
leg = plt.gca().get_legend()
ltext = leg.get_texts()
plt.setp(ltext, fontsize='small')
plt.grid(True)
plt.savefig("integrated_mass.png", bbox_inches='tight',
pad_inches=0.03, dpi=300)
def doPlot(f, name, i):
sp = f.sphere(f.center, f.width)
# Because ParticleProfilePlot doesn't exist, I will do the following trick.
p = yt.create_profile(sp, "particle_radius", "Masses",
weight_field=None, n_bins=50,
accumulation=True)
return p
...@@ -176,162 +176,3 @@ def doTemperaturePlot(f, name, i, fig, axes): ...@@ -176,162 +176,3 @@ def doTemperaturePlot(f, name, i, fig, axes):
# Add code name # Add code name
at = AnchoredText(name, loc=2, prop=dict(size=6), frameon=True) at = AnchoredText(name, loc=2, prop=dict(size=6), frameon=True)
plot.axes.add_artist(at) plot.axes.add_artist(at)
def doMassPlot(f, name, i, fig, axes, parttype):
"""
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
parttype: str
The name of the particle type to use
"""
if parttype == "stars":
if name == "GEAR":
parttype = "PartType1"
else:
parttype = "PartType4"
direction = "x"
field = (parttype, "Masses")
# compute the projection
p = yt.ParticleProjectionPlot(f, direction, field, center=f.center,
width=f.width)
# # Compute the limits
p.set_unit("Masses", "Msun")
data = p.to_fits_data()["Masses"].data
global limits_mass
if limits_mass is None:
dmin = np.nanmin(data)
if np.nanmin(data) == 0:
dmin = 1e-6 * np.nanmax(data)
else:
dmin *= 0.5
dmax = np.nanmax(data) * 2
limits_mass = (dmin, dmax)
else:
if limits_mass[0] > np.nanmin(data):
print("WARNING: data below min", np.nanmin(data))
if limits_mass[1] < np.nanmax(data):
print("WARNING: data above max", np.nanmax(data))
# 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)
def doMetalsPlot(f, name, i, fig, axes):
"""
Generate a metal 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", "Metallicity")
# compute the projection
p = yt.ProjectionPlot(f, direction, field, center=f.center,
width=f.width, buff_size=(800, 800),
weight_field=("PartType0", "Density"))
# Compute the limits
# p.set_unit("metallicity", "g / cm**2")
data = p.to_fits_data()["Metallicity"].data
global limits_metals
if limits_metals is None:
limits_metals = (data.min(), data.max())
if limits_metals[0] > data.min():
print("WARNING: data below min", data.min())
if limits_metals[1] < data.max():
print("WARNING: data above max", data.max())
# Adjust the plot
p.set_cmap(field=field, cmap=cmap_metals)
p.set_log(field, True)
if limits_metals[0] != limits_metals[1]:
p.set_zlim(field, limits_metals[0], limits_metals[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)
...@@ -144,134 +144,3 @@ def doAbundancesPlot(f, name, i): ...@@ -144,134 +144,3 @@ def doAbundancesPlot(f, name, i):
fe = getElement(f, sp, name, "Fe") fe = getElement(f, sp, name, "Fe")
return mg, fe return mg, fe
def saveIronDistributionPlot(profiles):
plt.figure(figsize=(8, 8))
plt.tight_layout(pad=0.4, w_pad=0.5, h_pad=0.5)
fe_max = max(p.max() for p in profiles[0])
fe_min = min(p.min() for p in profiles[0])
fe_min = np.log10((fe_min + 1e-10) / fe_sol_abund)
fe_max = np.log10(fe_max / fe_sol_abund)
for i, p in enumerate(profiles[0]):
fe = np.log10(p / fe_sol_abund)
plt.hist(fe, bins=20, range=[fe_min, fe_max], histtype="step",