#!/usr/bin/env python3 ############################################################################### # This file is part of SWIFT. # Copyright (c) 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 . # ############################################################################## import os import sys import argparse import matplotlib as mpl import numpy as np import swiftsimio import unyt from matplotlib import pyplot as plt # Parameters users should/may tweak snapshot_base = "output" # Snapshot basename plot_physical_quantities = True # Plot physical or comoving quantities time_first = 0 # Time of first snapshot def parse_args(): parser = argparse.ArgumentParser() parser.add_argument( "-z", "--redshift", help="Redshift domain to plot advection for", default="high" ) args = parser.parse_args() return args 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) snaplist = sorted(snaplist) if len(snaplist) == 0: print(f"No snapshots with base {snapshot_basename} found!") sys.exit(1) return snaplist def plot_param_over_time(snapshot_list, param="energy density", redshift_domain="high"): print(f"Now plotting {param} over time") # Grab number of photon groups data = swiftsimio.load(snapshot_list[0]) meta = data.metadata ngroups = int(meta.subgrid_scheme["PhotonGroupNumber"][0]) # Number of rows and columns nrows = 1 + int(plot_physical_quantities) ncols = ngroups # Create figure and axes fig, axs = plt.subplots(nrows, ncols, figsize=(5.04 * ncols, 5.4 * nrows), dpi=200) # Iterate over all photon groups for n in range(ngroups): # Arrays to keep track of plot_param and scale factor plot_param = [[], []] scale_factor = [] analytic_exponent = [0, 0] # Functions to convert between scale factor and redshift a2z = lambda a: 1 / a - 1 z2a = lambda z: 1 / (z + 1) for file in snapshot_list: data = swiftsimio.load(file) meta = data.metadata # Read comoving variables energy = getattr(data.gas.photon_energies, f"group{n+1}") mass = data.gas.masses rho = data.gas.densities volume = mass / rho energy_density = energy / volume if plot_physical_quantities: # The SWIFT cosmology module assumes 3-dimensional lengths and volumes, # so multiply by a to get the correct relations physical_energy_density = ( energy_density.to_physical() * meta.scale_factor ) physical_energy = energy.to_physical() if param == "energy density": plot_param[1].append( 1 * physical_energy_density.sum() / physical_energy_density.shape[0] ) analytic_exponent[1] = -2.0 elif param == "total energy": plot_param[1].append(1 * physical_energy.sum()) analytic_exponent[1] = 0.0 if param == "energy density": plot_param[0].append(1 * energy_density.sum() / energy_density.shape[0]) analytic_exponent[0] = 0.0 elif param == "total energy": plot_param[0].append(1 * energy.sum()) analytic_exponent[0] = 0.0 scale_factor.append(meta.scale_factor) if param == "energy density": titles = ["Comoving energy density", "Physical energy density $\\times a$"] ylabel = "Average energy density" figname = "output_energy_density_over_time" elif param == "total energy": titles = ["Comoving total energy", "Physical total energy"] ylabel = "Total energy" figname = "output_total_energy_over_time" # Analytic scale factor analytic_scale_factor = np.linspace(min(scale_factor), max(scale_factor), 1000) for i in range(nrows): ax = axs[i, n] ax.scatter(scale_factor, plot_param[i], label="Simulation") # Analytic scale factor relation analytic = analytic_scale_factor ** analytic_exponent[i] # Scale solution to correct offset analytic = analytic / analytic[0] * plot_param[i][0] ax.plot( analytic_scale_factor, analytic, c="r", label=f"Analytic solution $\propto a^{{{analytic_exponent[i]}}}$", ) ax.legend() ax.set_title(titles[i] + f" group {n+1}") ax.set_xlabel("Scale factor") secax = ax.secondary_xaxis("top", functions=(a2z, z2a)) secax.set_xlabel("Redshift") ax.yaxis.get_offset_text().set_position((-0.05, 1)) if analytic_exponent[i] == 0.0: ax.set_ylim(plot_param[i][0] * 0.95, plot_param[i][0] * 1.05) ylabel_scale = "" else: ylabel_scale = "$\\times a$" if n == 0: units = plot_param[i][0].units.latex_representation() ax.set_ylabel(f"{ylabel} [${units}$] {ylabel_scale}") plt.tight_layout() plt.savefig(f"{figname}-{redshift_domain}.png") plt.close() if __name__ in ("__main__"): # Get command line args args = parse_args() domain = args.redshift.lower() if domain in ("low", "l", "low_redshift", "low redshift", "low-redshift"): redshift_domain = "low_redshift" elif domain in ( "medium", "m", "medium_redshift", "medium redshift", "medium-redshift", ): redshift_domain = "medium_redshift" elif domain in ("high", "h", "high_redshift", "high redshift", "high-redshift"): redshift_domain = "high_redshift" else: print("Redshift domain not recognised!") sys.exit(1) snaplist = get_snapshot_list(snapshot_base + f"_{redshift_domain}") for param in ["energy density", "total energy"]: plot_param_over_time(snaplist, param, redshift_domain)