Skip to content

SWIFTsim
SWIFTsim
Commit 7000d664 authored by Loic Hausammann's avatar Loic Hausammann
Browse files
Options

GEAR: small improvement to the example

parent d897e2f9
Hide whitespace changes
Inline Side-by-side
Showing with 525 additions and 82 deletions
examples/GEAR/ZoomIn/add_fields.py
View file @ 7000d664
......@@ -3,6 +3,7 @@
import numpy as np
import yt.utilities.physical_constants as constants
from yt.mods import YTArray
from yt import derived_field
def addMassDeposit(f):
......@@ -65,3 +66,25 @@ def addTemperature(f):
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")
examples/GEAR/ZoomIn/cleanupSwift.py
View file @ 7000d664
......@@ -11,7 +11,7 @@ out = sys.argv[-1]
copyfile(filename, out)
f = File(out)
f = File(out, "a")
# read cosmological parameters
a = f["Cosmology"].attrs["Scale-factor"][0]
......@@ -23,12 +23,6 @@ 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
......@@ -43,11 +37,12 @@ for i in range(NPartType):
fields = [
("Coordinates", "Coordinates", h),
("Masses", "Masses", h),
("Velocities", "Velocities", a**0.5),
("Velocities", "Velocities", 1. / a**0.5),
("Density", "Densities", 1. / h**2),
("Entropies", "Entropies", 1.),
("InternalEnergy", "InternalEnergies", 1. / a**2),
("SmoothingLength", "SmoothingLengths", h)
("SmoothingLength", "SmoothingLengths", h),
("Metals", "SmoothedElementAbundances", 1.)
]
# create links
......@@ -67,7 +62,7 @@ cosmo = f["Cosmology"].attrs
head = f["Header"].attrs
head["Redshift"] = float(cosmo["Redshift"])
head["OmegaLambda"] = cosmo["Omega_lambda"]
head["Omega0"] = cosmo["Omega_b"]
head["Omega0"] = cosmo["Omega_m"]
head["HubbleParam"] = cosmo["h"][0]
head["Time"] = float(a)
......
examples/GEAR/ZoomIn/make_image.py
View file @ 7000d664
#!/usr/bin/env python3
import yt
from yt.units import Msun, kpc, s, km
import unyt
import sys
import matplotlib
import projection_plot
import velocity_plot
import phase_plot
import halo_distribution
import metal_plot
import add_fields
import star_plot
import profile_plot
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import AxesGrid
from matplotlib.offsetbox import AnchoredText
matplotlib.use('Agg')
# Parameters
......@@ -23,14 +21,29 @@ swift = "swift/snapshot_%04i.hdf5" % snap
gear = "gear/snapshot_%04i.hdf5" % snap
do_dmo = False
do_hydro = False
do_stars = True
do_feedback = True
do_plot = {
"projection_density": False,
"projection_temperature": False,
"projection_mass": False,
"halo_distribution": False,
"phase_1d": False,
"phase_2d": False,
"velocity": True
# DMO
"projection_mass": do_dmo,
"velocity": do_dmo,
"halo_distribution": False, # very slow
"profile": do_dmo,
# hydro
"projection_density": True, #do_hydro,
"projection_temperature": do_hydro,
"phase_1d_density": do_hydro,
"phase_1d_temperature": do_hydro,
"phase_2d": do_hydro,
"metals": do_hydro,
# stars
"SFR": do_stars,
"projection_mass_stars": do_stars,
# feedback
"abundances": do_feedback,
"projection_metals": do_feedback,
}
# Generate the figures
......@@ -39,7 +52,9 @@ figsize = (100, 100)
figures = {
"projection_density": 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)
}
......@@ -63,18 +78,31 @@ axes = {
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(
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)
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 = {
"phase_1d": ([], []),
"phase_1d_density": ([], []),
"phase_1d_temperature": ([], []),
"halo_distribution": ([], []),
"velocity": ([], []),
"metals": ([], []),
"SFR": ([], []),
"profile": ([], []),
"abundances": ([], []),
}
names = []
......@@ -88,13 +116,25 @@ def savePlot():
projection_plot.savePlot(figures["projection_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["phase_1d"]:
data["phase_1d"][1].extend(names)
phase_plot.save1DPlot(data["phase_1d"])
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"]:
phase_plot.save2DPlot(figures["phase_2d"])
......@@ -108,13 +148,38 @@ def savePlot():
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"]:
data["SFR"][1].extend(names)
star_plot.saveSFRPlot(data["SFR"])
if do_plot["profile"]:
data["profile"][1].extend(names)
profile_plot.savePlot(data["profile"])
def doPlot(filename, i, name):
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
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:
......@@ -135,17 +200,31 @@ def doPlot(filename, i, name):
f, name, i, figures["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"]:
add_fields.addMassDeposit(f)
projection_plot.doMassPlot(
f, name, i, figures["projection_mass"],
axes["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"]:
p = phase_plot.do1DPlot(f, name, i)
data["phase_1d"][0].append(p)
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
if do_plot["phase_2d"]:
......@@ -162,7 +241,26 @@ def doPlot(filename, i, name):
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"]:
p = star_plot.doSFRPlot(f, name, i)
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"]:
p = star_plot.doAbundancesPlot(f, name, i)
data["abundances"][0].append(p)
return center
doPlot(swift, 0, "SWIFT")
doPlot(gear, 1, "GEAR")
center = doPlot(gear, 1, "GEAR", center=None)
doPlot(swift, 0, "SWIFT", center=center)
savePlot()
examples/GEAR/ZoomIn/phase_plot.py
View file @ 7000d664
......@@ -4,10 +4,12 @@ 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
width = 5 * kpc
limits_temperature = (1e1 * K, 1e7 * K)
limits_density = (1e-10 * amu / cm**3, 1e3 * amu / cm**3)
width = 10 * kpc
limits_temperature = (1e1 * K, 1e5 * K)
limits_density = (1e-7 * amu / cm**3, 1e7 * amu / cm**3)
limits_mass = (1e3 * Msun, 1e7 * Msun)
......@@ -16,27 +18,54 @@ def save2DPlot(fig):
pad_inches=0.03, dpi=300)
def save1DPlot(profiles):
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{d}M\mathrm{/dlog}\/\mathrm{\rho}\/\mathrm{(M_{\odot})}$", 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.png", bbox_inches='tight', pad_inches=0.03, dpi=300)
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):
......@@ -77,12 +106,35 @@ def do2DPlot(f, name, i, fig, axes):
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)
def do1DPlotDensity(f, name, i):
sp = f.sphere(f.center, 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 do1DPlot(f, name, i):
sp = f.sphere("max", width)
def do1DPlotTemperature(f, name, i):
sp = f.sphere(f.center, 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=50,
accumulation=False)
p = yt.create_profile(sp, ("PartType0", "Temperature"),
("PartType0", "Masses"), weight_field=None,
n_bins=40, accumulation=False)
return p
examples/GEAR/ZoomIn/profile_plot.py 0 → 100644
View file @ 7000d664
#!/usr/bin/env python3
import yt
from yt.units import Msun, kpc
import matplotlib.pyplot as plt
width = 20 * kpc
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, 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
examples/GEAR/ZoomIn/projection_plot.py
View file @ 7000d664
......@@ -3,17 +3,17 @@
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"
center = (1441.9954833984375000,
1666.1169433593750000,
1891.6248779296875000)
small_width = 250 * kpc
large_width = 3000 * kpc
limits_density = (1e-10 * g / cm**2, 1e-2 * g / cm**2)
cmap_metals = "coolwarm"
small_width = 400 * kpc
large_width = 400 * kpc
limits_density = (1e-10 * g / cm**2, 1e2 * g / cm**2)
limits_temperature = (10 * K, 1e5 * K)
limits_metals = None
limits_mass = None
......@@ -58,10 +58,11 @@ def doDensityPlot(f, name, i, fig, axes):
"""
direction = "x"
field = ("PartType0", "density")
a = f.scale_factor
# compute the projection
p = yt.ProjectionPlot(f, direction, field, center=center,
width=small_width, buff_size=(800, 800))
p = yt.ProjectionPlot(f, direction, field, center=f.center,
width=small_width * a, buff_size=(800, 800))
# Compute the limits
p.set_unit("density", "g / cm**2")
......@@ -130,11 +131,12 @@ def doTemperaturePlot(f, name, i, fig, axes):
"""
direction = "x"
field = ("PartType0", "Temperature")
a = f.scale_factor
# compute the projection
p = yt.ProjectionPlot(f, direction, field, center=center,
p = yt.ProjectionPlot(f, direction, field, center=f.center,
weight_field="density",
width=small_width, buff_size=(800, 800))
width=small_width * a, buff_size=(800, 800))
# Compute the limits
p.set_unit("Temperature", "K")
......@@ -179,7 +181,7 @@ def doTemperaturePlot(f, name, i, fig, axes):
plot.axes.add_artist(at)
def doMassPlot(f, name, i, fig, axes):
def doMassPlot(f, name, i, fig, axes, parttype):
"""
Generate a mass projection (including dark matter) plot.
......@@ -200,36 +202,116 @@ def doMassPlot(f, name, i, fig, axes):
axes: list
The list of axes to use for the plot
parttype: str
The name of the particle type to use
"""
direction = "y"
field = ("all", "Masses")
width = large_width
if parttype == "stars":
width = small_width
if name == "GEAR":
parttype = "PartType1"
else:
parttype = "PartType4"
direction = "x"
field = (parttype, "Masses")
a = f.scale_factor
# compute the projection
p = yt.ParticleProjectionPlot(f, direction, field, center=center,
width=large_width)
p = yt.ParticleProjectionPlot(f, direction, field, center=f.center,
width=width * a)
# # 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())
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])
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")
a = f.scale_factor
# compute the projection
p = yt.ProjectionPlot(f, direction, field, center=f.center,
width=small_width * a, 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)
#p.set_zlim(field, limits_metals[0], limits_metals[1])
# Draw it into the correct figure
plot = p.plots[field]
......@@ -247,6 +329,15 @@ def doMassPlot(f, name, i, fig, axes):
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)
examples/GEAR/ZoomIn/star_plot.py 0 → 100644
View file @ 7000d664
#!/usr/bin/env python3
import matplotlib.pyplot as plt
import numpy as np
from yt.units import kpc
from yt.utilities.cosmology import Cosmology
from multiprocessing import Pool
from functools import partial
from h5py import File