Add test for metal advection

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:

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()
 

examples/ChemistryTests/MetalAdvectionTestEAGLE/plotAdvectedMetals.py

+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!")
+

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

examples/ChemistryTests/MetalAdvectionTestEAGLE/runSanityChecksAdvectedMetals.py

+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