""" Plots the mean temperature. """ import matplotlib.pyplot as plt from swiftsimio import load from unyt import Myr, K, mh, cm, erg import numpy as np def get_data_dump(metadata): """ Gets a big data dump from the SWIFT metadata """ try: viscosity = metadata.viscosity_info except: viscosity = "No info" try: diffusion = metadata.diffusion_info except: diffusion = "No info" output = ( "$\\bf{SWIFT}$\n" + metadata.code_info + "\n\n" + "$\\bf{Compiler}$\n" + metadata.compiler_info + "\n\n" + "$\\bf{Hydrodynamics}$\n" + metadata.hydro_info + "\n\n" + "$\\bf{Cooling}$\n" + metadata.subgrid_scheme["Cooling Model"].decode("utf-8") + "\n\n" + "$\\bf{Viscosity}$\n" + viscosity + "\n\n" + "$\\bf{Diffusion}$\n" + diffusion ) return output def setup_axes(): """ Sets up the axes and returns fig, ax """ fig, ax = plt.subplots(nrows=2, ncols=2, figsize=(8, 8)) ax = ax.flatten() ax[0].semilogy() ax[2].semilogy() ax[1].set_xlabel("Simulation time elapsed [Myr]") ax[2].set_xlabel("Simulation time elapsed [Myr]") ax[0].set_ylabel("Temperature of Universe [K]") ax[1].set_ylabel("Physical Density of Universe [$n_H$ cm$^{-3}$]") ax[2].set_ylabel("Physical Energy [erg]") for a in ax[:-1]: a.set_xlim(0, 100) return fig, ax def get_data(handle: float, n_snaps: int): """ Returns the elapsed simulation time, temperature, and density """ t0 = 0.0 times = [] temps = [] densities = [] energies = [] radiated_energies = [] for snap in range(n_snaps): try: data = load(f"data/{handle}_{snap:04d}.hdf5") if snap == 0: t0 = data.metadata.t.to(Myr).value times.append(data.metadata.t.to(Myr).value - t0) temps.append(np.mean(data.gas.temperatures.to(K).value)) densities.append( np.mean(data.gas.densities.to(mh / cm ** 3).value) / (data.metadata.scale_factor ** 3) ) try: energies.append( np.mean( (data.gas.internal_energies * data.gas.masses) .astype(np.float64) .to(erg) .value ) * data.metadata.scale_factor ** (2) ) radiated_energies.append( np.mean(data.gas.radiated_energy.astype(np.float64).to(erg).value) ) except AttributeError: continue except OSError: continue return times, temps, densities, energies, radiated_energies def get_n_steps(timesteps_filename: str) -> int: """ Reads the number of steps from the timesteps file. """ data = np.genfromtxt(timesteps_filename) return int(data[-1][0]) def plot_single_data( handle: str, n_snaps: int, label: str, ax: plt.axes, n_steps: int = 0, run: int = 0 ): """ Takes the a single simulation and plots it on the axes. """ data = get_data(handle, n_snaps) ax[0].plot(data[0], data[1], label=label, marker="o", ms=2, c=f"C{run}") ax[1].plot( data[0], data[2], label=f"Steps: {n_steps}", marker="o", ms=2, c=f"C{run}" ) if run == 0: label_energy = "Particle Energy" label_radiated = "Radiated Energy" label_sum = "Sum" else: label_energy = "" label_radiated = "" label_sum = "" ax[2].plot(data[0], data[3], label=label_energy, ls="dotted", color=f"C{run}") # ax[2].plot(data[0], data[4], label=label_radiated, ls="dashed", C=f"C{run}") # ax[2].plot( # data[0], [x + y for x, y in zip(*data[3:5])], label=label_sum, C=f"C{run}" # ) return def make_plot(handles, names, timestep_names, n_snaps=100): """ Makes the whole plot and returns the fig, ax objects. """ fig, ax = setup_axes() run = 0 for handle, name, timestep_name in zip(handles, names, timestep_names): try: n_steps = get_n_steps(timestep_name) except: n_steps = "Unknown" plot_single_data(handle, n_snaps, name, ax, n_steps=n_steps, run=run) run += 1 ax[0].legend() ax[1].legend() ax[2].legend() info_axis = ax[-1] try: info = get_data_dump(load(f"data/{handles[0]}_0000.hdf5").metadata) info_axis.text( 0.5, 0.45, info, ha="center", va="center", fontsize=7, transform=info_axis.transAxes, ) except OSError: pass info_axis.axis("off") return fig, ax if __name__ == "__main__": """ Plot everything! """ handles = [ "redshift_dependence_no_z", "redshift_dependence_low_z", "redshift_dependence_high_z", ] names = ["No Cosmology", "Low Redshift ($z=0.01$)", "High Redshift ($z=1$)"] timestep_names = ["timesteps_no.txt", "timesteps_low.txt", "timesteps_high.txt"] fig, ax = make_plot(handles, names, timestep_names) fig.tight_layout(pad=0.5) fig.savefig("redshift_dependence.png", dpi=300)