"""Demonstration script to plot the orbit of one galaxy in HYDRO and DM.
It demonstrates reading data from the pre-compiled galaxy tables, matching
subhaloes between the HYDRO and DM-ONLY versions of a simulation, and
working with the interpolated high-time-resolution orbit tables.
"""
# Import required packages
import numpy as np
import hydrangea as hy
import matplotlib.pyplot as plt
import matplotlib.colors as colors # For truncated grey-scale color map
from pdb import set_trace
# Set script parameters
sim_index = 0 # Which simulation do we want?
galaxy_id = 123 # Galaxy ID (in HYDRO) to plot
plotloc = 'orbit_plot.png' # Where to save the output plot?
# Set up the simulation object for the run we're working with, to easily
# get the relevant file paths
sim_hydro = hy.Simulation(index=sim_index)
sim_dm_only = hy.Simulation(index=sim_index, sim_type='DM')
# Prepare the plot
fig = plt.figure(figsize=(4*0.8/0.6, 4))
ax1 = fig.add_axes((0.15, 0.15, 0.65, 0.8))
# ------------------------------------------------------------------
# Part I: Find the corresponding galaxy ID in the DM-only simulation
# ------------------------------------------------------------------
# Find the corresponding galaxy ID in the DM simulation. For this, we find
# the snapshot in which the galaxy has its maximum total mass, and then
# match its subhalo
mtot_hydro = hy.hdf5.read_data(sim_hydro.fgt_loc, 'Msub',
read_index=galaxy_id)
snap_max = np.argmax(mtot_hydro)
print(f"Galaxy {galaxy_id} reaches its max subhalo mass in snap {snap_max}.")
# Now we find the subhalo index of the galaxy in this snapshot
shi_hydro_max = hy.hdf5.read_data(
sim_hydro.fgt_loc, 'SHI', read_index=galaxy_id)[snap_max]
# And now we find the matching DM-only subhalo in the same snapshot
shi_dm_max = hy.hdf5.read_data(sim_hydro.sh_extra_loc,
f'Snapshot_{snap_max:03d}/MatchInDM',
read_index=shi_hydro_max)
# And find the DM-only galaxy ID
galaxy_id_dmo = hy.hdf5.read_data(sim_dm_only.spider_loc,
f'Subhalo/Snapshot_{snap_max:03d}/Galaxy',
read_index=shi_dm_max)
print(f"Galaxy {galaxy_id} in HYDRO matched to {galaxy_id_dmo} in DM-only")
# ---------------------------------------------------------------
# Part II: Extract the orbits of the galaxies in both simulations
# ---------------------------------------------------------------
# Find the central of the galaxy at z = 0 in both simulations
central_id_hydro = hy.hdf5.read_data(sim_hydro.fgt_loc, 'CenGal',
read_index=galaxy_id)[29]
central_id_dmo = hy.hdf5.read_data(sim_dm_only.fgt_loc, 'CenGal',
read_index=galaxy_id_dmo)[29]
# We now extract the high-time-resolution orbits of all four galaxies
def get_galaxy_paths(galaxy_id, sim, load_times=False):
"""Extract the galaxy paths for a specified galaxy in a simulation."""
# First, we must look up the galaxy in the interpolation list
interp_id = hy.hdf5.read_data(sim.interpolation_loc, 'GalaxyRevIndex',
read_index=galaxy_id)
# If the galaxy is not in the interpolation list, quit
if interp_id < 0:
print(f"Could not locate galaxy {galaxy_id} in interpolation list.")
return
# Extract the interpolated orbit
orbit_pos = hy.hdf5.read_data(sim.interpolation_loc,
'InterpolatedPositions',
read_index=interp_id)
# If needed, also load the interpolation times
if not load_times:
return orbit_pos
orbit_times = hy.hdf5.read_data(sim.interpolation_loc,
'InterpolationTimes')
return orbit_pos, orbit_times
orbit_hydro, orbit_times = get_galaxy_paths(galaxy_id, sim_hydro,
load_times=True)
orbit_hydro -= get_galaxy_paths(central_id_hydro, sim_hydro)
orbit_dmo = (get_galaxy_paths(galaxy_id_dmo, sim_dm_only)
- get_galaxy_paths(central_id_dmo, sim_dm_only))
# -------------------------
# Part III: Plot the orbits
# -------------------------
# Create a truncated grey color map (Greys_r becomes white towards 1)
# [see https://stackoverflow.com/questions/18926031/]
cmap_greys = colors.LinearSegmentedColormap.from_list(
'trunc(Greys_r, 0.0, 0.8)', plt.cm.Greys_r(np.linspace(0.0, 0.8, 1000)))
plt.sca(ax1)
sc_dm = plt.scatter(orbit_dmo[0, :], orbit_dmo[1, :], c=orbit_times,
edgecolor='none', vmin=0, vmax=14.0, cmap=cmap_greys,
s=10)
sc_hy = plt.scatter(orbit_hydro[0, :], orbit_hydro[1, :], c=orbit_times,
edgecolor='none', vmin=0, vmax=14.0, cmap=plt.cm.viridis,
s=10)
# Some embellishments on the plot
ax1.set_xlim((-0.8, 0.8))
ax1.set_ylim((-0.8, 0.8))
ax1.set_xlabel(r'$\Delta x$ [pMpc]')
ax1.set_ylabel(r'$\Delta y$ [pMpc]')
# Add colour bars on the side
ax2 = fig.add_axes([0.81, 0.15, 0.03, 0.8])
ax2.set_xticks([])
ax2.set_yticks([])
cbar_dm = plt.colorbar(sc_dm, cax=ax2, orientation='vertical')
cbar_dm.set_ticks([])
ax3 = fig.add_axes([0.84, 0.15, 0.03, 0.8])
ax3.set_xticks([])
ax3.set_yticks([])
cbar_hy = plt.colorbar(sc_hy, cax=ax3, orientation='vertical')
fig.text(0.94, 0.55, 'Age of the Universe [Gyr]',
rotation=90.0, va='center', ha='left')
# Save the figure
plt.show
plt.savefig(plotloc, transparent=False, dpi=150)