#!/usr/bin/env python3 ############################################################################### # This file is part of SWIFT. # Copyright (c) 2022 Mladen Ivkovic (mladen.ivkovic@hotmail.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 ionizing species # ---------------------------------------------------- import os import sys import matplotlib as mpl import swiftsimio from matplotlib import pyplot as plt from matplotlib.colors import SymLogNorm from mpl_toolkits.axes_grid1 import make_axes_locatable import argparse # Parameters users should/may tweak # plot all groups and all photon quantities plot_all_data = True # fancy up the plots a bit? fancy = True # parameters for imshow plots imshow_kwargs = {"origin": "lower", "cmap": "brg"} # parameters for swiftsimio projections projection_kwargs = {"resolution": 512, "parallel": True} # ----------------------------------------------------------------------- mpl.rcParams["text.usetex"] = True # Read in cmdline arg: Are we plotting only one snapshot, or all? plot_all = False def parse_args(): parser = argparse.ArgumentParser() parser.add_argument( "-n", "--snapshot-number", help="Number of snapshot to plot", type=int ) parser.add_argument("-b", "--basename", help="Snapshot basename", default="output") 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 = [] if plot_all: 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) else: fname = snapshot_basename + "_" + str(snapnr).zfill(4) + ".hdf5" if not os.path.exists(fname): print("Didn't find file", fname) quit(1) snaplist.append(fname) return snaplist def set_colorbar(ax, im): """ Adapt the colorbar a bit for axis object and imshow instance """ divider = make_axes_locatable(ax) cax = divider.append_axes("right", size="5%", pad=0.05) plt.colorbar(im, cax=cax) return def plot_ionization(filename): """ Create and save the plot """ print("working on", filename) data = swiftsimio.load(filename) meta = data.metadata global imshow_kwargs imshow_kwargs["extent"] = [ 0.01 * meta.boxsize[0].v, 0.99 * meta.boxsize[0].v, 0.01 * meta.boxsize[1].v, 0.99 * meta.boxsize[1].v, ] mass_map = swiftsimio.visualisation.projection.project_gas( data, project="masses", **projection_kwargs ) data.gas.mXHI = data.gas.ion_mass_fractions.HI * data.gas.masses data.gas.mXHII = data.gas.ion_mass_fractions.HII * data.gas.masses data.gas.mXHeI = data.gas.ion_mass_fractions.HeI * data.gas.masses data.gas.mXHeII = data.gas.ion_mass_fractions.HeII * data.gas.masses data.gas.mXHeIII = data.gas.ion_mass_fractions.HeIII * data.gas.masses mass_weighted_XHImap = swiftsimio.visualisation.projection.project_gas( data, project="mXHI", **projection_kwargs ) mass_weighted_XHIImap = swiftsimio.visualisation.projection.project_gas( data, project="mXHII", **projection_kwargs ) mass_weighted_XHeImap = swiftsimio.visualisation.projection.project_gas( data, project="mXHeI", **projection_kwargs ) mass_weighted_XHeIImap = swiftsimio.visualisation.projection.project_gas( data, project="mXHeII", **projection_kwargs ) mass_weighted_XHeIIImap = swiftsimio.visualisation.projection.project_gas( data, project="mXHeIII", **projection_kwargs ) XHImap = mass_weighted_XHImap / mass_map XHIImap = mass_weighted_XHIImap / mass_map XHeImap = mass_weighted_XHeImap / mass_map XHeIImap = mass_weighted_XHeIImap / mass_map XHeIIImap = mass_weighted_XHeIIImap / mass_map fig = plt.figure(figsize=(20, 8), dpi=200) figname = filename[:-5] + "-mass-fractions.png" ax1 = fig.add_subplot(251) ax2 = fig.add_subplot(252) ax3 = fig.add_subplot(253) ax4 = fig.add_subplot(254) ax5 = fig.add_subplot(255) ax6 = fig.add_subplot(256) ax7 = fig.add_subplot(257) ax8 = fig.add_subplot(258) ax9 = fig.add_subplot(259) ax10 = fig.add_subplot(2, 5, 10) # im0 = ax0.imshow(mass_map.T, **imshow_kwargs) # set_colorbar(ax0, im0) # ax0.set_title("Mass") imshow_kwargs_ions = imshow_kwargs.copy() imshow_kwargs_ions["norm"] = SymLogNorm(vmin=0.0, vmax=1.0, linthresh=1e-3, base=10) im1 = ax1.imshow(XHImap.T, **imshow_kwargs_ions) set_colorbar(ax1, im1) ax1.set_title("HI Mass Fraction") im2 = ax2.imshow(XHIImap.T, **imshow_kwargs_ions) set_colorbar(ax2, im2) ax2.set_title("HII Mass Fraction") im3 = ax3.imshow(XHeImap.T, **imshow_kwargs_ions) set_colorbar(ax3, im3) ax3.set_title("HeI Mass Fraction") im4 = ax4.imshow(XHeIImap.T, **imshow_kwargs_ions) set_colorbar(ax4, im4) ax4.set_title("HeII Mass Fraction") im5 = ax5.imshow(XHeIIImap.T, **imshow_kwargs_ions) set_colorbar(ax5, im5) ax5.set_title("HeIII Mass Fraction") im6 = ax6.imshow(mass_weighted_XHImap.T, **imshow_kwargs) set_colorbar(ax6, im6) ax6.set_title("HI Mass") im7 = ax7.imshow(mass_weighted_XHIImap.T, **imshow_kwargs) set_colorbar(ax7, im7) ax7.set_title("HII Mass") im8 = ax8.imshow(mass_weighted_XHeImap.T, **imshow_kwargs) set_colorbar(ax8, im8) ax8.set_title("HeI Mass") im9 = ax9.imshow(mass_weighted_XHeIImap.T, **imshow_kwargs) set_colorbar(ax9, im9) ax9.set_title("HeII Mass") im10 = ax10.imshow(mass_weighted_XHeIIImap.T, **imshow_kwargs) set_colorbar(ax10, im10) ax10.set_title("HeIII Mass") for ax in [ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8, ax9, ax10]: ax.set_xlabel("[kpc]") ax.set_ylabel("[kpc]") title = filename.replace("_", r"\_") # exception handle underscore for latex if meta.cosmology is not None: title += ", $z$ = {0:.2e}".format(meta.z) title += ", $t$ = {0:.2e}".format(meta.time.to("Myr")) fig.suptitle(title) plt.tight_layout() plt.savefig(figname) plt.close() return if __name__ == "__main__": # Get command line arguments args = parse_args() if args.snapshot_number: plot_all = False snapnr = int(args.snapshot_numbeR) else: plot_all = True snapshot_base = args.basename snaplist = get_snapshot_list(snapshot_base) for f in snaplist: plot_ionization(f)