#!/usr/bin/env python3 ############################################################################### # This file is part of SWIFT. # Copyright (c) 2023 Yolan Uyttenhove (yolan.uyttenhove@ugent.be) # # 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 math import matplotlib.pyplot as plt from matplotlib.cm import ScalarMappable from matplotlib.colors import Normalize import numpy as np import unyt import swiftsimio from swiftsimio.visualisation import project_gas from pathlib import Path from makeIC import VELOCITY def plot_single(ax_mass, ax_fraction, name, title, data, mass_map, kwargs_inner): mass_weighted_map = project_gas(data, project=name, **kwargs_inner["projection"]) ax_mass.imshow(mass_weighted_map.in_cgs().value.T, **kwargs_inner["imshow_mass"]) ax_mass.set_title(title) ax_fraction.imshow( (mass_weighted_map / mass_map).T, **kwargs_inner["imshow_fraction"] ) ax_fraction.set_title(title) ax_mass.axis("off") ax_fraction.axis("off") def plot_all(fname, savename): data = swiftsimio.load(fname) # Shift coordinates to starting position velocity = 0.5 * math.sqrt(2) * VELOCITY * np.array([1, 1, 0]) * unyt.cm / unyt.s data.gas.coordinates -= data.metadata.time * velocity # Add mass weighted element mass fractions to gas dataset masses = data.gas.masses element_mass_fractions = data.gas.element_mass_fractions data.gas.element_mass_H = element_mass_fractions.hydrogen * masses data.gas.element_mass_He = element_mass_fractions.helium * masses data.gas.element_mass_C = element_mass_fractions.carbon * masses data.gas.element_mass_N = element_mass_fractions.nitrogen * masses data.gas.element_mass_O = element_mass_fractions.oxygen * masses data.gas.element_mass_Ne = element_mass_fractions.neon * masses data.gas.element_mass_Mg = element_mass_fractions.magnesium * masses data.gas.element_mass_Si = element_mass_fractions.silicon * masses data.gas.element_mass_Fe = element_mass_fractions.iron * masses data.gas.total_metal_mass = data.gas.metal_mass_fractions * masses # Create necessary figures and axes fig = plt.figure(layout="constrained", figsize=(16, 8)) fig.suptitle( f"Profiles shifted to starting position after t={data.metadata.time:.2f}", fontsize=14, ) fig_ratios, fig_masses = fig.subfigures(2, 1) fig_ratios.suptitle("Mass ratio of elements") fig_masses.suptitle("Surface density in elements") axes_ratios = fig_ratios.subplots(2, 5, sharex=True, sharey=True) axes_masses = fig_masses.subplots(2, 5, sharex=True, sharey=True) # parameters for swiftsimio projections projection_kwargs = { "region": np.array([0, 2, 0, 1, 0, 1]) * unyt.cm, "resolution": 500, "parallel": True, } # Parameters for imshow norm_ratios = Normalize(0, 0.95) norm_masses = Normalize(0, 0.9) imshow_fraction_kwargs = dict(norm=norm_ratios, cmap="rainbow") imshow_mass_kwargs = dict(norm=norm_masses, cmap="turbo") # Plot the quantities mass_map = project_gas(data, project="masses", **projection_kwargs) # Parameters for plotting: plotting_kwargs = dict( data=data, mass_map=mass_map, kwargs_inner=dict( projection=projection_kwargs, imshow_mass=imshow_mass_kwargs, imshow_fraction=imshow_fraction_kwargs, ), ) plot_single( axes_masses[0][0], axes_ratios[0][0], "total_metal_mass", "All metals", **plotting_kwargs, ) plot_single( axes_masses[0][1], axes_ratios[0][1], "element_mass_H", "Hydrogen", **plotting_kwargs, ) plot_single( axes_masses[0][2], axes_ratios[0][2], "element_mass_He", "Helium", **plotting_kwargs, ) plot_single( axes_masses[0][3], axes_ratios[0][3], "element_mass_C", "Carbon", **plotting_kwargs, ) plot_single( axes_masses[0][4], axes_ratios[0][4], "element_mass_N", "Nitrogen", **plotting_kwargs, ) plot_single( axes_masses[1][0], axes_ratios[1][0], "element_mass_O", "Oxygen", **plotting_kwargs, ) plot_single( axes_masses[1][1], axes_ratios[1][1], "element_mass_Ne", "Neon", **plotting_kwargs, ) plot_single( axes_masses[1][2], axes_ratios[1][2], "element_mass_Mg", "Magnesium", **plotting_kwargs, ) plot_single( axes_masses[1][3], axes_ratios[1][3], "element_mass_Si", "Silicon", **plotting_kwargs, ) plot_single( axes_masses[1][4], axes_ratios[1][4], "element_mass_Fe", "Iron", **plotting_kwargs, ) # Add Colorbars cb_masses = fig_masses.colorbar( ScalarMappable(**imshow_mass_kwargs), shrink=0.75, pad=0.01, ax=axes_masses ) cb_ratios = fig_ratios.colorbar( ScalarMappable(**imshow_fraction_kwargs), shrink=0.75, pad=0.01, ax=axes_ratios ) cb_masses.ax.set_ylabel("Surface density (g/cm^2)", rotation=270, labelpad=15) cb_ratios.ax.set_ylabel("Mass ratio", rotation=270, labelpad=15) # Save output fig.savefig(savename, dpi=300) if __name__ == "__main__": cwd = Path(__file__).parent print("Plotting metals for output_0000.hdf5...") plot_all(cwd / "output_0000.hdf5", savename=cwd / "metal_advection_0000.png") print("Plotting metals for output_0001.hdf5...") plot_all(cwd / "output_0001.hdf5", savename=cwd / "metal_advection_0001.png") print("Done!")