Merge branch 'GEAR_mass_fluxes' into 'master'

Chemistry: Metal fluxes for GEAR and AGORA (and QLA)

See merge request !1853

examples/ChemistryTests/MetalAdvectionTestGEAR/README

+Metal advection test for the GEAR/AGORA chemistry schemes
+
+In some schemes, like the meshless (gizmo) scheme, particles exchange 
+masses via fluxes. This example tests whether the metals are correctly
+advected with those mass fluxes.
+
+To run this test with GEAR, compile with:
+    --with-hydro-dimension=2 --with-hydro=gizmo-mfv --with-riemann-solver=hllc --with-chemistry=GEAR_X
+with X the desired number of elements. The number of elements must be also be set in the makeIC.py script
+and must be the equal to one used in the compilation flag. The default value is 4.
+
+To run this test with AGORA, compile with:
+    --with-hydro-dimension=2 --with-hydro=gizmo-mfv --with-riemann-solver=hllc --with-chemistry=AGORA
+NOTE: The AGORA scheme only traces Fe mass and total metal mass, so the number of elements in makeIC.py must be set
+to the value of 2.
+
+Due to the nature of this test, not much mass will be exchanged when running in Lagrangian mode,
+and hence not much advection will happen.
+To be able to see the effect of the advection, the hydrodynamics must be run in Eulerian mode.
+E.g. for gizmo-mvf: uncomment `#define GIZMO_FIX_PARTICLES` in src/const.h.
+
+Expected output when running in Eulerian mode is that the profiles should maintain their shape,
+but edges will be blurred due to numerical diffusion.

examples/ChemistryTests/MetalAdvectionTestGEAR/advect_metals.yml

+MetaData:
+  run_name: advect_metals_GEAR
+
+# Define the system of units to use internally. 
+InternalUnitSystem:
+  UnitMass_in_cgs:     1   # Grams
+  UnitLength_in_cgs:   1   # Centimeters
+  UnitVelocity_in_cgs: 1   # Centimeters per second
+  UnitCurrent_in_cgs:  1   # Amperes
+  UnitTemp_in_cgs:     1   # Kelvin
+
+# Parameters governing the time integration
+TimeIntegration:
+  time_begin: 0.    # The starting time of the simulation (in internal units).
+  time_end:   0.1   # 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).
+
+
+# Parameters governing the snapshots
+Snapshots:
+  basename:            output # Common part of the name of output files
+  time_first:          0.     # Time of the first output (in internal units)
+  delta_time:          0.1    # Time between snapshots
+
+# Parameters governing the conserved quantities statistics
+Statistics:
+  time_first:          0.
+  delta_time:          1e-1   # Time between statistics output
+
+# Parameters for the hydrodynamics scheme
+SPH:
+  resolution_eta:        1.2348   # Target smoothing length in units of the mean inter-particle separation (1.2348 == 48Ngbs with the cubic spline kernel).
+  CFL_condition:         0.1      # Courant-Friedrich-Levy condition for time integration.
+
+# Parameters related to the initial conditions
+InitialConditions:
+  file_name:  ./advect_metals.hdf5     # The file to read
+  periodic:   1                        # periodic ICs.
+
+GEARChemistry:
+  initial_metallicity: -1         # Negative values mean that the metallicity will be read from the ICs
+
+AGORAChemistry:
+  initial_metallicity: -1         # Negative values mean that the metallicity will be read from the ICs
+  scale_initial_metallicity: 1    # scale the initial metallicity using solar_abundance_Metals? (needed to prevent crash in writing of metadata)

examples/ChemistryTests/MetalAdvectionTestGEAR/getGlass.sh

+#!/bin/bash
+wget http://virgodb.cosma.dur.ac.uk/swift-webstorage/ICs/glassPlane_64.hdf5

examples/ChemistryTests/MetalAdvectionTestGEAR/makeIC.py

+#!/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 <http://www.gnu.org/licenses/>.
+#
+##############################################################################
+import math
+
+# ---------------------------------------------------------------------
+# Test the diffusion/advection of metals by creating regions with/without
+# initial metallicities. Run with EAGLE chemistry.
+# ---------------------------------------------------------------------
+
+import numpy as np
+import h5py
+
+GAMMA = 5 / 3
+RHO = 1
+P = 1
+VELOCITY = 2.5
+# Set ELEMENT_COUNT to 2 for AGORA runs, or for GEAR compile swift with `--with-chemistry=GEAR_X` with X equal to
+# ELEMENT_COUNT.
+ELEMENT_COUNT = 4
+
+outputfilename = "advect_metals.hdf5"
+
+
+def get_mask(boxsize, pos, element_idx):
+    x = boxsize[0]
+    right_half_mask = pos[:, 0] > x / 3
+
+    if element_idx == ELEMENT_COUNT - 1:
+        return right_half_mask
+    else:
+        periods = 2 * (element_idx + 2)
+        periods_mask = (pos[:, 0] * periods // x) % 2 == 0
+        return right_half_mask & periods_mask
+
+
+def get_abundance(element_idx):
+    if element_idx == ELEMENT_COUNT - 1:
+        return 0.2
+    else:
+        return 0.1 * 0.5 ** element_idx
+
+
+def get_element_abundances_metallicity(pos, boxsize):
+    element_abundances = []
+    for i in range(ELEMENT_COUNT):
+        element_abundances.append(np.where(get_mask(boxsize, pos, i), get_abundance(i), 0))
+
+    return np.stack(element_abundances, axis=1)
+
+
+if __name__ == "__main__":
+    glass = h5py.File("glassPlane_64.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([2.0, 1.0, 1.0])
+    n_part = len(h)
+    ids = np.arange(n_part) + 1
+
+    # Setup other particle quantities
+    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.5 * math.sqrt(2) * VELOCITY * np.array([1., 1., 0.])
+    internal_energy = P / (rho * (GAMMA - 1))
+
+    # Setup metallicities
+    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"] = 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]
+    grp.attrs["Time"] = 0.0
+    grp.attrs["NumFilesPerSnapshot"] = 1
+    grp.attrs["MassTable"] = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
+    grp.attrs["Flag_Entropy_ICs"] = 0
+    grp.attrs["Dimension"] = 2
+
+    # Units
+    grp = file.create_group("/Units")
+    grp.attrs["Unit length in cgs (U_L)"] = 1.0
+    grp.attrs["Unit mass in cgs (U_M)"] = 1.0
+    grp.attrs["Unit time in cgs (U_t)"] = 1.0
+    grp.attrs["Unit current in cgs (U_I)"] = 1.0
+    grp.attrs["Unit temperature in cgs (U_T)"] = 1.0
+
+    # Particle group
+    grp = file.create_group("/PartType0")
+    grp.create_dataset("Coordinates", data=pos, dtype="d")
+    grp.create_dataset("Velocities", data=velocities, dtype="f")
+    grp.create_dataset("Masses", data=masses, dtype="f")
+    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("MetalMassFraction", data=metallicities, dtype="f")
+
+    file.close()
+
+    # close up, and we're done!
+    file.close()
+
+    if ELEMENT_COUNT == 2:
+        print(f"Make sure swift was compiled with `--with-chemistry=GEAR_2` or `--with-chemistry=AGORA`")
+    else:
+        print(f"Make sure swift was compiled with `--with-chemistry=GEAR_{ELEMENT_COUNT}`")

examples/ChemistryTests/MetalAdvectionTestGEAR/plotAdvectedMetals.py

+#!/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 <http://www.gnu.org/licenses/>.
+#
+##############################################################################
+
+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=.21, linthresh=thresh, base=10)
+        norm_masses = SymLogNorm(vmin=0, vmax=.25, linthresh=thresh, base=10)
+    else:
+        norm_ratios = Normalize(vmin=0, vmax=.21)
+        norm_masses = Normalize(vmin=0, vmax=.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!")

examples/ChemistryTests/MetalAdvectionTestGEAR/run.sh

+#!/bin/bash
+
+# make run.sh fail if a subcommand fails
+set -e
+set -o pipefail
+
+if [ ! -e glassPlane_64.hdf5 ]
+then
+    echo "Fetching initial glass file ..."
+    ./getGlass.sh
+fi
+
+# Always generate ICs in case ELEMENT_COUNT has been changed
+echo "Generating ICs"
+python3 makeIC.py
+
+# Run SWIFT (must be compiled with --with-chemistry=GEAR_X or --with-chemistry=AGORA)
+../../../swift \
+    --hydro --threads=4 advect_metals.yml 2>&1 | tee output.log
+
+python3 runSanityChecksAdvectedMetals.py
+python3 plotAdvectedMetals.py

examples/ChemistryTests/MetalAdvectionTestGEAR/runSanityChecksAdvectedMetals.py

+#!/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 <http://www.gnu.org/licenses/>.
+#
+##############################################################################
+
+from pathlib import Path
+
+import numpy as np
+import swiftsimio
+
+from makeIC import ELEMENT_COUNT
+
+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("total metal mass conservation")
+def test_total_metal_mass_conservation(data_start, data_end):
+    def metal_mass(data):
+        columns = getattr(data.gas.metal_mass_fractions, "named_columns", None)
+        if columns is not None:
+            return sum(data.gas.masses * getattr(data.gas.metal_mass_fractions, c) for c in columns)
+        else:
+            return data.gas.masses.reshape(-1, 1) * 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_idx):
+    columns = getattr(data.gas.metal_mass_fractions, "named_columns", None)
+    if columns is not None:
+        return data.gas.masses * getattr(data.gas.metal_mass_fractions, columns[element_idx])
+    else:
+        return data.gas.masses * data.gas.metal_mass_fractions[:, element_idx]
+
+
+@test("element-wise mass conservation")
+def test_element_wise_mass_conservation(data_start, data_end):
+    for i in range(ELEMENT_COUNT):
+        assert approx_equal(
+            np.sum(element_mass(data_start, i)),
+            np.sum(element_mass(data_end, i))
+        )
+
+
+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_total_metal_mass_conservation(start, end)
+    test_element_wise_mass_conservation(start, end)
+
+    print("Done!")

src/chemistry/AGORA/chemistry_additions.h

+/*******************************************************************************
+ * 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 <http://www.gnu.org/licenses/>.
+ *
+ ******************************************************************************/
+
+#ifndef SWIFT_CHEMISTRY_AGORA_ADDITIONS_H
+#define SWIFT_CHEMISTRY_AGORA_ADDITIONS_H
+
+/**
+ * @brief Resets the metal mass fluxes for schemes that use them.
+ *
+ * @param p The particle to act upon.
+ */
+__attribute__((always_inline)) INLINE static void chemistry_reset_mass_fluxes(
+    struct part* restrict p) {
+  for (int i = 0; i < AGORA_CHEMISTRY_ELEMENT_COUNT; i++) {
+    p->chemistry_data.metal_mass_fluxes[i] = 0.f;
+  }
+}
+
+/**
+ * @brief Extra operations done during the kick. This needs to be
+ * done before the particle mass is updated in the hydro_kick_extra.
+ *
+ * @param p Particle to act upon.
+ * @param dt_therm Thermal energy time-step @f$\frac{dt}{a^2}@f$.
+ * @param dt_grav Gravity time-step @f$\frac{dt}{a}@f$.
+ * @param dt_hydro Hydro acceleration time-step
+ * @f$\frac{dt}{a^{3(\gamma{}-1)}}@f$.
+ * @param dt_kick_corr Gravity correction time-step @f$adt@f$.
+ * @param cosmo Cosmology.
+ * @param hydro_props Additional hydro properties.
+ */
+__attribute__((always_inline)) INLINE static void chemistry_kick_extra(
+    struct part* p, float dt_therm, float dt_grav, float dt_hydro,
+    float dt_kick_corr, const struct cosmology* cosmo,
+    const struct hydro_props* hydro_props) {
+  /* For hydro schemes that exchange mass fluxes between the particles,
+   * we want to advect the metals. */
+  if (p->flux.dt > 0.) {
+
+    /* Check for vacuum? */
+    if (p->conserved.mass + p->flux.mass <= 0.) {
+      for (int i = 0; i < AGORA_CHEMISTRY_ELEMENT_COUNT; i++) {
+        p->chemistry_data.metal_mass[i] = 0.;
+        p->chemistry_data.smoothed_metal_mass_fraction[i] = 0.;
+      }
+      chemistry_reset_mass_fluxes(p);
+      /* Nothing left to do */
+      return;
+    }
+
+    /* apply the metal mass fluxes and reset them */
+    const double* metal_fluxes = p->chemistry_data.metal_mass_fluxes;
+    for (int i = 0; i < AGORA_CHEMISTRY_ELEMENT_COUNT; i++) {
+      p->chemistry_data.metal_mass[i] =
+          fmax(p->chemistry_data.metal_mass[i] + metal_fluxes[i], 0.);
+    }
+
+#ifdef SWIFT_DEBUG_CHECKS
+    const double iron_mass = p->chemistry_data.metal_mass[0];
+    const double total_metal_mass =
+        p->chemistry_data.metal_mass[AGORA_CHEMISTRY_ELEMENT_COUNT - 1];
+    if (iron_mass > total_metal_mass)
+      error("Iron mass grew larger than total metal mass!");
+    if (total_metal_mass > (p->conserved.mass + p->flux.mass) * (1. + 1.e-4))
+      error("Total metal mass grew larger than total particle mass!");
+#endif
+    chemistry_reset_mass_fluxes(p);
+  }
+}
+
+/**
+ * @brief update metal mass fluxes between two interacting particles during
+ * hydro_iact_(non)sym(...) calls.
+ *
+ * Metals are advected. I.e. a particle loses metals according to its own
+ * metal mass fractions and gains mass according to the neighboring particle's
+ * mass fractions.
+ *
+ * @param pi first interacting particle
+ * @param pj second interacting particle
+ * @param mass_flux the mass flux between these two particles.
+ * @param flux_dt the time-step over which the fluxes are exchanged
+ * @param mode 0: non-symmetric interaction, update i only. 1: symmetric
+ * interaction.
+ **/
+__attribute__((always_inline)) INLINE static void runner_iact_chemistry_fluxes(
+    struct part* restrict pi, struct part* restrict pj, float mass_flux,
+    float flux_dt, int mode) {
+
+  const double mass_flux_integrated = mass_flux * flux_dt;
+  const float mi = pi->conserved.mass;
+  const float mi_inv = mi > 0 ? 1.f / mi : 0.f;
+  const float mj = pj->conserved.mass;
+  const float mj_inv = mj > 0 ? 1.f / mj : 0.f;
+
+  /* Convention: a positive mass flux means that pi is losing said mass and pj
+   * is gaining it. */
+  if (mass_flux > 0.f) {
+    /* pi is losing mass */
+    for (int i = 0; i < AGORA_CHEMISTRY_ELEMENT_COUNT; i++) {
+      pi->chemistry_data.metal_mass_fluxes[i] -=
+          mass_flux_integrated * pi->chemistry_data.metal_mass[i] * mi_inv;
+    }
+  } else {
+    /* pi is gaining mass: */
+    for (int i = 0; i < AGORA_CHEMISTRY_ELEMENT_COUNT; i++) {
+      pi->chemistry_data.metal_mass_fluxes[i] -=
+          mass_flux_integrated * pj->chemistry_data.metal_mass[i] * mj_inv;
+    }
+  }
+
+  /* update pj as well, even if it is inactive (flux.dt < 0.) */
+  if (mode == 1 || pj->flux.dt < 0.f) {
+    if (mass_flux > 0.f) {
+      /* pj is gaining mass */
+      for (int i = 0; i < AGORA_CHEMISTRY_ELEMENT_COUNT; i++) {
+        pj->chemistry_data.metal_mass_fluxes[i] +=
+            mass_flux_integrated * pi->chemistry_data.metal_mass[i] * mi_inv;
+      }
+    } else {
+      /* pj is losing mass */
+      for (int i = 0; i < AGORA_CHEMISTRY_ELEMENT_COUNT; i++) {
+        pj->chemistry_data.metal_mass_fluxes[i] +=
+            mass_flux_integrated * pj->chemistry_data.metal_mass[i] * mj_inv;
+      }
+    }
+  }
+}
+
+#endif  // SWIFT_CHEMISTRY_AGORA_ADDITIONS_H

src/chemistry/AGORA/chemistry_struct.h

@@ -46,6 +46,11 @@ struct chemistry_part_data {
   /*! Total mass of element in a particle. */
   double metal_mass[AGORA_CHEMISTRY_ELEMENT_COUNT];
 
+#ifdef HYDRO_DOES_MASS_FLUX
+  /*! Mass fluxes of the metals in a given element */
+  double metal_mass_fluxes[AGORA_CHEMISTRY_ELEMENT_COUNT];
+#endif
+
   /*! Smoothed fraction of the particle mass in a given element */
   double smoothed_metal_mass_fraction[AGORA_CHEMISTRY_ELEMENT_COUNT];
 };

src/chemistry/EAGLE/chemistry_additions.h

@@ -17,8 +17,8 @@
  *
  ******************************************************************************/
 
-#ifndef SWIFT_EAGLE_CHEMISTRY_ADDITIONS_H
-#define SWIFT_EAGLE_CHEMISTRY_ADDITIONS_H
+#ifndef SWIFT_CHEMISTRY_EAGLE_ADDITIONS_H
+#define SWIFT_CHEMISTRY_EAGLE_ADDITIONS_H
 
 /**
  * @brief Resets the metal mass fluxes for schemes that use them.
@@ -180,4 +180,4 @@ __attribute__((always_inline)) INLINE static void runner_iact_chemistry_fluxes(
   }
 }
 
-#endif  // SWIFT_EAGLE_CHEMISTRY_ADDITIONS_H
+#endif  // SWIFT_CHEMISTRY_EAGLE_ADDITIONS_H

src/chemistry/GEAR/chemistry_additions.h

+/*******************************************************************************
+ * 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 <http://www.gnu.org/licenses/>.
+ *
+ ******************************************************************************/
+
+#ifndef SWIFT_CHEMISTRY_GEAR_ADDITIONS_H
+#define SWIFT_CHEMISTRY_GEAR_ADDITIONS_H
+
+/**
+ * @brief Resets the metal mass fluxes for schemes that use them.
+ *
+ * @param p The particle to act upon.
+ */
+__attribute__((always_inline)) INLINE static void chemistry_reset_mass_fluxes(
+    struct part* restrict p) {
+  for (int i = 0; i < GEAR_CHEMISTRY_ELEMENT_COUNT; i++) {
+    p->chemistry_data.metal_mass_fluxes[i] = 0.f;
+  }
+}
+
+/**
+ * @brief Extra operations done during the kick. This needs to be
+ * done before the particle mass is updated in the hydro_kick_extra.
+ *
+ * @param p Particle to act upon.
+ * @param dt_therm Thermal energy time-step @f$\frac{dt}{a^2}@f$.
+ * @param dt_grav Gravity time-step @f$\frac{dt}{a}@f$.
+ * @param dt_hydro Hydro acceleration time-step
+ * @f$\frac{dt}{a^{3(\gamma{}-1)}}@f$.
+ * @param dt_kick_corr Gravity correction time-step @f$adt@f$.
+ * @param cosmo Cosmology.
+ * @param hydro_props Additional hydro properties.
+ */
+__attribute__((always_inline)) INLINE static void chemistry_kick_extra(
+    struct part* p, float dt_therm, float dt_grav, float dt_hydro,
+    float dt_kick_corr, const struct cosmology* cosmo,
+    const struct hydro_props* hydro_props) {
+  /* For hydro schemes that exchange mass fluxes between the particles,
+   * we want to advect the metals. */
+  if (p->flux.dt > 0.) {
+
+    /* Check for vacuum? */
+    if (p->conserved.mass + p->flux.mass <= 0.) {
+      for (int i = 0; i < GEAR_CHEMISTRY_ELEMENT_COUNT; i++) {
+        p->chemistry_data.metal_mass[i] = 0.;
+        p->chemistry_data.smoothed_metal_mass_fraction[i] = 0.;
+      }
+      chemistry_reset_mass_fluxes(p);
+      /* Nothing left to do */
+      return;
+    }
+
+    /* apply the metal mass fluxes and reset them */
+    const double* metal_fluxes = p->chemistry_data.metal_mass_fluxes;
+    for (int i = 0; i < GEAR_CHEMISTRY_ELEMENT_COUNT; i++) {
+      p->chemistry_data.metal_mass[i] =
+          fmax(p->chemistry_data.metal_mass[i] + metal_fluxes[i], 0.);
+    }
+
+#ifdef SWIFT_DEBUG_CHECKS
+    double sum = 0.;
+    for (int i = 0; i < GEAR_CHEMISTRY_ELEMENT_COUNT - 1; i++) {
+      sum += p->chemistry_data.metal_mass[i];
+    }
+    const double total_metal_mass =
+        p->chemistry_data.metal_mass[GEAR_CHEMISTRY_ELEMENT_COUNT - 1];
+    if (sum > total_metal_mass)
+      error(
+          "Sum of element-wise metal masses grew larger than total metal "
+          "mass!");
+    if (total_metal_mass > p->conserved.mass + p->flux.mass)
+      error("Total metal mass grew larger than the particle mass!");
+#endif
+    chemistry_reset_mass_fluxes(p);
+  }
+}
+
+/**
+ * @brief update metal mass fluxes between two interacting particles during
+ * hydro_iact_(non)sym(...) calls.
+ *
+ * Metals are advected. I.e. a particle loses metals according to its own
+ * metal mass fractions and gains mass according to the neighboring particle's
+ * mass fractions.
+ *
+ * @param pi first interacting particle
+ * @param pj second interacting particle
+ * @param mass_flux the mass flux between these two particles.
+ * @param flux_dt the time-step over which the fluxes are exchanged
+ * @param mode 0: non-symmetric interaction, update i only. 1: symmetric
+ * interaction.
+ **/
+__attribute__((always_inline)) INLINE static void runner_iact_chemistry_fluxes(
+    struct part* restrict pi, struct part* restrict pj, float mass_flux,
+    float flux_dt, int mode) {
+
+  const double mass_flux_integrated = mass_flux * flux_dt;
+  const float mi = pi->conserved.mass;
+  const float mi_inv = mi > 0 ? 1.f / mi : 0.f;
+  const float mj = pj->conserved.mass;
+  const float mj_inv = mj > 0 ? 1.f / mj : 0.f;
+
+  /* Convention: a positive mass flux means that pi is losing said mass and pj
+   * is gaining it. */
+  if (mass_flux > 0.f) {
+    /* pi is losing mass */
+    for (int i = 0; i < GEAR_CHEMISTRY_ELEMENT_COUNT; i++) {
+      pi->chemistry_data.metal_mass_fluxes[i] -=
+          mass_flux_integrated * pi->chemistry_data.metal_mass[i] * mi_inv;
+    }
+  } else {
+    /* pi is gaining mass: */
+    for (int i = 0; i < GEAR_CHEMISTRY_ELEMENT_COUNT; i++) {
+      pi->chemistry_data.metal_mass_fluxes[i] -=
+          mass_flux_integrated * pj->chemistry_data.metal_mass[i] * mj_inv;
+    }
+  }
+
+  /* update pj as well, even if it is inactive (flux.dt < 0.) */
+  if (mode == 1 || pj->flux.dt < 0.f) {
+    if (mass_flux > 0.f) {
+      /* pj is gaining mass */
+      for (int i = 0; i < GEAR_CHEMISTRY_ELEMENT_COUNT; i++) {
+        pj->chemistry_data.metal_mass_fluxes[i] +=
+            mass_flux_integrated * pi->chemistry_data.metal_mass[i] * mi_inv;
+      }
+    } else {
+      /* pj is losing mass */
+      for (int i = 0; i < GEAR_CHEMISTRY_ELEMENT_COUNT; i++) {
+        pj->chemistry_data.metal_mass_fluxes[i] +=
+            mass_flux_integrated * pj->chemistry_data.metal_mass[i] * mj_inv;
+      }
+    }
+  }
+}
+
+#endif  // SWIFT_CHEMISTRY_GEAR_ADDITIONS_H

src/chemistry/GEAR/chemistry_struct.h

@@ -44,6 +44,11 @@ struct chemistry_part_data {
   /*! Total mass of element in a particle. */
   double metal_mass[GEAR_CHEMISTRY_ELEMENT_COUNT];
 
+#ifdef HYDRO_DOES_MASS_FLUX
+  /*! Mass fluxes of the metals in a given element */
+  double metal_mass_fluxes[GEAR_CHEMISTRY_ELEMENT_COUNT];
+#endif
+
   /*! Smoothed fraction of the particle mass in a given element */
   double smoothed_metal_mass_fraction[GEAR_CHEMISTRY_ELEMENT_COUNT];
 };

src/chemistry/QLA/chemistry_additions.h

+/*******************************************************************************
+ * 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 <http://www.gnu.org/licenses/>.
+ *
+ ******************************************************************************/
+
+#ifndef SWIFT_CHEMISTRY_QLA_ADDITIONS_H
+#define SWIFT_CHEMISTRY_QLA_ADDITIONS_H
+
+/**
+ * @brief Extra operations done during the kick. This needs to be
+ * done before the particle mass is updated in the hydro_kick_extra.
+ *
+ * @param p Particle to act upon.
+ * @param dt_therm Thermal energy time-step @f$\frac{dt}{a^2}@f$.
+ * @param dt_grav Gravity time-step @f$\frac{dt}{a}@f$.
+ * @param dt_hydro Hydro acceleration time-step
+ * @f$\frac{dt}{a^{3(\gamma{}-1)}}@f$.
+ * @param dt_kick_corr Gravity correction time-step @f$adt@f$.
+ * @param cosmo Cosmology.
+ * @param hydro_props Additional hydro properties.
+ */
+__attribute__((always_inline)) INLINE static void chemistry_kick_extra(
+    struct part* p, float dt_therm, float dt_grav, float dt_hydro,
+    float dt_kick_corr, const struct cosmology* cosmo,
+    const struct hydro_props* hydro_props) {}
+
+/**
+ * @brief update metal mass fluxes between two interacting particles during
+ * hydro_iact_(non)sym(...) calls.
+ *
+ * @param pi first interacting particle
+ * @param pj second interacting particle
+ * @param mass_flux the mass flux between these two particles.
+ * @param flux_dt the time-step over which the fluxes are exchanged
+ * @param mode 0: non-symmetric interaction, update i only. 1: symmetric
+ * interaction.
+ **/
+__attribute__((always_inline)) INLINE static void runner_iact_chemistry_fluxes(
+    struct part* restrict pi, struct part* restrict pj, float mass_flux,
+    float flux_dt, int mode) {}
+
+#endif  // SWIFT_CHEMISTRY_QLA_ADDITIONS_H

src/chemistry/none/chemistry_additions.h

@@ -17,8 +17,8 @@
  *
  ******************************************************************************/
 
-#ifndef SWIFT_NONE_CHEMISTRY_ADDITIONS_H
-#define SWIFT_NONE_CHEMISTRY_ADDITIONS_H
+#ifndef SWIFT_CHEMISTRY_NONE_ADDITIONS_H
+#define SWIFT_CHEMISTRY_NONE_ADDITIONS_H
 
 /**
  * @brief Extra operations done during the kick. This needs to be
@@ -53,4 +53,4 @@ __attribute__((always_inline)) INLINE static void runner_iact_chemistry_fluxes(
     struct part* restrict pi, struct part* restrict pj, float mass_flux,
     float flux_dt, int mode) {}
 
-#endif  // SWIFT_NONE_CHEMISTRY_ADDITIONS_H
+#endif  // SWIFT_CHEMISTRY_NONE_ADDITIONS_H

src/chemistry_additions.h

@@ -32,10 +32,16 @@
 
 #ifdef HYDRO_DOES_MASS_FLUX
 /* Import the right chemistry definition */
-#if defined(CHEMISTRY_NONE)
-#include "./chemistry/none/chemistry_additions.h"
+#if defined(CHEMISTRY_AGORA)
+#include "./chemistry/AGORA/chemistry_additions.h"
 #elif defined(CHEMISTRY_EAGLE)
 #include "./chemistry/EAGLE/chemistry_additions.h"
+#elif defined(CHEMISTRY_GEAR)
+#include "./chemistry/GEAR/chemistry_additions.h"
+#elif defined(CHEMISTRY_NONE)
+#include "./chemistry/none/chemistry_additions.h"
+#elif defined(CHEMISTRY_NONE)
+#include "./chemistry/QLA/chemistry_additions.h"
 #else
 #error "Metal advection unimpmlemented for selected chemistry scheme!"
 #endif