diff --git a/examples/ChemistryTests/MetalAdvectionTestEAGLE/advect_metals.yml b/examples/ChemistryTests/MetalAdvectionTestEAGLE/advect_metals.yml
index 352dd55f2e147effead945ec8cb1698d1ad82351..4c6c184e5d5c50a567413efe92c9056e0db84ce6 100644
--- a/examples/ChemistryTests/MetalAdvectionTestEAGLE/advect_metals.yml
+++ b/examples/ChemistryTests/MetalAdvectionTestEAGLE/advect_metals.yml
@@ -12,7 +12,7 @@ InternalUnitSystem:
# Parameters governing the time integration
TimeIntegration:
time_begin: 0. # The starting time of the simulation (in internal units).
- time_end: 0.5 # The end time of the simulation (in internal units).
+ time_end: 0.25 # The end time of the simulation (in internal units).
dt_min: 1e-9 # The minimal time-step size of the simulation (in internal units).
dt_max: 1e-1 # The maximal time-step size of the simulation (in internal units).
@@ -21,12 +21,12 @@ TimeIntegration:
Snapshots:
basename: output # Common part of the name of output files
time_first: 0. # Time of the first output (in internal units)
- delta_time: 5e-2 # Time between snapshots
+ delta_time: 0.25 # Time between snapshots
# Parameters governing the conserved quantities statistics
Statistics:
time_first: 0.
- delta_time: 5e-2 # Time between statistics output
+ delta_time: 1e-1 # Time between statistics output
# Parameters for the hydrodynamics scheme
SPH:
diff --git a/examples/ChemistryTests/MetalAdvectionTestEAGLE/makeIC.py b/examples/ChemistryTests/MetalAdvectionTestEAGLE/makeIC.py
index 21c12b25f5ac21c40e66f83c083002e36dfb2561..aa6022c4d5098f638ccd38c39d41f83ee0b63ac2 100755
--- a/examples/ChemistryTests/MetalAdvectionTestEAGLE/makeIC.py
+++ b/examples/ChemistryTests/MetalAdvectionTestEAGLE/makeIC.py
@@ -17,7 +17,7 @@
# along with this program. If not, see <http://www.gnu.org/licenses/>.
#
##############################################################################
-
+import math
# ---------------------------------------------------------------------
# Test the diffusion/advection of metals by creating regions with/without
@@ -26,47 +26,80 @@
import numpy as np
import h5py
+import matplotlib.pyplot as plt
GAMMA = 5 / 3
RHO = 1
P = 1
VELOCITY = 1
-BOX_SIZE = 1
+ELEMENT_COUNT = 9
outputfilename = "advect_metals.hdf5"
+
+def get_masks(boxsize, pos):
+ masks = dict()
+ x = boxsize[0]
+ y = boxsize[1]
+
+ right_half_mask = pos[:, 0] > x / 2
+ masks["TotalMetallicity"] = right_half_mask
+ masks["He"] = pos[:, 0] * 2 % x > x / 2
+ masks["C"] = right_half_mask & (pos[:, 0] < 0.75 * x)
+ masks["N"] = right_half_mask & (pos[:, 0] > 0.75 * x)
+ masks["O"] = right_half_mask & (pos[:, 1] > 0.5 * y)
+ masks["Ne"] = right_half_mask & (pos[:, 1] < 0.5 * y)
+ masks["Mg"] = right_half_mask & (pos[:, 0] < 0.75 * x) & (pos[:, 1] > 0.5 * y)
+ masks["Si"] = right_half_mask & (pos[:, 0] > 0.75 * x) & (pos[:, 1] > 0.5 * y)
+ masks["Fe"] = right_half_mask & (pos[:, 0] < 0.75 * x) & (pos[:, 1] < 0.5 * y)
+
+ return masks
+
+
+def get_element_abundances_metallicity(pos, boxsize):
+ elements = ["He", "C", "N", "O", "Ne", "Mg", "Si", "Fe"]
+ abundances = [0.2, 0.2, 0.2, 0.2, 0.2, 0.1, 0.1, 0.1]
+
+ masks = get_masks(boxsize, pos)
+ total_metallicity = np.where(masks["TotalMetallicity"], 0.7, 0.1)
+ element_abundances = [np.where(masks[k], v, 0) for k, v in zip(elements, abundances)]
+
+ # Finally add H so that H + He + TotalMetallicity = 1
+ return np.stack([1 - element_abundances[0] - total_metallicity, ] + element_abundances, axis=1), total_metallicity
+
+
if __name__ == "__main__":
glass = h5py.File("glassPlane_128.hdf5", "r")
parts = glass["PartType0"]
pos = parts["Coordinates"][:]
+ pos = np.concatenate([pos, pos + np.array([1, 0, 0])])
h = parts["SmoothingLength"][:]
+ h = np.concatenate([h, h])
glass.close()
# Set up metadata
- boxsize = np.array([1.0, 1.0, 0.0])
+ boxsize = np.array([2.0, 1.0, 1.0])
n_part = len(h)
ids = np.arange(n_part) + 1
# Setup other particle quantities
- masses = RHO * BOX_SIZE ** 3 / n_part * np.ones(n_part)
+ rho = RHO * np.ones_like(h)
+ rho[pos[:, 1] < 0.5 * boxsize[1]] *= 0.5
+ masses = rho * np.prod(boxsize) / n_part
velocities = np.zeros((n_part, 3))
- velocities[:, 0] = VELOCITY
- internal_energy = P / (RHO * (GAMMA - 1)) * np.ones(n_part)
+ velocities[:, :] = 0.5 * math.sqrt(2) * VELOCITY * np.array([1., 1., 0.])
+ internal_energy = P / (rho * (GAMMA - 1))
# Setup metallicities
- metallicities = np.zeros(n_part)
- # Mask for middle square
- mask = ((1 / 3 * BOX_SIZE < pos[:, 0]) & (pos[:, 0] < 2 / 3 * BOX_SIZE) &
- (1 / 3 * BOX_SIZE < pos[:, 1]) & (pos[:, 1] < 2 / 3 * BOX_SIZE))
- metallicities[mask] = 1
+ element_abundances, metallicities = get_element_abundances_metallicity(pos, boxsize)
# Now open the file and write the data.
file = h5py.File(outputfilename, "w")
# Header
grp = file.create_group("/Header")
- grp.attrs["BoxSize"] = [1.0, 1.0, 1.0]
+ grp.attrs["BoxSize"] = boxsize
grp.attrs["NumPart_Total"] = [n_part, 0, 0, 0, 0, 0]
grp.attrs["NumPart_Total_HighWord"] = [0, 0, 0, 0, 0, 0]
grp.attrs["NumPart_ThisFile"] = [n_part, 0, 0, 0, 0, 0]
@@ -92,9 +125,8 @@ if __name__ == "__main__":
grp.create_dataset("SmoothingLength", data=h, dtype="f")
grp.create_dataset("InternalEnergy", data=internal_energy, dtype="f")
grp.create_dataset("ParticleIDs", data=ids, dtype="L")
- grp.create_dataset("Metallicity", data=metallicities, dtype="L")
-
- # TODO: add ElementAbundance arrays and IronMassFracFromSNIa (see EAGLE/chemistry_io.h)
+ grp.create_dataset("Metallicity", data=metallicities, dtype="f")
+ grp.create_dataset("ElementAbundance", data=element_abundances, dtype="f")
file.close()
diff --git a/examples/ChemistryTests/MetalAdvectionTestEAGLE/plotAdvectedMetals.py b/examples/ChemistryTests/MetalAdvectionTestEAGLE/plotAdvectedMetals.py
index 61b539cfca0b7f3ffab4ce9512adb9835a62e234..a9e3c4857323e9066df45472ab92833f9e45b02e 100644
--- a/examples/ChemistryTests/MetalAdvectionTestEAGLE/plotAdvectedMetals.py
+++ b/examples/ChemistryTests/MetalAdvectionTestEAGLE/plotAdvectedMetals.py
@@ -1,40 +1,109 @@
+import math
+
import matplotlib.pyplot as plt
+from matplotlib.cm import ScalarMappable
+from matplotlib.colors import Normalize
import sys
-
import numpy as np
+import unyt
+
import swiftsimio
from swiftsimio.visualisation import project_gas
-import scipy.interpolate as si
from pathlib import Path
-def project_nn(x, y, data, data_xy):
- xv, yv = np.meshgrid(x, y)
- return si.griddata(data_xy, data, (xv, yv), method="nearest")
+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")
-if __name__ == "__main__":
- try:
- n = sys.argv[1]
- except IndexError:
- n = 2
- cwd = Path(__file__).parent
+def plot_all(fname, savename):
+ data = swiftsimio.load(fname)
- fname = cwd / f"output_{n:04}.hdf5"
- # fname = cwd / f"IC.hdf5"
+ # Shift coordinates to starting position
+ velocity = 0.5 * math.sqrt(2) * np.array([1, 1, 0]) * unyt.cm / unyt.s
+ data.gas.coordinates -= data.metadata.time * velocity
- data = swiftsimio.load(fname)
+ # 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 = {"resolution": 1024, "parallel": True}
+ projection_kwargs = {
+ "region": np.array([0, 2, 0, 1, 0, 1]) * unyt.cm,
+ "resolution": 500,
+ "parallel": True
+ }
+ # Parameters for imshow
+ norm_ratios = Normalize(0, .95)
+ norm_masses = Normalize(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)
- mass_weighted_metal_map = project_gas(data, project="metal_mass_fractions", **projection_kwargs)
- metal_map = mass_weighted_metal_map / mass_map
-
- res = 1000
- plt.imshow(metal_map.T)
- plt.axis("off")
- plt.tight_layout()
- plt.savefig("metal_advection.png", dpi=300)
- plt.show()
\ No newline at end of file
+
+ # 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!")
+
diff --git a/examples/ChemistryTests/MetalAdvectionTestEAGLE/run.sh b/examples/ChemistryTests/MetalAdvectionTestEAGLE/run.sh
index 910f64576d1bc2ebcb973fd5126793d97d9f5506..7acc2cc921781beb9dde4549051e0489266ecd35 100755
--- a/examples/ChemistryTests/MetalAdvectionTestEAGLE/run.sh
+++ b/examples/ChemistryTests/MetalAdvectionTestEAGLE/run.sh
@@ -19,4 +19,5 @@ fi
../../../swift \
--hydro --threads=4 advect_metals.yml 2>&1 | tee output.log
+python3 runSanityChecksAdvectedMetals.py
python3 plotAdvectedMetals.py
diff --git a/examples/ChemistryTests/MetalAdvectionTestEAGLE/runSanityChecksAdvectedMetals.py b/examples/ChemistryTests/MetalAdvectionTestEAGLE/runSanityChecksAdvectedMetals.py
new file mode 100644
index 0000000000000000000000000000000000000000..a1d8cfa45558d712aaddfb14eff7059ccee303db
--- /dev/null
+++ b/examples/ChemistryTests/MetalAdvectionTestEAGLE/runSanityChecksAdvectedMetals.py
@@ -0,0 +1,137 @@
+from pathlib import Path
+
+import numpy as np
+import swiftsimio
+
+
+def test(name):
+ def test_decorator(test_func):
+ def test_runner(*args, **kwargs):
+ try:
+ test_func(*args, **kwargs)
+ except Exception as e:
+ print(f"\tTest {name} failed!")
+ raise e
+ else:
+ print(f"\tTest {name} \u2705")
+ return test_runner
+ return test_decorator
+
+
+def approx_equal(a, b, threshold=0.001):
+ return 2 * abs(a - b) / (abs(a) + abs(b)) < threshold
+
+
+@test("mass fractions sum to 1")
+def test_sum_mass_fractions(data):
+ metal_fractions = data.gas.metal_mass_fractions
+ h_fraction = data.gas.element_mass_fractions.hydrogen
+ he_fraction = data.gas.element_mass_fractions.helium
+ total = metal_fractions + he_fraction + h_fraction
+
+ assert np.all(approx_equal(total, 1))
+
+@test("total metal mass conservation")
+def test_total_metal_mass_conservation(data_start, data_end):
+ def metal_mass(data):
+ return data.gas.masses * data.gas.metal_mass_fractions
+
+ assert approx_equal(np.sum(metal_mass(data_start)), np.sum(metal_mass(data_end)))
+
+
+def element_mass(data, element_name):
+ return data.gas.masses * getattr(data.gas.element_mass_fractions, element_name)
+
+
+@test("hydrogen mass conservation")
+def test_h_mass_conservation(data_start, data_end):
+ assert approx_equal(
+ np.sum(element_mass(data_start, "hydrogen")),
+ np.sum(element_mass(data_end, "hydrogen"))
+ )
+
+
+@test("helium mass conservation")
+def test_he_mass_conservation(data_start, data_end):
+ assert approx_equal(
+ np.sum(element_mass(data_start, "helium")),
+ np.sum(element_mass(data_end, "helium"))
+ )
+
+
+@test("carbon mass conservation")
+def test_c_mass_conservation(data_start, data_end):
+ assert approx_equal(
+ np.sum(element_mass(data_start, "carbon")),
+ np.sum(element_mass(data_end, "carbon"))
+ )
+
+
+@test("nitrogen mass conservation")
+def test_n_mass_conservation(data_start, data_end):
+ assert approx_equal(
+ np.sum(element_mass(data_start, "nitrogen")),
+ np.sum(element_mass(data_end, "nitrogen"))
+ )
+
+
+@test("oxygen mass conservation")
+def test_o_mass_conservation(data_start, data_end):
+ assert approx_equal(
+ np.sum(element_mass(data_start, "oxygen")),
+ np.sum(element_mass(data_end, "oxygen"))
+ )
+
+
+@test("neon mass conservation")
+def test_ne_mass_conservation(data_start, data_end):
+ assert approx_equal(
+ np.sum(element_mass(data_start, "neon")),
+ np.sum(element_mass(data_end, "neon"))
+ )
+
+
+@test("magnesium mass conservation")
+def test_mg_mass_conservation(data_start, data_end):
+ assert approx_equal(
+ np.sum(element_mass(data_start, "magnesium")),
+ np.sum(element_mass(data_end, "magnesium"))
+ )
+
+
+@test("silicon mass conservation")
+def test_si_mass_conservation(data_start, data_end):
+ assert approx_equal(
+ np.sum(element_mass(data_start, "silicon")),
+ np.sum(element_mass(data_end, "silicon"))
+ )
+
+
+@test("iron mass conservation")
+def test_fe_mass_conservation(data_start, data_end):
+ assert approx_equal(
+ np.sum(element_mass(data_start, "iron")),
+ np.sum(element_mass(data_end, "iron"))
+ )
+
+
+if __name__ == "__main__":
+ print("Running sanity checks...")
+
+ cwd = Path(__file__).parent
+ start = swiftsimio.load(cwd / "output_0000.hdf5")
+ end = swiftsimio.load(cwd / "output_0001.hdf5")
+
+ test_sum_mass_fractions(end)
+ test_total_metal_mass_conservation(start, end)
+ test_h_mass_conservation(start, end)
+ test_he_mass_conservation(start, end)
+ test_c_mass_conservation(start, end)
+ test_n_mass_conservation(start, end)
+ test_o_mass_conservation(start, end)
+ test_ne_mass_conservation(start, end)
+ test_mg_mass_conservation(start, end)
+ test_si_mass_conservation(start, end)
+ test_fe_mass_conservation(start, end)
+
+ print("Done!")
\ No newline at end of file