#!/usr/bin/env python3 ############################################################################### # This file is part of SWIFT. # Copyright (c) 2022 Mladen Ivkovic (mladen.ivkovic@hotmail.com) # 2024 Stan Verhoeve (s06verhoeve@gmail.com) # This program is free software: you can redistribute it and/or modify # it under the terms of the GNU Lesser General Public License as published # by the Free Software Foundation, either version 3 of the License, or # (at your option) any later version. # # This program is distributed in the hope that it will be useful, # but WITHOUT ANY WARRANTY; without even the implied warranty of # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the # GNU General Public License for more details. # # You should have received a copy of the GNU Lesser General Public License # along with this program. If not, see . # ############################################################################## # ------------------------------------------- # Plot the gas temperature, mean molecular # weight, and mass fractions # ------------------------------------------- import copy import os import numpy as np import swiftsimio import unyt from matplotlib import pyplot as plt # arguments for plots of results plotkwargs = {"alpha": 0.5} # arguments for plot of references referenceplotkwargs = {"color": "grey", "lw": 4, "alpha": 0.6} # arguments for legends legendprops = {"size": 8} # snapshot basenames snapshot_base = "output" # Reference file name reference_file = "CosmoRTCoolingTestReference.txt" # Plot time on x-axis plot_time = False # ----------------------------------------------------------------------- energy_units = unyt.Msun * unyt.kpc ** 2 / unyt.kyr ** 2 mass_units = unyt.Msun length_units = unyt.kpc def mean_molecular_weight(XH0, XHp, XHe0, XHep, XHepp): """ Determines the mean molecular weight for given mass fractions of hydrogen: XH0 H+: XHp He: XHe0 He+: XHep He++: XHepp returns: mu: mean molecular weight [in atomic mass units] NOTE: to get the actual mean mass, you still need to multiply it by m_u, as is tradition in the formulae """ # 1/mu = sum_j X_j / A_j * (1 + E_j) # A_H = 1, E_H = 0 # A_Hp = 1, E_Hp = 1 # A_He = 4, E_He = 0 # A_Hep = 4, E_Hep = 1 # A_Hepp = 4, E_Hepp = 2 one_over_mu = XH0 + 2 * XHp + 0.25 * XHe0 + 0.5 * XHep + 0.75 * XHepp return 1.0 / one_over_mu def gas_temperature(u, mu, gamma): """ Compute the gas temperature given the specific internal energy u and the mean molecular weight mu """ # Using u = 1 / (gamma - 1) * p / rho # and p = N/V * kT = rho / (mu * m_u) * kT T = u * (gamma - 1) * mu * unyt.atomic_mass_unit / unyt.boltzmann_constant return T def get_snapshot_list(snapshot_basename="output"): """ Find the snapshot(s) that are to be plotted and return their names as list """ snaplist = [] dirlist = os.listdir() for f in dirlist: if f.startswith(snapshot_basename) and f.endswith("hdf5"): snaplist.append(f) if len(snaplist) == 0: raise FileNotFoundError( "Didn't find any snapshots with basename '" + snapshot_basename + "'" ) snaplist = sorted(snaplist) return snaplist def get_ion_mass_fractions(swiftsimio_loaded_data): """ Returns the ion mass fractions according to the used scheme. swiftsimio_loaded_data: the swiftsimio.load() object """ data = swiftsimio_loaded_data meta = data.metadata gas = data.gas with_rt = True scheme = None try: scheme = str(meta.subgrid_scheme["RT Scheme"].decode("utf-8")) except KeyError: # allow to read in solutions with only cooling, without RT with_rt = False if with_rt: if scheme.startswith("GEAR M1closure"): imf = data.gas.ion_mass_fractions elif scheme.startswith("SPH M1closure"): # atomic mass mamu = { "e": 0.0, "HI": 1.0, "HII": 1.0, "HeI": 4.0, "HeII": 4.0, "HeIII": 4.0, } mass_function_hydrogen = data.gas.rt_element_mass_fractions.hydrogen imf = copy.deepcopy(data.gas.rt_species_abundances) named_columns = data.gas.rt_species_abundances.named_columns for column in named_columns: # abundance is in n_X/n_H unit. We convert it to mass fraction by multipling mass fraction of H mass_function = ( getattr(data.gas.rt_species_abundances, column) * mass_function_hydrogen * mamu[column] ) setattr(imf, column, mass_function) else: raise ValueError("Unknown scheme", scheme) else: # try to find solutions for cooling only runs imf = { "HI": gas.hi[:], "HII": gas.hii[:], "HeI": gas.he_i[:], "HeII": gas.he_ii[:], "HeIII": gas.he_iii[:], } return imf def get_snapshot_data(snaplist): """ Extract the relevant data from the list of snapshots. Returns: numpy arrays of: time temperatures mean molecular weights mass fractions """ nsnaps = len(snaplist) firstdata = swiftsimio.load(snaplist[0]) with_rt = True try: ngroups = int(firstdata.metadata.subgrid_scheme["PhotonGroupNumber"][0]) except KeyError: # allow to read in solutions with only cooling, without RT with_rt = False # Read internal units unit_system = firstdata.metadata.units mass_units = unit_system.mass length_units = unit_system.length time_units = unit_system.time # Units derived from base velocity_units = length_units / time_units energy_units = mass_units * velocity_units ** 2 density_units = mass_units / length_units ** 3 scale_factors = np.zeros(nsnaps) times = np.zeros(nsnaps) * unyt.Myr temperatures = np.zeros(nsnaps) * unyt.K mean_molecular_weights = np.zeros(nsnaps) mass_fractions = np.zeros((nsnaps, 5)) internal_energies = np.zeros(nsnaps) * energy_units specific_internal_energies = np.zeros(nsnaps) * energy_units / mass_units densities = np.zeros(nsnaps) * density_units if with_rt: photon_energies = np.zeros((ngroups, nsnaps)) * energy_units else: photon_energies = None for i, snap in enumerate(snaplist): data = swiftsimio.load(snap) gamma = data.gas.metadata.gas_gamma[0] time = data.metadata.time.copy() scale_factor = data.metadata.scale_factor.copy() gas = data.gas u = gas.internal_energies[:].to(energy_units / mass_units) u.convert_to_physical() rho = gas.densities rho.convert_to_physical() rho.to(density_units) rho.convert_to_physical() masses = gas.masses[:].to(mass_units) masses.convert_to_physical() imf = get_ion_mass_fractions(data) mu = mean_molecular_weight(imf.HI, imf.HII, imf.HeI, imf.HeII, imf.HeIII) mu.convert_to_physical() T = gas_temperature(u, mu, gamma).to("K") um = u.to(energy_units / mass_units) * masses um.to(energy_units) scale_factors[i] = scale_factor times[i] = time.to("Myr") temperatures[i] = np.mean(T) mean_molecular_weights[i] = np.mean(mu) internal_energies[i] = np.mean(um) specific_internal_energies[i] = u.to(energy_units / mass_units).sum() / len(u) densities[i] = rho.sum() / len(rho) mass_fractions[i, 0] = np.mean(imf.HI) mass_fractions[i, 1] = np.mean(imf.HII) mass_fractions[i, 2] = np.mean(imf.HeI) mass_fractions[i, 3] = np.mean(imf.HeII) mass_fractions[i, 4] = np.mean(imf.HeIII) if with_rt: for g in range(ngroups): en = getattr(data.gas.photon_energies, "group" + str(g + 1)) en = en[:].to_physical().to(energy_units) photon_energies[g, i] = en.sum() / en.shape[0] return ( scale_factors, times, temperatures, mean_molecular_weights, mass_fractions, internal_energies, photon_energies, specific_internal_energies, densities, ) def get_reference_data(reference_file="CosmoRTCoolingTestReference.txt"): # Grab lines containing units with open(reference_file) as f: # Skip first line f.readline() mass_units_line = f.readline() length_units_line = f.readline() velocity_units_line = f.readline() # Extract numbers from lines mass_units = float(mass_units_line[18:-4]) * unyt.g length_units = float(length_units_line[20:-5]) * unyt.cm velocity_units = float(velocity_units_line[22:-7]) * unyt.cm / unyt.s # Derived units energy_units = mass_units * velocity_units ** 2 density_units = mass_units / length_units ** 3 # Read in data refdata = np.loadtxt(reference_file, dtype=np.float64) a_ref = refdata[:, 1] t_ref = refdata[:, 3] * unyt.yr dt_ref = refdata[:, 4] * unyt.yr T_ref = refdata[:, 5] * unyt.K mu_ref = refdata[:, 6] rho_ref = refdata[:, 7] * density_units rhoHI_ref = refdata[:, 8] * density_units rhoHII_ref = refdata[:, 9] * density_units rhoHeI_ref = refdata[:, 10] * density_units rhoHeII_ref = refdata[:, 11] * density_units rhoHeIII_ref = refdata[:, 12] * density_units rhoe_ref = refdata[:, 13] * density_units u_ref = refdata[:, -1] * energy_units / mass_units mass_fraction_ref = np.empty((t_ref.shape[0], 5)) mass_fraction_ref[:, 0] = rhoHI_ref / rho_ref mass_fraction_ref[:, 1] = rhoHII_ref / rho_ref mass_fraction_ref[:, 2] = rhoHeI_ref / rho_ref mass_fraction_ref[:, 3] = rhoHeII_ref / rho_ref mass_fraction_ref[:, 4] = rhoHeIII_ref / rho_ref return ( a_ref, t_ref, T_ref, mu_ref, mass_fraction_ref, dt_ref, rho_ref, rhoHI_ref, rhoHII_ref, rhoHeI_ref, rhoHeII_ref, rhoHeIII_ref, rhoe_ref, u_ref, ) if __name__ == "__main__": # Get list of shapshots snaplist = get_snapshot_list(snapshot_base) # Read snapshot data a, t, T, mu, mass_fraction, u, photon_energies, us, rho = get_snapshot_data( snaplist ) with_rt = photon_energies is not None if with_rt: ngroups = photon_energies.shape[0] # Read reference solution data a_ref, t_ref, T_ref, mu_ref, mass_fraction_ref, *__ = get_reference_data( reference_file ) # Convert t_ref to Myr t_ref.convert_to_units("Myr") # Translate snapshot times to start at t=0 t -= t[0] # Grab x-coordinate if plot_time: xcoords = t xcoords_ref = t_ref xlabel = "Time [Myr]" xscale = "log" xlims = (0.1, max(t)) else: xcoords = a xcoords_ref = a_ref xlabel = "Scale factor [1]" xscale = "linear" xlims = (min(a), max(a)) # ------------------ # Plot figures # ------------------ fig = plt.figure(figsize=(8, 8), dpi=300) ax1 = fig.add_subplot(2, 2, 1) ax2 = fig.add_subplot(2, 2, 2) ax3 = fig.add_subplot(2, 2, 3) ax4 = fig.add_subplot(2, 2, 4) ax1.set_xscale(xscale) ax2.set_xscale(xscale) ax3.set_xscale(xscale) ax4.set_xscale(xscale) ax1.set_xlim(*xlims) ax2.set_xlim(*xlims) ax3.set_xlim(*xlims) ax4.set_xlim(*xlims) ax1.semilogy(xcoords_ref, T_ref, label="reference", **referenceplotkwargs) ax1.semilogy(xcoords, T, label="obtained results") ax1.set_yscale("log") ax1.set_xlabel(xlabel) ax1.set_ylabel("gas temperature [K]") ax1.legend(prop=legendprops) ax1.grid() ax2.plot(xcoords_ref, mu_ref, label="reference", **referenceplotkwargs) ax2.plot(xcoords, mu, label="obtained results") ax2.set_xlabel(xlabel) ax2.set_ylabel("mean molecular weight") ax2.legend(prop=legendprops) ax2.grid() total_mass_fraction = np.sum(mass_fraction, axis=1) ax3.plot(xcoords, total_mass_fraction, "k", label="total", ls="-") ax3.plot( xcoords_ref[1:], mass_fraction_ref[1:, 0], label="reference", **referenceplotkwargs, zorder=0, ) ax3.plot(xcoords_ref[1:], mass_fraction_ref[1:, 1], **referenceplotkwargs, zorder=0) ax3.plot(xcoords_ref[1:], mass_fraction_ref[1:, 2], **referenceplotkwargs, zorder=0) ax3.plot(xcoords_ref[1:], mass_fraction_ref[1:, 3], **referenceplotkwargs, zorder=0) ax3.plot(xcoords_ref[1:], mass_fraction_ref[1:, 4], **referenceplotkwargs, zorder=0) ax3.plot(xcoords, mass_fraction[:, 0], label="HI", ls=":", **plotkwargs, zorder=1) ax3.plot(xcoords, mass_fraction[:, 1], label="HII", ls="-.", **plotkwargs, zorder=1) ax3.plot(xcoords, mass_fraction[:, 2], label="HeI", ls=":", **plotkwargs, zorder=1) ax3.plot( xcoords, mass_fraction[:, 3], label="HeII", ls="-.", **plotkwargs, zorder=1 ) ax3.plot( xcoords, mass_fraction[:, 4], label="HeIII", ls="--", **plotkwargs, zorder=1 ) ax3.legend(loc="upper right", prop=legendprops) ax3.set_xlabel(xlabel) ax3.set_ylabel("gas mass fractions [1]") ax3.grid() if with_rt: u.convert_to_units(energy_units) photon_energies.convert_to_units(energy_units) tot_energy = u + np.sum(photon_energies, axis=0) ax4.plot( xcoords, tot_energy, label=f"total energy budget", color="k", ls="--", **plotkwargs, ) for g in range(ngroups): ax4.plot( xcoords, photon_energies[g, :], label=f"photon energies group {g + 1}", **plotkwargs, ) ax4.plot(xcoords, u, label="gas internal energy", **plotkwargs) ax4.set_xlabel(xlabel) ax4.set_ylabel( r"energy budget [$" + u.units.latex_representation() + "$]", usetex=True ) ax4.legend(prop=legendprops) ax4.grid() plt.tight_layout() # plt.show() plt.savefig("cooling_test.png")