#!/usr/bin/env python3 ############################################################################### # This file is part of SWIFT. # Copyright (c) 2024 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 SymLogNorm, Normalize import numpy as np import unyt import swiftsimio from swiftsimio.visualisation import project_gas from pathlib import Path from makeIC import VELOCITY, ELEMENT_COUNT 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).value.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.metal_mass_fractions columns = getattr(element_mass_fractions, "named_columns", None) for i in range(ELEMENT_COUNT): if columns is None: data.gas.__setattr__( f"element_mass_sp{i}", element_mass_fractions[:, i] * masses ) else: data.gas.__setattr__( f"element_mass_sp{i}", getattr(element_mass_fractions, columns[i]) * masses, ) # Create necessary figures and axes fig = plt.figure(layout="constrained", figsize=(8, 2 * ELEMENT_COUNT + 1)) fig.suptitle( f"Profiles shifted to starting position after t={data.metadata.time:.2f}", fontsize=14, ) fig_ratios, fig_masses = fig.subfigures(1, 2) fig_ratios.suptitle("Mass ratio of elements") fig_masses.suptitle("Surface density in elements") axes_ratios = fig_ratios.subplots(ELEMENT_COUNT, 1, sharex=True, sharey=True) axes_masses = fig_masses.subplots(ELEMENT_COUNT, 1, 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 if ELEMENT_COUNT > 5: thresh = 10 ** math.floor(math.log10(0.5 ** (ELEMENT_COUNT - 2))) norm_ratios = SymLogNorm(vmin=0, vmax=0.21, linthresh=thresh, base=10) norm_masses = SymLogNorm(vmin=0, vmax=0.25, linthresh=thresh, base=10) else: norm_ratios = Normalize(vmin=0, vmax=0.21) norm_masses = Normalize(vmin=0, vmax=0.25) 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, ), ) if columns is None: columns = [f"Species {i + 1}" for i in range(ELEMENT_COUNT)] for i in range(ELEMENT_COUNT): plot_single( axes_masses[i], axes_ratios[i], f"element_mass_sp{i}", columns[i], **plotting_kwargs, ) # Add Colorbars cb_masses = fig_masses.colorbar( ScalarMappable(**imshow_mass_kwargs), orientation="horizontal", shrink=0.75, pad=0.01, ax=axes_masses, ) cb_ratios = fig_ratios.colorbar( ScalarMappable(**imshow_fraction_kwargs), orientation="horizontal", shrink=0.75, pad=0.01, ax=axes_ratios, ) cb_masses.ax.set_xlabel("Surface density (g/cm^2)") cb_ratios.ax.set_xlabel("Mass ratio") # 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!")