Commit fe17b4e5 authored by Josh Borrow's avatar Josh Borrow Committed by Matthieu Schaller
Browse files

Santa barbara analysis scripts

parent 0131d5ec
......@@ -14,3 +14,8 @@ particles will be generated at startup.
MD5 check-sum of the ICS:
ba9ab4f00a70d39fa601a4a59984b343 SantaBarbara.hdf5
You can use the script run_velociraptor.sh to also run a basic 3D FoF
with VELOCIraptor on your output data. You will need to set the
VELOCIRAPTOR_PATH environment variable to tell us where the stf-gas
binary lives.
"""
Makes an image of the Santa Barbara cluster.
Requires py-sphviewer.
Invoke as follows:
python3 makeImage.py <name of hdf5 file> \
<number of particle type (i.e. 0 or 1)> \
<colour map to use (default viridis)> \
<text color (default white)> \
<image resolution (default 2048x2048)>
"""
import numpy as np
import matplotlib.pyplot as plt
import h5py
import matplotlib
from sphviewer.tools import QuickView
from matplotlib.patches import Rectangle
from typing import Tuple
from collections import namedtuple
# Set up our simulation data collection to keep stuff together
SimulationData = namedtuple(
"SimulationData",
["coordinates", "masses", "sph_name", "dark_matter_mass", "swift_name", "boxsize"],
)
def latex_float(f):
"""
Taken from:
https://stackoverflow.com/questions/13490292/format-number-using-latex-notation-in-python.
Formats a float to LaTeX style.
"""
float_str = "{0:.2g}".format(f)
if "e" in float_str:
base, exponent = float_str.split("e")
return r"{0} \times 10^{{{1}}}".format(base, int(exponent))
else:
return float_str
def read_data_from_file(filename: str, part_type=0) -> SimulationData:
"""
Reads the relevant data from the HDF5 file.
"""
part_type_name = f"PartType{part_type}"
with h5py.File(filename, "r") as file:
coordinates, boxsize = move_box(file[f"{part_type_name}/Coordinates"][...].T)
masses = file[f"{part_type_name}/Masses"][...]
sph_name = file["HydroScheme"].attrs["Scheme"].decode("utf-8")
unit_mass = (
float(file["Units"].attrs["Unit mass in cgs (U_M)"]) / 2e33
) # in M_sun
dark_matter_mass = float(file["PartType1/Masses"][0]) * unit_mass
code_revision = file["Code"].attrs["Git Revision"].decode("utf-8")
swift_name = f"SWIFT {code_revision}"
data = SimulationData(
coordinates=coordinates,
masses=masses,
sph_name=sph_name,
dark_matter_mass=dark_matter_mass,
swift_name=swift_name,
boxsize=boxsize,
)
return data
def move_box(coordinates: np.ndarray) -> np.ndarray:
"""
Takes the coordinates and moves them in the x-y plane. This moves them 20
code units to the left/right to ensure that the zoomed-out version of the
cluster image is nicely shown
"""
boxsize = np.max(coordinates[0])
coordinates[0] -= 20
coordinates[1] -= 20
coordinates[0] %= boxsize
coordinates[1] %= boxsize
return coordinates, boxsize
def generate_views(data: SimulationData, res=2048) -> Tuple[np.ndarray]:
"""
Generates the views on the data from py-sphviewer.
Returns the overall image for the whole box and then a zoomed region.
"""
qv_all = QuickView(
data.coordinates,
data.masses,
r="infinity",
plot=False,
xsize=res,
ysize=res,
logscale=False,
p=0,
np=48,
)
zoomed_res = (res * 6) // 10
mask = np.logical_and(
np.logical_and(
data.coordinates[0] > (data.boxsize/2-4-20),
data.coordinates[0] < (data.boxsize/2+6-20)
),
np.logical_and(
data.coordinates[1] > (data.boxsize/2-3.5-20),
data.coordinates[1] < (data.boxsize/2+6.5-20)
)
)
qv_zoomed = QuickView(
data.coordinates.T[mask].T,
data.masses[mask],
r="infinity",
plot=False,
xsize=zoomed_res,
ysize=zoomed_res,
logscale=False,
np=48,
)
return qv_all.get_image(), qv_zoomed.get_image()
def create_plot(data: SimulationData, res=2048, cmap="viridis", text_color="white"):
"""
Creates a figure and axes object and returns them for you to do with what you wish.
"""
img_all, img_zoomed = generate_views(data, res)
fig, ax = plt.subplots(figsize=(8, 8))
# Set up in "image" mode
ax.axis("off")
fig.subplots_adjust(0, 0, 1, 1)
ax.imshow(
np.log10(img_all + np.min(img_all[img_all != 0])),
origin="lower",
extent=[-1, 1, -1, 1],
cmap=cmap,
)
lower_left = [(-24 / (0.5 * data.boxsize)), (-23.5 / (0.5 * data.boxsize))]
zoom_rect = Rectangle(
lower_left,
10 / (0.5 * data.boxsize),
10 / (0.5 * data.boxsize),
linewidth=2,
edgecolor=text_color,
facecolor="none",
)
ax.add_patch(zoom_rect)
# Remove ticks as we want "image mode"
ax2 = fig.add_axes([0.35, 0.35, 0.6, 0.6], frame_on=True, xticks=[], yticks=[])
ax2.imshow(
np.log10(img_zoomed + np.min(img_zoomed[img_zoomed != 0])),
origin="lower",
extent=[-1, 1, -1, 1],
cmap=cmap,
)
# This ugly hack sets the box around the subfigure to be white
for child in ax2.get_children():
if isinstance(child, matplotlib.spines.Spine):
child.set_color(text_color)
child.set_linewidth(2)
# Draw lines between boxes
# Bottom Right
ax.plot(
[(-14 / (0.5 * data.boxsize)), 0.9],
[(-23.5 / (0.5 * data.boxsize)), -0.3],
lw=2,
color=text_color,
)
# Top Left
ax.plot(
[(-24 / (0.5 * data.boxsize)), -0.3],
[(-13.5 / (0.5 * data.boxsize)), 0.9],
lw=2,
color=text_color,
)
ax.text(0.95, -0.95, data.swift_name, color=text_color, ha="right")
formatted_dark_matter_mass = latex_float(data.dark_matter_mass)
ax.text(
-0.95,
0.95,
rf"M$_{{\rm DM}} = {formatted_dark_matter_mass}$ M$_\odot$",
color=text_color,
va="top",
)
ax.text(
-0.95,
-0.95,
data.sph_name + "\n" + r"Santa Barbara Cluster (re-ran from Frenk+ 1999)",
color=text_color,
)
return fig, ax
if __name__ == "__main__":
import sys
try:
filename = sys.argv[1]
except IndexError:
filename = "santabarbara_0153.hdf5"
try:
part_type = int(sys.argv[2])
except IndexError:
part_type = 0
try:
cmap = sys.argv[3]
except IndexError:
cmap = "viridis"
try:
text_color = sys.argv[4]
except IndexError:
text_color = "white"
try:
res = int(sys.argv[5])
except IndexError:
res = 2048
# Read in the data from file
try:
data = read_data_from_file(filename, part_type)
except IndexError:
# Must be a dark matter only run
part_type = 1
data = read_data_from_file(filename, part_type)
# Make the plot
fig, ax = create_plot(data, res, cmap, text_color)
fig.savefig(
f"SantaBarbara_{data.sph_name[:8]}_{cmap}_PartType{part_type}_res{res}.png",
dpi=res // 8,
)
"""
Plots the "solution" (i.e. some profiles) for the Santa Barbara cluster.
Invoke as follows:
python3 plotSolution.py <snapshot number> <catalogue directory> <number of bins (optional)>
"""
import matplotlib.pyplot as plt
import numpy as np
import h5py
from collections import namedtuple
from typing import Tuple
try:
import makeImage
create_images = True
except:
create_images = False
# Simulation data
SimulationParticleData = namedtuple(
"SimulationData", ["gas", "dark_matter", "metadata"]
)
ParticleData = namedtuple(
"ParticleData", ["coordinates", "radii", "masses", "densities", "energies"]
)
MetaData = namedtuple("MetaData", ["header", "code", "hydroscheme"])
HaloData = namedtuple("HaloData", ["c", "Rvir", "Mvir", "center"])
def get_energies(handle: h5py.File):
"""
Gets the energies with the correct units.
"""
u = handle["PartType0/InternalEnergy"][:]
unit_length_in_cgs = handle["/Units"].attrs["Unit length in cgs (U_L)"]
unit_mass_in_cgs = handle["/Units"].attrs["Unit mass in cgs (U_M)"]
unit_time_in_cgs = handle["/Units"].attrs["Unit time in cgs (U_t)"]
gas_gamma = handle["/HydroScheme"].attrs["Adiabatic index"][0]
a = handle["/Cosmology"].attrs["Scale-factor"][0]
unit_length_in_si = 0.01 * unit_length_in_cgs
unit_mass_in_si = 0.001 * unit_mass_in_cgs
unit_time_in_si = unit_time_in_cgs
u *= unit_length_in_si ** 2 / unit_time_in_si ** 2
u /= a ** (3 * (gas_gamma - 1.))
return u
def load_data(filename: str, center: np.array) -> SimulationParticleData:
"""
Loads the relevant data for making the profiles, as well as some metadata
for the plot.
Center is the center of the SB cluster and is used to calculate the radial
distances to the particles.
"""
with h5py.File(filename, "r") as file:
gas_handle = file["PartType0"]
dm_handle = file["PartType1"]
gas_data = ParticleData(
coordinates=gas_handle["Coordinates"][...],
radii=get_radial_distances(gas_handle["Coordinates"][...], center),
masses=gas_handle["Masses"][...],
energies=get_energies(file),
densities=gas_handle["Density"][...],
)
dm_data = ParticleData(
coordinates=dm_handle["Coordinates"][...],
radii=get_radial_distances(dm_handle["Coordinates"][...], center),
masses=dm_handle["Masses"][...],
energies=None,
densities=None,
)
metadata = MetaData(
header=dict(file["Header"].attrs),
code=dict(file["Code"].attrs),
hydroscheme=dict(file["HydroScheme"].attrs),
)
simulation_data = SimulationParticleData(
gas=gas_data, dark_matter=dm_data, metadata=metadata
)
return simulation_data
def get_halo_data(catalogue_filename: str) -> HaloData:
"""
Gets the halo center of the largest halo (i.e. the SB cluster).
You will want the .properties file, probably
halo/santabarbara.properties
that is given by VELOCIraptor.
"""
with h5py.File(catalogue_filename, "r") as file:
x = file["Xc"][0]
y = file["Yc"][0]
z = file["Zc"][0]
Mvir = file["Mass_200crit"][0]
Rvir = file["R_200crit"][0]
c = file["cNFW"][0]
return HaloData(c=c, Rvir=Rvir, Mvir=Mvir, center=np.array([x, y, z]))
def get_radial_distances(coordinates: np.ndarray, center: np.array) -> np.array:
"""
Gets the radial distances for all particles.
"""
dx = coordinates - center
return np.sqrt(np.sum(dx * dx, axis=1))
def get_radial_density_profile(radii, masses, bins: int) -> Tuple[np.ndarray]:
"""
Gets the radial gas density profile, after generating similar bins to those
used in similar works.
"""
bins = np.logspace(-2, 1, bins)
histogram, bin_edges = np.histogram(a=radii, weights=masses, bins=bins)
volumes = np.array(
[
(4. * np.pi / 3.) * (r_outer ** 3 - r_inner ** 3)
for r_outer, r_inner in zip(bin_edges[1:], bin_edges[:-1])
]
)
return histogram / volumes, bin_edges # densities
def mu(T, H_frac, T_trans):
"""
Get the molecular weight as a function of temperature.
"""
if T > T_trans:
return 4. / (8. - 5. * (1. - H_frac))
else:
return 4. / (1. + 3. * H_frac)
def T(u, metadata: MetaData):
"""
Temperature of primordial gas.
"""
gas_gamma = metadata.hydroscheme["Adiabatic index"][0]
H_frac = metadata.hydroscheme["Hydrogen mass fraction"][0]
T_trans = metadata.hydroscheme["Hydrogen ionization transition temperature"][0]
k_in_J_K = 1.38064852e-23
mH_in_kg = 1.6737236e-27
T_over_mu = (gas_gamma - 1.) * u * mH_in_kg / k_in_J_K
ret = np.ones(np.size(u)) * T_trans
# Enough energy to be ionized?
mask_ionized = T_over_mu > (T_trans + 1) / mu(T_trans + 1, H_frac, T_trans)
if np.sum(mask_ionized) > 0:
ret[mask_ionized] = T_over_mu[mask_ionized] * mu(T_trans * 10, H_frac, T_trans)
# Neutral gas?
mask_neutral = T_over_mu < (T_trans - 1) / mu((T_trans - 1), H_frac, T_trans)
if np.sum(mask_neutral) > 0:
ret[mask_neutral] = T_over_mu[mask_neutral] * mu(0, H_frac, T_trans)
return ret
def get_radial_temperature_profile(
data: SimulationParticleData, bins: int
) -> np.ndarray:
"""
Gets the radial gas density profile, after generating similar bins to those
used in similar works.
"""
temperatures = T(data.gas.energies, data.metadata)
radii = data.gas.radii
bins = np.logspace(-2, 1, bins)
histogram, _ = np.histogram(a=radii, weights=temperatures, bins=bins)
counts, _ = np.histogram(a=radii, weights=np.ones_like(radii), bins=bins)
return histogram / counts # need to get mean value in bin
def get_radial_entropy_profile(data: SimulationParticleData, bins: int) -> np.ndarray:
"""
Gets the radial gas density profile, after generating similar bins to those
used in similar works.
"""
gas_gamma = data.metadata.hydroscheme["Adiabatic index"][0]
gamma_minus_one = gas_gamma - 1.0
entropies = (
data.gas.energies * (gamma_minus_one) / data.gas.densities ** gamma_minus_one
)
print("Warning: Current entropy profile assumes all gas is ionised")
radii = data.gas.radii
bins = np.logspace(-2, 1, bins)
histogram, _ = np.histogram(a=radii, weights=entropies, bins=bins)
counts, _ = np.histogram(a=radii, weights=np.ones_like(radii), bins=bins)
return histogram / counts # need to get mean value in bin
def nfw(R, halo_data: HaloData):
"""
NFW profile at radius R.
"""
R_s = halo_data.Rvir / halo_data.c
rho_0 = (4 * np.pi * R_s ** 3) / (halo_data.Mvir)
rho_0 *= np.log(1 + halo_data.c) - halo_data.c / (halo_data.c + 1)
rho_0 = 1.0 / rho_0
ratio = R / R_s
return rho_0 / (ratio * (1 + ratio) ** 2)
def create_plot(
data: SimulationParticleData,
halo_data: HaloData,
bins: int,
create_images: bool,
image_data: np.ndarray,
):
"""
Creates the figure and axes objects and plots the data on them.
"""
fig, axes = plt.subplots(2, 3, figsize=(12, 8))
gas_density, bin_edges = get_radial_density_profile(
data.gas.radii, data.gas.masses, bins=bins
)
dm_density, _ = get_radial_density_profile(
data.dark_matter.radii, data.dark_matter.masses, bins=bins
)
temperature = get_radial_temperature_profile(data, bins=bins)
entropy = get_radial_entropy_profile(data, bins=bins)
bin_centers = [0.5 * (x + y) for x, y in zip(bin_edges[:-1], bin_edges[1:])]
nfw_R = np.logspace(-2, 1, bins * 100)
nfw_rho = nfw(nfw_R, halo_data)
axes[0][0].loglog()
axes[0][0].plot(nfw_R, 0.1 * nfw_rho, ls="dashed", color="grey")
axes[0][0].scatter(bin_centers, gas_density)
axes[0][0].set_ylabel(r"$\rho_{\rm gas} (R)$ [$10^{10}$ M$_\odot$ Mpc$^{-3}$]")
axes[0][0].set_xlabel(r"R [Mpc]")
axes[0][0].set_xlim(0.01, 10)
axes[0][1].semilogx()
axes[0][1].scatter(bin_centers, np.log(entropy))
axes[0][1].set_ylabel(
r"Entropy $\log(A$ [K ($10^{10}$ M$_\odot$)$^{2/3}$ Mpc$^{-2}$])"
)
axes[0][1].set_xlabel(r"R [Mpc]")
axes[0][1].set_xlim(0.01, 10)
if create_images:
axes[0][2].imshow(np.log10(image_data))
axes[0][2].set_xticks([])
axes[0][2].set_yticks([])
axes[1][0].loglog()
axes[1][0].scatter(bin_centers, temperature)
axes[1][0].set_ylabel(r"$T_{\rm gas} (R)$ [K]")
axes[1][0].set_xlabel(r"R [Mpc]")
axes[1][0].set_xlim(0.01, 10)
axes[1][1].loglog()
axes[1][1].scatter(bin_centers, dm_density)
axes[1][1].plot(nfw_R, 0.9 * nfw_rho, ls="dashed", color="grey")
axes[1][1].set_ylabel(r"$\rho_{\rm DM} (R)$ [$10^{10}$ M$_\odot$ Mpc$^{-3}$]")
axes[1][1].set_xlabel(r"R [Mpc]")
axes[1][1].set_xlim(0.01, 10)
axes[1][1].text(
0.02,
5,
"$c_{{vir}} = {:2.2f}$\n$R_{{vir}} = {:2.2f}$ Mpc\n$M_{{vir}} = {:2.2f}$ $10^{{10}}$ M$_\odot$".format(
halo_data.c, halo_data.Rvir, halo_data.Mvir
),
va="bottom",
ha="left",
)
axes[1][2].text(
-0.49,
0.7,
"Santa Barbara with $\\gamma={:2.2f}$ in 3D".format(
data.metadata.hydroscheme["Adiabatic index"][0]
),
)
scheme_list = data.metadata.hydroscheme["Scheme"].decode("utf-8").split(" ")
i = 4
while i < len(scheme_list):
scheme_list.insert(i, "\n")