diff --git a/configure.ac b/configure.ac
index 3c5e43b579e8cd6065d962294f535467f8fd1441..8202e6c405e97dde2e8a788b2ecfebde0ddc378f 100644
--- a/configure.ac
+++ b/configure.ac
@@ -425,6 +425,18 @@ elif test "$fixed_boundary_particles" != "no"; then
    AC_DEFINE_UNQUOTED([SWIFT_FIXED_BOUNDARY_PARTICLES], [$enableval] ,[Particles with smaller ID than this will be considered as boundaries.])
 fi
 
+# Check if fixed entropy is on for settling planetary initial conditions
+AC_ARG_ENABLE([planetary-fixed-entropy],
+   [AS_HELP_STRING([--enable-planetary-fixed-entropy],
+     [Force entropies to stay fixed for settling planetary initial conditions @<:@yes/no@:>@]
+   )],
+   [planetary_fixed_entropy="$enableval"],
+   [planetary_fixed_entropy="no"]
+)
+if test "$planetary_fixed_entropy" = "yes"; then
+   AC_DEFINE([PLANETARY_FIXED_ENTROPY],1,[Enable planetary fixed entropy])
+fi
+
 # Check whether we have any of the ARM v8.1 tick timers
 AX_ASM_ARM_PMCCNTR
 AX_ASM_ARM_CNTVCT
@@ -1244,8 +1256,6 @@ AC_ARG_WITH([velociraptor],
     [with_velociraptor="no"]
 )
 if test "x$with_velociraptor" != "xno"; then
-   AC_PROG_FC
-   AC_FC_LIBRARY_LDFLAGS
    if test "x$with_velociraptor" != "xyes" -a "x$with_velociraptor" != "x"; then
       VELOCIRAPTOR_LIBS="-L$with_velociraptor -lvelociraptor -lmpi -lstdc++ -lhdf5"
       CFLAGS="$CFLAGS -fopenmp"
@@ -1890,6 +1900,60 @@ case "$with_riemann" in
       AC_MSG_ERROR([Unknown Riemann solver: $with_riemann])
    ;;
 esac
+
+#  chemistry function
+AC_ARG_WITH([chemistry],
+   [AS_HELP_STRING([--with-chemistry=<function>],
+      [chemistry function @<:@none, GEAR_*, GEAR_DIFFUSION_*, QLA, EAGLE default: none@:>@
+      For GEAR, you need to provide the number of elements (e.g. GEAR_10)]
+   )],
+   [with_chemistry="$withval"],
+   [with_chemistry="none"]
+)
+
+if test "$with_subgrid" != "none"; then
+   if test "$with_chemistry" != "none"; then
+      AC_MSG_ERROR([Cannot provide with-subgrid and with-chemistry together])
+   else
+      with_chemistry="$with_subgrid_chemistry"
+   fi
+fi
+
+case "$with_chemistry" in
+   none)
+      AC_DEFINE([CHEMISTRY_NONE], [1], [No chemistry function])
+   ;;
+   GEAR_DIFFUSION_*)
+      AC_DEFINE([CHEMISTRY_GEAR_DIFFUSION], [1], [Chemistry taken from the GEAR model including the metal diffusion])
+      number_element=${with_chemistry:15}
+      AC_DEFINE_UNQUOTED([GEAR_CHEMISTRY_ELEMENT_COUNT], [$number_element], [Number of element to follow])
+   ;;
+   GEAR_*)
+      AC_DEFINE([CHEMISTRY_GEAR], [1], [Chemistry taken from the GEAR model])
+      number_element=${with_chemistry:5}
+      AC_DEFINE_UNQUOTED([GEAR_CHEMISTRY_ELEMENT_COUNT], [$number_element], [Number of element to follow])
+   ;;
+   QLA)
+      AC_DEFINE([CHEMISTRY_QLA], [1], [Chemistry taken from the Quick-Lyman-alpha model])
+   ;;
+   EAGLE)
+      AC_DEFINE([CHEMISTRY_EAGLE], [1], [Chemistry taken from the EAGLE model])
+   ;;
+   *)
+      AC_MSG_ERROR([Unknown chemistry function: $with_chemistry])
+   ;;
+esac
+
+if test "$with_chemistry" != "none"; then
+   if test "$enable_hand_vec" == "yes"; then
+      if test "$enable_vec" == "yes"; then
+         if test "$with_hydro" == "gadget2"; then
+            AC_MSG_ERROR([Cannot run with hand vectorisation and chemistry yet. Please use --disable-hand-vec])
+         fi
+      fi
+   fi
+fi
+
 #  Cooling function
 AC_ARG_WITH([cooling],
    [AS_HELP_STRING([--with-cooling=<model>],
@@ -1918,9 +1982,6 @@ case "$with_cooling" in
    const-lambda)
       AC_DEFINE([COOLING_CONST_LAMBDA], [1], [Const Lambda cooling function])
    ;;
-   compton)
-      AC_DEFINE([COOLING_COMPTON], [1], [Compton cooling off the CMB])
-   ;;
    grackle_*)
       AC_DEFINE([COOLING_GRACKLE], [1], [Cooling via the grackle library])
       primordial_chemistry=${with_cooling:8}
@@ -1951,51 +2012,9 @@ case "$with_cooling" in
    ;;
 esac
 
-#  chemistry function
-AC_ARG_WITH([chemistry],
-   [AS_HELP_STRING([--with-chemistry=<function>],
-      [chemistry function @<:@none, GEAR_*, QLA, EAGLE default: none@:>@
-      For GEAR, you need to provide the number of elements (e.g. GEAR_10)]
-   )],
-   [with_chemistry="$withval"],
-   [with_chemistry="none"]
-)
-
-if test "$with_subgrid" != "none"; then
-   if test "$with_chemistry" != "none"; then
-      AC_MSG_ERROR([Cannot provide with-subgrid and with-chemistry together])
-   else
-      with_chemistry="$with_subgrid_chemistry"
-   fi
-fi
-
-case "$with_chemistry" in
-   none)
-      AC_DEFINE([CHEMISTRY_NONE], [1], [No chemistry function])
-   ;;
-   GEAR_*)
-      AC_DEFINE([CHEMISTRY_GEAR], [1], [Chemistry taken from the GEAR model])
-      number_element=${with_chemistry:5}
-      AC_DEFINE_UNQUOTED([GEAR_CHEMISTRY_ELEMENT_COUNT], [$number_element], [Number of element to follow])
-   ;;
-   QLA)
-      AC_DEFINE([CHEMISTRY_QLA], [1], [Chemistry taken from the Quick-Lyman-alpha model])
-   ;;
-   EAGLE)
-      AC_DEFINE([CHEMISTRY_EAGLE], [1], [Chemistry taken from the EAGLE model])
-   ;;
-   *)
-      AC_MSG_ERROR([Unknown chemistry function: $with_chemistry])
-   ;;
-esac
-
-if test "$with_chemistry" != "none"; then
-   if test "$enable_hand_vec" == "yes"; then
-      if test "$enable_vec" == "yes"; then
-         if test "$with_hydro" == "gadget2"; then
-            AC_MSG_ERROR([Cannot run with hand vectorisation and chemistry yet. Please use --disable-hand-vec])
-         fi
-      fi
+if test "${with_cooling:0:5}" = "EAGLE" || test "${with_cooling:0:7}" = "COLIBRE"; then
+   if test "$with_chemistry" = "none"; then
+      AC_MSG_ERROR([Cannot run with EAGLE or COLIBRE cooling without chemistry. Please pick a chemistry model])
    fi
 fi
 
@@ -2346,22 +2365,22 @@ DX_INIT_DOXYGEN(SWIFT, doc/Doxyfile, doc/)
 AM_CONDITIONAL([HAVE_DOXYGEN], [test "$ac_cv_path_ac_pt_DX_DOXYGEN" != ""])
 
 # Check if using QLA cooling
-AM_CONDITIONAL([HAVEQLACOOLING], [test "${with_cooling:0:3}" == "QLA"])
+AM_CONDITIONAL([HAVEQLACOOLING], [test "${with_cooling:0:3}" = "QLA"])
 
 # Check if using EAGLE cooling
-AM_CONDITIONAL([HAVEEAGLECOOLING], [test "${with_cooling:0:5}" == "EAGLE"])
+AM_CONDITIONAL([HAVEEAGLECOOLING], [test "${with_cooling:0:5}" = "EAGLE"])
 
 # Check if using COLIBRE cooling
-AM_CONDITIONAL([HAVECOLIBRECOOLING], [test "${with_cooling:0:7}" == "COLIBRE"])
+AM_CONDITIONAL([HAVECOLIBRECOOLING], [test "${with_cooling:0:7}" = "COLIBRE"])
 
 # Check if using EAGLE feedback
-AM_CONDITIONAL([HAVEEAGLEFEEDBACK], [test "$with_feedback" == "EAGLE"])
+AM_CONDITIONAL([HAVEEAGLEFEEDBACK], [test "$with_feedback" = "EAGLE"])
 
 # check if using grackle cooling
-AM_CONDITIONAL([HAVEGRACKLECOOLING], [test "${with_cooling:0:7}" == "grackle"])
+AM_CONDITIONAL([HAVEGRACKLECOOLING], [test "${with_cooling:0:7}" = "grackle"])
 
 # check if using gear feedback
-AM_CONDITIONAL([HAVEGEARFEEDBACK], [test "$with_feedback" == "GEAR"])
+AM_CONDITIONAL([HAVEGEARFEEDBACK], [test "$with_feedback" = "GEAR"])
 
 # Handle .in files.
 AC_CONFIG_FILES([Makefile src/Makefile examples/Makefile examples/Cooling/CoolingRates/Makefile doc/Makefile doc/Doxyfile tests/Makefile])
@@ -2459,6 +2478,7 @@ AC_MSG_RESULT([
    Custom icbrtf               : $enable_custom_icbrtf
    Boundary particles          : $boundary_particles
    Fixed boundary particles    : $fixed_boundary_particles
+   Planetary fixed entropy     : $planetary_fixed_entropy
 
    Particle Logger      : $with_logger
    Python enabled       : $have_python
diff --git a/doc/RTD/source/HydroSchemes/hydro.dot b/doc/RTD/source/HydroSchemes/hydro.dot
new file mode 100644
index 0000000000000000000000000000000000000000..85a6bac17be88935e40f22c79dc1bb3f774db197
--- /dev/null
+++ b/doc/RTD/source/HydroSchemes/hydro.dot
@@ -0,0 +1,158 @@
+digraph task_dep {
+	 # Header
+	 compound=true;
+	 ratio=1.41;
+	 node[nodesep=0.15, fontsize=30, penwidth=3.];
+	 edge[fontsize=0, penwidth=0.5];
+	 ranksep=0.8;
+
+	 # Special tasks
+	 sort[color=blue3];
+	 self_density[color=blue3];
+	 self_gradient[color=blue3];
+	 self_force[color=blue3];
+	 self_stars_density[color=darkorange1];
+	 pair_density[color=blue3];
+	 pair_gradient[color=blue3];
+	 pair_force[color=blue3];
+	 pair_limiter[color=black];
+	 pair_stars_density[color=darkorange1];
+	 sub_self_density[color=blue3];
+	 sub_self_gradient[color=blue3];
+	 sub_self_force[color=blue3];
+	 sub_self_limiter[color=black];
+	 sub_self_stars_density[color=darkorange1];
+	 sub_pair_density[color=blue3];
+	 sub_pair_gradient[color=blue3];
+	 sub_pair_force[color=blue3];
+	 sub_pair_limiter[color=black];
+	 sub_pair_stars_density[color=darkorange1];
+	 ghost_in[style=filled,fillcolor=grey90,color=blue3];
+	 ghost[color=blue3];
+	 ghost_out[style=filled,fillcolor=grey90,color=blue3];
+	 extra_ghost[color=blue3];
+	 drift_part[color=blue3];
+	 end_hydro_force[color=blue3];
+	 send_gradient[shape=diamond,style=filled,fillcolor=azure,color=blue3];
+	 send_xv[shape=diamond,style=filled,fillcolor=azure,color=blue3];
+	 send_rho[shape=diamond,style=filled,fillcolor=azure,color=blue3];
+	 recv_gradient[shape=diamond,style=filled,fillcolor=azure,color=blue3];
+	 recv_tend_part[shape=diamond,style=filled,fillcolor=azure,color=blue3];
+	 recv_xv[shape=diamond,style=filled,fillcolor=azure,color=blue3];
+	 recv_rho[shape=diamond,style=filled,fillcolor=azure,color=blue3];
+	 cooling_in[style=filled,fillcolor=grey90,color=blue3];
+	 
+	 # Clusters
+	 subgraph clusterDensity {
+		 label="";
+		 bgcolor="grey99";
+		 pair_density;
+		 self_density;
+		 sub_pair_density;
+		 sub_self_density;
+	 };
+
+	 subgraph clusterForce {
+		 label="";
+		 bgcolor="grey99";
+		 pair_force;
+		 self_force;
+		 sub_pair_force;
+		 sub_self_force;
+	 };
+
+	 subgraph clusterGradient {
+		 label="";
+		 bgcolor="grey99";
+		 pair_gradient;
+		 self_gradient;
+		 sub_pair_gradient;
+		 sub_self_gradient;
+	 };
+
+	 subgraph clusterStarsDensity {
+		 label="";
+		 bgcolor="grey99";
+		 pair_stars_density;
+		 self_stars_density;
+		 sub_pair_stars_density;
+		 sub_self_stars_density;
+	 };
+
+	 subgraph clusterTimestep_limiter {
+		 label="";
+		 bgcolor="grey99";
+		 pair_limiter;
+		 self_limiter;
+		 sub_pair_limiter;
+		 sub_self_limiter;
+	 };
+
+
+	 # Dependencies
+	 sort->pair_density[fontcolor=blue3,color=blue3]
+	 sort->pair_stars_density[fontcolor=blue3,color=blue3]
+	 sort->recv_rho[fontcolor=blue3,color=blue3]
+	 sort->sub_self_density[fontcolor=blue3,color=blue3]
+	 sort->sub_self_stars_density[fontcolor=blue3,color=blue3]
+	 sort->sub_pair_density[fontcolor=blue3,color=blue3]
+	 sort->sub_pair_stars_density[fontcolor=blue3,color=blue3]
+	 self_density->ghost_in[fontcolor=blue3,color=blue3]
+	 self_gradient->extra_ghost[fontcolor=blue3,color=blue3]
+	 self_force->end_hydro_force[fontcolor=blue3,color=blue3]
+	 pair_density->ghost_in[fontcolor=blue3,color=blue3]
+	 pair_density->recv_rho[fontcolor=blue3,color=blue3]
+	 pair_gradient->extra_ghost[fontcolor=blue3,color=blue3]
+	 pair_gradient->recv_gradient[fontcolor=blue3,color=blue3]
+	 pair_force->end_hydro_force[fontcolor=blue3,color=blue3]
+	 pair_force->recv_tend_part[fontcolor=blue3,color=blue3]
+	 sub_self_density->ghost_in[fontcolor=blue3,color=blue3]
+	 sub_self_gradient->extra_ghost[fontcolor=blue3,color=blue3]
+	 sub_self_force->end_hydro_force[fontcolor=blue3,color=blue3]
+	 sub_pair_density->ghost_in[fontcolor=blue3,color=blue3]
+	 sub_pair_density->recv_rho[fontcolor=blue3,color=blue3]
+	 sub_pair_gradient->extra_ghost[fontcolor=blue3,color=blue3]
+	 sub_pair_gradient->recv_gradient[fontcolor=blue3,color=blue3]
+	 sub_pair_force->end_hydro_force[fontcolor=blue3,color=blue3]
+	 sub_pair_force->recv_tend_part[fontcolor=blue3,color=blue3]
+	 ghost_in->ghost[fontcolor=blue3,color=blue3]
+	 ghost->ghost_out[fontcolor=blue3,color=blue3]
+	 ghost_out->self_gradient[fontcolor=blue3,color=blue3]
+	 ghost_out->pair_gradient[fontcolor=blue3,color=blue3]
+	 ghost_out->send_rho[fontcolor=blue3,color=blue3]
+	 ghost_out->sub_self_gradient[fontcolor=blue3,color=blue3]
+	 ghost_out->sub_pair_gradient[fontcolor=blue3,color=blue3]
+	 extra_ghost->self_force[fontcolor=blue3,color=blue3]
+	 extra_ghost->pair_force[fontcolor=blue3,color=blue3]
+	 extra_ghost->send_gradient[fontcolor=blue3,color=blue3]
+	 extra_ghost->sub_self_force[fontcolor=blue3,color=blue3]
+	 extra_ghost->sub_pair_force[fontcolor=blue3,color=blue3]
+	 drift_part->self_density[fontcolor=blue3,color=blue3]
+	 drift_part->self_stars_density[fontcolor=blue3,color=blue3]
+	 drift_part->self_limiter[fontcolor=blue3,color=blue3]
+	 drift_part->pair_density[fontcolor=blue3,color=blue3]
+	 drift_part->pair_stars_density[fontcolor=blue3,color=blue3]
+	 drift_part->pair_limiter[fontcolor=blue3,color=blue3]
+	 drift_part->sort[fontcolor=blue3,color=blue3]
+	 drift_part->send_rho[fontcolor=blue3,color=blue3]
+	 drift_part->send_xv[fontcolor=blue3,color=blue3]
+	 drift_part->sub_self_density[fontcolor=blue3,color=blue3]
+	 drift_part->sub_self_stars_density[fontcolor=blue3,color=blue3]
+	 drift_part->sub_self_limiter[fontcolor=blue3,color=blue3]
+	 drift_part->sub_pair_density[fontcolor=blue3,color=blue3]
+	 drift_part->sub_pair_stars_density[fontcolor=blue3,color=blue3]
+	 drift_part->sub_pair_limiter[fontcolor=blue3,color=blue3]
+	 end_hydro_force->cooling_in[fontcolor=blue3,color=blue3]
+	 send_gradient->end_hydro_force[fontcolor=blue3,color=blue3]
+	 send_xv->ghost_in[fontcolor=blue3,color=blue3]
+	 send_rho->extra_ghost[fontcolor=blue3,color=blue3]
+	 recv_gradient->pair_force[fontcolor=blue3,color=blue3]
+	 recv_gradient->sub_pair_force[fontcolor=blue3,color=blue3]
+	 recv_xv->sort[fontcolor=blue3,color=blue3]
+	 recv_xv->pair_density[fontcolor=blue3,color=blue3]
+	 recv_xv->sub_pair_density[fontcolor=blue3,color=blue3]
+	 recv_rho->pair_gradient[fontcolor=blue3,color=blue3]
+	 recv_rho->pair_stars_density[fontcolor=blue3,color=blue3]
+	 recv_rho->sub_pair_gradient[fontcolor=blue3,color=blue3]
+	 recv_rho->sub_pair_stars_density[fontcolor=blue3,color=blue3]
+}
\ No newline at end of file
diff --git a/doc/RTD/source/HydroSchemes/hydro.png b/doc/RTD/source/HydroSchemes/hydro.png
new file mode 100644
index 0000000000000000000000000000000000000000..0095655a662819bf205c141a71f85c413bd49996
Binary files /dev/null and b/doc/RTD/source/HydroSchemes/hydro.png differ
diff --git a/doc/RTD/source/HydroSchemes/index.rst b/doc/RTD/source/HydroSchemes/index.rst
index d276b56abf3b6b77c1c9af569c6d37143a117c06..530712ff0026ce4d558ecf2b4c5130408003dac1 100644
--- a/doc/RTD/source/HydroSchemes/index.rst
+++ b/doc/RTD/source/HydroSchemes/index.rst
@@ -9,6 +9,24 @@ Hydrodynamics Schemes
 This section of the documentation includes information on the hydrodynamics
 schemes available in SWIFT, as well as how to implement your own.
 
+Depending on the scheme used, the algorithm will need either 2
+(e.g. GADGET-2) or 3 (e.g. GIZMO and SPHENIX) interaction loops.
+Here we show the task dependencies for the hydrodynamics assuming 3 loops.
+In case the case of a 2 loop scheme, SWIFT removes the gradient loop and the extra ghost.
+
+.. figure:: hydro.png
+    :width: 400px
+    :align: center
+    :figclass: align-center
+    :alt: Task dependencies for the hydrodynamics.
+
+    This figure shows the task dependencies for the hydrodynamics assuming a scheme with the gradient loop.
+    The first tasks to be executed are at top the (without any incoming links) and then in the order of the links
+    until the last tasks without any outgoing links.
+    For the hydrodynamics tasks (in blue), the rectangles represent (from top to bottom) the density, gradient and force loops.
+    As this graph was created manually, the task dependencies might not reflect a real run depending on the physics simulated.
+    This was done with SWIFT v0.9.0.
+
 .. toctree::
    :maxdepth: 2
    :caption: Contents:
diff --git a/doc/RTD/source/ParameterFiles/parameter_description.rst b/doc/RTD/source/ParameterFiles/parameter_description.rst
index 9a34a26ae5670bb3ce0e673cd1dab45d9748f6b6..cb6248cb66a2e56b1c9d744d5ee78bf5092f6456 100644
--- a/doc/RTD/source/ParameterFiles/parameter_description.rst
+++ b/doc/RTD/source/ParameterFiles/parameter_description.rst
@@ -781,6 +781,19 @@ the SHUFFLE filter is also applied to get higher compression rates. Note that up
 until HDF5 1.10.x this option is not available when using the MPI-parallel
 version of the i/o routines.
 
+Users can run a program after a snapshot is dumped to disk using the following
+parameters:
+
+* Use the extra command after snapshot creation: ``run_on_dump`` (default :``0``)
+* Command to run after snapshot creation: ``dump_command`` (default: nothing)
+
+These are particularly useful should you wish to submit a job for postprocessing
+the snapshot after it has just been created. Your script will be invoked with
+two parameters, the snapshot base-name, and the snapshot number that was just
+output as a zero-padded integer. For example, if the base-name is "eagle" and
+snapshot 7 was just dumped, with ``dump_command`` set to ``./postprocess.sh``,
+then SWIFT will run ``./postprocess.sh eagle 0007``.
+
 Finally, it is possible to specify a different system of units for the snapshots
 than the one that was used internally by SWIFT. The format is identical to the
 one described above (See the :ref:`Parameters_units` section) and read:
diff --git a/doc/RTD/source/Planetary/index.rst b/doc/RTD/source/Planetary/index.rst
index 7b0a227e1bf387d383cfe42fa63d21c57ec87281..5cc45fa0e51fbfc38e42e7cf111a74ccfe9f835e 100644
--- a/doc/RTD/source/Planetary/index.rst
+++ b/doc/RTD/source/Planetary/index.rst
@@ -26,7 +26,7 @@ These allow for multiple materials to be used,
 chosen from the several available equations of state.
 
 .. toctree::
-   :maxdepth: 2
+   :maxdepth: 1
    :caption: More information:
    
    Hydro Scheme <hydro_scheme>
diff --git a/doc/RTD/source/Planetary/initial_conditions.rst b/doc/RTD/source/Planetary/initial_conditions.rst
index d4c20f1fa1074c57fa3864f6859980e5f1d832b0..fd348a9fd134fca3a35be4a7eeb16669ea56534a 100755
--- a/doc/RTD/source/Planetary/initial_conditions.rst
+++ b/doc/RTD/source/Planetary/initial_conditions.rst
@@ -20,4 +20,19 @@ Ruiz-Bonilla et al. (2020).
 They are available with documentation and examples at 
 https://github.com/srbonilla/WoMa and https://github.com/jkeger/seagen,
 or can be installed directly with ``pip``
-(https://pypi.org/project/woma/, https://pypi.org/project/seagen/).
\ No newline at end of file
+(https://pypi.org/project/woma/, https://pypi.org/project/seagen/).
+
+
+Settling initial conditions with fixed entropies
+------------------------------------------------
+
+If the particles' equations of state include specific entropies, 
+and the initial conditions file includes specific entropies for each particle
+(in ``PartType0/Entropies``), 
+then configuring SWIFT with ``--enable-planetary-fixed-entropy``
+will override the internal energy of each particle each step such that its 
+specific entropy remains constant. 
+
+This should be used with caution, but may be a convenient way to maintain an 
+entropy profile while initial conditions settle to equilibrium with their 
+slightly different SPH densities.
diff --git a/doc/RTD/source/Snapshots/index.rst b/doc/RTD/source/Snapshots/index.rst
index f3e81c6410da14f859eb8843e91f7f45889ddfd1..ae82ab48c0da639ef96ed803f6828b94373c5e37 100644
--- a/doc/RTD/source/Snapshots/index.rst
+++ b/doc/RTD/source/Snapshots/index.rst
@@ -24,10 +24,18 @@ that is always set to the string ``SWIFT``, which can be used to identify
 SWIFT-generated snapshots and hence make use of all the extensions to the file
 format described below.
 
-The most important quantity of the header is the array ``NumPart_ThisFile``
-which contains the number of particles of each type in this snapshot. This is an
-array of 6 numbers; one for each of the 5 supported types and a dummy "type 3"
-field only used for compatibility reasons but always containing a zero.
+The most important quantity of the header is the array ``NumPart_Total`` which
+contains the number of particles of each type in this snapshot. This is an array
+of 6 numbers; one for each of the supported types. The field
+``NumPart_ThisFile`` contains the number of particles in this sub-snapshot file
+when the user asked for distributed snapshots (see :ref:`Parameters_snapshots`);
+otherwise it contains the same information as ``NumPart_Total``. The field
+``NumFilesPerSnapshot`` specifies the number of sub-snapshot files (always 1
+unless a distributed snapshot was asked).
+
+The field ``InitialMassTable`` contains the *mean* initial mass of each of the
+particle types present in the initial conditions. This can be used as estimator
+of the mass resolution of the run. The masses are expressed in internal units.
 
 The ``RunName`` field contains the name of the simulation that was specified as
 the ``run_name`` in the :ref:`Parameters_meta_data` section of the YAML
diff --git a/doc/RTD/source/Task/current_dependencies.rst b/doc/RTD/source/Task/current_dependencies.rst
new file mode 100644
index 0000000000000000000000000000000000000000..2b6d9031ec0d597e97b649fb04a86b3b80f268fe
--- /dev/null
+++ b/doc/RTD/source/Task/current_dependencies.rst
@@ -0,0 +1,56 @@
+.. Current task dependencies
+   Loic Hausammann, 2020
+
+
+.. _current_dependencies:
+
+
+
+Current Task Dependencies
+=========================
+
+In order to compute the physics in the correct order, SWIFT uses dependencies in between the tasks.
+In :ref:`Analysis_Tools`, we describe our tool to generate a graph of the dependencies but,
+unfortunately, the graphs tend to be too large.
+Therefore in this section we show some smaller graphs created by hand (thus they are not necessarily reflecting an actual run depending on the physics simulated).
+The task in the form of ellipses are computational tasks while the diamond are communications.
+The colors are picked depending on the type of physics computed (blue for hydro, yellow for stars, red for gravity and black for the rest).
+The gray tasks are implicit tasks that do not compute anything but are useful to simplify the task dependencies.
+The first graph shows the full graph (without AGN and sink particles) but with some tasks collapsed into a single meta-task
+(hydrodynamics, gravity and stellar feedback):
+
+.. figure:: reduced.png
+    :width: 400px
+    :align: center
+    :figclass: align-center
+    :alt: Task dependencies for a run with gravity, hydro, cooling, star formation and stellar feedback.
+
+    This figure shows the task dependencies for a run with gravity, hydro, cooling, star formation and stellar feedback.
+    The tasks for the limiter, hydrodynamics, stellar feedback and gravity are collapsed into a single meta-task.
+    The other graphs are showing the details of the collapsed tasks except for the limiter that is done in the same way than the density loop.
+    The first tasks to be executed are at the top (without any incoming links) and then follow the order of the links
+    until the last tasks without any outgoing links.
+    This was done with SWIFT v0.9.0.
+
+As the hydrodynamics are described in :ref:`hydro`, we are only showing the gravity and stellar feedback here:
+
+.. figure:: grav.png
+    :width: 400px
+    :align: center
+    :figclass: align-center
+    :alt: Task dependencies for the gravity.
+
+    This figure shows the task dependencies for the gravity.
+    The rectangle represents the computation of the forces through either the direct computation or the multipole expansion.
+    This was done with SWIFT v0.9.0.
+
+.. figure:: stars.png
+    :width: 400px
+    :align: center
+    :figclass: align-center
+    :alt: Task dependencies for the stellar feedback.
+
+    This figure shows the task dependencies for the stellar feedback.
+    The first rectangle groups the tasks that compute the smoothing length of the stars and
+    the second one the tasks that deposit the energy into the surrounding gas.
+    This was done with SWIFT v0.9.0.
diff --git a/doc/RTD/source/Task/grav.dot b/doc/RTD/source/Task/grav.dot
new file mode 100644
index 0000000000000000000000000000000000000000..5ac1e7be4cef144faf36ea1f43ed94eafca5847b
--- /dev/null
+++ b/doc/RTD/source/Task/grav.dot
@@ -0,0 +1,59 @@
+digraph task_dep {
+	 # Header
+	 compound=true;
+	 ratio=1.41;
+	 node[nodesep=0.15, fontsize=30, penwidth=3.];
+	 edge[fontsize=0, penwidth=0.5];
+	 ranksep=0.8;
+
+	 # Special tasks
+	 self_grav[color=red3];
+	 pair_grav[color=red3];
+	 init_grav[color=red3];
+	 init_grav_out[style=filled,fillcolor=grey90,color=red3];
+	 drift_gpart[color=red3];
+	 drift_gpart_out[style=filled,fillcolor=grey90,color=red3];
+	 kick2[color=black];
+	 send_gpart[shape=diamond,style=filled,fillcolor=azure,color=red3];
+	 recv_gpart[shape=diamond,style=filled,fillcolor=azure,color=red3];
+	 grav_long_range[color=red3];
+	 grav_mm[color=red3];
+	 grav_down_in[style=filled,fillcolor=grey90,color=red3];
+	 grav_down[color=red3];
+	 grav_end_force[color=red3];
+	 recv_tend_gpart[shape=diamond,style=filled,fillcolor=azure,color=red3];
+
+	 subgraph clusterGravity {
+		 label="";
+		 bgcolor="grey99";
+		 grav_long_range;
+		 grav_mm;
+		 pair_grav;
+		 self_grav;
+	 };
+
+
+	 # Dependencies
+	 self_grav->grav_down_in[fontcolor=red3,color=red3]
+	 pair_grav->grav_down_in[fontcolor=red3,color=red3]
+	 pair_grav->recv_tend_gpart[fontcolor=red3,color=red3]
+	 init_grav->grav_long_range[fontcolor=red3,color=red3]
+	 init_grav->init_grav_out[fontcolor=red3,color=red3]
+	 init_grav_out->self_grav[fontcolor=red3,color=red3]
+	 init_grav_out->pair_grav[fontcolor=red3,color=red3]
+	 init_grav_out->init_grav_out[fontcolor=red3,color=red3]
+	 init_grav_out->grav_mm[fontcolor=red3,color=red3]
+	 drift_gpart->drift_gpart_out[fontcolor=red3,color=red3]
+	 drift_gpart->send_gpart[fontcolor=red3,color=red3]
+	 drift_gpart_out->self_grav[fontcolor=red3,color=red3]
+	 drift_gpart_out->pair_grav[fontcolor=red3,color=red3]
+	 drift_gpart_out->drift_gpart_out[fontcolor=red3,color=red3]
+	 send_gpart->grav_down[fontcolor=red3,color=red3]
+	 recv_gpart->pair_grav[fontcolor=red3,color=red3]
+	 grav_long_range->grav_down[fontcolor=red3,color=red3]
+	 grav_mm->grav_down_in[fontcolor=red3,color=red3]
+	 grav_down_in->grav_down[fontcolor=red3,color=red3]
+	 grav_down_in->grav_down_in[fontcolor=red3,color=red3]
+	 grav_down->grav_end_force[fontcolor=red3,color=red3]
+	 grav_end_force->kick2[fontcolor=red3,color=red3]
+}
\ No newline at end of file
diff --git a/doc/RTD/source/Task/grav.png b/doc/RTD/source/Task/grav.png
new file mode 100644
index 0000000000000000000000000000000000000000..c72bc4a72106168c9a498e17c9de1e052ffbd2dd
Binary files /dev/null and b/doc/RTD/source/Task/grav.png differ
diff --git a/doc/RTD/source/Task/index.rst b/doc/RTD/source/Task/index.rst
index 2f64f8eb9ed31abc709a0d2b01f8d88a251119ba..98442108894c355fc14c44d0bca3161d33eaeb67 100644
--- a/doc/RTD/source/Task/index.rst
+++ b/doc/RTD/source/Task/index.rst
@@ -17,6 +17,7 @@ Everything is described in :ref:`Analysis_Tools`.
    :maxdepth: 2
    :caption: Contents:
 
+   current_dependencies
    adding_your_own
    adding_your_own_neighbour_loop
    adding_to_dependency_plotting_tool
diff --git a/doc/RTD/source/Task/reduced.dot b/doc/RTD/source/Task/reduced.dot
new file mode 100644
index 0000000000000000000000000000000000000000..00f2d268042652de3fc16577cd72ce5e777ab9a8
--- /dev/null
+++ b/doc/RTD/source/Task/reduced.dot
@@ -0,0 +1,83 @@
+digraph task_dep {
+	 # Header
+	 compound=true;
+	 node[nodesep=0.1, fontsize=20, penwidth=3.];
+	 edge[fontsize=0, penwidth=0.5];
+	 ranksep=0.8;
+
+	 # Special tasks
+	 hydro[color=blue3,shape=folder];
+	 limiter[color=black,shape=folder];
+	 grav[color=red3,shape=folder];
+	 stars[color=darkorange1,shape=folder];
+	 drift_part[color=blue3];
+	 drift_spart[color=darkorange1];
+	 drift_gpart[color=red3];
+	 kick2[color=black];
+	 timestep[color=black];
+	 timestep_limiter[color=black];
+	 timestep_sync[color=black];
+	 recv_limiter[shape=diamond,style=filled,fillcolor=azure,color=black];
+	 recv_tend_part[shape=diamond,style=filled,fillcolor=azure,color=blue3];
+	 recv_sf_count[shape=diamond,style=filled,fillcolor=azure,color=darkorange1];
+	 cooling[color=blue3];
+	 cooling_in[style=filled,fillcolor=grey90,color=blue3];
+	 cooling_out[style=filled,fillcolor=grey90,color=blue3];
+	 star_formation[color=blue3];
+	 kick1[color=black];
+	 recv_tend_gpart[shape=diamond,style=filled,fillcolor=azure,color=red3];
+	 recv_tend_spart[shape=diamond,style=filled,fillcolor=azure,color=darkorange1];
+	 send_tend_part[shape=diamond,style=filled,fillcolor=azure,color=blue3,rank=min];
+	 send_limiter[shape=diamond,style=filled,fillcolor=azure,color=black,rank=min];
+	 send_tend_spart[shape=diamond,style=filled,fillcolor=azure,color=darkorange1,rank=min];
+	 send_tend_gpart[shape=diamond,style=filled,fillcolor=azure,color=red3,rank=min];
+	 send_sf_count[shape=diamond,style=filled,fillcolor=azure,color=darkorange1];
+
+	 # Dependencies
+	 hydro->stars[fontcolor=blue3,color=blue3]
+	 limiter->kick1[fontcolor=black,color=black]
+	 limiter->timestep_limiter[fontcolor=black,color=black]
+	 stars->timestep_sync[fontcolor=darkorange1,color=darkorange1]
+	 hydro->recv_tend_part[fontcolor=blue3,color=blue3]
+	 grav->recv_tend_gpart[fontcolor=red3,color=red3]
+	 stars->recv_tend_spart[fontcolor=darkorange1,color=darkorange1]
+	 stars->recv_tend_spart[fontcolor=darkorange1,color=darkorange1]
+	 drift_part->hydro[fontcolor=blue3,color=blue3]
+	 drift_part->stars[fontcolor=blue3,color=blue3]
+	 drift_part->limiter[fontcolor=blue3,color=blue3]
+	 drift_spart->kick2[fontcolor=darkorange1,color=darkorange1]
+	 drift_spart->stars[fontcolor=darkorange1,color=darkorange1]
+	 drift_gpart->grav[fontcolor=red3,color=red3]
+	 hydro->cooling_in[fontcolor=blue3,color=blue3]
+	 kick2->timestep[fontcolor=black,color=black]
+	 kick2->stars[fontcolor=black,color=black]
+	 kick2->star_formation[fontcolor=black,color=black]
+	 timestep->kick1[fontcolor=black,color=black]
+	 timestep->timestep_limiter[fontcolor=black,color=black]
+	 timestep->timestep_sync[fontcolor=black,color=black]
+	 timestep->limiter[fontcolor=black,color=black]
+	 timestep->send_tend_part[fontcolor=black,color=black]
+	 timestep->send_limiter[fontcolor=black,color=black]
+	 timestep->send_tend_spart[fontcolor=black,color=black]
+	 timestep->send_tend_gpart[fontcolor=black,color=black]
+	 timestep_limiter->kick1[fontcolor=black,color=black]
+	 timestep_limiter->timestep_sync[fontcolor=black,color=black]
+	 timestep_sync->kick1[fontcolor=black,color=black]
+	 recv_limiter->limiter[fontcolor=black,color=black]
+	 recv_tend_part->limiter[fontcolor=blue3,color=blue3]
+	  recv_sf_count->stars[fontcolor=darkorange1,color=darkorange1]
+	 grav->kick2[fontcolor=red3,color=red3]
+	 cooling->cooling_out[fontcolor=blue3,color=blue3]
+	 cooling_in->cooling[fontcolor=blue3,color=blue3]
+	 cooling_out->kick2[fontcolor=blue3,color=blue3]
+	 star_formation->timestep[fontcolor=blue3,color=blue3]
+	 star_formation->stars[fontcolor=blue3,color=blue3]
+	 star_formation->send_sf_count[fontcolor=blue3,color=blue3]
+	 stars->timestep[fontcolor=darkorange1,color=darkorange1]
+
+   # style
+   timestep_limiter->send_tend_part[style=invis];
+   timestep_limiter->send_tend_gpart[style=invis];
+   timestep_limiter->send_tend_spart[style=invis];
+   timestep_limiter->send_limiter[style=invis];
+}
\ No newline at end of file
diff --git a/doc/RTD/source/Task/reduced.png b/doc/RTD/source/Task/reduced.png
new file mode 100644
index 0000000000000000000000000000000000000000..b032fc7c8c5c74c6de3be3bf6b19e18eaa537772
Binary files /dev/null and b/doc/RTD/source/Task/reduced.png differ
diff --git a/doc/RTD/source/Task/stars.dot b/doc/RTD/source/Task/stars.dot
new file mode 100644
index 0000000000000000000000000000000000000000..ff17d9eb67a616807ed46cc2d5fb8ed9cb09ec6e
--- /dev/null
+++ b/doc/RTD/source/Task/stars.dot
@@ -0,0 +1,113 @@
+digraph task_dep {
+	 # Header
+	 compound=true;
+	 ratio=1.41;
+	 node[nodesep=0.15, fontsize=30, penwidth=3.];
+	 edge[fontsize=0, penwidth=0.5];
+	 ranksep=0.8;
+
+	 # Special tasks
+	 sort[color=blue3];
+	 self_stars_density[color=darkorange1];
+	 self_stars_feedback[color=darkorange1];
+	 pair_stars_density[color=darkorange1];
+	 pair_stars_feedback[color=darkorange1];
+	 sub_self_stars_density[color=darkorange1];
+	 sub_self_stars_feedback[color=darkorange1];
+	 sub_pair_stars_density[color=darkorange1];
+	 sub_pair_stars_feedback[color=darkorange1];
+	 drift_part[color=blue3];
+	 drift_spart[color=darkorange1];
+	 kick2[color=black];
+	 timestep[color=black];
+	 timestep_sync[color=black];
+	 send_spart[shape=diamond,style=filled,fillcolor=azure,color=darkorange1];
+	 recv_rho[shape=diamond,style=filled,fillcolor=azure,color=blue3];
+	 recv_spart[shape=diamond,style=filled,fillcolor=azure,color=darkorange1];
+	 recv_sf_count[shape=diamond,style=filled,fillcolor=azure,color=darkorange1];
+	 star_formation[color=blue3];
+	 stars_in[style=filled,fillcolor=grey90,color=darkorange1];
+	 stars_out[style=filled,fillcolor=grey90,color=darkorange1];
+	 stars_ghost[color=darkorange1];
+	 stars_sort[color=darkorange1];
+	 stars_resort[color=darkorange1];
+	 recv_tend_spart[shape=diamond,style=filled,fillcolor=azure,color=darkorange1];
+	 send_sf_count[shape=diamond,style=filled,fillcolor=azure,color=darkorange1];
+
+	 subgraph clusterStarsDensity {
+		 label="";
+		 bgcolor="grey99";
+		 pair_stars_density;
+		 self_stars_density;
+		 sub_pair_stars_density;
+		 sub_self_stars_density;
+	 };
+
+	 subgraph clusterStarsFeedback {
+		 label="";
+		 bgcolor="grey99";
+		 pair_stars_feedback;
+		 self_stars_feedback;
+		 sub_pair_stars_feedback;
+		 sub_self_stars_feedback;
+	 };
+
+
+	 # Dependencies
+	 sort->pair_stars_density[fontcolor=blue3,color=blue3]
+	 sort->sub_self_stars_density[fontcolor=blue3,color=blue3]
+	 sort->sub_pair_stars_density[fontcolor=blue3,color=blue3]
+	 self_stars_density->stars_ghost[fontcolor=darkorange1,color=darkorange1]
+	 self_stars_feedback->stars_out[fontcolor=darkorange1,color=darkorange1]
+	 self_stars_feedback->timestep_sync[fontcolor=darkorange1,color=darkorange1]
+	 pair_stars_density->stars_ghost[fontcolor=darkorange1,color=darkorange1]
+	 pair_stars_density->recv_spart[fontcolor=darkorange1,color=darkorange1]
+	 pair_stars_feedback->stars_out[fontcolor=darkorange1,color=darkorange1]
+	 pair_stars_feedback->timestep_sync[fontcolor=darkorange1,color=darkorange1]
+	 pair_stars_feedback->recv_tend_spart[fontcolor=darkorange1,color=darkorange1]
+	 sub_self_stars_density->stars_ghost[fontcolor=darkorange1,color=darkorange1]
+	 sub_self_stars_feedback->stars_out[fontcolor=darkorange1,color=darkorange1]
+	 sub_self_stars_feedback->timestep_sync[fontcolor=darkorange1,color=darkorange1]
+	 sub_pair_stars_density->stars_ghost[fontcolor=darkorange1,color=darkorange1]
+	 sub_pair_stars_density->recv_spart[fontcolor=darkorange1,color=darkorange1]
+	 sub_pair_stars_feedback->stars_out[fontcolor=darkorange1,color=darkorange1]
+	 sub_pair_stars_feedback->timestep_sync[fontcolor=darkorange1,color=darkorange1]
+	 sub_pair_stars_feedback->recv_tend_spart[fontcolor=darkorange1,color=darkorange1]
+	 drift_part->self_stars_density[fontcolor=blue3,color=blue3]
+	 drift_part->pair_stars_density[fontcolor=blue3,color=blue3]
+	 drift_part->sub_self_stars_density[fontcolor=blue3,color=blue3]
+	 drift_part->sub_pair_stars_density[fontcolor=blue3,color=blue3]
+	 drift_spart->kick2[fontcolor=darkorange1,color=darkorange1]
+	 drift_spart->self_stars_density[fontcolor=darkorange1,color=darkorange1]
+	 drift_spart->pair_stars_density[fontcolor=darkorange1,color=darkorange1]
+	 drift_spart->stars_sort[fontcolor=darkorange1,color=darkorange1]
+	 drift_spart->send_spart[fontcolor=darkorange1,color=darkorange1]
+	 drift_spart->sub_self_stars_density[fontcolor=darkorange1,color=darkorange1]
+	 drift_spart->sub_pair_stars_density[fontcolor=darkorange1,color=darkorange1]
+	 kick2->stars_in[fontcolor=black,color=black]
+	 kick2->star_formation[fontcolor=black,color=black]
+	 send_spart->stars_out[fontcolor=darkorange1,color=darkorange1]
+	 recv_rho->pair_stars_density[fontcolor=blue3,color=blue3]
+	 recv_spart->stars_sort[fontcolor=darkorange1,color=darkorange1]
+	 recv_spart->pair_stars_feedback[fontcolor=darkorange1,color=darkorange1]
+	 recv_spart->sub_pair_stars_feedback[fontcolor=darkorange1,color=darkorange1]
+	 recv_sf_count->recv_spart[fontcolor=darkorange1,color=darkorange1]
+	 star_formation->stars_resort[fontcolor=blue3,color=blue3]
+	 star_formation->send_sf_count[fontcolor=blue3,color=blue3]
+	 stars_in->self_stars_density[fontcolor=darkorange1,color=darkorange1]
+	 stars_in->pair_stars_density[fontcolor=darkorange1,color=darkorange1]
+	 stars_in->sub_self_stars_density[fontcolor=darkorange1,color=darkorange1]
+	 stars_in->sub_pair_stars_density[fontcolor=darkorange1,color=darkorange1]
+	 stars_out->timestep[fontcolor=darkorange1,color=darkorange1]
+	 stars_ghost->self_stars_feedback[fontcolor=darkorange1,color=darkorange1]
+	 stars_ghost->pair_stars_feedback[fontcolor=darkorange1,color=darkorange1]
+	 stars_ghost->send_spart[fontcolor=darkorange1,color=darkorange1]
+	 stars_ghost->sub_self_stars_feedback[fontcolor=darkorange1,color=darkorange1]
+	 stars_ghost->sub_pair_stars_feedback[fontcolor=darkorange1,color=darkorange1]
+	 stars_sort->pair_stars_density[fontcolor=darkorange1,color=darkorange1]
+	 stars_sort->pair_stars_feedback[fontcolor=darkorange1,color=darkorange1]
+	 stars_sort->sub_self_stars_density[fontcolor=darkorange1,color=darkorange1]
+	 stars_sort->sub_pair_stars_density[fontcolor=darkorange1,color=darkorange1]
+	 stars_sort->sub_pair_stars_feedback[fontcolor=darkorange1,color=darkorange1]
+	 stars_resort->stars_in[fontcolor=darkorange1,color=darkorange1]
+}
\ No newline at end of file
diff --git a/doc/RTD/source/Task/stars.png b/doc/RTD/source/Task/stars.png
new file mode 100644
index 0000000000000000000000000000000000000000..affd9a26d16289593d6ca6a56f2e8ba941586f16
Binary files /dev/null and b/doc/RTD/source/Task/stars.png differ
diff --git a/examples/EAGLE_low_z/EAGLE_12/eagle_12.yml b/examples/EAGLE_low_z/EAGLE_12/eagle_12.yml
index 4db4228d3ec8aa01ec608f66d58805a4993e5071..9e6e5658822746a164fc2d85d5deff563b52e402 100644
--- a/examples/EAGLE_low_z/EAGLE_12/eagle_12.yml
+++ b/examples/EAGLE_low_z/EAGLE_12/eagle_12.yml
@@ -194,8 +194,8 @@ EAGLEFeedback:
 # EAGLE AGN model
 EAGLEAGN:
   subgrid_seed_mass_Msun:             1.5e5      # Black hole subgrid mass at creation time in solar masses.
-  multi_phase_bondi:                  0          # Compute Bondi rates per neighbour particle?
-  subgrid_bondi:                      0          # Compute Bondi rates using the subgrid extrapolation of the gas properties around the BH?
+  use_multi_phase_bondi:              0          # Compute Bondi rates per neighbour particle?
+  use_subgrid_bondi:                  0          # Compute Bondi rates using the subgrid extrapolation of the gas properties around the BH?
   with_angmom_limiter:                1          # Are we applying the Rosas-Guevara et al. (2015) viscous time-scale reduction term?
   viscous_alpha:                      1e6        # Normalisation constant of the viscous time-scale in the accretion reduction term
   with_boost_factor:                  0          # Are we using the model from Booth & Schaye (2009)?
diff --git a/examples/EAGLE_low_z/EAGLE_25/eagle_25.yml b/examples/EAGLE_low_z/EAGLE_25/eagle_25.yml
index f37a52ff03093f085115f30322cb579877d87ca9..022474e6da69351ac174310a75dabdedd5815a77 100644
--- a/examples/EAGLE_low_z/EAGLE_25/eagle_25.yml
+++ b/examples/EAGLE_low_z/EAGLE_25/eagle_25.yml
@@ -202,8 +202,8 @@ EAGLEFeedback:
 # EAGLE AGN model
 EAGLEAGN:
   subgrid_seed_mass_Msun:             1.5e5      # Black hole subgrid mass at creation time in solar masses.
-  multi_phase_bondi:                  0          # Compute Bondi rates per neighbour particle?
-  subgrid_bondi:                      0          # Compute Bondi rates using the subgrid extrapolation of the gas properties around the BH?
+  use_multi_phase_bondi:              0          # Compute Bondi rates per neighbour particle?
+  use_subgrid_bondi:                  0          # Compute Bondi rates using the subgrid extrapolation of the gas properties around the BH?
   with_angmom_limiter:                1          # Are we applying the Rosas-Guevara et al. (2015) viscous time-scale reduction term?
   viscous_alpha:                      1e6        # Normalisation constant of the viscous time-scale in the accretion reduction term
   with_boost_factor:                  0          # Are we using the model from Booth & Schaye (2009)?
diff --git a/examples/EAGLE_low_z/EAGLE_50/eagle_50.yml b/examples/EAGLE_low_z/EAGLE_50/eagle_50.yml
index e1ce6362e88a110a615f617bf568ec62d544fea2..ea5bc25942e9c6253c9c8192ff6dfb51cf610e88 100644
--- a/examples/EAGLE_low_z/EAGLE_50/eagle_50.yml
+++ b/examples/EAGLE_low_z/EAGLE_50/eagle_50.yml
@@ -193,8 +193,8 @@ EAGLEFeedback:
 # EAGLE AGN model
 EAGLEAGN:
   subgrid_seed_mass_Msun:             1.5e5      # Black hole subgrid mass at creation time in solar masses.
-  multi_phase_bondi:                  0          # Compute Bondi rates per neighbour particle?
-  subgrid_bondi:                      0          # Compute Bondi rates using the subgrid extrapolation of the gas properties around the BH?
+  use_multi_phase_bondi:              0          # Compute Bondi rates per neighbour particle?
+  use_subgrid_bondi:                  0          # Compute Bondi rates using the subgrid extrapolation of the gas properties around the BH?
   with_angmom_limiter:                1          # Are we applying the Rosas-Guevara et al. (2015) viscous time-scale reduction term?
   viscous_alpha:                      1e6        # Normalisation constant of the viscous time-scale in the accretion reduction term
   with_boost_factor:                  0          # Are we using the model from Booth & Schaye (2009)?
diff --git a/examples/EAGLE_low_z/EAGLE_6/eagle_6.yml b/examples/EAGLE_low_z/EAGLE_6/eagle_6.yml
index 3cb8c288ca6b508efe76be8ec1a9767b5c2f8c7f..1bb9b94b9341e4f285080d41523d7f0191792572 100644
--- a/examples/EAGLE_low_z/EAGLE_6/eagle_6.yml
+++ b/examples/EAGLE_low_z/EAGLE_6/eagle_6.yml
@@ -205,8 +205,8 @@ EAGLEFeedback:
 # EAGLE AGN model
 EAGLEAGN:
   subgrid_seed_mass_Msun:             1.5e5      # Black hole subgrid mass at creation time in solar masses.
-  multi_phase_bondi:                  0          # Compute Bondi rates per neighbour particle?
-  subgrid_bondi:                      0          # Compute Bondi rates using the subgrid extrapolation of the gas properties around the BH?
+  use_multi_phase_bondi:              0          # Compute Bondi rates per neighbour particle?
+  use_subgrid_bondi:                  0          # Compute Bondi rates using the subgrid extrapolation of the gas properties around the BH?
   with_angmom_limiter:                1          # Are we applying the Rosas-Guevara et al. (2015) viscous time-scale reduction term?
   viscous_alpha:                      1e6        # Normalisation constant of the viscous time-scale in the accretion reduction term
   with_boost_factor:                  0          # Are we using the model from Booth & Schaye (2009)?
diff --git a/examples/GEAR/AgoraDisk/agora_disk.yml b/examples/GEAR/AgoraDisk/agora_disk.yml
index e85fe1e9a8bf75a4df0451d959f9b4692ebc1d7f..18b9b618026dbc4a4cf26ff046041365990293d7 100644
--- a/examples/GEAR/AgoraDisk/agora_disk.yml
+++ b/examples/GEAR/AgoraDisk/agora_disk.yml
@@ -88,7 +88,7 @@ GrackleCooling:
 GEARStarFormation:
   star_formation_efficiency: 0.01   # star formation efficiency (c_*)
   maximal_temperature:  1e10         # Upper limit to the temperature of a star forming particle
-  n_stars_per_particle: 4
+  n_stars_per_particle: 1
   min_mass_frac: 0.5
 
 GEARPressureFloor:
@@ -107,5 +107,5 @@ GEARChemistry:
   scale_initial_metallicity: 1
 
 Restarts:
-  delta_hours:        72        # (Optional) decimal hours between dumps of restart files.
+  delta_hours:   1.        # (Optional) decimal hours between dumps of restart files.
 
diff --git a/examples/GEAR/AgoraDisk/plotSolution.py b/examples/GEAR/AgoraDisk/plotSolution.py
index b3a831b2d902c901f3e06c968c7fef28d4db3429..c17c67aebae9b0b699fa854a50bb6601893d85ed 100644
--- a/examples/GEAR/AgoraDisk/plotSolution.py
+++ b/examples/GEAR/AgoraDisk/plotSolution.py
@@ -717,10 +717,10 @@ for time in range(len(times)):
                                 pf.add_field(("gas", "metallicity"), function=_metallicity_3, force_override=True, display_name="Metallicity", particle_type=False, take_log=True, units="")
                         elif codes[code] == 'SWIFT': # "Metals" in SWIFT is 10-species field ([:,9] is the total metal fraction), so requires a change in _vector_fields in frontends/gadget/io.py: added ("Metals", 10)
                                 def _metallicity_2(field, data):
-                                        if len(data[PartType_Gas_to_use, "SmoothedElementAbundances"].shape) == 1:
-                                                return data[PartType_Gas_to_use, "SmoothedElementAbundances"]
+                                        if len(data[PartType_Gas_to_use, "SmoothedMetalMassFractions"].shape) == 1:
+                                                return data[PartType_Gas_to_use, "SmoothedMetalMassFractions"]
                                         else:
-                                                return data[PartType_Gas_to_use, "SmoothedElementAbundances"][:,9].in_units("") # in_units("") turned out to be crucial!; otherwise code_metallicity will be used and it will mess things up
+                                                return data[PartType_Gas_to_use, "SmoothedMetalMassFractions"][:,9].in_units("") # in_units("") turned out to be crucial!; otherwise code_metallicity will be used and it will mess things up
                                 # We are creating ("Gas", "Metallicity") here, different from ("Gas", "metallicity") which is auto-generated by yt but doesn't work properly
                                 pf.add_field((PartType_Gas_to_use, MetallicityType_to_use), function=_metallicity_2, display_name="Metallicity", particle_type=True, take_log=True, units="")
                                 # Also creating smoothed field following an example in yt-project.org/docs/dev/cookbook/calculating_information.html; use hardcoded num_neighbors as in frontends/gadget/fields.py
diff --git a/examples/GEAR/ZoomIn/README b/examples/GEAR/ZoomIn/README
index cffc275f2ae1046156d392f8725a7b542c80471a..b41ebabb38ea4478519069ad2e0b0e339afe88fd 100644
--- a/examples/GEAR/ZoomIn/README
+++ b/examples/GEAR/ZoomIn/README
@@ -6,11 +6,10 @@ The cosmology is taken from Planck 2015.
 
 The initial conditions have been cleaned to contain only the required
 fields. The ICs have been created for Gadget and the positions and box
-size are hence expressed in h-full units (e.g. box size of 32 / h Mpc).
+size are hence expressed in h-full units (e.g. box size of 3.4 / h Mpc).
 Similarly, the peculiar velocitites contain an extra sqrt(a) factor. 
 
-We will use SWIFT to cancel the h- and a-factors from the ICs. Gas
-particles will be generated at startup.
+We will use SWIFT to cancel the h- and a-factors from the ICs.
 
 MD5 check-sum of the ICS: 
 9aafe154438478ed435e88664c1c5dba zoom_in.hdf5
diff --git a/examples/GEAR/ZoomIn/snaplist.txt b/examples/GEAR/ZoomIn/snaplist.txt
new file mode 100644
index 0000000000000000000000000000000000000000..a14ce398716247b094a5cbe3c1730bab582b50ed
--- /dev/null
+++ b/examples/GEAR/ZoomIn/snaplist.txt
@@ -0,0 +1,1012 @@
+# Scale Factor
+0.0141
+0.0146484375
+0.0156250000
+0.0166015625
+0.0175781250
+0.0185546875
+0.0195312500
+0.0205078125
+0.0214843750
+0.0224609375
+0.0234375000
+0.0244140625
+0.0253906250
+0.0263671875
+0.0273437500
+0.0283203125
+0.0292968750
+0.0302734375
+0.0312500000
+0.0322265625
+0.0332031250
+0.0341796875
+0.0351562500
+0.0361328125
+0.0371093750
+0.0380859375
+0.0390625000
+0.0400390625
+0.0410156250
+0.0419921875
+0.0429687500
+0.0439453125
+0.0449218750
+0.0458984375
+0.0468750000
+0.0478515625
+0.0488281250
+0.0498046875
+0.0507812500
+0.0517578125
+0.0527343750
+0.0537109375
+0.0546875000
+0.0556640625
+0.0566406250
+0.0576171875
+0.0585937500
+0.0595703125
+0.0605468750
+0.0615234375
+0.0625000000
+0.0634765625
+0.0644531250
+0.0654296875
+0.0664062500
+0.0673828125
+0.0683593750
+0.0693359375
+0.0703125000
+0.0712890625
+0.0722656250
+0.0732421875
+0.0742187500
+0.0751953125
+0.0761718750
+0.0771484375
+0.0781250000
+0.0791015625
+0.0800781250
+0.0810546875
+0.0820312500
+0.0830078125
+0.0839843750
+0.0849609375
+0.0859375000
+0.0869140625
+0.0878906250
+0.0888671875
+0.0898437500
+0.0908203125
+0.0917968750
+0.0927734375
+0.0937500000
+0.0947265625
+0.0957031250
+0.0966796875
+0.0976562500
+0.0986328125
+0.0996093750
+0.1005859375
+0.1015625000
+0.1025390625
+0.1035156250
+0.1044921875
+0.1054687500
+0.1064453125
+0.1074218750
+0.1083984375
+0.1093750000
+0.1103515625
+0.1113281250
+0.1123046875
+0.1132812500
+0.1142578125
+0.1152343750
+0.1162109375
+0.1171875000
+0.1181640625
+0.1191406250
+0.1201171875
+0.1210937500
+0.1220703125
+0.1230468750
+0.1240234375
+0.1250000000
+0.1259765625
+0.1269531250
+0.1279296875
+0.1289062500
+0.1298828125
+0.1308593750
+0.1318359375
+0.1328125000
+0.1337890625
+0.1347656250
+0.1357421875
+0.1367187500
+0.1376953125
+0.1386718750
+0.1396484375
+0.1406250000
+0.1416015625
+0.1425781250
+0.1435546875
+0.1445312500
+0.1455078125
+0.1464843750
+0.1474609375
+0.1484375000
+0.1494140625
+0.1503906250
+0.1513671875
+0.1523437500
+0.1533203125
+0.1542968750
+0.1552734375
+0.1562500000
+0.1572265625
+0.1582031250
+0.1591796875
+0.1601562500
+0.1611328125
+0.1621093750
+0.1630859375
+0.1640625000
+0.1650390625
+0.1660156250
+0.1669921875
+0.1679687500
+0.1689453125
+0.1699218750
+0.1708984375
+0.1718750000
+0.1728515625
+0.1738281250
+0.1748046875
+0.1757812500
+0.1767578125
+0.1777343750
+0.1787109375
+0.1796875000
+0.1806640625
+0.1816406250
+0.1826171875
+0.1835937500
+0.1845703125
+0.1855468750
+0.1865234375
+0.1875000000
+0.1884765625
+0.1894531250
+0.1904296875
+0.1914062500
+0.1923828125
+0.1933593750
+0.1943359375
+0.1953125000
+0.1962890625
+0.1972656250
+0.1982421875
+0.1992187500
+0.2001953125
+0.2011718750
+0.2021484375
+0.2031250000
+0.2041015625
+0.2050781250
+0.2060546875
+0.2070312500
+0.2080078125
+0.2089843750
+0.2099609375
+0.2109375000
+0.2119140625
+0.2128906250
+0.2138671875
+0.2148437500
+0.2158203125
+0.2167968750
+0.2177734375
+0.2187500000
+0.2197265625
+0.2207031250
+0.2216796875
+0.2226562500
+0.2236328125
+0.2246093750
+0.2255859375
+0.2265625000
+0.2275390625
+0.2285156250
+0.2294921875
+0.2304687500
+0.2314453125
+0.2324218750
+0.2333984375
+0.2343750000
+0.2353515625
+0.2363281250
+0.2373046875
+0.2382812500
+0.2392578125
+0.2402343750
+0.2412109375
+0.2421875000
+0.2431640625
+0.2441406250
+0.2451171875
+0.2460937500
+0.2470703125
+0.2480468750
+0.2490234375
+0.2500000000
+0.2509765625
+0.2519531250
+0.2529296875
+0.2539062500
+0.2548828125
+0.2558593750
+0.2568359375
+0.2578125000
+0.2587890625
+0.2597656250
+0.2607421875
+0.2617187500
+0.2626953125
+0.2636718750
+0.2646484375
+0.2656250000
+0.2666015625
+0.2675781250
+0.2685546875
+0.2695312500
+0.2705078125
+0.2714843750
+0.2724609375
+0.2734375000
+0.2744140625
+0.2753906250
+0.2763671875
+0.2773437500
+0.2783203125
+0.2792968750
+0.2802734375
+0.2812500000
+0.2822265625
+0.2832031250
+0.2841796875
+0.2851562500
+0.2861328125
+0.2871093750
+0.2880859375
+0.2890625000
+0.2900390625
+0.2910156250
+0.2919921875
+0.2929687500
+0.2939453125
+0.2949218750
+0.2958984375
+0.2968750000
+0.2978515625
+0.2988281250
+0.2998046875
+0.3007812500
+0.3017578125
+0.3027343750
+0.3037109375
+0.3046875000
+0.3056640625
+0.3066406250
+0.3076171875
+0.3085937500
+0.3095703125
+0.3105468750
+0.3115234375
+0.3125000000
+0.3134765625
+0.3144531250
+0.3154296875
+0.3164062500
+0.3173828125
+0.3183593750
+0.3193359375
+0.3203125000
+0.3212890625
+0.3222656250
+0.3232421875
+0.3242187500
+0.3251953125
+0.3261718750
+0.3271484375
+0.3281250000
+0.3291015625
+0.3300781250
+0.3310546875
+0.3320312500
+0.3330078125
+0.3339843750
+0.3349609375
+0.3359375000
+0.3369140625
+0.3378906250
+0.3388671875
+0.3398437500
+0.3408203125
+0.3417968750
+0.3427734375
+0.3437500000
+0.3447265625
+0.3457031250
+0.3466796875
+0.3476562500
+0.3486328125
+0.3496093750
+0.3505859375
+0.3515625000
+0.3525390625
+0.3535156250
+0.3544921875
+0.3554687500
+0.3564453125
+0.3574218750
+0.3583984375
+0.3593750000
+0.3603515625
+0.3613281250
+0.3623046875
+0.3632812500
+0.3642578125
+0.3652343750
+0.3662109375
+0.3671875000
+0.3681640625
+0.3691406250
+0.3701171875
+0.3710937500
+0.3720703125
+0.3730468750
+0.3740234375
+0.3750000000
+0.3759765625
+0.3769531250
+0.3779296875
+0.3789062500
+0.3798828125
+0.3808593750
+0.3818359375
+0.3828125000
+0.3837890625
+0.3847656250
+0.3857421875
+0.3867187500
+0.3876953125
+0.3886718750
+0.3896484375
+0.3906250000
+0.3916015625
+0.3925781250
+0.3935546875
+0.3945312500
+0.3955078125
+0.3964843750
+0.3974609375
+0.3984375000
+0.3994140625
+0.4003906250
+0.4013671875
+0.4023437500
+0.4033203125
+0.4042968750
+0.4052734375
+0.4062500000
+0.4072265625
+0.4082031250
+0.4091796875
+0.4101562500
+0.4111328125
+0.4121093750
+0.4130859375
+0.4140625000
+0.4150390625
+0.4160156250
+0.4169921875
+0.4179687500
+0.4189453125
+0.4199218750
+0.4208984375
+0.4218750000
+0.4228515625
+0.4238281250
+0.4248046875
+0.4257812500
+0.4267578125
+0.4277343750
+0.4287109375
+0.4296875000
+0.4306640625
+0.4316406250
+0.4326171875
+0.4335937500
+0.4345703125
+0.4355468750
+0.4365234375
+0.4375000000
+0.4384765625
+0.4394531250
+0.4404296875
+0.4414062500
+0.4423828125
+0.4433593750
+0.4443359375
+0.4453125000
+0.4462890625
+0.4472656250
+0.4482421875
+0.4492187500
+0.4501953125
+0.4511718750
+0.4521484375
+0.4531250000
+0.4541015625
+0.4550781250
+0.4560546875
+0.4570312500
+0.4580078125
+0.4589843750
+0.4599609375
+0.4609375000
+0.4619140625
+0.4628906250
+0.4638671875
+0.4648437500
+0.4658203125
+0.4667968750
+0.4677734375
+0.4687500000
+0.4697265625
+0.4707031250
+0.4716796875
+0.4726562500
+0.4736328125
+0.4746093750
+0.4755859375
+0.4765625000
+0.4775390625
+0.4785156250
+0.4794921875
+0.4804687500
+0.4814453125
+0.4824218750
+0.4833984375
+0.4843750000
+0.4853515625
+0.4863281250
+0.4873046875
+0.4882812500
+0.4892578125
+0.4902343750
+0.4912109375
+0.4921875000
+0.4931640625
+0.4941406250
+0.4951171875
+0.4960937500
+0.4970703125
+0.4980468750
+0.4990234375
+0.5000000000
+0.5009765625
+0.5019531250
+0.5029296875
+0.5039062500
+0.5048828125
+0.5058593750
+0.5068359375
+0.5078125000
+0.5087890625
+0.5097656250
+0.5107421875
+0.5117187500
+0.5126953125
+0.5136718750
+0.5146484375
+0.5156250000
+0.5166015625
+0.5175781250
+0.5185546875
+0.5195312500
+0.5205078125
+0.5214843750
+0.5224609375
+0.5234375000
+0.5244140625
+0.5253906250
+0.5263671875
+0.5273437500
+0.5283203125
+0.5292968750
+0.5302734375
+0.5312500000
+0.5322265625
+0.5332031250
+0.5341796875
+0.5351562500
+0.5361328125
+0.5371093750
+0.5380859375
+0.5390625000
+0.5400390625
+0.5410156250
+0.5419921875
+0.5429687500
+0.5439453125
+0.5449218750
+0.5458984375
+0.5468750000
+0.5478515625
+0.5488281250
+0.5498046875
+0.5507812500
+0.5517578125
+0.5527343750
+0.5537109375
+0.5546875000
+0.5556640625
+0.5566406250
+0.5576171875
+0.5585937500
+0.5595703125
+0.5605468750
+0.5615234375
+0.5625000000
+0.5634765625
+0.5644531250
+0.5654296875
+0.5664062500
+0.5673828125
+0.5683593750
+0.5693359375
+0.5703125000
+0.5712890625
+0.5722656250
+0.5732421875
+0.5742187500
+0.5751953125
+0.5761718750
+0.5771484375
+0.5781250000
+0.5791015625
+0.5800781250
+0.5810546875
+0.5820312500
+0.5830078125
+0.5839843750
+0.5849609375
+0.5859375000
+0.5869140625
+0.5878906250
+0.5888671875
+0.5898437500
+0.5908203125
+0.5917968750
+0.5927734375
+0.5937500000
+0.5947265625
+0.5957031250
+0.5966796875
+0.5976562500
+0.5986328125
+0.5996093750
+0.6005859375
+0.6015625000
+0.6025390625
+0.6035156250
+0.6044921875
+0.6054687500
+0.6064453125
+0.6074218750
+0.6083984375
+0.6093750000
+0.6103515625
+0.6113281250
+0.6123046875
+0.6132812500
+0.6142578125
+0.6152343750
+0.6162109375
+0.6171875000
+0.6181640625
+0.6191406250
+0.6201171875
+0.6210937500
+0.6220703125
+0.6230468750
+0.6240234375
+0.6250000000
+0.6259765625
+0.6269531250
+0.6279296875
+0.6289062500
+0.6298828125
+0.6308593750
+0.6318359375
+0.6328125000
+0.6337890625
+0.6347656250
+0.6357421875
+0.6367187500
+0.6376953125
+0.6386718750
+0.6396484375
+0.6406250000
+0.6416015625
+0.6425781250
+0.6435546875
+0.6445312500
+0.6455078125
+0.6464843750
+0.6474609375
+0.6484375000
+0.6494140625
+0.6503906250
+0.6513671875
+0.6523437500
+0.6533203125
+0.6542968750
+0.6552734375
+0.6562500000
+0.6572265625
+0.6582031250
+0.6591796875
+0.6601562500
+0.6611328125
+0.6621093750
+0.6630859375
+0.6640625000
+0.6650390625
+0.6660156250
+0.6669921875
+0.6679687500
+0.6689453125
+0.6699218750
+0.6708984375
+0.6718750000
+0.6728515625
+0.6738281250
+0.6748046875
+0.6757812500
+0.6767578125
+0.6777343750
+0.6787109375
+0.6796875000
+0.6806640625
+0.6816406250
+0.6826171875
+0.6835937500
+0.6845703125
+0.6855468750
+0.6865234375
+0.6875000000
+0.6884765625
+0.6894531250
+0.6904296875
+0.6914062500
+0.6923828125
+0.6933593750
+0.6943359375
+0.6953125000
+0.6962890625
+0.6972656250
+0.6982421875
+0.6992187500
+0.7001953125
+0.7011718750
+0.7021484375
+0.7031250000
+0.7041015625
+0.7050781250
+0.7060546875
+0.7070312500
+0.7080078125
+0.7089843750
+0.7099609375
+0.7109375000
+0.7119140625
+0.7128906250
+0.7138671875
+0.7148437500
+0.7158203125
+0.7167968750
+0.7177734375
+0.7187500000
+0.7197265625
+0.7207031250
+0.7216796875
+0.7226562500
+0.7236328125
+0.7246093750
+0.7255859375
+0.7265625000
+0.7275390625
+0.7285156250
+0.7294921875
+0.7304687500
+0.7314453125
+0.7324218750
+0.7333984375
+0.7343750000
+0.7353515625
+0.7363281250
+0.7373046875
+0.7382812500
+0.7392578125
+0.7402343750
+0.7412109375
+0.7421875000
+0.7431640625
+0.7441406250
+0.7451171875
+0.7460937500
+0.7470703125
+0.7480468750
+0.7490234375
+0.7500000000
+0.7509765625
+0.7519531250
+0.7529296875
+0.7539062500
+0.7548828125
+0.7558593750
+0.7568359375
+0.7578125000
+0.7587890625
+0.7597656250
+0.7607421875
+0.7617187500
+0.7626953125
+0.7636718750
+0.7646484375
+0.7656250000
+0.7666015625
+0.7675781250
+0.7685546875
+0.7695312500
+0.7705078125
+0.7714843750
+0.7724609375
+0.7734375000
+0.7744140625
+0.7753906250
+0.7763671875
+0.7773437500
+0.7783203125
+0.7792968750
+0.7802734375
+0.7812500000
+0.7822265625
+0.7832031250
+0.7841796875
+0.7851562500
+0.7861328125
+0.7871093750
+0.7880859375
+0.7890625000
+0.7900390625
+0.7910156250
+0.7919921875
+0.7929687500
+0.7939453125
+0.7949218750
+0.7958984375
+0.7968750000
+0.7978515625
+0.7988281250
+0.7998046875
+0.8007812500
+0.8017578125
+0.8027343750
+0.8037109375
+0.8046875000
+0.8056640625
+0.8066406250
+0.8076171875
+0.8085937500
+0.8095703125
+0.8105468750
+0.8115234375
+0.8125000000
+0.8134765625
+0.8144531250
+0.8154296875
+0.8164062500
+0.8173828125
+0.8183593750
+0.8193359375
+0.8203125000
+0.8212890625
+0.8222656250
+0.8232421875
+0.8242187500
+0.8251953125
+0.8261718750
+0.8271484375
+0.8281250000
+0.8291015625
+0.8300781250
+0.8310546875
+0.8320312500
+0.8330078125
+0.8339843750
+0.8349609375
+0.8359375000
+0.8369140625
+0.8378906250
+0.8388671875
+0.8398437500
+0.8408203125
+0.8417968750
+0.8427734375
+0.8437500000
+0.8447265625
+0.8457031250
+0.8466796875
+0.8476562500
+0.8486328125
+0.8496093750
+0.8505859375
+0.8515625000
+0.8525390625
+0.8535156250
+0.8544921875
+0.8554687500
+0.8564453125
+0.8574218750
+0.8583984375
+0.8593750000
+0.8603515625
+0.8613281250
+0.8623046875
+0.8632812500
+0.8642578125
+0.8652343750
+0.8662109375
+0.8671875000
+0.8681640625
+0.8691406250
+0.8701171875
+0.8710937500
+0.8720703125
+0.8730468750
+0.8740234375
+0.8750000000
+0.8759765625
+0.8769531250
+0.8779296875
+0.8789062500
+0.8798828125
+0.8808593750
+0.8818359375
+0.8828125000
+0.8837890625
+0.8847656250
+0.8857421875
+0.8867187500
+0.8876953125
+0.8886718750
+0.8896484375
+0.8906250000
+0.8916015625
+0.8925781250
+0.8935546875
+0.8945312500
+0.8955078125
+0.8964843750
+0.8974609375
+0.8984375000
+0.8994140625
+0.9003906250
+0.9013671875
+0.9023437500
+0.9033203125
+0.9042968750
+0.9052734375
+0.9062500000
+0.9072265625
+0.9082031250
+0.9091796875
+0.9101562500
+0.9111328125
+0.9121093750
+0.9130859375
+0.9140625000
+0.9150390625
+0.9160156250
+0.9169921875
+0.9179687500
+0.9189453125
+0.9199218750
+0.9208984375
+0.9218750000
+0.9228515625
+0.9238281250
+0.9248046875
+0.9257812500
+0.9267578125
+0.9277343750
+0.9287109375
+0.9296875000
+0.9306640625
+0.9316406250
+0.9326171875
+0.9335937500
+0.9345703125
+0.9355468750
+0.9365234375
+0.9375000000
+0.9384765625
+0.9394531250
+0.9404296875
+0.9414062500
+0.9423828125
+0.9433593750
+0.9443359375
+0.9453125000
+0.9462890625
+0.9472656250
+0.9482421875
+0.9492187500
+0.9501953125
+0.9511718750
+0.9521484375
+0.9531250000
+0.9541015625
+0.9550781250
+0.9560546875
+0.9570312500
+0.9580078125
+0.9589843750
+0.9599609375
+0.9609375000
+0.9619140625
+0.9628906250
+0.9638671875
+0.9648437500
+0.9658203125
+0.9667968750
+0.9677734375
+0.9687500000
+0.9697265625
+0.9707031250
+0.9716796875
+0.9726562500
+0.9736328125
+0.9746093750
+0.9755859375
+0.9765625000
+0.9775390625
+0.9785156250
+0.9794921875
+0.9804687500
+0.9814453125
+0.9824218750
+0.9833984375
+0.9843750000
+0.9853515625
+0.9863281250
+0.9873046875
+0.9882812500
+0.9892578125
+0.9902343750
+0.9912109375
+0.9921875000
+0.9931640625
+0.9941406250
+0.9951171875
+0.9960937500
+0.9970703125
+0.9980468750
+0.9990234375
+1.0000000000
diff --git a/examples/GEAR/ZoomIn/zoom_in.yml b/examples/GEAR/ZoomIn/zoom_in.yml
index d17e285b8024e8366292322d54a5d63135347973..5dd2d508e3ab400a01b26697c5715bfb71a11ef5 100644
--- a/examples/GEAR/ZoomIn/zoom_in.yml
+++ b/examples/GEAR/ZoomIn/zoom_in.yml
@@ -1,7 +1,7 @@
 # Define the system of units to use internally. 
 InternalUnitSystem:
   UnitMass_in_cgs:     1.98841e43    # 10^10 M_sun in grams
-  UnitLength_in_cgs:   3.08567758e24 # Mpc in centimeters
+  UnitLength_in_cgs:   3.08567758e21 # kpc in centimeters
   UnitVelocity_in_cgs: 1e5           # km/s in centimeters per second
   UnitCurrent_in_cgs:  1             # Amperes
   UnitTemp_in_cgs:     1             # Kelvin
@@ -52,10 +52,10 @@ Gravity:
   MAC:                           geometric  # Choice of mulitpole acceptance criterion: 'adaptive' OR 'geometric'.
   epsilon_fmm:                   0.001     # Tolerance parameter for the adaptive multipole acceptance criterion.
   theta_cr:                      0.7       # Opening angle for the purely gemoetric criterion.
-  comoving_DM_softening:     0.1278 # Comoving softening length (in internal units).
-  max_physical_DM_softening: 0.03365    # Physical softening length (in internal units).
-  comoving_baryon_softening:     0.03365 # Comoving softening length (in internal units).
-  max_physical_baryon_softening: 0.00673    # Physical softening length (in internal units).
+  comoving_DM_softening:     0.2823 # Comoving softening length (in internal units).
+  max_physical_DM_softening: 0.07429    # Physical softening length (in internal units).
+  comoving_baryon_softening:     0.07429 # Comoving softening length (in internal units).
+  max_physical_baryon_softening: 0.01485    # Physical softening length (in internal units).
   softening_ratio_background:    0.0285714      # Fraction of the mean inter-particle separation to use as Plummer-equivalent softening for the background DM particles.
   mesh_side_length:       128        # Number of cells along each axis for the periodic gravity mesh.
   
@@ -101,7 +101,7 @@ GEARPressureFloor:
   jeans_factor: 10
 
 GEARFeedback:
-  supernovae_energy_erg: 0.135e51
+  supernovae_energy_erg: 0.1e51
   yields_table: chemistry-AGB+OMgSFeZnSrYBaEu-16072013.h5
   discrete_yields: 1
   yields_table_first_stars: chemistry-PopIII.hdf5          # Table containing the yields of the first stars.
diff --git a/examples/SubgridTests/CosmologicalStellarEvolution/run.sh b/examples/SubgridTests/CosmologicalStellarEvolution/run.sh
index 637a100567f691d883b95f78c4513e0640c1932a..ec267259e95368dcbda1282a8c625e5fd2647e40 100755
--- a/examples/SubgridTests/CosmologicalStellarEvolution/run.sh
+++ b/examples/SubgridTests/CosmologicalStellarEvolution/run.sh
@@ -26,22 +26,22 @@ then
     ./getSolutions.sh
 fi
 
-../../swift  --feedback --stars --hydro --cosmology --external-gravity --threads=4 stellar_evolution.yml -P EAGLEChemistry:init_abundance_metal:0.08 2>&1 | tee output_0p08.log
+../../swift  --temperature --feedback --stars --hydro --cosmology --external-gravity --threads=4 stellar_evolution.yml -P EAGLEChemistry:init_abundance_metal:0.08 -P EAGLEChemistry:init_abundance_Hydrogen:0.71 -P EAGLEChemistry:init_abundance_Helium:0.21 2>&1 | tee output_0p08.log
 
-python plot_box_evolution.py
+python3 plot_box_evolution.py
 
-../../swift  --feedback --stars --hydro --cosmology --external-gravity --threads=4 stellar_evolution.yml -P EAGLEChemistry:init_abundance_metal:0.04 2>&1 | tee output_0p04.log
+../../swift  --temperature --feedback --stars --hydro --cosmology --external-gravity --threads=4 stellar_evolution.yml -P EAGLEChemistry:init_abundance_metal:0.04 -P EAGLEChemistry:init_abundance_Hydrogen:0.74 -P EAGLEChemistry:init_abundance_Helium:0.23 2>&1 | tee output_0p04.log
 
-python plot_box_evolution.py
+python3 plot_box_evolution.py
 
-../../swift  --feedback --stars --hydro --cosmology --external-gravity --threads=4 stellar_evolution.yml -P EAGLEChemistry:init_abundance_metal:0.01 2>&1 | tee output_0p01.log
+../../swift  --temperature --feedback --stars --hydro --cosmology --external-gravity --threads=4 stellar_evolution.yml -P EAGLEChemistry:init_abundance_metal:0.01 2>&1 | tee output_0p01.log
 
-python plot_box_evolution.py
+python3 plot_box_evolution.py
 
-../../swift  --feedback --stars --hydro --cosmology --external-gravity --threads=4 stellar_evolution.yml -P EAGLEChemistry:init_abundance_metal:0.001 2>&1 | tee output_0p001.log
+../../swift  --temperature --feedback --stars --hydro --cosmology --external-gravity --threads=4 stellar_evolution.yml -P EAGLEChemistry:init_abundance_metal:0.001 2>&1 | tee output_0p001.log
 
-python plot_box_evolution.py
+python3 plot_box_evolution.py
 
-../../swift  --feedback --stars --hydro --cosmology --external-gravity --threads=4 stellar_evolution.yml -P EAGLEChemistry:init_abundance_metal:0.0001 2>&1 | tee output_0p0001.log
+../../swift  --temperature --feedback --stars --hydro --cosmology --external-gravity --threads=4 stellar_evolution.yml -P EAGLEChemistry:init_abundance_metal:0.0001 2>&1 | tee output_0p0001.log
 
-python plot_box_evolution.py
+python3 plot_box_evolution.py
diff --git a/examples/SubgridTests/CosmologicalStellarEvolution/stellar_evolution.yml b/examples/SubgridTests/CosmologicalStellarEvolution/stellar_evolution.yml
index e6b24b5194ed2dd40515ebf22292939d483b263f..032ac26df052fb8f11c13d20bb29dd197b4af304 100644
--- a/examples/SubgridTests/CosmologicalStellarEvolution/stellar_evolution.yml
+++ b/examples/SubgridTests/CosmologicalStellarEvolution/stellar_evolution.yml
@@ -96,6 +96,7 @@ EAGLEFeedback:
   IMF_max_mass_Msun:                  100.0             # Maximal stellar mass considered for the Chabrier IMF in solar masses.
   SNII_min_mass_Msun:                   6.0             # Minimal mass considered for SNII stars in solar masses.
   SNII_max_mass_Msun:                 100.0             # Maximal mass considered for SNII stars in solar masses.
+  SNII_feedback_model:                  Random          # Feedback modes: Random, Isotropic, MinimumDistance, MinimumDensity
   SNII_sampled_delay:                   0               # Sample the SNII lifetimes to do feedback.
   SNII_wind_delay_Gyr:                  0.03            # Time in Gyr between a star's birth and the SNII thermal feedback event when not sampling.
   SNII_delta_T_K:                       3.16228e7       # Change in temperature to apply to the gas particle in a SNII thermal feedback event in Kelvin.
@@ -106,7 +107,8 @@ EAGLEFeedback:
   SNII_energy_fraction_n_0_H_p_cm3:     0.67            # Pivot point for the birth density dependance of the SNII energy fraction in cm^-3.
   SNII_energy_fraction_n_Z:             0.8686          # Power-law for the metallicity dependance of the SNII energy fraction.
   SNII_energy_fraction_n_n:             0.8686          # Power-law for the birth density dependance of the SNII energy fraction.
-  SNII_energy_fraction_use_birth_props: 0               # Are we using the density and metallicity at birth to compute f_E or at feedback time?
+  SNII_energy_fraction_use_birth_density: 0             # Are we using the density and metallicity at birth to compute f_E or at feedback time?
+  SNII_energy_fraction_use_birth_metallicity: 0         # Are we using the density and metallicity at birth to compute f_E or at feedback time?
   SNIa_DTD:                             Exponential     # Use the EAGLE-Ref SNIa DTD.
   SNIa_DTD_delay_Gyr:                   0.04            # Age of the most massive SNIa in Gyr.
   SNIa_DTD_exp_timescale_Gyr:           2.0             # Time-scale of the exponential decay of the SNIa rates in Gyr.
diff --git a/examples/SubgridTests/SmoothedMetallicity/makeIC.py b/examples/SubgridTests/SmoothedMetallicity/makeIC.py
index 542b4c5911c942015d16595f42e73ca8978d20da..ce5a2b857763aff6bf690995353d98f84da53dbc 100644
--- a/examples/SubgridTests/SmoothedMetallicity/makeIC.py
+++ b/examples/SubgridTests/SmoothedMetallicity/makeIC.py
@@ -1,40 +1,45 @@
 ###############################################################################
 # This file is part of SWIFT.
 # Copyright (c) 2016 Matthieu Schaller (matthieu.schaller@durham.ac.uk)
-# 
+#
 # 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 h5py
 import numpy as np
 
-# Generates a swift IC file for the Smoothed Metallicity test in a periodic cubic box
+# Generates a swift IC file for the Smoothed Metallicity test
+# in a periodic cubic box
 
 # Parameters
 gamma = 5./3.      # Gas adiabatic index
 rho0 = 1.          # Background density
-P0 = 1.e-6         # Background pressure
-Nelem = 9          # Gear: 9, EAGLE: 9
+P0 = 1e-6          # Background pressure
+Nelem = 10         # Gear: 10, EAGLE: 9
 low_metal = -6     # Low iron fraction
-high_metal = -5    # high iron fraction
-sigma_metal = 0.1  # relative standard deviation for the metallicities
+high_metal = -5.5  # high iron fraction
+max_shift = 1      # Shift between the different elements
+sigma_metal = 0.2  # relative standard deviation for the metallicities
 fileName = "smoothed_metallicity.hdf5"
 
 # shift all metals in order to obtain nicer plots
-low_metal = [low_metal] * Nelem + np.linspace(0, 3, Nelem)
-high_metal = [high_metal] * Nelem + np.linspace(0, 3, Nelem)
+low_metal = [low_metal] * Nelem + np.linspace(0, max_shift, Nelem)
+low_metal = 10**low_metal
+
+high_metal = [high_metal] * Nelem + np.linspace(0, max_shift, Nelem)
+high_metal = 10**high_metal
 
 # ---------------------------------------------------
 glass = h5py.File("glassCube_32.hdf5", "r")
@@ -52,21 +57,24 @@ ids = np.linspace(1, numPart, numPart)
 m = np.zeros(numPart)
 u = np.zeros(numPart)
 r = np.zeros(numPart)
-Z = np.zeros((numPart, Nelem))
+mass_frac = np.zeros((numPart, Nelem))
 
-r = np.sqrt((pos[:, 0] - 0.5)**2 + (pos[:, 1] - 0.5)**2 + (pos[:, 2] - 0.5)**2)
 m[:] = rho0 * vol / numPart
 u[:] = P0 / (rho0 * (gamma - 1))
 
 # set metallicities
 select = pos[:, 0] < 0.5
 nber = sum(select)
-Z[select, :] = low_metal * (1 + np.random.normal(loc=0., scale=sigma_metal,
-                                                 size=(nber, Nelem)))
+mass_frac[select, :] = low_metal * (
+    1 + np.random.normal(loc=0., scale=sigma_metal,
+                         size=(nber, Nelem)))
 nber = numPart - nber
-Z[np.logical_not(select), :] = high_metal * (1 + np.random.normal(
+mass_frac[~select, :] = high_metal * (1 + np.random.normal(
     loc=0., scale=sigma_metal, size=(nber, Nelem)))
 
+v[select, 2] = 1
+v[~select, 2] = -1
+
 # --------------------------------------------------
 
 # File
@@ -100,7 +108,8 @@ grp.create_dataset('Masses', data=m, dtype='f')
 grp.create_dataset('SmoothingLength', data=h, dtype='f')
 grp.create_dataset('InternalEnergy', data=u, dtype='f')
 grp.create_dataset('ParticleIDs', data=ids, dtype='L')
-grp.create_dataset('ElementAbundance', data=Z, dtype='f')
+grp.create_dataset('MetalMassFraction', data=mass_frac, dtype='d')
+grp.create_dataset('ElementAbundance', data=mass_frac, dtype='d')
 
 
 file.close()
diff --git a/examples/SubgridTests/SmoothedMetallicity/plotSolution.py b/examples/SubgridTests/SmoothedMetallicity/plotSolution.py
index 068fe5378e19c34ee8a68398f4e0ed096d0982e0..ab5e18d7d9d9bb15b154693e61c6fb93bb8bcc79 100644
--- a/examples/SubgridTests/SmoothedMetallicity/plotSolution.py
+++ b/examples/SubgridTests/SmoothedMetallicity/plotSolution.py
@@ -30,14 +30,17 @@ matplotlib.use("Agg")
 import matplotlib.pyplot as plt
 
 # Parameters
-low_metal = -6  # low metal abundance
-high_metal = -5  # High metal abundance
-sigma_metal = 0.1  # relative standard deviation for Z
+low_metal = -6     # low metal abundance
+high_metal = -5.5  # High metal abundance
+sigma_metal = 0.2  # relative standard deviation for Z
+max_shift = 1      # Shift between the different elements
 
-Nelem = 9
+Nelem = 10
 # shift all metals in order to obtain nicer plots
-low_metal = [low_metal] * Nelem + np.linspace(0, 3, Nelem)
-high_metal = [high_metal] * Nelem + np.linspace(0, 3, Nelem)
+low_metal = [low_metal] * Nelem + np.linspace(0, max_shift, Nelem)
+low_metal = 10**low_metal
+high_metal = [high_metal] * Nelem + np.linspace(0, max_shift, Nelem)
+high_metal = 10**high_metal
 
 # ---------------------------------------------------------------
 # Don't touch anything after this.
@@ -80,12 +83,20 @@ kernel = sim["/HydroScheme"].attrs["Kernel function"]
 neighbours = sim["/HydroScheme"].attrs["Kernel target N_ngb"]
 eta = sim["/HydroScheme"].attrs["Kernel eta"]
 chemistry = sim["/SubgridScheme"].attrs["Chemistry Model"]
+chemistry = str(chemistry)
 git = sim["Code"].attrs["Git Revision"]
 
 pos = sim["/PartType0/Coordinates"][:, :]
 d = pos[:, 0] - boxSize / 2
-smooth_metal = sim["/PartType0/SmoothedElementMassFractions"][:, :]
-metal = sim["/PartType0/ElementMassFractions"][:, :]
+gear_name = "/PartType0/MetalMassFractions"
+eagle_name = "/PartType0/ElementAbundances"
+if gear_name in sim:
+    smooth_metal = sim["/PartType0/SmoothedMetalMassFractions"][:, :]
+    metal = sim[gear_name][:, :]
+else:
+    smooth_metal = sim["/PartType0/SmoothedElementAbundances"][:, :]
+    metal = sim[eagle_name][:, :]
+
 h = sim["/PartType0/SmoothingLengths"][:]
 h = np.mean(h)
 
@@ -165,15 +176,15 @@ plt.figure()
 # Metallicity --------------------------------
 plt.subplot(221)
 for e in range(Nelem):
-    plt.plot(metal[:, e], smooth_metal[:, e], ".", ms=0.5, alpha=0.2)
+    plt.loglog(metal[:, e], smooth_metal[:, e], ".", ms=0.5, alpha=0.2)
 
 xmin, xmax = metal.min(), metal.max()
 ymin, ymax = smooth_metal.min(), smooth_metal.max()
 x = max(xmin, ymin)
 y = min(xmax, ymax)
-plt.plot([x, y], [x, y], "--k", lw=1.0)
-plt.xlabel("${\\rm{Metallicity}}~Z_\\textrm{part}$", labelpad=0)
-plt.ylabel("${\\rm{Smoothed~Metallicity}}~Z_\\textrm{sm}$", labelpad=0)
+plt.loglog([x, y], [x, y], "--k", lw=1.0)
+plt.xlabel("${\\rm{Metal~mass~fraction}}$", labelpad=0)
+plt.ylabel("${\\rm{Smoothed~metal~mass~fraction}}$", labelpad=0)
 
 # Metallicity --------------------------------
 e = 0
@@ -183,10 +194,11 @@ plt.plot(d_a, sol[:, e], "--", color="b", alpha=0.8, lw=1.2)
 plt.fill_between(
     d_a, sig[:, e, 0], sig[:, e, 1], facecolor="b", interpolate=True, alpha=0.5
 )
+plt.yscale("log")
 plt.xlabel("${\\rm{Distance}}~r$", labelpad=0)
-plt.ylabel("${\\rm{Smoothed~Metallicity}}~Z_\\textrm{sm}$", labelpad=0)
+plt.ylabel("${\\rm{Smoothed~metal~mass~fraction}}$", labelpad=0)
 plt.xlim(-0.5, 0.5)
-plt.ylim(low_metal[e] - 1, high_metal[e] + 1)
+# plt.ylim(low_metal[e] - 1, high_metal[e] + 1)
 
 # Information -------------------------------------
 plt.subplot(222, frameon=False)
diff --git a/examples/SubgridTests/SmoothedMetallicity/smoothed_metallicity.yml b/examples/SubgridTests/SmoothedMetallicity/smoothed_metallicity.yml
index f6841c6bd0744b4bbeacbe136a126b4ed5631f6f..aeb725235f59f9b0d2c2126d3af26ab3ec7e50da 100644
--- a/examples/SubgridTests/SmoothedMetallicity/smoothed_metallicity.yml
+++ b/examples/SubgridTests/SmoothedMetallicity/smoothed_metallicity.yml
@@ -33,3 +33,7 @@ InitialConditions:
   file_name:  ./smoothed_metallicity.hdf5          # The file to read
   periodic:   1
 
+GEARChemistry:
+  scale_initial_metallicity: 0 # Scale the initial metallicity with solar abundances
+  initial_metallicity: -1 # Read from the snapshot
+  diffusion_coefficient: 1e-3 # Coefficient for the diffusion (see Shen et al. 2010)
diff --git a/examples/SubgridTests/StellarEvolution/plot_box_evolution.py b/examples/SubgridTests/StellarEvolution/plot_box_evolution.py
index aa1f07f98ecede241423bc72a3e0cf7d76963979..eb555bc3c87cfda5f6766ce630254ecafd605d7a 100644
--- a/examples/SubgridTests/StellarEvolution/plot_box_evolution.py
+++ b/examples/SubgridTests/StellarEvolution/plot_box_evolution.py
@@ -1,29 +1,30 @@
 ###############################################################################
- # This file is part of SWIFT.
- # Copyright (c) 2018 Matthieu Schaller (matthieu.schaller@durham.ac.uk)
- # 
- # 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/>.
- # 
- ##############################################################################
-
-# Script used to plot time evolution of gas particle properties. Intended to 
+# This file is part of SWIFT.
+# Copyright (c) 2018 Matthieu Schaller (matthieu.schaller@durham.ac.uk)
+#
+# 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/>.
+#
+##############################################################################
+
+# Script used to plot time evolution of gas particle properties. Intended to
 # compare result of feedback due to one star placed in centre of uniform box
 # of gas with output from EAGLE feedback test. Can also use as input output
 # from SWIFT feedback test (tests/testFeedback) with the appropriate change
 # to filepath.
 
 import matplotlib
+
 matplotlib.use("Agg")
 from pylab import *
 from scipy import stats
@@ -33,31 +34,34 @@ import glob
 import os.path
 
 # Plot parameters
-params = {'axes.labelsize': 10,
-'axes.titlesize': 10,
-'font.size': 12,
-'legend.fontsize': 12,
-'xtick.labelsize': 10,
-'ytick.labelsize': 10,
-'text.usetex': True,
- 'figure.figsize' : (9.90,6.45),
-'figure.subplot.left'    : 0.05,
-'figure.subplot.right'   : 0.995,
-'figure.subplot.bottom'  : 0.06,
-'figure.subplot.top'     : 0.92,
-'figure.subplot.wspace'  : 0.25,
-'figure.subplot.hspace'  : 0.2,
-'lines.markersize' : 6,
-'lines.linewidth' : 3.,
-'text.latex.unicode': True
+params = {
+    "axes.labelsize": 10,
+    "axes.titlesize": 10,
+    "font.size": 12,
+    "legend.fontsize": 12,
+    "xtick.labelsize": 10,
+    "ytick.labelsize": 10,
+    "text.usetex": True,
+    "figure.figsize": (9.90, 6.45),
+    "figure.subplot.left": 0.05,
+    "figure.subplot.right": 0.995,
+    "figure.subplot.bottom": 0.06,
+    "figure.subplot.top": 0.92,
+    "figure.subplot.wspace": 0.25,
+    "figure.subplot.hspace": 0.2,
+    "lines.markersize": 6,
+    "lines.linewidth": 3.0,
+    "text.latex.unicode": True,
 }
 rcParams.update(params)
-rc('font',**{'family':'sans-serif','sans-serif':['Times']})
+rc("font", **{"family": "sans-serif", "sans-serif": ["Times"]})
 
 
 # Number of snapshots and elements
-newest_snap_name = max(glob.glob('stellar_evolution_*.hdf5'))#, key=os.path.getctime)
-n_snapshots = int(newest_snap_name.replace('stellar_evolution_','').replace('.hdf5','')) + 1
+newest_snap_name = max(glob.glob("stellar_evolution_*.hdf5"))  # , key=os.path.getctime)
+n_snapshots = (
+    int(newest_snap_name.replace("stellar_evolution_", "").replace(".hdf5", "")) + 1
+)
 n_elements = 9
 
 # Read the simulation data
@@ -72,7 +76,9 @@ git = sim["Code"].attrs["Git Revision"]
 stellar_mass = sim["/PartType4/Masses"][0]
 E_SNII_cgs = double(sim["/Parameters"].attrs["EAGLEFeedback:SNII_energy_erg"])
 E_SNIa_cgs = double(sim["/Parameters"].attrs["EAGLEFeedback:SNIa_energy_erg"])
-ejecta_vel_cgs = double(sim["/Parameters"].attrs["EAGLEFeedback:AGB_ejecta_velocity_km_p_s"]) * 1e5
+ejecta_vel_cgs = (
+    double(sim["/Parameters"].attrs["EAGLEFeedback:AGB_ejecta_velocity_km_p_s"]) * 1e5
+)
 
 # Units
 unit_length_in_cgs = sim["/Units"].attrs["Unit length in cgs (U_L)"]
@@ -81,11 +87,11 @@ unit_time_in_cgs = sim["/Units"].attrs["Unit time in cgs (U_t)"]
 unit_temp_in_cgs = sim["/Units"].attrs["Unit temperature in cgs (U_T)"]
 unit_vel_in_cgs = unit_length_in_cgs / unit_time_in_cgs
 unit_energy_in_cgs = unit_mass_in_cgs * unit_vel_in_cgs * unit_vel_in_cgs
-unit_density_in_cgs = unit_mass_in_cgs*unit_length_in_cgs**-3
-unit_pressure_in_cgs = unit_mass_in_cgs/unit_length_in_cgs*unit_time_in_cgs**-2
-unit_int_energy_in_cgs = unit_energy_in_cgs/unit_mass_in_cgs
-unit_entropy_in_cgs = unit_energy_in_cgs/unit_temp_in_cgs
-Gyr_in_cgs = 1e9 * 365. * 24 * 3600.
+unit_density_in_cgs = unit_mass_in_cgs * unit_length_in_cgs ** -3
+unit_pressure_in_cgs = unit_mass_in_cgs / unit_length_in_cgs * unit_time_in_cgs ** -2
+unit_int_energy_in_cgs = unit_energy_in_cgs / unit_mass_in_cgs
+unit_entropy_in_cgs = unit_energy_in_cgs / unit_temp_in_cgs
+Gyr_in_cgs = 1e9 * 365.0 * 24 * 3600.0
 Msun_in_cgs = 1.98848e33
 
 # Declare arrays to store SWIFT data
@@ -98,148 +104,247 @@ swift_box_gas_metal_mass = zeros(n_snapshots)
 swift_box_gas_metal_mass_AGB = zeros(n_snapshots)
 swift_box_gas_metal_mass_SNII = zeros(n_snapshots)
 swift_box_gas_metal_mass_SNIa = zeros(n_snapshots)
-swift_element_mass = zeros((n_snapshots,n_elements))
+swift_element_mass = zeros((n_snapshots, n_elements))
 swift_internal_energy = zeros(n_snapshots)
 swift_kinetic_energy = zeros(n_snapshots)
 swift_total_energy = zeros(n_snapshots)
-swift_mean_u_start = 0.
+swift_mean_u_start = 0.0
 t = zeros(n_snapshots)
 
 # Read data from snapshots
 for i in range(n_snapshots):
-        #print("reading snapshot "+str(i))
-        sim = h5py.File("stellar_evolution_%04d.hdf5"%i, "r")
-        t[i] = sim["/Header"].attrs["Time"][0]
+    # print("reading snapshot "+str(i))
+    sim = h5py.File("stellar_evolution_%04d.hdf5" % i, "r")
+    t[i] = sim["/Header"].attrs["Time"][0]
+
+    masses = sim["/PartType0/Masses"][:]
+    swift_box_gas_mass[i] = np.sum(masses)
 
-        masses = sim["/PartType0/Masses"][:]
-        swift_box_gas_mass[i] = np.sum(masses)
+    AGB_mass = sim["/PartType0/MassesFromAGB"][:]
+    SNII_mass = sim["/PartType0/MassesFromSNII"][:]
+    SNIa_mass = sim["/PartType0/MassesFromSNIa"][:]
 
-        AGB_mass = sim["/PartType0/MassesFromAGB"][:]
-        SNII_mass = sim["/PartType0/MassesFromSNII"][:]
-        SNIa_mass = sim["/PartType0/MassesFromSNIa"][:]
+    swift_box_gas_mass_AGB[i] = np.sum(AGB_mass)
+    swift_box_gas_mass_SNII[i] = np.sum(SNII_mass)
+    swift_box_gas_mass_SNIa[i] = np.sum(SNIa_mass)
 
-        swift_box_gas_mass_AGB[i] = np.sum(AGB_mass)
-        swift_box_gas_mass_SNII[i] = np.sum(SNII_mass)
-        swift_box_gas_mass_SNIa[i] = np.sum(SNIa_mass)
+    Z_star = sim["/PartType4/MetalMassFractions"][0]
+    star_masses = sim["/PartType4/Masses"][:]
+    swift_box_star_mass[i] = np.sum(star_masses)
 
-        Z_star = sim["/PartType4/MetalMassFractions"][0]
-        star_masses = sim["/PartType4/Masses"][:]
-        swift_box_star_mass[i] = np.sum(star_masses)
+    metallicities = sim["/PartType0/MetalMassFractions"][:]
+    swift_box_gas_metal_mass[i] = np.sum(metallicities * masses)
 
-        metallicities = sim["/PartType0/MetalMassFractions"][:]
-        swift_box_gas_metal_mass[i] = np.sum(metallicities * masses)
+    AGB_Z_frac = sim["/PartType0/MetalMassFractionsFromAGB"][:]
+    SNII_Z_frac = sim["/PartType0/MetalMassFractionsFromSNII"][:]
+    SNIa_Z_frac = sim["/PartType0/MetalMassFractionsFromSNIa"][:]
 
-        AGB_Z_frac = sim["/PartType0/MetalMassFractionsFromAGB"][:]
-        SNII_Z_frac = sim["/PartType0/MetalMassFractionsFromSNII"][:]
-        SNIa_Z_frac = sim["/PartType0/MetalMassFractionsFromSNIa"][:]
+    swift_box_gas_metal_mass_AGB[i] = np.sum(AGB_Z_frac * masses)
+    swift_box_gas_metal_mass_SNII[i] = np.sum(SNII_Z_frac * masses)
+    swift_box_gas_metal_mass_SNIa[i] = np.sum(SNIa_Z_frac * masses)
 
-        swift_box_gas_metal_mass_AGB[i] = np.sum(AGB_Z_frac * masses)
-        swift_box_gas_metal_mass_SNII[i] = np.sum(SNII_Z_frac * masses)
-        swift_box_gas_metal_mass_SNIa[i] = np.sum(SNIa_Z_frac * masses)
+    element_abundances = sim["/PartType0/ElementMassFractions"][:][:]
+    for j in range(n_elements):
+        swift_element_mass[i, j] = np.sum(element_abundances[:, j] * masses)
 
-        element_abundances = sim["/PartType0/ElementMassFractions"][:][:]
-        for j in range(n_elements):
-                swift_element_mass[i,j] = np.sum(element_abundances[:,j] * masses)
+    v = sim["/PartType0/Velocities"][:, :]
+    v2 = v[:, 0] ** 2 + v[:, 1] ** 2 + v[:, 2] ** 2
+    u = sim["/PartType0/InternalEnergies"][:]
+    swift_internal_energy[i] = np.sum(masses * u)
+    swift_kinetic_energy[i] = np.sum(0.5 * masses * v2)
+    swift_total_energy[i] = swift_kinetic_energy[i] + swift_internal_energy[i]
 
-        v = sim["/PartType0/Velocities"][:,:]
-        v2 = v[:,0]**2 + v[:,1]**2 + v[:,2]**2
-        u = sim["/PartType0/InternalEnergies"][:]
-        swift_internal_energy[i] = np.sum(masses * u)
-        swift_kinetic_energy[i] = np.sum(0.5 * masses * v2)
-        swift_total_energy[i] = swift_kinetic_energy[i] + swift_internal_energy[i]
+    if i == 0:
+        swift_mean_u_start = np.mean(u)
 
-        if i == 0:
-                swift_mean_u_start = np.mean(u)
-        
-        sim.close()
+    sim.close()
 
 # Read expected yields from EAGLE. Choose which file to use based on metallicity used when
 # running SWIFT (can be specified in yml file)
-filename = "./StellarEvolutionSolution/Z_%.4f/StellarEvolutionTotal.txt"%Z_star
+filename = "./StellarEvolutionSolution/Z_%.4f/StellarEvolutionTotal.txt" % Z_star
 
 # Read EAGLE test output
 data = loadtxt(filename)
-eagle_time_Gyr = data[:,0]
-eagle_total_mass = data[:,1] * stellar_mass / Msun_in_cgs * unit_mass_in_cgs
-eagle_total_metal_mass = data[:,2] * stellar_mass / Msun_in_cgs * unit_mass_in_cgs
-eagle_total_element_mass = data[:, 3:3+n_elements] * stellar_mass / Msun_in_cgs * unit_mass_in_cgs
-
-eagle_energy_from_mass_cgs = eagle_total_mass * Msun_in_cgs * swift_mean_u_start * unit_int_energy_in_cgs
-eagle_energy_ejecta_cgs = 0.5 * (eagle_total_mass * Msun_in_cgs) * ejecta_vel_cgs**2 
+eagle_time_Gyr = data[:, 0]
+eagle_total_mass = data[:, 1] * stellar_mass / Msun_in_cgs * unit_mass_in_cgs
+eagle_total_metal_mass = data[:, 2] * stellar_mass / Msun_in_cgs * unit_mass_in_cgs
+eagle_total_element_mass = (
+    data[:, 3 : 3 + n_elements] * stellar_mass / Msun_in_cgs * unit_mass_in_cgs
+)
+
+eagle_energy_from_mass_cgs = (
+    eagle_total_mass * Msun_in_cgs * swift_mean_u_start * unit_int_energy_in_cgs
+)
+eagle_energy_ejecta_cgs = 0.5 * (eagle_total_mass * Msun_in_cgs) * ejecta_vel_cgs ** 2
 
 # Read the mass per channel
-filename = "./StellarEvolutionSolution/Z_%.4f/StellarEvolutionAGB.txt"%Z_star
+filename = "./StellarEvolutionSolution/Z_%.4f/StellarEvolutionAGB.txt" % Z_star
 data = loadtxt(filename)
-eagle_total_mass_AGB = data[:,1] * stellar_mass / Msun_in_cgs * unit_mass_in_cgs
-eagle_total_metal_mass_AGB = data[:,2] * stellar_mass / Msun_in_cgs * unit_mass_in_cgs
+eagle_total_mass_AGB = data[:, 1] * stellar_mass / Msun_in_cgs * unit_mass_in_cgs
+eagle_total_metal_mass_AGB = data[:, 2] * stellar_mass / Msun_in_cgs * unit_mass_in_cgs
 
-filename = "./StellarEvolutionSolution/Z_%.4f/StellarEvolutionII.txt"%Z_star
+filename = "./StellarEvolutionSolution/Z_%.4f/StellarEvolutionII.txt" % Z_star
 data = loadtxt(filename)
-eagle_total_mass_SNII = data[:,1] * stellar_mass / Msun_in_cgs * unit_mass_in_cgs
-eagle_total_metal_mass_SNII = data[:,2] * stellar_mass / Msun_in_cgs * unit_mass_in_cgs
+eagle_total_mass_SNII = data[:, 1] * stellar_mass / Msun_in_cgs * unit_mass_in_cgs
+eagle_total_metal_mass_SNII = data[:, 2] * stellar_mass / Msun_in_cgs * unit_mass_in_cgs
 
-filename = "./StellarEvolutionSolution/Z_%.4f/StellarEvolutionIa.txt"%Z_star
+filename = "./StellarEvolutionSolution/Z_%.4f/StellarEvolutionIa.txt" % Z_star
 data = loadtxt(filename)
-eagle_total_mass_SNIa = data[:,1] * stellar_mass / Msun_in_cgs * unit_mass_in_cgs
-eagle_total_metal_mass_SNIa = data[:,2] * stellar_mass / Msun_in_cgs * unit_mass_in_cgs
+eagle_total_mass_SNIa = data[:, 1] * stellar_mass / Msun_in_cgs * unit_mass_in_cgs
+eagle_total_metal_mass_SNIa = data[:, 2] * stellar_mass / Msun_in_cgs * unit_mass_in_cgs
 
 
 # Plot the interesting quantities
 figure()
 
-suptitle("Star metallicity Z = %.4f"%Z_star)
+suptitle("Star metallicity Z = %.4f" % Z_star)
 
 # Box gas mass --------------------------------
 subplot(221)
-plot(t[1:] * unit_time_in_cgs / Gyr_in_cgs, (swift_box_gas_mass[1:] - swift_box_gas_mass[0])* unit_mass_in_cgs / Msun_in_cgs, linewidth=0.5, color='k', label='Total')
-plot(t * unit_time_in_cgs / Gyr_in_cgs, swift_box_gas_mass_AGB * unit_mass_in_cgs / Msun_in_cgs, linewidth=0.5, color='C0', label='AGB')
-plot(t * unit_time_in_cgs / Gyr_in_cgs, swift_box_gas_mass_SNII * unit_mass_in_cgs / Msun_in_cgs, linewidth=0.5, color='C1', label='SNII')
-plot(t * unit_time_in_cgs / Gyr_in_cgs, swift_box_gas_mass_SNIa * unit_mass_in_cgs / Msun_in_cgs, linewidth=0.5, color='C2', label='SNIa')
-plot(eagle_time_Gyr[1:],eagle_total_mass[:-1],linewidth=0.5, color='k', ls='--')
-plot(eagle_time_Gyr[1:],eagle_total_mass_AGB[:-1],linewidth=0.5, color='C0', ls='--')
-plot(eagle_time_Gyr[1:],eagle_total_mass_SNII[:-1],linewidth=0.5, color='C1', ls='--')
-plot(eagle_time_Gyr[1:],eagle_total_mass_SNIa[:-1],linewidth=0.5, color='C2', ls='--')
+plot(
+    t[1:] * unit_time_in_cgs / Gyr_in_cgs,
+    (swift_box_gas_mass[1:] - swift_box_gas_mass[0]) * unit_mass_in_cgs / Msun_in_cgs,
+    linewidth=0.5,
+    color="k",
+    label="Total",
+)
+plot(
+    t * unit_time_in_cgs / Gyr_in_cgs,
+    swift_box_gas_mass_AGB * unit_mass_in_cgs / Msun_in_cgs,
+    linewidth=0.5,
+    color="C0",
+    label="AGB",
+)
+plot(
+    t * unit_time_in_cgs / Gyr_in_cgs,
+    swift_box_gas_mass_SNII * unit_mass_in_cgs / Msun_in_cgs,
+    linewidth=0.5,
+    color="C1",
+    label="SNII",
+)
+plot(
+    t * unit_time_in_cgs / Gyr_in_cgs,
+    swift_box_gas_mass_SNIa * unit_mass_in_cgs / Msun_in_cgs,
+    linewidth=0.5,
+    color="C2",
+    label="SNIa",
+)
+plot(eagle_time_Gyr[1:], eagle_total_mass[:-1], linewidth=0.5, color="k", ls="--")
+plot(eagle_time_Gyr[1:], eagle_total_mass_AGB[:-1], linewidth=0.5, color="C0", ls="--")
+plot(eagle_time_Gyr[1:], eagle_total_mass_SNII[:-1], linewidth=0.5, color="C1", ls="--")
+plot(eagle_time_Gyr[1:], eagle_total_mass_SNIa[:-1], linewidth=0.5, color="C2", ls="--")
 legend(loc="lower right", ncol=2, fontsize=8)
 xlabel("${\\rm Time~[Gyr]}$", labelpad=0)
 ylabel("Change in total gas particle mass ${[\\rm M_\\odot]}$", labelpad=2)
-ticklabel_format(style='sci', axis='y', scilimits=(0,0))
+ticklabel_format(style="sci", axis="y", scilimits=(0, 0))
 
 # Box star mass --------------------------------
 subplot(222)
-plot(t * unit_time_in_cgs / Gyr_in_cgs, (swift_box_star_mass)* unit_mass_in_cgs / Msun_in_cgs, linewidth=0.5, color='k', label='SWIFT')
-plot(eagle_time_Gyr[1:], swift_box_star_mass[0] * unit_mass_in_cgs / Msun_in_cgs - eagle_total_mass[:-1],linewidth=0.5,color='k',label='EAGLE test', ls='--')
+plot(
+    t * unit_time_in_cgs / Gyr_in_cgs,
+    (swift_box_star_mass) * unit_mass_in_cgs / Msun_in_cgs,
+    linewidth=0.5,
+    color="k",
+    label="SWIFT",
+)
+plot(
+    eagle_time_Gyr[1:],
+    swift_box_star_mass[0] * unit_mass_in_cgs / Msun_in_cgs - eagle_total_mass[:-1],
+    linewidth=0.5,
+    color="k",
+    label="EAGLE test",
+    ls="--",
+)
 xlabel("${\\rm Time~[Gyr]}$", labelpad=0)
 ylabel("Change in total star particle mass ${[\\rm M_\\odot]}$", labelpad=2)
-ticklabel_format(style='sci', axis='y', scilimits=(0,0))
+ticklabel_format(style="sci", axis="y", scilimits=(0, 0))
 legend()
 
 # Box gas element  mass --------------------------------
-colours = ['k','r','g','b','c','y','m','skyblue','plum']
-element_names = ['H','He','C','N','O','Ne','Mg','Si','Fe']
+colours = ["k", "r", "g", "b", "c", "y", "m", "skyblue", "plum"]
+element_names = ["H", "He", "C", "N", "O", "Ne", "Mg", "Si", "Fe"]
 subplot(223)
 for j in range(n_elements):
-        plot(t[1:] * unit_time_in_cgs / Gyr_in_cgs, (swift_element_mass[1:,j] - swift_element_mass[0,j]) * unit_mass_in_cgs / Msun_in_cgs, linewidth=0.5, color=colours[j], ms=0.5, label=element_names[j])
-        plot(eagle_time_Gyr[1:],eagle_total_element_mass[:-1,j],linewidth=1,color=colours[j],linestyle='--')
+    plot(
+        t[1:] * unit_time_in_cgs / Gyr_in_cgs,
+        (swift_element_mass[1:, j] - swift_element_mass[0, j])
+        * unit_mass_in_cgs
+        / Msun_in_cgs,
+        linewidth=0.5,
+        color=colours[j],
+        ms=0.5,
+        label=element_names[j],
+    )
+    plot(
+        eagle_time_Gyr[1:],
+        eagle_total_element_mass[:-1, j],
+        linewidth=1,
+        color=colours[j],
+        linestyle="--",
+    )
 xlabel("${\\rm Time~[Gyr]}$", labelpad=0)
 ylabel("Change in element mass of gas particles ${[\\rm M_\\odot]}$", labelpad=2)
 xscale("log")
 yscale("log")
-legend(bbox_to_anchor=(1.005, 1.), ncol=1, fontsize=8, handlelength=1)
+legend(bbox_to_anchor=(1.005, 1.0), ncol=1, fontsize=8, handlelength=1)
 
 # Box gas metal mass --------------------------------
 subplot(224)
-plot(t[1:] * unit_time_in_cgs / Gyr_in_cgs, (swift_box_gas_metal_mass[1:] - swift_box_gas_metal_mass[0])* unit_mass_in_cgs / Msun_in_cgs, linewidth=0.5, color='k', label='Total')
-plot(t * unit_time_in_cgs / Gyr_in_cgs, swift_box_gas_metal_mass_AGB * unit_mass_in_cgs / Msun_in_cgs, linewidth=0.5, color='C0', label='AGB')
-plot(t * unit_time_in_cgs / Gyr_in_cgs, swift_box_gas_metal_mass_SNII * unit_mass_in_cgs / Msun_in_cgs, linewidth=0.5, color='C1', label='SNII')
-plot(t * unit_time_in_cgs / Gyr_in_cgs, swift_box_gas_metal_mass_SNIa * unit_mass_in_cgs / Msun_in_cgs, linewidth=0.5, color='C2', label='SNIa')
-plot(eagle_time_Gyr[1:],eagle_total_metal_mass[:-1],linewidth=0.5,color='k', ls='--')
-plot(eagle_time_Gyr[1:],eagle_total_metal_mass_AGB[:-1],linewidth=0.5, color='C0', ls='--')
-plot(eagle_time_Gyr[1:],eagle_total_metal_mass_SNII[:-1],linewidth=0.5, color='C1', ls='--')
-plot(eagle_time_Gyr[1:],eagle_total_metal_mass_SNIa[:-1],linewidth=0.5, color='C2', ls='--')
+plot(
+    t[1:] * unit_time_in_cgs / Gyr_in_cgs,
+    (swift_box_gas_metal_mass[1:] - swift_box_gas_metal_mass[0])
+    * unit_mass_in_cgs
+    / Msun_in_cgs,
+    linewidth=0.5,
+    color="k",
+    label="Total",
+)
+plot(
+    t * unit_time_in_cgs / Gyr_in_cgs,
+    swift_box_gas_metal_mass_AGB * unit_mass_in_cgs / Msun_in_cgs,
+    linewidth=0.5,
+    color="C0",
+    label="AGB",
+)
+plot(
+    t * unit_time_in_cgs / Gyr_in_cgs,
+    swift_box_gas_metal_mass_SNII * unit_mass_in_cgs / Msun_in_cgs,
+    linewidth=0.5,
+    color="C1",
+    label="SNII",
+)
+plot(
+    t * unit_time_in_cgs / Gyr_in_cgs,
+    swift_box_gas_metal_mass_SNIa * unit_mass_in_cgs / Msun_in_cgs,
+    linewidth=0.5,
+    color="C2",
+    label="SNIa",
+)
+plot(eagle_time_Gyr[1:], eagle_total_metal_mass[:-1], linewidth=0.5, color="k", ls="--")
+plot(
+    eagle_time_Gyr[1:],
+    eagle_total_metal_mass_AGB[:-1],
+    linewidth=0.5,
+    color="C0",
+    ls="--",
+)
+plot(
+    eagle_time_Gyr[1:],
+    eagle_total_metal_mass_SNII[:-1],
+    linewidth=0.5,
+    color="C1",
+    ls="--",
+)
+plot(
+    eagle_time_Gyr[1:],
+    eagle_total_metal_mass_SNIa[:-1],
+    linewidth=0.5,
+    color="C2",
+    ls="--",
+)
 legend(loc="center right", ncol=2, fontsize=8)
 xlabel("${\\rm Time~[Gyr]}$", labelpad=0)
 ylabel("Change in total metal mass of gas particles ${[\\rm M_\\odot]}$", labelpad=2)
-ticklabel_format(style='sci', axis='y', scilimits=(0,0))
-
-savefig("box_evolution_Z_%.4f.png"%(Z_star), dpi=200)
+ticklabel_format(style="sci", axis="y", scilimits=(0, 0))
 
+savefig("box_evolution_Z_%.4f.png" % (Z_star), dpi=200)
diff --git a/examples/SubgridTests/StellarEvolution/run.sh b/examples/SubgridTests/StellarEvolution/run.sh
index 032083db5055dadac8c1bdcdc6cad3e2e03498a1..28509b8fe52a39e2a2d4049ee62813747a72b0b1 100755
--- a/examples/SubgridTests/StellarEvolution/run.sh
+++ b/examples/SubgridTests/StellarEvolution/run.sh
@@ -28,20 +28,20 @@ fi
 
 ../../swift  --temperature --feedback --stars --hydro --external-gravity --threads=4 stellar_evolution.yml -P EAGLEChemistry:init_abundance_metal:0.08 -P EAGLEChemistry:init_abundance_Hydrogen:0.71 -P EAGLEChemistry:init_abundance_Helium:0.21 2>&1 | tee output_0p08.log
 
-python plot_box_evolution.py
+python3 plot_box_evolution.py
 
 ../../swift  --temperature --feedback --stars --hydro --external-gravity --threads=4 stellar_evolution.yml -P EAGLEChemistry:init_abundance_metal:0.04 -P EAGLEChemistry:init_abundance_Hydrogen:0.74 -P EAGLEChemistry:init_abundance_Helium:0.23 2>&1 | tee output_0p04.log
 
-python plot_box_evolution.py
+python3 plot_box_evolution.py
 
 ../../swift  --temperature --feedback --stars --hydro --external-gravity --threads=4 stellar_evolution.yml -P EAGLEChemistry:init_abundance_metal:0.01 2>&1 | tee output_0p01.log
 
-python plot_box_evolution.py
+python3 plot_box_evolution.py
 
 ../../swift  --temperature --feedback --stars --hydro --external-gravity --threads=4 stellar_evolution.yml -P EAGLEChemistry:init_abundance_metal:0.001 2>&1 | tee output_0p001.log
 
-python plot_box_evolution.py
+python3 plot_box_evolution.py
 
 ../../swift  --temperature --feedback --stars --hydro --external-gravity --threads=4 stellar_evolution.yml -P EAGLEChemistry:init_abundance_metal:0.0001 2>&1 | tee output_0p0001.log
 
-python plot_box_evolution.py
+python3 plot_box_evolution.py
diff --git a/examples/SubgridTests/StellarEvolution/stellar_evolution.yml b/examples/SubgridTests/StellarEvolution/stellar_evolution.yml
index 9836aa5931dddb5d7f8e5036d246ab9ec51e70b1..86bc8e234ebba2af6581f74c90538a539c48126e 100644
--- a/examples/SubgridTests/StellarEvolution/stellar_evolution.yml
+++ b/examples/SubgridTests/StellarEvolution/stellar_evolution.yml
@@ -90,6 +90,7 @@ EAGLEFeedback:
   IMF_max_mass_Msun:                  100.0             # Maximal stellar mass considered for the Chabrier IMF in solar masses.
   SNII_min_mass_Msun:                   6.0             # Minimal mass considered for SNII stars in solar masses.
   SNII_max_mass_Msun:                 100.0             # Maximal mass considered for SNII stars in solar masses.
+  SNII_feedback_model:                  Random          # Feedback modes: Random, Isotropic, MinimumDistance, MinimumDensity
   SNII_sampled_delay:                   0               # Sample the SNII lifetimes to do feedback.
   SNII_wind_delay_Gyr:                  0.03            # Time in Gyr between a star's birth and the SNII thermal feedback event when not sampling.
   SNII_delta_T_K:                       3.16228e7       # Change in temperature to apply to the gas particle in a SNII thermal feedback event in Kelvin.
@@ -100,7 +101,8 @@ EAGLEFeedback:
   SNII_energy_fraction_n_0_H_p_cm3:     0.67            # Pivot point for the birth density dependance of the SNII energy fraction in cm^-3.
   SNII_energy_fraction_n_Z:             0.8686          # Power-law for the metallicity dependance of the SNII energy fraction.
   SNII_energy_fraction_n_n:             0.8686          # Power-law for the birth density dependance of the SNII energy fraction.
-  SNII_energy_fraction_use_birth_props: 0               # Are we using the density and metallicity at birth to compute f_E or at feedback time?
+  SNII_energy_fraction_use_birth_density: 0             # Are we using the density and metallicity at birth to compute f_E or at feedback time?
+  SNII_energy_fraction_use_birth_metallicity: 0         # Are we using the density and metallicity at birth to compute f_E or at feedback time?
   SNIa_DTD:                             Exponential     # Use the EAGLE-Ref SNIa DTD.
   SNIa_DTD_delay_Gyr:                   0.04            # Age of the most massive SNIa in Gyr.
   SNIa_DTD_exp_timescale_Gyr:           2.0             # Time-scale of the exponential decay of the SNIa rates in Gyr.
diff --git a/examples/main.c b/examples/main.c
index 4c9db982e43982c6d9788100e8bb2c96552d60af..d31e4fe358d0ed2fc3e49f986c76e19270795fbc 100644
--- a/examples/main.c
+++ b/examples/main.c
@@ -1559,8 +1559,8 @@ int main(int argc, char *argv[]) {
     /* Dump MPI requests if collected. */
 #if defined(SWIFT_MPIUSE_REPORTS) && defined(WITH_MPI)
     {
-      char dumpfile[40];
-      snprintf(dumpfile, 40, "mpiuse_report-rank%d-step%d.dat", engine_rank,
+      char dumpfile[80];
+      snprintf(dumpfile, 80, "mpiuse_report-rank%d-step%d.dat", engine_rank,
                j + 1);
       mpiuse_log_dump(dumpfile, e.tic_step);
     }
@@ -1569,12 +1569,12 @@ int main(int argc, char *argv[]) {
 #ifdef SWIFT_DEBUG_THREADPOOL
     /* Dump the task data using the given frequency. */
     if (dump_threadpool && (dump_threadpool == 1 || j % dump_threadpool == 1)) {
-      char dumpfile[40];
+      char dumpfile[80];
 #ifdef WITH_MPI
-      snprintf(dumpfile, 40, "threadpool_info-rank%d-step%d.dat", engine_rank,
+      snprintf(dumpfile, 80, "threadpool_info-rank%d-step%d.dat", engine_rank,
                j + 1);
 #else
-      snprintf(dumpfile, 40, "threadpool_info-step%d.dat", j + 1);
+      snprintf(dumpfile, 80, "threadpool_info-step%d.dat", j + 1);
 #endif  // WITH_MPI
       threadpool_dump_log(&e.threadpool, dumpfile, 1);
     } else {
diff --git a/examples/parameter_example.yml b/examples/parameter_example.yml
index 66bcfefc8d0eb1b9e603bec0c39ffd8272ed80d2..a4af19afc9019190665994b454e80270039abde2 100644
--- a/examples/parameter_example.yml
+++ b/examples/parameter_example.yml
@@ -158,6 +158,9 @@ Snapshots:
   output_list:         snaplist.txt # (Optional) File containing the output times (see documentation in "Parameter File" section)
   select_output_on:    0  # (Optional) Enable the output selection behaviour
   select_output:       selectoutput.yml # (Optional) File containing information to select outputs with (see documentation in the "Output Selection" section)
+  run_on_dump:         0 # (Optional) Run the dump_command each time that a snapshot is dumped?
+  dump_command:        ./submit_velociraptor.sh # (Optional) Command to run each time that a snapshot is dumped.
+  
 
 # Parameters governing the logger snapshot system
 Logger:
@@ -454,6 +457,7 @@ EAGLEChemistry:
 GEARChemistry:
   initial_metallicity: 1         # Initial metallicity of the gas (mass fraction)
   scale_initial_metallicity: 1   # Should we scale the initial metallicity with the solar one?
+  diffusion_coefficient: 1e-3    # Coefficient for the diffusion (see Shen et al. 2010; differs by gamma^2 due to h^2).
 
 
 # Parameters related to star formation models  -----------------------------------------------
diff --git a/src/Makefile.am b/src/Makefile.am
index f1558a1504b352752f57c1b1946e3ab380ace12d..cdcdb04295aec71c2c68ea081f870e393724e7b9 100644
--- a/src/Makefile.am
+++ b/src/Makefile.am
@@ -139,7 +139,7 @@ nobase_noinst_HEADERS += gravity_iact.h kernel_long_gravity.h vector.h accumulat
 nobase_noinst_HEADERS += runner_doiact_nosort.h runner_doiact_hydro.h runner_doiact_stars.h runner_doiact_black_holes.h runner_doiact_grav.h 
 nobase_noinst_HEADERS += runner_doiact_functions_hydro.h runner_doiact_functions_stars.h runner_doiact_functions_black_holes.h 
 nobase_noinst_HEADERS += runner_doiact_functions_limiter.h runner_doiact_limiter.h units.h intrinsics.h minmax.h 
-nobase_noinst_HEADERS += runner_doiact_rt.h runner_doiact_functions_rt.h runner_doiact_sinks.h
+nobase_noinst_HEADERS += runner_doiact_rt.h runner_doiact_functions_rt.h runner_doiact_sinks.h runner_doiact_functions_sinks.h
 nobase_noinst_HEADERS += kick.h timestep.h drift.h adiabatic_index.h io_properties.h dimension.h part_type.h periodic.h memswap.h 
 nobase_noinst_HEADERS += timestep_limiter.h timestep_limiter_iact.h timestep_sync.h timestep_sync_part.h timestep_limiter_struct.h 
 nobase_noinst_HEADERS += dump.h logger.h sign.h logger_io.h hashmap.h gravity.h gravity_io.h gravity_logger.h  gravity_cache.h output_options.h
@@ -260,8 +260,6 @@ nobase_noinst_HEADERS += star_formation/GEAR/star_formation_logger.h star_format
 nobase_noinst_HEADERS += star_formation/none/star_formation_logger.h star_formation/none/star_formation_logger_struct.h 
 nobase_noinst_HEADERS += cooling/none/cooling.h cooling/none/cooling_struct.h 
 nobase_noinst_HEADERS += cooling/none/cooling_io.h cooling/none/cooling_properties.h 
-nobase_noinst_HEADERS += cooling/Compton/cooling.h cooling/Compton/cooling_struct.h 
-nobase_noinst_HEADERS += cooling/Compton/cooling_io.h cooling/Compton/cooling_properties.h 
 nobase_noinst_HEADERS += cooling/const_du/cooling.h cooling/const_du/cooling_struct.h 
 nobase_noinst_HEADERS += cooling/const_du/cooling_io.h cooling/const_du/cooling_properties.h 
 nobase_noinst_HEADERS += cooling/const_lambda/cooling.h cooling/const_lambda/cooling_struct.h 
@@ -286,6 +284,10 @@ nobase_noinst_HEADERS += chemistry/GEAR/chemistry.h
 nobase_noinst_HEADERS += chemistry/GEAR/chemistry_io.h 
 nobase_noinst_HEADERS += chemistry/GEAR/chemistry_struct.h 
 nobase_noinst_HEADERS += chemistry/GEAR/chemistry_iact.h 
+nobase_noinst_HEADERS += chemistry/GEAR_DIFFUSION/chemistry.h
+nobase_noinst_HEADERS += chemistry/GEAR_DIFFUSION/chemistry_io.h
+nobase_noinst_HEADERS += chemistry/GEAR_DIFFUSION/chemistry_struct.h
+nobase_noinst_HEADERS += chemistry/GEAR_DIFFUSION/chemistry_iact.h
 nobase_noinst_HEADERS += chemistry/EAGLE/chemistry.h 
 nobase_noinst_HEADERS += chemistry/EAGLE/chemistry_io.h 
 nobase_noinst_HEADERS += chemistry/EAGLE/chemistry_struct.h
@@ -324,8 +326,7 @@ nobase_noinst_HEADERS += pressure_floor/GEAR/pressure_floor_iact.h pressure_floo
 nobase_noinst_HEADERS += pressure_floor/GEAR/pressure_floor_struct.h pressure_floor/none/pressure_floor_struct.h 
 nobase_noinst_HEADERS += sink/Default/sink.h sink/Default/sink_io.h sink/Default/sink_part.h sink/Default/sink_properties.h 
 nobase_noinst_HEADERS += sink/Default/sink_iact.h
-nobase_noinst_HEADERS += sink.h sink_io.h sink_properties.h runner_doiact_functions_sinks.h
-
+nobase_noinst_HEADERS += sink.h sink_io.h sink_properties.h
 
 # Sources and special flags for the gravity library
 libgrav_la_SOURCES = runner_doiact_grav.c
diff --git a/src/cell.c b/src/cell.c
index 8bd26b32343f2307673c5f0d9c3abfd68172c187..66f7453667e202494515e0eb45ab7028e75a8303 100644
--- a/src/cell.c
+++ b/src/cell.c
@@ -1386,12 +1386,14 @@ void cell_check_sort_flags(const struct cell *c) {
   const int do_stars_sub_sort = cell_get_flag(c, cell_flag_do_stars_sub_sort);
 
   if (do_hydro_sub_sort)
-    error("cell %d has a hydro sub_sort flag set. Node=%d depth=%d maxdepth=%d",
-          c->cellID, c->nodeID, c->depth, c->maxdepth);
+    error(
+        "cell %lld has a hydro sub_sort flag set. Node=%d depth=%d maxdepth=%d",
+        c->cellID, c->nodeID, c->depth, c->maxdepth);
 
   if (do_stars_sub_sort)
-    error("cell %d has a stars sub_sort flag set. Node=%d depth=%d maxdepth=%d",
-          c->cellID, c->nodeID, c->depth, c->maxdepth);
+    error(
+        "cell %lld has a stars sub_sort flag set. Node=%d depth=%d maxdepth=%d",
+        c->cellID, c->nodeID, c->depth, c->maxdepth);
 
   if (c->split) {
     for (int k = 0; k < 8; ++k) {
diff --git a/src/cell.h b/src/cell.h
index d80d11b8229032795a64b2c862d3dc0bc310de70..b6d9bb9ad66ca30bbd55839d59f40986d404e224 100644
--- a/src/cell.h
+++ b/src/cell.h
@@ -62,6 +62,11 @@ struct scheduler;
 /* Global variables. */
 extern int cell_next_tag;
 
+/*! Counter for cell IDs (when exceeding max values for uniqueness) */
+#if defined(SWIFT_DEBUG_CHECKS) || defined(SWIFT_CELL_GRAPH)
+extern long long last_cell_id;
+#endif
+
 /* Struct to temporarily buffer the particle locations and bin id. */
 struct cell_buff {
   double x[3];
@@ -225,7 +230,7 @@ struct pcell {
 
 #ifdef SWIFT_DEBUG_CHECKS
   /* Cell ID (for debugging) */
-  int cellID;
+  long long cellID;
 #endif
 
 } SWIFT_STRUCT_ALIGN;
@@ -452,7 +457,7 @@ struct cell {
 
 #if defined(SWIFT_DEBUG_CHECKS) || defined(SWIFT_CELL_GRAPH)
   /* Cell ID (for debugging) */
-  int cellID;
+  long long cellID;
 #endif
 
 #ifdef SWIFT_DEBUG_CHECKS
@@ -1279,4 +1284,95 @@ __attribute__((always_inline)) INLINE static struct task *cell_get_recv(
 #endif
 }
 
+/**
+ * @brief Generate the cell ID for top level cells. Only used for debugging.
+ *
+ * Cell IDs are stored in the long long `cell->cellID`. Top level cells get
+ * their index according to their location on the top level grid, and are
+ * marked with a minus sign.
+ * We have 15 bits set aside in `cell->cellID` for the top level cells. Hence
+ * if we have more that 32^3 top level cells, the cell IDs won't be guaranteed
+ * to be unique. Top level cells will still be recognizable by the minus sign.
+ */
+__attribute__((always_inline)) INLINE void cell_assign_top_level_cell_index(
+    struct cell *c, int cdim[3], double dim[3], double width[3]) {
+
+#if defined(SWIFT_DEBUG_CHECKS) || defined(SWIFT_CELL_GRAPH)
+  if (c->depth != 0) {
+    error("assigning top level cell index to cell with depth > 0");
+  } else {
+    if (cdim[0] * cdim[1] * cdim[2] > 32 * 32 * 32) {
+      /* print warning only once */
+      if (last_cell_id == 1) {
+        message(
+            "Warning: Got %d x %d x %d top level cells."
+            "Cell IDs are only guaranteed to be unique if count is < 32^3",
+            cdim[0], cdim[1], cdim[2]);
+      }
+      c->cellID = -last_cell_id;
+      atomic_inc(&last_cell_id);
+    }
+
+    int i = (int)(c->loc[0] / width[0]);
+    int j = (int)(c->loc[1] / width[1]);
+    int k = (int)(c->loc[2] / width[2]);
+    c->cellID = -(long long)(cell_getid(cdim, i, j, k) + 1);
+  }
+#endif
+}
+
+/**
+ * @brief Generate the cell ID for progeny cells. Only used for debugging.
+ *
+ * Cell IDs are stored in the long long `cell->cellID`.
+ * We have 15 bits set aside in `cell->cellID` for the top level cells, and
+ * one for a minus sign to mark top level cells. The remaining 48 bits are
+ * for all other cells. Each progeny cell gets a unique ID by inheriting
+ * its parent ID and adding 3 bits on the right side, which are set according
+ * to the progeny's location within its parent cell. Hence we can store up to
+ * 16 levels of depth uniquely.
+ * If the depth exceeds 16, we use the old scheme where we just add up a
+ * counter. This gives us 32^3 new unique cell IDs, previously reserved for
+ * top level cells, but the IDs won't be thread safe and will vary each run.
+ * After the 32^3 cells are filled, we reach degeneracy.
+ */
+__attribute__((always_inline)) INLINE void cell_assign_cell_index(
+    struct cell *c, const struct cell *parent) {
+
+#if defined(SWIFT_DEBUG_CHECKS) || defined(SWIFT_CELL_GRAPH)
+  if (c->depth == 0) {
+    error("assigning progeny cell index to top level cell.");
+  } else if (c->depth > 16 || last_cell_id > 1) {
+    /* last_cell_id > 1 => too many top level cells for clever IDs */
+    /* print warning only once */
+    if (last_cell_id == 1) {
+      message(
+          "Warning: Got depth %d > 16."
+          "IDs are only guaranteed unique if depth <= 16",
+          c->depth);
+    }
+    c->cellID = last_cell_id;
+    atomic_inc(&last_cell_id);
+  } else {
+    /* we're good to go for unique IDs */
+    /* first inherit the parent's ID and mark it as not top-level*/
+    long long child_id = llabs(parent->cellID);
+
+    /* make place for new bits */
+    /* parent's ID needs to be leading bits, so 000 children still
+     * change the value of the cellID */
+    child_id <<= 3;
+
+    /* get progeny index in parent cell */
+    if (c->loc[0] > parent->loc[0]) child_id |= 1LL;
+    if (c->loc[1] > parent->loc[1]) child_id |= 2LL;
+    if (c->loc[2] > parent->loc[2]) child_id |= 4LL;
+
+    /* add progeny index to cell index */
+    c->cellID = child_id;
+  }
+
+#endif
+}
+
 #endif /* SWIFT_CELL_H */
diff --git a/src/chemistry.h b/src/chemistry.h
index c6113cbbc64e3d1dcbd4c0cea0fdef25b1aaa75f..45db33d1f746294af017c123d7cbf02ac818b995 100644
--- a/src/chemistry.h
+++ b/src/chemistry.h
@@ -35,6 +35,9 @@
 #elif defined(CHEMISTRY_GEAR)
 #include "./chemistry/GEAR/chemistry.h"
 #include "./chemistry/GEAR/chemistry_iact.h"
+#elif defined(CHEMISTRY_GEAR_DIFFUSION)
+#include "./chemistry/GEAR_DIFFUSION/chemistry.h"
+#include "./chemistry/GEAR_DIFFUSION/chemistry_iact.h"
 #elif defined(CHEMISTRY_QLA)
 #include "./chemistry/QLA/chemistry.h"
 #include "./chemistry/QLA/chemistry_iact.h"
diff --git a/src/chemistry/EAGLE/chemistry.h b/src/chemistry/EAGLE/chemistry.h
index cab8916bce17fb5beda5f01ef679bf2fdbdaf0ad..f4e109eda5ac14ad8b5e496d76a3f89cb294c89e 100644
--- a/src/chemistry/EAGLE/chemistry.h
+++ b/src/chemistry/EAGLE/chemistry.h
@@ -294,9 +294,13 @@ static INLINE void chemistry_print_backend(
  *
  * @param p The particle to act upon.
  * @param cosmo The current cosmological model.
+ * @param with_cosmology Are we running with the cosmology?
+ * @param time Current time of the simulation.
+ * @param dt Time step (in physical units).
  */
 __attribute__((always_inline)) INLINE static void chemistry_end_force(
-    struct part* restrict p, const struct cosmology* cosmo) {}
+    struct part* restrict p, const struct cosmology* cosmo,
+    const int with_cosmology, const double time, const double dt) {}
 
 /**
  * @brief Computes the chemistry-related time-step constraint.
diff --git a/src/chemistry/EAGLE/chemistry_iact.h b/src/chemistry/EAGLE/chemistry_iact.h
index 573291637d66e39a973e662edec084d3a4687050..a6f0be62c7182162bd8127aa8ec5e56e25922c76 100644
--- a/src/chemistry/EAGLE/chemistry_iact.h
+++ b/src/chemistry/EAGLE/chemistry_iact.h
@@ -130,4 +130,54 @@ __attribute__((always_inline)) INLINE static void runner_iact_nonsym_chemistry(
       mj * chj->iron_mass_fraction_from_SNIa * wi;
 }
 
+/**
+ * @brief do metal diffusion computation in the <FORCE LOOP>
+ * (symmetric version)
+ *
+ * @param r2 Comoving square distance between the two particles.
+ * @param dx Comoving vector separating both particles (pi - pj).
+ * @param hi Comoving smoothing-length of particle i.
+ * @param hj Comoving smoothing-length of particle j.
+ * @param pi First particle.
+ * @param pj Second particle.
+ * @param a Current scale factor.
+ * @param H Current Hubble parameter.
+ * @param time_base The time base used in order to convert integer to float
+ * time.
+ * @param ti_current The current time (in integer)
+ * @param cosmo The #cosmology.
+ * @param with_cosmology Are we running with cosmology?
+ *
+ */
+__attribute__((always_inline)) INLINE static void runner_iact_diffusion(
+    float r2, const float *dx, float hi, float hj, struct part *restrict pi,
+    struct part *restrict pj, float a, float H, float time_base,
+    integertime_t t_current, const struct cosmology *cosmo,
+    const int with_cosmology) {}
+
+/**
+ * @brief do metal diffusion computation in the <FORCE LOOP>
+ * (nonsymmetric version)
+ *
+ * @param r2 Comoving square distance between the two particles.
+ * @param dx Comoving vector separating both particles (pi - pj).
+ * @param hi Comoving smoothing-length of particle i.
+ * @param hj Comoving smoothing-length of particle j.
+ * @param pi First particle.
+ * @param pj Second particle.
+ * @param a Current scale factor.
+ * @param H Current Hubble parameter.
+ * @param time_base The time base used in order to convert integer to float
+ * time.
+ * @param ti_current The current time (in integer)
+ * @param cosmo The #cosmology.
+ * @param with_cosmology Are we running with cosmology?
+ *
+ */
+__attribute__((always_inline)) INLINE static void runner_iact_nonsym_diffusion(
+    float r2, const float *dx, float hi, float hj, struct part *restrict pi,
+    struct part *restrict pj, float a, float H, float time_base,
+    integertime_t t_current, const struct cosmology *cosmo,
+    const int with_cosmology) {}
+
 #endif /* SWIFT_EAGLE_CHEMISTRY_IACT_H */
diff --git a/src/chemistry/EAGLE/chemistry_struct.h b/src/chemistry/EAGLE/chemistry_struct.h
index 9e99ea7bd8e34bf103deba3b31cb7b841f869ec3..491bc444d7f201f9a9a389c6a04d2ba04e927914 100644
--- a/src/chemistry/EAGLE/chemistry_struct.h
+++ b/src/chemistry/EAGLE/chemistry_struct.h
@@ -89,6 +89,8 @@ struct chemistry_part_data {
   float smoothed_iron_mass_fraction_from_SNIa;
 };
 
+#define chemistry_spart_data chemistry_part_data
+
 /**
  * @brief Chemical abundances traced by the #bpart in the EAGLE model.
  */
diff --git a/src/chemistry/GEAR/chemistry.h b/src/chemistry/GEAR/chemistry.h
index 754d1636ddf9ecf30931e3ecf2db7ad3025822b0..9e47c0edb926e19a6de708b02a12a77e668bce7c 100644
--- a/src/chemistry/GEAR/chemistry.h
+++ b/src/chemistry/GEAR/chemistry.h
@@ -145,6 +145,12 @@ static INLINE void chemistry_init_backend(struct swift_params* parameter_file,
   const float initial_metallicity = parser_get_param_float(
       parameter_file, "GEARChemistry:initial_metallicity");
 
+  if (initial_metallicity < 0) {
+    message("Setting the initial metallicity from the snapshot.");
+  } else {
+    message("Setting the initial metallicity from the parameter file.");
+  }
+
   /* Set the initial metallicities */
   for (int i = 0; i < GEAR_CHEMISTRY_ELEMENT_COUNT; i++) {
     data->initial_metallicities[i] = initial_metallicity;
@@ -177,10 +183,6 @@ __attribute__((always_inline)) INLINE static void chemistry_init_part(
   for (int i = 0; i < GEAR_CHEMISTRY_ELEMENT_COUNT; i++) {
     /* Reset the smoothed metallicity */
     cpd->smoothed_metal_mass_fraction[i] = 0.f;
-
-    /* Convert the total mass into mass fraction */
-    /* Now the metal mass is not available anymore */
-    cpd->metal_mass_fraction[i] = cpd->metal_mass[i] / p->mass;
   }
 }
 
@@ -204,21 +206,15 @@ __attribute__((always_inline)) INLINE static void chemistry_end_density(
   const float h = p->h;
   const float h_inv = 1.0f / h;                       /* 1/h */
   const float factor = pow_dimension(h_inv) / p->rho; /* 1 / h^d * rho */
-  const float m = p->mass;
 
   struct chemistry_part_data* cpd = &p->chemistry_data;
 
   for (int i = 0; i < GEAR_CHEMISTRY_ELEMENT_COUNT; i++) {
     /* Final operation on the density (add self-contribution). */
-    cpd->smoothed_metal_mass_fraction[i] +=
-        m * cpd->metal_mass_fraction[i] * kernel_root;
+    cpd->smoothed_metal_mass_fraction[i] += cpd->metal_mass[i] * kernel_root;
 
     /* Finish the calculation by inserting the missing h-factors */
     cpd->smoothed_metal_mass_fraction[i] *= factor;
-
-    /* Convert the mass fraction into a total mass */
-    /* Now the metal mass fraction is not available anymore */
-    cpd->metal_mass[i] = m * cpd->metal_mass_fraction[i];
   }
 }
 
@@ -229,7 +225,8 @@ __attribute__((always_inline)) INLINE static void chemistry_end_density(
  * @param cosmo The current cosmological model.
  */
 __attribute__((always_inline)) INLINE static void chemistry_end_force(
-    struct part* restrict p, const struct cosmology* cosmo) {}
+    struct part* restrict p, const struct cosmology* cosmo,
+    const int with_cosmology, const double time, const double dt) {}
 
 /**
  * @brief Sets all particle fields to sensible values when the #part has 0 ngbs.
@@ -248,9 +245,7 @@ chemistry_part_has_no_neighbours(struct part* restrict p,
   /* Set the smoothed fractions with the non smoothed fractions */
   for (int i = 0; i < GEAR_CHEMISTRY_ELEMENT_COUNT; i++) {
     p->chemistry_data.smoothed_metal_mass_fraction[i] =
-        p->chemistry_data.metal_mass_fraction[i];
-    p->chemistry_data.metal_mass[i] =
-        p->chemistry_data.metal_mass_fraction[i] * p->mass;
+        p->chemistry_data.metal_mass[i] / hydro_get_mass(p);
   }
 }
 
@@ -294,7 +289,14 @@ __attribute__((always_inline)) INLINE static void chemistry_first_init_part(
     struct xpart* restrict xp) {
 
   for (int i = 0; i < GEAR_CHEMISTRY_ELEMENT_COUNT; i++) {
-    p->chemistry_data.metal_mass[i] = data->initial_metallicities[i] * p->mass;
+    if (data->initial_metallicities[i] < 0) {
+      /* Use the value from the IC. We are reading the metal mass fraction. */
+      p->chemistry_data.metal_mass[i] *= hydro_get_mass(p);
+    } else {
+      /* Use the value from the parameter file */
+      p->chemistry_data.metal_mass[i] =
+          data->initial_metallicities[i] * hydro_get_mass(p);
+    }
   }
 
   chemistry_init_part(p, data);
@@ -405,7 +407,7 @@ __attribute__((always_inline)) INLINE static void chemistry_split_part(
  */
 __attribute__((always_inline)) INLINE static float const*
 chemistry_get_metal_mass_fraction_for_feedback(const struct part* restrict p) {
-
+  error("Not implemented");
   return NULL;
 }
 
@@ -420,6 +422,7 @@ chemistry_get_metal_mass_fraction_for_feedback(const struct part* restrict p) {
 __attribute__((always_inline)) INLINE static float
 chemistry_get_total_metal_mass_fraction_for_feedback(
     const struct part* restrict p) {
+  error("Not implemented");
   return 0.f;
 }
 
@@ -537,7 +540,7 @@ chemistry_get_star_total_metal_mass_for_stats(const struct spart* restrict sp) {
  */
 __attribute__((always_inline)) INLINE static float
 chemistry_get_bh_total_metal_mass_for_stats(const struct bpart* restrict bp) {
-
+  error("Not implemented");
   return 0.f;
 }
 
diff --git a/src/chemistry/GEAR/chemistry_iact.h b/src/chemistry/GEAR/chemistry_iact.h
index ff49af4abb141ee69dea48bbaf99f6b38db4ce45..e2ea733d44c0855ebe2bdc38f85a1fab2fa6ce45 100644
--- a/src/chemistry/GEAR/chemistry_iact.h
+++ b/src/chemistry/GEAR/chemistry_iact.h
@@ -49,30 +49,24 @@ __attribute__((always_inline)) INLINE static void runner_iact_chemistry(
   struct chemistry_part_data *chi = &pi->chemistry_data;
   struct chemistry_part_data *chj = &pj->chemistry_data;
 
-  float wi, wi_dx;
-  float wj, wj_dx;
-
-  /* Get the masses. */
-  const float mi = pi->mass;
-  const float mj = pj->mass;
+  float wi;
+  float wj;
 
   /* Get r */
   const float r = sqrtf(r2);
 
   /* Compute the kernel function for pi */
   const float ui = r / hi;
-  kernel_deval(ui, &wi, &wi_dx);
+  kernel_eval(ui, &wi);
 
   /* Compute the kernel function for pj */
   const float uj = r / hj;
-  kernel_deval(uj, &wj, &wj_dx);
+  kernel_eval(uj, &wj);
 
   /* Compute contribution to the smooth metallicity */
   for (int i = 0; i < GEAR_CHEMISTRY_ELEMENT_COUNT; i++) {
-    chi->smoothed_metal_mass_fraction[i] +=
-        mj * chj->metal_mass_fraction[i] * wi;
-    chj->smoothed_metal_mass_fraction[i] +=
-        mi * chi->metal_mass_fraction[i] * wj;
+    chi->smoothed_metal_mass_fraction[i] += chj->metal_mass[i] * wi;
+    chj->smoothed_metal_mass_fraction[i] += chi->metal_mass[i] * wj;
   }
 }
 
@@ -96,23 +90,69 @@ __attribute__((always_inline)) INLINE static void runner_iact_nonsym_chemistry(
   struct chemistry_part_data *chi = &pi->chemistry_data;
   const struct chemistry_part_data *chj = &pj->chemistry_data;
 
-  float wi, wi_dx;
-
-  /* Get the masses. */
-  const float mj = pj->mass;
+  float wi;
 
   /* Get r */
   const float r = sqrtf(r2);
 
   /* Compute the kernel function for pi */
   const float ui = r / hi;
-  kernel_deval(ui, &wi, &wi_dx);
+  kernel_eval(ui, &wi);
 
   /* Compute contribution to the smooth metallicity */
   for (int i = 0; i < GEAR_CHEMISTRY_ELEMENT_COUNT; i++) {
-    chi->smoothed_metal_mass_fraction[i] +=
-        mj * chj->metal_mass_fraction[i] * wi;
+    chi->smoothed_metal_mass_fraction[i] += chj->metal_mass[i] * wi;
   }
 }
 
+/**
+ * @brief do metal diffusion computation in the <FORCE LOOP>
+ * (symmetric version)
+ *
+ * @param r2 Comoving square distance between the two particles.
+ * @param dx Comoving vector separating both particles (pi - pj).
+ * @param hi Comoving smoothing-length of particle i.
+ * @param hj Comoving smoothing-length of particle j.
+ * @param pi First particle.
+ * @param pj Second particle.
+ * @param a Current scale factor.
+ * @param H Current Hubble parameter.
+ * @param time_base The time base used in order to convert integer to float
+ * time.
+ * @param ti_current The current time (in integer)
+ * @param cosmo The #cosmology.
+ * @param with_cosmology Are we running with cosmology?
+ *
+ */
+__attribute__((always_inline)) INLINE static void runner_iact_diffusion(
+    float r2, const float *dx, float hi, float hj, struct part *restrict pi,
+    struct part *restrict pj, float a, float H, float time_base,
+    integertime_t t_current, const struct cosmology *cosmo,
+    const int with_cosmology) {}
+
+/**
+ * @brief do metal diffusion computation in the <FORCE LOOP>
+ * (nonsymmetric version)
+ *
+ * @param r2 Comoving square distance between the two particles.
+ * @param dx Comoving vector separating both particles (pi - pj).
+ * @param hi Comoving smoothing-length of particle i.
+ * @param hj Comoving smoothing-length of particle j.
+ * @param pi First particle.
+ * @param pj Second particle.
+ * @param a Current scale factor.
+ * @param H Current Hubble parameter.
+ * @param time_base The time base used in order to convert integer to float
+ * time.
+ * @param ti_current The current time (in integer)
+ * @param cosmo The #cosmology.
+ * @param with_cosmology Are we running with cosmology?
+ *
+ */
+__attribute__((always_inline)) INLINE static void runner_iact_nonsym_diffusion(
+    float r2, const float *dx, float hi, float hj, struct part *restrict pi,
+    struct part *restrict pj, float a, float H, float time_base,
+    integertime_t t_current, const struct cosmology *cosmo,
+    const int with_cosmology) {}
+
 #endif /* SWIFT_GEAR_CHEMISTRY_IACT_H */
diff --git a/src/chemistry/GEAR/chemistry_io.h b/src/chemistry/GEAR/chemistry_io.h
index 8112bfb13b3203d75905167f3e75c8b06f05fb83..b485329a36733226113343b942b50f28578dc0ad 100644
--- a/src/chemistry/GEAR/chemistry_io.h
+++ b/src/chemistry/GEAR/chemistry_io.h
@@ -42,8 +42,8 @@ INLINE static int chemistry_read_particles(struct part* parts,
 
   /* List what we want to read */
   list[0] = io_make_input_field(
-      "ElementAbundance", DOUBLE, GEAR_CHEMISTRY_ELEMENT_COUNT, OPTIONAL,
-      UNIT_CONV_NO_UNITS, parts, chemistry_data.metal_mass_fraction);
+      "MetalMassFraction", DOUBLE, GEAR_CHEMISTRY_ELEMENT_COUNT, OPTIONAL,
+      UNIT_CONV_NO_UNITS, parts, chemistry_data.metal_mass);
 
   return 1;
 }
@@ -53,7 +53,7 @@ INLINE static void convert_gas_metals(const struct engine* e,
                                       const struct xpart* xp, double* ret) {
 
   for (int i = 0; i < GEAR_CHEMISTRY_ELEMENT_COUNT; i++) {
-    ret[0] = p->chemistry_data.metal_mass[i] / hydro_get_mass(p);
+    ret[i] = p->chemistry_data.metal_mass[i] / hydro_get_mass(p);
   }
 }
 
@@ -74,13 +74,13 @@ INLINE static int chemistry_write_particles(const struct part* parts,
 
   /* List what we want to write */
   list[0] = io_make_output_field(
-      "SmoothedElementAbundances", DOUBLE, GEAR_CHEMISTRY_ELEMENT_COUNT,
+      "SmoothedMetalMassFractions", DOUBLE, GEAR_CHEMISTRY_ELEMENT_COUNT,
       UNIT_CONV_NO_UNITS, 0.f, parts,
       chemistry_data.smoothed_metal_mass_fraction,
-      "Element abundances smoothed over the neighbors");
+      "Mass fraction of each element smoothed over the neighbors");
 
   list[1] = io_make_output_field_convert_part(
-      "ElementAbundances", DOUBLE, GEAR_CHEMISTRY_ELEMENT_COUNT,
+      "MetalMassFractions", DOUBLE, GEAR_CHEMISTRY_ELEMENT_COUNT,
       UNIT_CONV_NO_UNITS, 0.f, parts, xparts, convert_gas_metals,
       "Mass fraction of each element");
 
@@ -100,7 +100,7 @@ INLINE static int chemistry_write_sparticles(const struct spart* sparts,
 
   /* List what we want to write */
   list[0] = io_make_output_field(
-      "ElementAbundances", DOUBLE, GEAR_CHEMISTRY_ELEMENT_COUNT,
+      "MetalMassFractions", DOUBLE, GEAR_CHEMISTRY_ELEMENT_COUNT,
       UNIT_CONV_NO_UNITS, 0.f, sparts, chemistry_data.metal_mass_fraction,
       "Mass fraction of each element");
 
diff --git a/src/chemistry/GEAR/chemistry_struct.h b/src/chemistry/GEAR/chemistry_struct.h
index 1332deea87e74a874c3b0a76339d88f8edcef13a..9ad831b27d84409c131fd5aa8df1a4204dbb73ef 100644
--- a/src/chemistry/GEAR/chemistry_struct.h
+++ b/src/chemistry/GEAR/chemistry_struct.h
@@ -24,7 +24,7 @@
  */
 struct chemistry_global_data {
 
-  /* Initial metallicity Z */
+  /* Initial mass fraction */
   double initial_metallicities[GEAR_CHEMISTRY_ELEMENT_COUNT];
 };
 
@@ -33,15 +33,8 @@ struct chemistry_global_data {
  */
 struct chemistry_part_data {
 
-  union {
-    /*! Fraction of the particle mass in a given element.
-      This field is available only during the density hydro loop. */
-    double metal_mass_fraction[GEAR_CHEMISTRY_ELEMENT_COUNT];
-
-    /*! Total mass of element in a particle.
-      This field is available only outside the density hydro loop. */
-    double metal_mass[GEAR_CHEMISTRY_ELEMENT_COUNT];
-  };
+  /*! Total mass of element in a particle. */
+  double metal_mass[GEAR_CHEMISTRY_ELEMENT_COUNT];
 
   /*! Smoothed fraction of the particle mass in a given element */
   double smoothed_metal_mass_fraction[GEAR_CHEMISTRY_ELEMENT_COUNT];
diff --git a/src/chemistry/GEAR_DIFFUSION/chemistry.h b/src/chemistry/GEAR_DIFFUSION/chemistry.h
new file mode 100644
index 0000000000000000000000000000000000000000..80cbcbc4e40909e07a6727dd7b8cc6fe72926dca
--- /dev/null
+++ b/src/chemistry/GEAR_DIFFUSION/chemistry.h
@@ -0,0 +1,595 @@
+/*******************************************************************************
+ * This file is part of SWIFT.
+ * Copyright (c) 2020 Loic Hausammann (loic.hausammann@epfl.ch)
+ *
+ * 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_DIFFUSION_H
+#define SWIFT_CHEMISTRY_GEAR_DIFFUSION_H
+
+/**
+ * @file src/chemistry/GEAR_DIFFUSION/chemistry.h
+ * Follows Shen et al. 2010
+ */
+
+/* Some standard headers. */
+#include <float.h>
+#include <math.h>
+#include <string.h>
+
+/* Local includes. */
+#include "chemistry_struct.h"
+#include "error.h"
+#include "hydro.h"
+#include "parser.h"
+#include "part.h"
+#include "physical_constants.h"
+#include "units.h"
+
+/**
+ * @brief Copies the chemistry properties of the gas particle over to the
+ * star particle.
+ *
+ * @param p the gas particles.
+ * @param xp the additional properties of the gas particles.
+ * @param sp the new created star particle with its properties.
+ */
+INLINE static void chemistry_copy_star_formation_properties(
+    struct part* p, const struct xpart* xp, struct spart* sp) {
+
+  float mass = hydro_get_mass(p);
+
+  /* Store the chemistry struct in the star particle */
+  for (int i = 0; i < GEAR_CHEMISTRY_ELEMENT_COUNT; i++) {
+    sp->chemistry_data.metal_mass_fraction[i] =
+        p->chemistry_data.smoothed_metal_mass_fraction[i];
+
+    /* Remove the metals taken by the star. */
+    p->chemistry_data.metal_mass[i] *= mass / (mass + sp->mass);
+  }
+}
+
+/**
+ * @brief Prints the properties of the chemistry model to stdout.
+ *
+ * @brief The #chemistry_global_data containing information about the current
+ * model.
+ */
+static INLINE void chemistry_print_backend(
+    const struct chemistry_global_data* data) {
+
+  message("Chemistry function is 'Gear with diffusion'.");
+}
+
+/**
+ * @brief Read the solar abundances and scale with them the initial
+ * metallicities.
+ *
+ * @param parameter_file The parsed parameter file.
+ * @param data The properties to initialise.
+ */
+static INLINE void chemistry_scale_initial_metallicities(
+    struct swift_params* parameter_file, struct chemistry_global_data* data) {
+#ifdef HAVE_HDF5
+
+  /* Get the yields table */
+  char filename[DESCRIPTION_BUFFER_SIZE];
+  parser_get_param_string(parameter_file, "GEARFeedback:yields_table",
+                          filename);
+
+  /* Open file. */
+  hid_t file_id = H5Fopen(filename, H5F_ACC_RDONLY, H5P_DEFAULT);
+  if (file_id < 0) error("unable to open file %s.\n", filename);
+
+  /* Open group. */
+  hid_t group_id = H5Gopen(file_id, "Data", H5P_DEFAULT);
+  if (group_id < 0) error("unable to open group Data.\n");
+
+  /* Read the data */
+  float* sol_ab = (float*)malloc(sizeof(float) * GEAR_CHEMISTRY_ELEMENT_COUNT);
+  io_read_array_attribute(group_id, "SolarMassAbundances", FLOAT, sol_ab,
+                          GEAR_CHEMISTRY_ELEMENT_COUNT);
+
+  /* Close group */
+  hid_t status = H5Gclose(group_id);
+  if (status < 0) error("error closing group.");
+
+  /* Close file */
+  status = H5Fclose(file_id);
+  if (status < 0) error("error closing file.");
+
+  /* Scale the initial metallicities */
+  char txt[DESCRIPTION_BUFFER_SIZE] = "Scaling initial metallicities by:";
+  for (int i = 0; i < GEAR_CHEMISTRY_ELEMENT_COUNT; i++) {
+    data->initial_metallicities[i] *= sol_ab[i];
+    char tmp[10];
+    sprintf(tmp, " %.2g", sol_ab[i]);
+    strcat(txt, tmp);
+  }
+
+  if (engine_rank == 0) {
+    message("%s", txt);
+  }
+#else
+  error("Cannot scale the solar abundances without HDF5");
+#endif
+}
+
+/**
+ * @brief Initialises the chemistry properties.
+ *
+ * Nothing to do here.
+ *
+ * @param parameter_file The parsed parameter file.
+ * @param us The current internal system of units.
+ * @param phys_const The physical constants in internal units.
+ * @param data The properties to initialise.
+ */
+static INLINE void chemistry_init_backend(struct swift_params* parameter_file,
+                                          const struct unit_system* us,
+                                          const struct phys_const* phys_const,
+                                          struct chemistry_global_data* data) {
+
+  /* read parameters */
+  const float initial_metallicity = parser_get_param_float(
+      parameter_file, "GEARChemistry:initial_metallicity");
+
+  if (initial_metallicity < 0) {
+    message("Setting the initial metallicity from the snapshot.");
+  } else {
+    message("Setting the initial metallicity from the parameter file.");
+  }
+
+  /* Set the initial metallicities */
+  for (int i = 0; i < GEAR_CHEMISTRY_ELEMENT_COUNT; i++) {
+    data->initial_metallicities[i] = initial_metallicity;
+  }
+
+  /* Read the diffusion coefficient */
+  data->C = parser_get_param_float(parameter_file,
+                                   "GEARChemistry:diffusion_coefficient");
+
+  /* Check if need to scale the initial metallicity */
+  const int scale_metallicity = parser_get_opt_param_int(
+      parameter_file, "GEARChemistry:scale_initial_metallicity", 0);
+
+  /* Scale the metallicities if required */
+  if (scale_metallicity) {
+    chemistry_scale_initial_metallicities(parameter_file, data);
+  }
+}
+
+/**
+ * @brief Prepares a particle for the smooth metal calculation.
+ *
+ * Zeroes all the relevant arrays in preparation for the sums taking place in
+ * the various smooth metallicity tasks
+ *
+ * @param p The particle to act upon
+ * @param cd #chemistry_global_data containing chemistry informations.
+ */
+__attribute__((always_inline)) INLINE static void chemistry_init_part(
+    struct part* restrict p, const struct chemistry_global_data* cd) {
+
+  struct chemistry_part_data* cpd = &p->chemistry_data;
+
+  for (int i = 0; i < GEAR_CHEMISTRY_ELEMENT_COUNT; i++) {
+    /* Reset the smoothed metallicity */
+    cpd->smoothed_metal_mass_fraction[i] = 0.f;
+
+    /* Reset the diffusion equation */
+    cpd->metal_mass_dt[i] = 0;
+  }
+
+  /* Reset the shear tensor */
+  for (int i = 0; i < 3; i++) {
+    cpd->S[i][0] = 0;
+    cpd->S[i][1] = 0;
+    cpd->S[i][2] = 0;
+  }
+
+  /* Reset the diffusion. */
+  cpd->diff_coef = 0;
+}
+
+/**
+ * @brief Finishes the smooth metal calculation.
+ *
+ * Multiplies the smoothed metallicity and number of neighbours by the
+ * appropiate constants and add the self-contribution term.
+ *
+ * This function requires the #hydro_end_density to have been called.
+ *
+ * @param p The particle to act upon.
+ * @param cd #chemistry_global_data containing chemistry informations.
+ * @param cosmo The current cosmological model.
+ */
+__attribute__((always_inline)) INLINE static void chemistry_end_density(
+    struct part* restrict p, const struct chemistry_global_data* cd,
+    const struct cosmology* cosmo) {
+
+  /* Some smoothing length multiples. */
+  const float h = p->h;
+  const float h_inv = 1.0f / h;                       /* 1/h */
+  const float h_inv_dim = pow_dimension(h_inv);       /* 1/h^d */
+  const float h_inv_dim_plus_one = h_inv_dim * h_inv; /* 1/h^(d+1) */
+  const float rho = hydro_get_comoving_density(p);
+  const float factor = h_inv_dim / rho; /* 1 / h^d * rho */
+  const float rho_inv = 1.0f / rho;     /* 1 / rho */
+
+  struct chemistry_part_data* cpd = &p->chemistry_data;
+
+  for (int i = 0; i < GEAR_CHEMISTRY_ELEMENT_COUNT; i++) {
+    /* Final operation on the density (add self-contribution). */
+    cpd->smoothed_metal_mass_fraction[i] += cpd->metal_mass[i] * kernel_root;
+
+    /* Finish the calculation by inserting the missing h-factors */
+    cpd->smoothed_metal_mass_fraction[i] *= factor;
+  }
+
+  /* convert the shear factor into physical */
+  const float factor_shear = h_inv_dim_plus_one * rho_inv * cosmo->a2_inv;
+  for (int k = 0; k < 3; k++) {
+    cpd->S[k][0] *= factor_shear;
+    cpd->S[k][1] *= factor_shear;
+    cpd->S[k][2] *= factor_shear;
+  }
+
+  /* Compute the trace over 3 and add the hubble flow. */
+  float trace_3 = 0;
+  for (int i = 0; i < 3; i++) {
+    cpd->S[i][i] += cosmo->H;
+    trace_3 += cpd->S[i][i];
+  }
+  trace_3 /= 3.;
+
+  float S_t[3][3];
+  for (int i = 0; i < 3; i++) {
+    /* Make the tensor symmetric. */
+    float avg = 0.5 * (cpd->S[i][0] + cpd->S[0][i]);
+    S_t[i][0] = avg;
+    S_t[0][i] = avg;
+
+    avg = 0.5 * (cpd->S[i][1] + cpd->S[1][i]);
+    S_t[i][1] = avg;
+    S_t[1][i] = avg;
+
+    avg = 0.5 * (cpd->S[i][2] + cpd->S[2][i]);
+    S_t[i][2] = avg;
+    S_t[2][i] = avg;
+
+    /* Remove the trace. */
+    S_t[i][i] -= trace_3;
+  }
+
+  /* Compute the norm. */
+  float norm = 0;
+  for (int i = 0; i < 3; i++) {
+    norm += S_t[i][0] * S_t[i][0];
+    norm += S_t[i][1] * S_t[i][1];
+    norm += S_t[i][2] * S_t[i][2];
+  }
+  norm = sqrt(norm);
+
+  /* Compute the diffusion coefficient in physical coordinates.
+   * The norm is already in physical coordinates.
+   * We do not include kernel_gamma on purpose. */
+  const float h_phys = cosmo->a * p->h;
+  cpd->diff_coef = cd->C * norm * h_phys * h_phys;
+}
+
+/**
+ * @brief Updates to the chemistry data after the hydro force loop.
+ *
+ * @param p The particle to act upon.
+ * @param cosmo The current cosmological model.
+ * @param with_cosmology Are we running with the cosmology?
+ * @param time Current time of the simulation.
+ * @param dt Time step (in physical units).
+ */
+__attribute__((always_inline)) INLINE static void chemistry_end_force(
+    struct part* restrict p, const struct cosmology* cosmo,
+    const int with_cosmology, const double time, const double dt) {
+  if (dt == 0) {
+    return;
+  }
+
+  struct chemistry_part_data* ch = &p->chemistry_data;
+  const float h_inv = cosmo->a / p->h;
+  const float h_inv_dim = pow_dimension(h_inv); /* 1/h^d */
+  /* Missing factors in iact. */
+  const float factor = h_inv_dim * h_inv;
+
+  const double sum = 0;
+  for (int i = 0; i < GEAR_CHEMISTRY_ELEMENT_COUNT; i++) {
+    ch->metal_mass[i] += ch->metal_mass_dt[i] * dt * factor;
+    /* Make sure that the metallicity is 0 <= x <= 1 */
+    if (ch->metal_mass[i] < 0 || ch->metal_mass[i] > hydro_get_mass(p)) {
+      error("Negative mass or mass fraction larger than 1.");
+    }
+    /* Make sure that we do not have more metals than the sum. */
+    if (i != GEAR_CHEMISTRY_ELEMENT_COUNT) {
+      sum += ch->metal_mass[i];
+    } else if (sum > ch->metal_mass[i]) {
+      error("Found more individual elements than the sum of all of them.");
+    }
+  }
+}
+
+/**
+ * @brief Sets all particle fields to sensible values when the #part has 0 ngbs.
+ *
+ * @param p The particle to act upon
+ * @param xp The extended particle data to act upon
+ * @param cd #chemistry_global_data containing chemistry informations.
+ * @param cosmo The current cosmological model.
+ */
+__attribute__((always_inline)) INLINE static void
+chemistry_part_has_no_neighbours(struct part* restrict p,
+                                 struct xpart* restrict xp,
+                                 const struct chemistry_global_data* cd,
+                                 const struct cosmology* cosmo) {
+
+  /* Set the smoothed fractions with the non smoothed fractions */
+  for (int i = 0; i < GEAR_CHEMISTRY_ELEMENT_COUNT; i++) {
+    p->chemistry_data.smoothed_metal_mass_fraction[i] =
+        p->chemistry_data.metal_mass[i] / hydro_get_mass(p);
+  }
+}
+
+/**
+ * @brief Computes the chemistry-related time-step constraint.
+ *
+ * No constraints in the GEAR model (no diffusion) --> FLT_MAX
+ *
+ * @param phys_const The physical constants in internal units.
+ * @param cosmo The current cosmological model.
+ * @param us The internal system of units.
+ * @param hydro_props The properties of the hydro scheme.
+ * @param cd The global properties of the chemistry scheme.
+ * @param p Pointer to the particle data.
+ */
+__attribute__((always_inline)) INLINE static float chemistry_timestep(
+    const struct phys_const* restrict phys_const,
+    const struct cosmology* restrict cosmo,
+    const struct unit_system* restrict us,
+    const struct hydro_props* hydro_props,
+    const struct chemistry_global_data* cd, const struct part* restrict p) {
+  return FLT_MAX;
+}
+
+/**
+ * @brief Sets the chemistry properties of the (x-)particles to a valid start
+ * state.
+ *
+ * @param phys_const The #phys_const.
+ * @param us The #unit_system.
+ * @param cosmo The #cosmology.
+ * @param data The global chemistry information.
+ * @param p Pointer to the particle data.
+ * @param xp Pointer to the extended particle data.
+ */
+__attribute__((always_inline)) INLINE static void chemistry_first_init_part(
+    const struct phys_const* restrict phys_const,
+    const struct unit_system* restrict us,
+    const struct cosmology* restrict cosmo,
+    const struct chemistry_global_data* data, struct part* restrict p,
+    struct xpart* restrict xp) {
+
+  for (int i = 0; i < GEAR_CHEMISTRY_ELEMENT_COUNT; i++) {
+    if (data->initial_metallicities[i] < 0) {
+      /* Use the value from the IC. We are reading the metal mass fraction. */
+      p->chemistry_data.metal_mass[i] *= hydro_get_mass(p);
+    } else {
+      /* Use the value from the parameter file */
+      p->chemistry_data.metal_mass[i] =
+          data->initial_metallicities[i] * hydro_get_mass(p);
+    }
+  }
+
+  chemistry_init_part(p, data);
+}
+
+/**
+ * @brief Sets the chemistry properties of the sparticles to a valid start
+ * state.
+ *
+ * @param data The global chemistry information.
+ * @param sp Pointer to the sparticle data.
+ */
+__attribute__((always_inline)) INLINE static void chemistry_first_init_spart(
+    const struct chemistry_global_data* data, struct spart* restrict sp) {
+
+  for (int i = 0; i < GEAR_CHEMISTRY_ELEMENT_COUNT; i++) {
+    sp->chemistry_data.metal_mass_fraction[i] =
+        data->initial_metallicities[i] * sp->mass;
+  }
+}
+
+/**
+ * @brief Initialise the chemistry properties of a black hole with
+ * the chemistry properties of the gas it is born from.
+ *
+ * Nothing to do here.
+ *
+ * @param bp_data The black hole data to initialise.
+ * @param p_data The gas data to use.
+ * @param gas_mass The mass of the gas particle.
+ */
+__attribute__((always_inline)) INLINE static void chemistry_bpart_from_part(
+    struct chemistry_bpart_data* bp_data,
+    const struct chemistry_part_data* p_data, const double gas_mass) {
+  error("Loic: to be implemented");
+}
+
+/**
+ * @brief Add the chemistry data of a gas particle to a black hole.
+ *
+ * Nothing to do here.
+ *
+ * @param bp_data The black hole data to add to.
+ * @param p_data The gas data to use.
+ * @param gas_mass The mass of the gas particle.
+ */
+__attribute__((always_inline)) INLINE static void chemistry_add_part_to_bpart(
+    struct chemistry_bpart_data* bp_data,
+    const struct chemistry_part_data* p_data, const double gas_mass) {
+  error("Loic: to be implemented");
+}
+
+/**
+ * @brief Add the chemistry data of a black hole to another one.
+ *
+ * Nothing to do here.
+ *
+ * @param bp_data The black hole data to add to.
+ * @param swallowed_data The black hole data to use.
+ */
+__attribute__((always_inline)) INLINE static void chemistry_add_bpart_to_bpart(
+    struct chemistry_bpart_data* bp_data,
+    const struct chemistry_bpart_data* swallowed_data) {
+  error("Loic: to be implemented");
+}
+
+/**
+ * @brief Split the metal content of a particle into n pieces
+ *
+ * @param p The #part.
+ * @param n The number of pieces to split into.
+ */
+__attribute__((always_inline)) INLINE static void chemistry_split_part(
+    struct part* p, const double n) {
+  error("Loic: to be implemented");
+}
+
+/**
+ * @brief Returns the total metallicity (metal mass fraction) of the
+ * star particle to be used in feedback/enrichment related routines.
+ *
+ * @param sp Pointer to the particle data.
+ */
+__attribute__((always_inline)) INLINE static double
+chemistry_get_total_metal_mass_fraction_for_feedback(
+    const struct spart* restrict sp) {
+
+  return sp->chemistry_data
+      .metal_mass_fraction[GEAR_CHEMISTRY_ELEMENT_COUNT - 1];
+}
+
+/**
+ * @brief Returns the abundances (metal mass fraction) of the
+ * star particle to be used in feedback/enrichment related routines.
+ *
+ * @param sp Pointer to the particle data.
+ */
+__attribute__((always_inline)) INLINE static double const*
+chemistry_get_metal_mass_fraction_for_feedback(
+    const struct spart* restrict sp) {
+
+  return sp->chemistry_data.metal_mass_fraction;
+}
+
+/**
+ * @brief Returns the total metallicity (metal mass fraction) of the
+ * gas particle to be used in cooling related routines.
+ *
+ * @param p Pointer to the particle data.
+ */
+__attribute__((always_inline)) INLINE static double
+chemistry_get_total_metal_mass_fraction_for_cooling(
+    const struct part* restrict p) {
+
+  return p->chemistry_data
+      .smoothed_metal_mass_fraction[GEAR_CHEMISTRY_ELEMENT_COUNT - 1];
+}
+
+/**
+ * @brief Returns the abundance array (metal mass fractions) of the
+ * gas particle to be used in cooling related routines.
+ *
+ * @param p Pointer to the particle data.
+ */
+__attribute__((always_inline)) INLINE static double const*
+chemistry_get_metal_mass_fraction_for_cooling(const struct part* restrict p) {
+
+  return p->chemistry_data.smoothed_metal_mass_fraction;
+}
+
+/**
+ * @brief Returns the total metallicity (metal mass fraction) of the
+ * gas particle to be used in star formation related routines.
+ *
+ * @param p Pointer to the particle data.
+ */
+__attribute__((always_inline)) INLINE static double
+chemistry_get_total_metal_mass_fraction_for_star_formation(
+    const struct part* restrict p) {
+
+  return p->chemistry_data
+      .smoothed_metal_mass_fraction[GEAR_CHEMISTRY_ELEMENT_COUNT - 1];
+}
+
+/**
+ * @brief Returns the abundance array (metal mass fractions) of the
+ * gas particle to be used in star formation related routines.
+ *
+ * @param p Pointer to the particle data.
+ */
+__attribute__((always_inline)) INLINE static double const*
+chemistry_get_metal_mass_fraction_for_star_formation(
+    const struct part* restrict p) {
+
+  return p->chemistry_data.smoothed_metal_mass_fraction;
+}
+
+/**
+ * @brief Returns the total metallicity (metal mass fraction) of the
+ * black hole particle to be used in the stats related routines.
+ *
+ * @param p Pointer to the particle data.
+ */
+__attribute__((always_inline)) INLINE static float
+chemistry_get_bh_total_metal_mass_for_stats(const struct bpart* restrict bp) {
+  error("Not implemented");
+  return 0.f;
+}
+
+/**
+ * @brief Returns the total metallicity (metal mass fraction) of the
+ * gas particle to be used in the stats related routines.
+ *
+ * @param p Pointer to the particle data.
+ */
+__attribute__((always_inline)) INLINE static float
+chemistry_get_total_metal_mass_for_stats(const struct part* restrict p) {
+
+  return p->chemistry_data.metal_mass[GEAR_CHEMISTRY_ELEMENT_COUNT - 1];
+}
+
+/**
+ * @brief Returns the total metallicity (metal mass fraction) of the
+ * star particle to be used in the stats related routines.
+ *
+ * @param p Pointer to the particle data.
+ */
+__attribute__((always_inline)) INLINE static float
+chemistry_get_star_total_metal_mass_for_stats(const struct spart* restrict sp) {
+
+  return sp->chemistry_data
+             .metal_mass_fraction[GEAR_CHEMISTRY_ELEMENT_COUNT - 1] *
+         sp->mass;
+}
+
+#endif /* SWIFT_CHEMISTRY_GEAR_DIFFUSION_H */
diff --git a/src/chemistry/GEAR_DIFFUSION/chemistry_iact.h b/src/chemistry/GEAR_DIFFUSION/chemistry_iact.h
new file mode 100644
index 0000000000000000000000000000000000000000..896307c011e2e501e03a076e8662bd9e7bdd101c
--- /dev/null
+++ b/src/chemistry/GEAR_DIFFUSION/chemistry_iact.h
@@ -0,0 +1,297 @@
+/*******************************************************************************
+ * This file is part of SWIFT.
+ * Copyright (c) 2020 Loic Hausammann (loic.hausammann@epfl.ch)
+ *
+ * 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_GEAR_DIFFUSION_CHEMISTRY_IACT_H
+#define SWIFT_GEAR_DIFFUSION_CHEMISTRY_IACT_H
+
+/**
+ * @file GEAR/chemistry_iact.h
+ * @brief Smooth metal interaction + diffusion functions following the GEAR
+ * version of smooth metalicity.
+ */
+
+/**
+ * @brief do chemistry computation after the runner_iact_density (symmetric
+ * version)
+ *
+ * @param r2 Comoving square distance between the two particles.
+ * @param dx Comoving vector separating both particles (pi - pj).
+ * @param hi Comoving smoothing-length of particle i.
+ * @param hj Comoving smoothing-length of particle j.
+ * @param pi First particle.
+ * @param pj Second particle.
+ * @param a Current scale factor.
+ * @param H Current Hubble parameter.
+ */
+__attribute__((always_inline)) INLINE static void runner_iact_chemistry(
+    float r2, const float *dx, float hi, float hj, struct part *restrict pi,
+    struct part *restrict pj, float a, float H) {
+
+  struct chemistry_part_data *chi = &pi->chemistry_data;
+  struct chemistry_part_data *chj = &pj->chemistry_data;
+
+  float wi, wi_dx;
+  float wj, wj_dx;
+
+  /* Get the masses. */
+  const float mj = hydro_get_mass(pj);
+  const float mi = hydro_get_mass(pi);
+
+  /* Get r */
+  const float r = sqrtf(r2);
+  const float r_inv = 1.f / r;
+
+  /* Compute the kernel function for pi */
+  const float ui = r / hi;
+  kernel_deval(ui, &wi, &wi_dx);
+
+  /* Compute the kernel function for pj */
+  const float uj = r / hj;
+  kernel_deval(uj, &wj, &wj_dx);
+
+  const float wi_dr = wi_dx * r_inv;
+  const float mj_wi_dr = mj * wi_dr;
+
+  const float wj_dr = wj_dx * r_inv;
+  const float mi_wj_dr = mi * wj_dr;
+
+  /* Compute contribution to the smooth metallicity */
+  for (int i = 0; i < GEAR_CHEMISTRY_ELEMENT_COUNT; i++) {
+    chi->smoothed_metal_mass_fraction[i] += chj->metal_mass[i] * wi;
+    chj->smoothed_metal_mass_fraction[i] += chi->metal_mass[i] * wj;
+  }
+
+  /* Compute the shear tensor */
+  for (int i = 0; i < 3; i++) {
+    chi->S[i][0] += (pj->v[0] - pi->v[0]) * dx[i] * mj_wi_dr;
+    chi->S[i][1] += (pj->v[1] - pi->v[1]) * dx[i] * mj_wi_dr;
+    chi->S[i][2] += (pj->v[2] - pi->v[2]) * dx[i] * mj_wi_dr;
+
+    chj->S[i][0] -= (pj->v[0] - pi->v[0]) * dx[i] * mi_wj_dr;
+    chj->S[i][1] -= (pj->v[1] - pi->v[1]) * dx[i] * mi_wj_dr;
+    chj->S[i][2] -= (pj->v[2] - pi->v[2]) * dx[i] * mi_wj_dr;
+  }
+}
+
+/**
+ * @brief do chemistry computation after the runner_iact_density (non symmetric
+ * version)
+ *
+ * @param r2 Comoving square distance between the two particles.
+ * @param dx Comoving vector separating both particles (pi - pj).
+ * @param hi Comoving smoothing-length of particle i.
+ * @param hj Comoving smoothing-length of particle j.
+ * @param pi First particle.
+ * @param pj Second particle (not updated).
+ * @param a Current scale factor.
+ * @param H Current Hubble parameter.
+ */
+__attribute__((always_inline)) INLINE static void runner_iact_nonsym_chemistry(
+    float r2, const float *dx, float hi, float hj, struct part *restrict pi,
+    const struct part *restrict pj, float a, float H) {
+
+  struct chemistry_part_data *chi = &pi->chemistry_data;
+  const struct chemistry_part_data *chj = &pj->chemistry_data;
+
+  float wi, wi_dx;
+
+  /* Get the masses. */
+  const float mj = hydro_get_mass(pj);
+
+  /* Get r */
+  const float r = sqrtf(r2);
+  const float r_inv = 1.f / r;
+
+  /* Compute the kernel function for pi */
+  const float ui = r / hi;
+  kernel_deval(ui, &wi, &wi_dx);
+
+  const float wi_dr = wi_dx * r_inv;
+  const float mj_wi_dr = mj * wi_dr;
+
+  /* Compute contribution to the smooth metallicity */
+  for (int i = 0; i < GEAR_CHEMISTRY_ELEMENT_COUNT; i++) {
+    chi->smoothed_metal_mass_fraction[i] += chj->metal_mass[i] * wi;
+  }
+
+  /* Compute the shear tensor */
+  for (int i = 0; i < 3; i++) {
+    chi->S[i][0] += (pj->v[0] - pi->v[0]) * dx[i] * mj_wi_dr;
+    chi->S[i][1] += (pj->v[1] - pi->v[1]) * dx[i] * mj_wi_dr;
+    chi->S[i][2] += (pj->v[2] - pi->v[2]) * dx[i] * mj_wi_dr;
+  }
+}
+
+/**
+ * @brief do metal diffusion computation in the <FORCE LOOP>
+ * (symmetric version)
+ *
+ * @param r2 Comoving square distance between the two particles.
+ * @param dx Comoving vector separating both particles (pi - pj).
+ * @param hi Comoving smoothing-length of particle i.
+ * @param hj Comoving smoothing-length of particle j.
+ * @param pi First particle.
+ * @param pj Second particle.
+ * @param a Current scale factor.
+ * @param H Current Hubble parameter.
+ * @param time_base The time base used in order to convert integer to float
+ * time.
+ * @param ti_current The current time (in integer)
+ * @param cosmo The #cosmology.
+ * @param with_cosmology Are we running with cosmology?
+ *
+ */
+__attribute__((always_inline)) INLINE static void runner_iact_diffusion(
+    float r2, const float *dx, float hi, float hj, struct part *restrict pi,
+    struct part *restrict pj, float a, float H, float time_base,
+    integertime_t t_current, const struct cosmology *cosmo,
+    const int with_cosmology) {
+
+  struct chemistry_part_data *chi = &pi->chemistry_data;
+  struct chemistry_part_data *chj = &pj->chemistry_data;
+
+  /* No need to diffuse if both particles are not diffusing. */
+  if (chj->diff_coef > 0 && chi->diff_coef > 0) {
+
+    /* Get mass */
+    const float mj = hydro_get_mass(pj);
+    const float mi = hydro_get_mass(pi);
+    const float rhoj = hydro_get_physical_density(pj, cosmo);
+    const float rhoi = hydro_get_physical_density(pi, cosmo);
+
+    float wi, wj, dwi_dx, dwj_dx;
+
+    /* Get r */
+    const float r = sqrtf(r2);
+
+    /* part j */
+    /* Get the kernel for hj */
+    const float hj_inv = 1.0f / hj;
+
+    /* Compute the kernel function for pj */
+    const float xj = r * hj_inv;
+    kernel_deval(xj, &wj, &dwj_dx);
+
+    /* part i */
+    /* Get the kernel for hi */
+    const float hi_inv = 1.0f / hi;
+
+    /* Compute the kernel function for pi */
+    const float xi = r * hi_inv;
+    kernel_deval(xi, &wi, &dwi_dx);
+
+    /* Get 1/r */
+    const float r_inv = 1.f / sqrtf(r2);
+
+    const float wi_dr = dwi_dx * r_inv;
+    const float wj_dr = dwj_dx * r_inv;
+
+    const float mj_dw_r = mj * wi_dr;
+    const float mi_dw_r = mi * wj_dr;
+
+    /* Compute the diffusion coefficient <D> / <rho> in physical units. */
+    const float coef = 2. * (chi->diff_coef + chj->diff_coef) / (rhoi + rhoj);
+
+    const float coef_i = coef * mj_dw_r;
+    const float coef_j = coef * mi_dw_r;
+
+    /* Compute the time derivative */
+    const float inv_mi = 1.f / hydro_get_mass(pi);
+    const float inv_mj = 1.f / hydro_get_mass(pj);
+    for (int i = 0; i < GEAR_CHEMISTRY_ELEMENT_COUNT; i++) {
+      const double mi_frac = chi->metal_mass[i] * inv_mi;
+      const double mj_frac = chj->metal_mass[i] * inv_mj;
+      const double dm = mi_frac - mj_frac;
+      chi->metal_mass_dt[i] += coef_i * dm;
+      chj->metal_mass_dt[i] -= coef_j * dm;
+    }
+  }
+}
+
+/**
+ * @brief do metal diffusion computation in the <FORCE LOOP>
+ * (nonsymmetric version)
+ *
+ * @param r2 Comoving square distance between the two particles.
+ * @param dx Comoving vector separating both particles (pi - pj).
+ * @param hi Comoving smoothing-length of particle i.
+ * @param hj Comoving smoothing-length of particle j.
+ * @param pi First particle.
+ * @param pj Second particle.
+ * @param a Current scale factor.
+ * @param H Current Hubble parameter.
+ * @param time_base The time base used in order to convert integer to float
+ * time.
+ * @param ti_current The current time (in integer)
+ * @param cosmo The #cosmology.
+ * @param with_cosmology Are we running with cosmology?
+ *
+ */
+__attribute__((always_inline)) INLINE static void runner_iact_nonsym_diffusion(
+    float r2, const float *dx, float hi, float hj, struct part *restrict pi,
+    struct part *restrict pj, float a, float H, float time_base,
+    integertime_t t_current, const struct cosmology *cosmo,
+    const int with_cosmology) {
+
+  struct chemistry_part_data *chi = &pi->chemistry_data;
+  struct chemistry_part_data *chj = &pj->chemistry_data;
+
+  if (chj->diff_coef > 0 && chi->diff_coef > 0) {
+
+    /* Get mass */
+    const float mj = hydro_get_mass(pj);
+    const float rhoj = hydro_get_physical_density(pj, cosmo);
+    const float rhoi = hydro_get_physical_density(pi, cosmo);
+
+    float wi, dwi_dx;
+
+    /* Get r */
+    const float r = sqrtf(r2);
+
+    /* part i */
+    /* Get the kernel for hi */
+    const float hi_inv = 1.0f / hi;
+
+    /* Compute the kernel function for pi */
+    const float xi = r * hi_inv;
+    kernel_deval(xi, &wi, &dwi_dx);
+
+    /* Get 1/r */
+    const float r_inv = 1.f / sqrtf(r2);
+    const float wi_dr = dwi_dx * r_inv;
+
+    const float mj_dw_r = mj * wi_dr;
+
+    /* Compute the diffusion coefficient <D> / <rho> in physical units. */
+    const float coef = 2. * (chi->diff_coef + chj->diff_coef) / (rhoi + rhoj);
+
+    const float coef_i = coef * mj_dw_r;
+
+    /* Compute the time derivative */
+    const float inv_mi = 1.f / hydro_get_mass(pi);
+    const float inv_mj = 1.f / hydro_get_mass(pj);
+    for (int i = 0; i < GEAR_CHEMISTRY_ELEMENT_COUNT; i++) {
+      const double mi_frac = chi->metal_mass[i] * inv_mi;
+      const double mj_frac = chj->metal_mass[i] * inv_mj;
+      const double dm = mi_frac - mj_frac;
+      chi->metal_mass_dt[i] += coef_i * dm;
+    }
+  }
+}
+
+#endif /* SWIFT_GEAR_DIFFUSION_CHEMISTRY_IACT_H */
diff --git a/src/chemistry/GEAR_DIFFUSION/chemistry_io.h b/src/chemistry/GEAR_DIFFUSION/chemistry_io.h
new file mode 100644
index 0000000000000000000000000000000000000000..29761c0efc759533b7cf53f105f7d6575ee9d790
--- /dev/null
+++ b/src/chemistry/GEAR_DIFFUSION/chemistry_io.h
@@ -0,0 +1,160 @@
+/*******************************************************************************
+ * This file is part of SWIFT.
+ * Copyright (c) 2020 Loic Hausammann (loic.hausammann@epfl.ch)
+ *
+ * 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_IO_GEAR_DIFFUSION_H
+#define SWIFT_CHEMISTRY_IO_GEAR_DIFFUSION_H
+
+#include "chemistry_struct.h"
+#include "engine.h"
+#include "error.h"
+#include "feedback.h"
+#include "io_properties.h"
+#include "parser.h"
+#include "part.h"
+#include "physical_constants.h"
+#include "units.h"
+
+/**
+ * @brief Specifies which particle fields to read from a dataset
+ *
+ * @param parts The particle array.
+ * @param list The list of i/o properties to read.
+ *
+ * @return Returns the number of fields to read.
+ */
+INLINE static int chemistry_read_particles(struct part* parts,
+                                           struct io_props* list) {
+
+  /* List what we want to read */
+  list[0] = io_make_input_field(
+      "MetalMassFraction", DOUBLE, GEAR_CHEMISTRY_ELEMENT_COUNT, OPTIONAL,
+      UNIT_CONV_NO_UNITS, parts, chemistry_data.metal_mass);
+
+  return 1;
+}
+
+INLINE static void convert_gas_metals(const struct engine* e,
+                                      const struct part* p,
+                                      const struct xpart* xp, double* ret) {
+
+  for (int i = 0; i < GEAR_CHEMISTRY_ELEMENT_COUNT; i++) {
+    ret[i] = p->chemistry_data.metal_mass[i] / hydro_get_mass(p);
+  }
+}
+
+/**
+ * @brief Specifies which particle fields to write to a dataset
+ *
+ * @param parts The particle array.
+ * @param xparts The extra particle array.
+ * @param list The list of i/o properties to write.
+ *
+ * @return Returns the number of fields to write.
+ */
+INLINE static int chemistry_write_particles(const struct part* parts,
+                                            const struct xpart* xparts,
+                                            struct io_props* list,
+                                            const int with_cosmology) {
+
+  /* List what we want to write */
+  list[0] = io_make_output_field(
+      "SmoothedMetalMassFractions", DOUBLE, GEAR_CHEMISTRY_ELEMENT_COUNT,
+      UNIT_CONV_NO_UNITS, 0.f, parts,
+      chemistry_data.smoothed_metal_mass_fraction,
+      "Mass fraction of each element smoothed over the neighbors");
+
+  list[1] = io_make_output_field_convert_part(
+      "MetalMassFractions", DOUBLE, GEAR_CHEMISTRY_ELEMENT_COUNT,
+      UNIT_CONV_NO_UNITS, 0.f, parts, xparts, convert_gas_metals,
+      "Mass fraction of each element");
+
+  return 2;
+}
+
+/**
+ * @brief Specifies which sparticle fields to write to a dataset
+ *
+ * @param sparts The sparticle array.
+ * @param list The list of i/o properties to write.
+ *
+ * @return Returns the number of fields to write.
+ */
+INLINE static int chemistry_write_sparticles(const struct spart* sparts,
+                                             struct io_props* list) {
+
+  /* List what we want to write */
+  list[0] = io_make_output_field(
+      "MetalMassFractions", DOUBLE, GEAR_CHEMISTRY_ELEMENT_COUNT,
+      UNIT_CONV_NO_UNITS, 0.f, sparts, chemistry_data.metal_mass_fraction,
+      "Mass fraction of each element");
+
+  return 1;
+}
+
+/**
+ * @brief Specifies which black hole particle fields to write to a dataset
+ *
+ * @param bparts The black hole particle array.
+ * @param list The list of i/o properties to write.
+ *
+ * @return Returns the number of fields to write.
+ */
+INLINE static int chemistry_write_bparticles(const struct bpart* bparts,
+                                             struct io_props* list) {
+
+  /* No fields to write here */
+  return 0;
+}
+
+#ifdef HAVE_HDF5
+
+/**
+ * @brief Writes the current model of SPH to the file
+ * @param h_grp The HDF5 group in which to write
+ * @param h_grp_columns The HDF5 group containing named columns
+ * @param e The #engine.
+ */
+INLINE static void chemistry_write_flavour(hid_t h_grp, hid_t h_grp_columns,
+                                           const struct engine* e) {
+
+  io_write_attribute_s(h_grp, "Chemistry Model", "GEAR with diffusion");
+  io_write_attribute_d(h_grp, "Chemistry element count",
+                       GEAR_CHEMISTRY_ELEMENT_COUNT);
+#ifdef FEEDBACK_GEAR
+  const char* element_names = e->feedback_props->stellar_model.elements_name;
+
+  /* Add to the named columns */
+  hsize_t dims[1] = {GEAR_CHEMISTRY_ELEMENT_COUNT};
+  hid_t type = H5Tcopy(H5T_C_S1);
+  H5Tset_size(type, GEAR_LABELS_SIZE);
+  hid_t space = H5Screate_simple(1, dims, NULL);
+  hid_t dset = H5Dcreate(h_grp_columns, "ElementAbundances", type, space,
+                         H5P_DEFAULT, H5P_DEFAULT, H5P_DEFAULT);
+  H5Dwrite(dset, type, H5S_ALL, H5S_ALL, H5P_DEFAULT, element_names);
+  H5Dclose(dset);
+  dset = H5Dcreate(h_grp_columns, "SmoothedElementAbundances", type, space,
+                   H5P_DEFAULT, H5P_DEFAULT, H5P_DEFAULT);
+  H5Dwrite(dset, type, H5S_ALL, H5S_ALL, H5P_DEFAULT, element_names);
+  H5Dclose(dset);
+
+  H5Tclose(type);
+#endif
+}
+#endif
+
+#endif /* SWIFT_CHEMISTRY_IO_GEAR_DIFFUSION_H */
diff --git a/src/chemistry/GEAR_DIFFUSION/chemistry_struct.h b/src/chemistry/GEAR_DIFFUSION/chemistry_struct.h
new file mode 100644
index 0000000000000000000000000000000000000000..eb8b1e255e4d4130ef843251dff34bd199e42295
--- /dev/null
+++ b/src/chemistry/GEAR_DIFFUSION/chemistry_struct.h
@@ -0,0 +1,70 @@
+/*******************************************************************************
+ * This file is part of SWIFT.
+ * Copyright (c) 2020 Loic Hausammann (loic.hausammann@epfl.ch)
+ *
+ * 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_STRUCT_GEAR_DIFFUSION_H
+#define SWIFT_CHEMISTRY_STRUCT_GEAR_DIFFUSION_H
+
+/**
+ * @brief Global chemical abundance information.
+ */
+struct chemistry_global_data {
+
+  /* Initial mass fraction */
+  double initial_metallicities[GEAR_CHEMISTRY_ELEMENT_COUNT];
+
+  /*! Diffusion coefficent (no unit) */
+  float C;
+};
+
+/**
+ * @brief Properties of the chemistry function for #part.
+ */
+struct chemistry_part_data {
+
+  /*! Total mass of element in a particle.
+    This field is available only outside the density hydro loop. */
+  double metal_mass[GEAR_CHEMISTRY_ELEMENT_COUNT];
+
+  /*! Smoothed fraction of the particle mass in a given element */
+  double smoothed_metal_mass_fraction[GEAR_CHEMISTRY_ELEMENT_COUNT];
+
+  /*! Diffusion coefficient */
+  float diff_coef;
+
+  /*! Variation of the metal mass */
+  double metal_mass_dt[GEAR_CHEMISTRY_ELEMENT_COUNT];
+
+  /*! shear tensor in internal and physical units. */
+  float S[3][3];
+};
+
+/**
+ * @brief Properties of the chemistry function for #spart.
+ */
+struct chemistry_spart_data {
+
+  /*! Fraction of the particle mass in a given element */
+  double metal_mass_fraction[GEAR_CHEMISTRY_ELEMENT_COUNT];
+};
+
+/**
+ * @brief Chemical abundances traced by the #bpart in the GEAR model.
+ */
+struct chemistry_bpart_data {};
+
+#endif /* SWIFT_CHEMISTRY_STRUCT_GEAR_DIFFUSION_H */
diff --git a/src/chemistry/QLA/chemistry.h b/src/chemistry/QLA/chemistry.h
index c9823214570a82649c4db77049e9d611c26e94d7..7d6ffba8facf8d1e3b1e6d2abafbd55a2f3c2ba8 100644
--- a/src/chemistry/QLA/chemistry.h
+++ b/src/chemistry/QLA/chemistry.h
@@ -94,11 +94,17 @@ __attribute__((always_inline)) INLINE static void chemistry_end_density(
 /**
  * @brief Updates to the chemistry data after the hydro force loop.
  *
+ * Nothing to do here.
+ *
  * @param p The particle to act upon.
  * @param cosmo The current cosmological model.
+ * @param with_cosmology Are we running with the cosmology?
+ * @param time Current time of the simulation.
+ * @param dt Time step (in physical units).
  */
 __attribute__((always_inline)) INLINE static void chemistry_end_force(
-    struct part* restrict p, const struct cosmology* cosmo) {}
+    struct part* restrict p, const struct cosmology* cosmo,
+    const int with_cosmology, const double time, const double dt) {}
 
 /**
  * @brief Computes the chemistry-related time-step constraint.
diff --git a/src/chemistry/QLA/chemistry_iact.h b/src/chemistry/QLA/chemistry_iact.h
index 89f0d4a3c8a2dd4fbde40b4525196d439ba66f85..1c3dcfed72fd31312f5b7bc337e599991c017afe 100644
--- a/src/chemistry/QLA/chemistry_iact.h
+++ b/src/chemistry/QLA/chemistry_iact.h
@@ -58,4 +58,54 @@ __attribute__((always_inline)) INLINE static void runner_iact_nonsym_chemistry(
     float r2, const float *dx, float hi, float hj, struct part *restrict pi,
     const struct part *restrict pj, float a, float H) {}
 
+/**
+ * @brief do metal diffusion computation in the <FORCE LOOP>
+ * (symmetric version)
+ *
+ * @param r2 Comoving square distance between the two particles.
+ * @param dx Comoving vector separating both particles (pi - pj).
+ * @param hi Comoving smoothing-length of particle i.
+ * @param hj Comoving smoothing-length of particle j.
+ * @param pi First particle.
+ * @param pj Second particle.
+ * @param a Current scale factor.
+ * @param H Current Hubble parameter.
+ * @param time_base The time base used in order to convert integer to float
+ * time.
+ * @param ti_current The current time (in integer)
+ * @param cosmo The #cosmology.
+ * @param with_cosmology Are we running with cosmology?
+ *
+ */
+__attribute__((always_inline)) INLINE static void runner_iact_diffusion(
+    float r2, const float *dx, float hi, float hj, struct part *restrict pi,
+    struct part *restrict pj, float a, float H, float time_base,
+    integertime_t t_current, const struct cosmology *cosmo,
+    const int with_cosmology) {}
+
+/**
+ * @brief do metal diffusion computation in the <FORCE LOOP>
+ * (nonsymmetric version)
+ *
+ * @param r2 Comoving square distance between the two particles.
+ * @param dx Comoving vector separating both particles (pi - pj).
+ * @param hi Comoving smoothing-length of particle i.
+ * @param hj Comoving smoothing-length of particle j.
+ * @param pi First particle.
+ * @param pj Second particle.
+ * @param a Current scale factor.
+ * @param H Current Hubble parameter.
+ * @param time_base The time base used in order to convert integer to float
+ * time.
+ * @param ti_current The current time (in integer)
+ * @param cosmo The #cosmology.
+ * @param with_cosmology Are we running with cosmology?
+ *
+ */
+__attribute__((always_inline)) INLINE static void runner_iact_nonsym_diffusion(
+    float r2, const float *dx, float hi, float hj, struct part *restrict pi,
+    struct part *restrict pj, float a, float H, float time_base,
+    integertime_t t_current, const struct cosmology *cosmo,
+    const int with_cosmology) {}
+
 #endif /* SWIFT_QLA_CHEMISTRY_IACT_H */
diff --git a/src/chemistry/QLA/chemistry_struct.h b/src/chemistry/QLA/chemistry_struct.h
index 97194433b1d1e33a5ae66d67f344ac48d58f71ce..ac5ac803773790704666bbab5ac38bd16c766575 100644
--- a/src/chemistry/QLA/chemistry_struct.h
+++ b/src/chemistry/QLA/chemistry_struct.h
@@ -43,6 +43,8 @@ struct chemistry_global_data {};
  */
 struct chemistry_part_data {};
 
+#define chemistry_spart_data chemistry_part_data
+
 /**
  * @brief Chemistry properties carried by the #bpart.
  *
diff --git a/src/chemistry/none/chemistry.h b/src/chemistry/none/chemistry.h
index 3cd78a2cc63af5caf62f439ad8df5cf05cc33683..fa6e0c892d1d36d472306bb228ba803cbdbe112c 100644
--- a/src/chemistry/none/chemistry.h
+++ b/src/chemistry/none/chemistry.h
@@ -101,11 +101,17 @@ __attribute__((always_inline)) INLINE static void chemistry_end_density(
 /**
  * @brief Updates to the chemistry data after the hydro force loop.
  *
+ * Nothing to do here.
+ *
  * @param p The particle to act upon.
  * @param cosmo The current cosmological model.
+ * @param with_cosmology Are we running with the cosmology?
+ * @param time Current time of the simulation.
+ * @param dt Time step (in physical units).
  */
 __attribute__((always_inline)) INLINE static void chemistry_end_force(
-    struct part* restrict p, const struct cosmology* cosmo) {}
+    struct part* restrict p, const struct cosmology* cosmo,
+    const int with_cosmology, const double time, const double dt) {}
 
 /**
  * @brief Computes the chemistry-related time-step constraint.
diff --git a/src/chemistry/none/chemistry_iact.h b/src/chemistry/none/chemistry_iact.h
index 52499c5e5cc3d5749904251cc967193d875579b9..13233f6a0f0d4657f49a51b2da5bc456bd829759 100644
--- a/src/chemistry/none/chemistry_iact.h
+++ b/src/chemistry/none/chemistry_iact.h
@@ -59,4 +59,54 @@ __attribute__((always_inline)) INLINE static void runner_iact_nonsym_chemistry(
     float r2, const float *dx, float hi, float hj, struct part *restrict pi,
     const struct part *restrict pj, float a, float H) {}
 
+/**
+ * @brief do metal diffusion computation in the <FORCE LOOP>
+ * (symmetric version)
+ *
+ * @param r2 Comoving square distance between the two particles.
+ * @param dx Comoving vector separating both particles (pi - pj).
+ * @param hi Comoving smoothing-length of particle i.
+ * @param hj Comoving smoothing-length of particle j.
+ * @param pi First particle.
+ * @param pj Second particle.
+ * @param a Current scale factor.
+ * @param H Current Hubble parameter.
+ * @param time_base The time base used in order to convert integer to float
+ * time.
+ * @param ti_current The current time (in integer)
+ * @param cosmo The #cosmology.
+ * @param with_cosmology Are we running with cosmology?
+ *
+ */
+__attribute__((always_inline)) INLINE static void runner_iact_diffusion(
+    float r2, const float *dx, float hi, float hj, struct part *restrict pi,
+    struct part *restrict pj, float a, float H, float time_base,
+    integertime_t t_current, const struct cosmology *cosmo,
+    const int with_cosmology) {}
+
+/**
+ * @brief do metal diffusion computation in the <FORCE LOOP>
+ * (nonsymmetric version)
+ *
+ * @param r2 Comoving square distance between the two particles.
+ * @param dx Comoving vector separating both particles (pi - pj).
+ * @param hi Comoving smoothing-length of particle i.
+ * @param hj Comoving smoothing-length of particle j.
+ * @param pi First particle.
+ * @param pj Second particle.
+ * @param a Current scale factor.
+ * @param H Current Hubble parameter.
+ * @param time_base The time base used in order to convert integer to float
+ * time.
+ * @param ti_current The current time (in integer)
+ * @param cosmo The #cosmology.
+ * @param with_cosmology Are we running with cosmology?
+ *
+ */
+__attribute__((always_inline)) INLINE static void runner_iact_nonsym_diffusion(
+    float r2, const float *dx, float hi, float hj, struct part *restrict pi,
+    struct part *restrict pj, float a, float H, float time_base,
+    integertime_t t_current, const struct cosmology *cosmo,
+    const int with_cosmology) {}
+
 #endif /* SWIFT_NONE_CHEMISTRY_IACT_H */
diff --git a/src/chemistry_io.h b/src/chemistry_io.h
index c94461c84bdf243990fa8f28f0049178e180d9ec..d7386f8c68a42b797a8d201e8a68329f00b4f6fc 100644
--- a/src/chemistry_io.h
+++ b/src/chemistry_io.h
@@ -27,6 +27,8 @@
 #include "./chemistry/none/chemistry_io.h"
 #elif defined(CHEMISTRY_GEAR)
 #include "./chemistry/GEAR/chemistry_io.h"
+#elif defined(CHEMISTRY_GEAR_DIFFUSION)
+#include "./chemistry/GEAR_DIFFUSION/chemistry_io.h"
 #elif defined(CHEMISTRY_QLA)
 #include "./chemistry/QLA/chemistry_io.h"
 #elif defined(CHEMISTRY_EAGLE)
diff --git a/src/chemistry_struct.h b/src/chemistry_struct.h
index 05074d8759be660876544aaa20cffec6cddb8bce..0b1cb59e592d9dbce8e77fb01cf45e1be846a55b 100644
--- a/src/chemistry_struct.h
+++ b/src/chemistry_struct.h
@@ -32,6 +32,8 @@
 #include "./chemistry/none/chemistry_struct.h"
 #elif defined(CHEMISTRY_GEAR)
 #include "./chemistry/GEAR/chemistry_struct.h"
+#elif defined(CHEMISTRY_GEAR_DIFFUSION)
+#include "./chemistry/GEAR_DIFFUSION/chemistry_struct.h"
 #elif defined(CHEMISTRY_QLA)
 #include "./chemistry/QLA/chemistry_struct.h"
 #elif defined(CHEMISTRY_EAGLE)
diff --git a/src/cooling/COLIBRE/cooling.c b/src/cooling/COLIBRE/cooling.c
index 0109baf3ea51e08a79a1a29d9ab0709b8f57a074..2073a5daec6700bb75e93c111006410751ab9f39 100644
--- a/src/cooling/COLIBRE/cooling.c
+++ b/src/cooling/COLIBRE/cooling.c
@@ -49,6 +49,7 @@
 #include "part.h"
 #include "physical_constants.h"
 #include "space.h"
+#include "star_formation.h"
 #include "units.h"
 
 /* Maximum number of iterations for
@@ -1148,7 +1149,7 @@ void cooling_Hydrogen_reionization(const struct cooling_function_data *cooling,
 
     if (part_is_inhibited(p, s->e)) continue;
 
-    if (xp->sf_data.SFR <= 0.) {
+    if (star_formation_get_SFR(p, xp) == 0.f) {
       const float old_u = hydro_get_physical_internal_energy(p, xp, cosmo);
       const float new_u = old_u + extra_heat;
 
diff --git a/src/cooling/Compton/cooling.h b/src/cooling/Compton/cooling.h
deleted file mode 100644
index 7b2530c04590e9257f96e8b15c78e0c2915c72a0..0000000000000000000000000000000000000000
--- a/src/cooling/Compton/cooling.h
+++ /dev/null
@@ -1,469 +0,0 @@
-/*******************************************************************************
- * This file is part of SWIFT.
- * Copyright (c) 2018 Matthieu Schaller (matthieu.schaller@durham.ac.uk)
- *
- * 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_COOLING_COMPTON_H
-#define SWIFT_COOLING_COMPTON_H
-
-/**
- * @file src/cooling/Compton/cooling.h
- * @brief Routines related to the "Compton" cooling function.
- *
- * This model compute the cooling rate from the Compton interaction with
- * the CMB photons.
- */
-
-/* Config parameters. */
-#include "../config.h"
-
-/* Some standard headers. */
-#include <float.h>
-#include <math.h>
-
-/* Local includes. */
-#include "cooling_properties.h"
-#include "entropy_floor.h"
-#include "error.h"
-#include "hydro.h"
-#include "parser.h"
-#include "part.h"
-#include "physical_constants.h"
-#include "units.h"
-
-/**
- * @brief Common operations performed on the cooling function at a
- * given time-step or redshift.
- *
- * @param cosmo The current cosmological model.
- * @param cooling The #cooling_function_data used in the run.
- * @param s The #space containing all the particles.
- */
-INLINE static void cooling_update(const struct cosmology* cosmo,
-                                  struct cooling_function_data* cooling,
-                                  struct space* s) {
-  // Add content if required.
-}
-
-/**
- * @brief Calculates du/dt in CGS units for a particle.
- *
- * @param cosmo The current cosmological model.
- * @param phys_const The physical constants in internal units.
- * @param hydro_props The properties of the hydro scheme.
- * @param cooling The #cooling_function_data used in the run.
- * @param z The current redshift.
- * @param u The current internal energy in internal units.
- * @param p Pointer to the particle data.
- * @return The change in energy per unit mass due to cooling for this particle
- * in cgs units [erg * g^-1 * s^-1].
- */
-__attribute__((always_inline)) INLINE static double Compton_cooling_rate_cgs(
-    const struct cosmology* cosmo, const struct phys_const* restrict phys_const,
-    const struct hydro_props* hydro_props,
-    const struct cooling_function_data* cooling, const double z, const double u,
-    const struct part* p) {
-
-  /* Get particle density */
-  const double rho = hydro_get_physical_density(p, cosmo);
-  const double rho_cgs = rho * cooling->conv_factor_density_to_cgs;
-
-  /* Powers of (1 + z) */
-  const double zp1 = z + 1.;
-  const double zp1p2 = zp1 * zp1;
-  const double zp1p4 = zp1p2 * zp1p2; /* (1 + z)^4 */
-
-  /* CMB temperature at this redshift */
-  const double T_CMB = cooling->const_T_CMB_0 * zp1;
-
-  /* Physical constants */
-  const double m_H = phys_const->const_proton_mass;
-  const double k_B = phys_const->const_boltzmann_k;
-
-  /* Gas properties */
-  const double T_transition = hydro_props->hydrogen_ionization_temperature;
-  const double mu_neutral = hydro_props->mu_neutral;
-  const double mu_ionised = hydro_props->mu_ionised;
-
-  /* Temperature over mean molecular weight */
-  const double T_over_mu = hydro_gamma_minus_one * u * m_H / k_B;
-
-  double T;
-
-  /* Are we above or below the HII -> HI transition? */
-  if (T_over_mu > (T_transition + 1.) / mu_ionised)
-    T = T_over_mu * mu_ionised;
-  else if (T_over_mu < (T_transition - 1.) / mu_neutral)
-    T = T_over_mu * mu_neutral;
-  else
-    T = T_transition;
-
-  /* Electron abundance */
-  double electron_abundance = 0.;  // MATTHIEU: To do: compute X_e
-
-  /* Temperature difference with the CMB */
-  const double delta_T = T - T_CMB;
-
-  /* Electron density */
-  const double electron_density_cgs =
-      rho_cgs * electron_abundance * cooling->proton_mass_cgs_inv;
-
-  /* Compton formula */
-  return cooling->const_Compton_rate_cgs * delta_T * zp1p4 *
-         electron_density_cgs / rho_cgs;
-}
-
-/**
- * @brief Apply the cooling function to a particle.
- *
- * @param phys_const The physical constants in internal units.
- * @param us The internal system of units.
- * @param cosmo The current cosmological model.
- * @param hydro_props The properties of the hydro scheme.
- * @param floor_props Properties of the entropy floor.
- * @param cooling The #cooling_function_data used in the run.
- * @param p Pointer to the particle data.
- * @param xp Pointer to the particle' extended data.
- * @param dt The time-step of this particle.
- * @param dt_therm The time-step operator used for thermal quantities.
- * @param time Time since Big Bang (or start of the simulation) in internal
- * units.
- */
-__attribute__((always_inline)) INLINE static void cooling_cool_part(
-    const struct phys_const* restrict phys_const,
-    const struct unit_system* restrict us,
-    const struct cosmology* restrict cosmo,
-    const struct hydro_props* hydro_props,
-    const struct entropy_floor_properties* floor_props,
-    const struct cooling_function_data* restrict cooling,
-    struct part* restrict p, struct xpart* restrict xp, const float dt,
-    const float dt_therm, const double time) {
-
-  /* Nothing to do here? */
-  if (dt == 0.) return;
-
-  /* Current energy */
-  const float u_old = hydro_get_physical_internal_energy(p, xp, cosmo);
-
-  /* Current du_dt in physical coordinates (internal units) */
-  const float hydro_du_dt = hydro_get_physical_internal_energy_dt(p, cosmo);
-
-  /* Calculate cooling du_dt (in cgs units) */
-  const double cooling_du_dt_cgs = Compton_cooling_rate_cgs(
-      cosmo, phys_const, hydro_props, cooling, cosmo->z, u_old, p);
-
-  /* Convert to internal units */
-  float cooling_du_dt =
-      cooling_du_dt_cgs * cooling->conv_factor_energy_rate_from_cgs;
-
-  /* Add cosmological term */
-  cooling_du_dt *= cosmo->a * cosmo->a;
-
-  float total_du_dt = hydro_du_dt + cooling_du_dt;
-
-  /* We now need to check that we are not going to go below any of the limits */
-
-  /* Limit imposed by the entropy floor */
-  const float A_floor = entropy_floor(p, cosmo, floor_props);
-  const float rho = hydro_get_physical_density(p, cosmo);
-  const float u_floor = gas_internal_energy_from_entropy(rho, A_floor);
-
-  /* Absolute minimum */
-  const float u_minimal = hydro_props->minimal_internal_energy;
-
-  /* Largest of both limits */
-  const float u_limit = max(u_minimal, u_floor);
-
-  /* First, check whether we may end up below the minimal energy after
-   * this step 1/2 kick + another 1/2 kick that could potentially be for
-   * a time-step twice as big. We hence check for 1.5 delta_t. */
-  if (u_old + total_du_dt * 1.5 * dt_therm < u_limit) {
-    total_du_dt = (u_limit - u_old) / (1.5f * dt_therm);
-  }
-
-  /* Second, check whether the energy used in the prediction could get negative.
-   * We need to check for the 1/2 dt kick followed by a full time-step drift
-   * that could potentially be for a time-step twice as big. We hence check
-   * for 2.5 delta_t but this time against 0 energy not the minimum */
-  if (u_old + total_du_dt * 2.5 * dt_therm < 0.) {
-    total_du_dt = -u_old / ((2.5f + 0.0001f) * dt_therm);
-  }
-
-  /* Update the internal energy time derivative */
-  hydro_set_physical_internal_energy_dt(p, cosmo, total_du_dt);
-
-  /* Store the radiated energy (assuming dt will not change) */
-  xp->cooling_data.radiated_energy +=
-      -hydro_get_mass(p) * (total_du_dt - hydro_du_dt) * dt_therm;
-}
-
-/**
- * @brief Computes the time-step due to cooling for this particle.
- *
- * We impose no time-step limit.
- *
- * @param cooling The #cooling_function_data used in the run.
- * @param phys_const The physical constants in internal units.
- * @param cosmo The current cosmological model.
- * @param hydro_props The properties of the hydro scheme.
- * @param us The internal system of units.
- * @param p Pointer to the particle data.
- * @param xp Pointer to the extended data of the particle.
- */
-__attribute__((always_inline)) INLINE static float cooling_timestep(
-    const struct cooling_function_data* restrict cooling,
-    const struct phys_const* restrict phys_const,
-    const struct cosmology* restrict cosmo,
-    const struct unit_system* restrict us,
-    const struct hydro_props* hydro_props, const struct part* restrict p,
-    const struct xpart* restrict xp) {
-
-  return FLT_MAX;
-}
-
-/**
- * @brief Sets the cooling properties of the (x-)particles to a valid start
- * state.
- *
- * Nothing to do here. Just set the radiated energy counter to 0.
- *
- * @param phys_const The physical constants in internal units.
- * @param cooling The properties of the cooling function.
- * @param us The internal system of units.
- * @param hydro_props The properties of the hydro scheme.
- * @param cosmo The current cosmological model.
- * @param p Pointer to the particle data.
- * @param xp Pointer to the extended particle data.
- */
-__attribute__((always_inline)) INLINE static void cooling_first_init_part(
-    const struct phys_const* restrict phys_const,
-    const struct unit_system* restrict us,
-    const struct hydro_props* hydro_props,
-    const struct cosmology* restrict cosmo,
-    const struct cooling_function_data* restrict cooling,
-    const struct part* restrict p, struct xpart* restrict xp) {
-
-  xp->cooling_data.radiated_energy = 0.f;
-}
-
-/**
- * @brief Compute the temperature of a #part based on the cooling function.
- *
- * @param phys_const #phys_const data structure.
- * @param hydro_props The properties of the hydro scheme.
- * @param us The internal system of units.
- * @param cosmo #cosmology data structure.
- * @param cooling #cooling_function_data struct.
- * @param p #part data.
- * @param xp Pointer to the #xpart data.
- */
-INLINE static float cooling_get_temperature(
-    const struct phys_const* restrict phys_const,
-    const struct hydro_props* restrict hydro_props,
-    const struct unit_system* restrict us,
-    const struct cosmology* restrict cosmo,
-    const struct cooling_function_data* restrict cooling,
-    const struct part* restrict p, const struct xpart* restrict xp) {
-
-  /* Physical constants */
-  const double m_H = phys_const->const_proton_mass;
-  const double k_B = phys_const->const_boltzmann_k;
-
-  /* Gas properties */
-  const double T_transition = hydro_props->hydrogen_ionization_temperature;
-  const double mu_neutral = hydro_props->mu_neutral;
-  const double mu_ionised = hydro_props->mu_ionised;
-
-  /* Particle temperature */
-  const double u = hydro_get_physical_internal_energy(p, xp, cosmo);
-
-  /* Temperature over mean molecular weight */
-  const double T_over_mu = hydro_gamma_minus_one * u * m_H / k_B;
-
-  /* Are we above or below the HII -> HI transition? */
-  if (T_over_mu > (T_transition + 1.) / mu_ionised)
-    return T_over_mu * mu_ionised;
-  else if (T_over_mu < (T_transition - 1.) / mu_neutral)
-    return T_over_mu * mu_neutral;
-  else
-    return T_transition;
-}
-
-/**
- * @brief Returns the subgrid temperature of a particle.
- *
- * This model has no subgrid quantity. We return an error.
- *
- * @param p The particle.
- * @param xp The extended particle data.
- */
-INLINE static float cooling_get_subgrid_temperature(const struct part* p,
-                                                    const struct xpart* xp) {
-  error("This cooling model does not use subgrid quantities!");
-  return -1.f;
-}
-
-/**
- * @brief Returns the subgrid density of a particle.
- *
- * This model has no subgrid quantity. We return an error.
- *
- * @param p The particle.
- * @param xp The extended particle data.
- */
-INLINE static float cooling_get_subgrid_density(const struct part* p,
-                                                const struct xpart* xp) {
-  error("This cooling model does not use subgrid quantities!");
-  return -1.f;
-}
-
-/**
- * @brief Returns the total radiated energy by this particle.
- *
- * @param xp The extended particle data
- */
-__attribute__((always_inline)) INLINE static float cooling_get_radiated_energy(
-    const struct xpart* restrict xp) {
-
-  return xp->cooling_data.radiated_energy;
-}
-
-/**
- * @brief Split the coolong content of a particle into n pieces
- *
- * @param p The #part.
- * @param xp The #xpart.
- * @param n The number of pieces to split into.
- */
-static INLINE void cooling_split_part(struct part* p, struct xpart* xp,
-                                      double n) {
-
-  xp->cooling_data.radiated_energy /= n;
-}
-
-/**
- * @brief Initialises the cooling properties.
- *
- * @param parameter_file The parsed parameter file.
- * @param us The current internal system of units.
- * @param phys_const The physical constants in internal units.
- * @param hydro_props The properties of the hydro scheme.
- * @param cooling The cooling properties to initialize
- */
-static INLINE void cooling_init_backend(struct swift_params* parameter_file,
-                                        const struct unit_system* us,
-                                        const struct phys_const* phys_const,
-                                        const struct hydro_props* hydro_props,
-                                        struct cooling_function_data* cooling) {
-
-  /* Some useful conversion values */
-  cooling->conv_factor_density_to_cgs =
-      units_cgs_conversion_factor(us, UNIT_CONV_DENSITY);
-  cooling->conv_factor_energy_to_cgs =
-      units_cgs_conversion_factor(us, UNIT_CONV_ENERGY_PER_UNIT_MASS);
-  cooling->conv_factor_energy_rate_from_cgs =
-      units_cgs_conversion_factor(us, UNIT_CONV_TIME) /
-      units_cgs_conversion_factor(us, UNIT_CONV_ENERGY_PER_UNIT_MASS);
-
-  /* Useful constants */
-  cooling->proton_mass_cgs_inv =
-      1. / (phys_const->const_proton_mass *
-            units_cgs_conversion_factor(us, UNIT_CONV_MASS));
-
-  /* Temperature of the CMB in CGS */
-  const double T_CMB_0 = phys_const->const_T_CMB_0 *
-                         units_cgs_conversion_factor(us, UNIT_CONV_TEMPERATURE);
-  cooling->const_T_CMB_0 = T_CMB_0; /* [K] */
-
-  /* Compute the coefficient at the front of the Compton cooling expression */
-  const double radiation_constant =
-      4. * phys_const->const_stefan_boltzmann / phys_const->const_speed_light_c;
-  const double compton_coefficient =
-      4. * radiation_constant * phys_const->const_thomson_cross_section *
-      phys_const->const_boltzmann_k /
-      (phys_const->const_electron_mass * phys_const->const_speed_light_c);
-  const float dimension_coefficient[5] = {1, 2, -3, 0, -5};
-
-  /* This should be ~1.0178085e-37 [g cm^2 s^-3 K^-5] */
-  const double compton_coefficient_cgs =
-      compton_coefficient *
-      units_general_cgs_conversion_factor(us, dimension_coefficient);
-
-  /* And now the Compton rate [g cm^2 s^-3 K^-1] == [erg s^-1 K^-1]*/
-  cooling->const_Compton_rate_cgs =
-      compton_coefficient_cgs * T_CMB_0 * T_CMB_0 * T_CMB_0 * T_CMB_0;
-}
-
-/**
- * @brief Restore cooling tables (if applicable) after
- * restart
- *
- * @param cooling the cooling_function_data structure
- * @param cosmo cosmology structure
- */
-static INLINE void cooling_restore_tables(struct cooling_function_data* cooling,
-                                          const struct cosmology* cosmo) {}
-
-/**
- * @brief Prints the properties of the cooling model to stdout.
- *
- * @param cooling The properties of the cooling function.
- */
-static INLINE void cooling_print_backend(
-    const struct cooling_function_data* cooling) {
-
-  message("Cooling function is 'Compton cooling'.");
-}
-
-/**
- * @brief Clean-up the memory allocated for the cooling routines
- *
- * @param cooling the cooling data structure.
- */
-static INLINE void cooling_clean(struct cooling_function_data* cooling) {}
-
-/**
- * @brief Write a cooling struct to the given FILE as a stream of bytes.
- *
- * Nothing to do beyond writing the structure from the stream.
- *
- * @param cooling the struct
- * @param stream the file stream
- */
-static INLINE void cooling_struct_dump(
-    const struct cooling_function_data* cooling, FILE* stream) {
-  restart_write_blocks((void*)cooling, sizeof(struct cooling_function_data), 1,
-                       stream, "cooling", "cooling function");
-}
-
-/**
- * @brief Restore a hydro_props struct from the given FILE as a stream of
- * bytes.
- *
- * Nothing to do beyond reading the structure from the stream.
- *
- * @param cooling the struct
- * @param stream the file stream
- * @param cosmo #cosmology structure
- */
-static INLINE void cooling_struct_restore(struct cooling_function_data* cooling,
-                                          FILE* stream,
-                                          const struct cosmology* cosmo) {
-  restart_read_blocks((void*)cooling, sizeof(struct cooling_function_data), 1,
-                      stream, NULL, "cooling function");
-}
-
-#endif /* SWIFT_COOLING_COMPTON_H */
diff --git a/src/cooling/Compton/cooling_io.h b/src/cooling/Compton/cooling_io.h
deleted file mode 100644
index e289544e2d40e920ae208ee8dedbfb40c0028a24..0000000000000000000000000000000000000000
--- a/src/cooling/Compton/cooling_io.h
+++ /dev/null
@@ -1,78 +0,0 @@
-/*******************************************************************************
- * This file is part of SWIFT.
- * Copyright (c) 2018 Matthieu Schaller (matthieu.schaller@durham.ac.uk)
- *
- * 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_COOLING_IO_COMPTON_H
-#define SWIFT_COOLING_IO_COMPTON_H
-
-/* Config parameters. */
-#include "../config.h"
-
-/* Local includes */
-#include "cooling.h"
-#include "engine.h"
-#include "io_properties.h"
-
-#ifdef HAVE_HDF5
-
-/**
- * @brief Writes the current model of cooling to the file
- *
- * @param h_grp The HDF5 group in which to write
- * @param h_grp_columns The HDF5 group containing named columns
- * @param cooling The #cooling_function_data
- */
-__attribute__((always_inline)) INLINE static void cooling_write_flavour(
-    hid_t h_grp, hid_t h_grp_columns,
-    const struct cooling_function_data* cooling) {
-
-  io_write_attribute_s(h_grp, "Cooling Model", "Compton cooling");
-  io_write_attribute_d(h_grp, "Compton rate [erg s^-1 K^-1]",
-                       cooling->const_Compton_rate_cgs);
-}
-#endif
-
-INLINE static void convert_part_T(const struct engine* e, const struct part* p,
-                                  const struct xpart* xp, float* ret) {
-
-  ret[0] = cooling_get_temperature(e->physical_constants, e->hydro_properties,
-                                   e->internal_units, e->cosmology,
-                                   e->cooling_func, p, xp);
-}
-
-/**
- * @brief Specifies which particle fields to write to a dataset
- *
- * @param parts The particle array.
- * @param xparts The extended particle array.
- * @param list The list of i/o properties to write.
- * @param cooling The #cooling_function_data
- *
- * @return Returns the number of fields to write.
- */
-__attribute__((always_inline)) INLINE static int cooling_write_particles(
-    const struct part* parts, const struct xpart* xparts, struct io_props* list,
-    const struct cooling_function_data* cooling) {
-
-  list[0] = io_make_output_field_convert_part(
-      "Temperatures", FLOAT, 1, UNIT_CONV_TEMPERATURE, 0.f, parts, xparts,
-      convert_part_T, "Temperatures of the gas particles");
-
-  return 1;
-}
-
-#endif /* SWIFT_COOLING_IO_COMPTON_H */
diff --git a/src/cooling/Compton/cooling_properties.h b/src/cooling/Compton/cooling_properties.h
deleted file mode 100644
index 79b40db4ff8d53709c7c6d3f677ca98803419c1e..0000000000000000000000000000000000000000
--- a/src/cooling/Compton/cooling_properties.h
+++ /dev/null
@@ -1,47 +0,0 @@
-/*******************************************************************************
- * This file is part of SWIFT.
- * Copyright (c) 2018 Matthieu Schaller (matthieu.schaller@durham.ac.uk)
- *
- * 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_COOLING_PROPERTIES_COMPTON_H
-#define SWIFT_COOLING_PROPERTIES_COMPTON_H
-
-/**
- * @brief Properties of the cooling function.
- */
-struct cooling_function_data {
-
-  /*! Compton rate in cgs [g cm^2 s^-3 K^-1] */
-  double const_Compton_rate_cgs;
-
-  /*! Temperature of the CMB at redshift 0 in cgs [K] */
-  double const_T_CMB_0;
-
-  /*! Conversion factor from internal units to cgs for density */
-  double conv_factor_density_to_cgs;
-
-  /*! Conversion factor from internal units to cgs for internal energy */
-  double conv_factor_energy_to_cgs;
-
-  /*! Conversion factor from internal units from cgs for internal energy
-   * derivative */
-  double conv_factor_energy_rate_from_cgs;
-
-  /*! Inverse of the proton mass in cgs units [g^-1] */
-  double proton_mass_cgs_inv;
-};
-
-#endif /* SWIFT_COOLING_PROPERTIES_COMPTON_H */
diff --git a/src/cooling/Compton/cooling_struct.h b/src/cooling/Compton/cooling_struct.h
deleted file mode 100644
index 4cd75d745dbb8b3d9e0ad8aca050aafe115d1a8f..0000000000000000000000000000000000000000
--- a/src/cooling/Compton/cooling_struct.h
+++ /dev/null
@@ -1,36 +0,0 @@
-/*******************************************************************************
- * This file is part of SWIFT.
- * Copyright (c) 2018 Matthieu Schaller (matthieu.schaller@durham.ac.uk)
- *
- * 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_COOLING_STRUCT_COMPTON_H
-#define SWIFT_COOLING_STRUCT_COMPTON_H
-
-/**
- * @brief Properties of the cooling stored in the #part data.
- */
-struct cooling_part_data {};
-
-/**
- * @brief Properties of the cooling stored in the particle data.
- */
-struct cooling_xpart_data {
-
-  /*! Energy radiated away by this particle since the start of the run */
-  float radiated_energy;
-};
-
-#endif /* SWIFT_COOLING_STRUCT_COMPTON_H */
diff --git a/src/cooling/grackle/cooling.c b/src/cooling/grackle/cooling.c
index 7d4f88a013a0dd7172d126faa2f8e215675a977b..84b9fa4457f5cce15ea0cdb61f530d8548c43ae0 100644
--- a/src/cooling/grackle/cooling.c
+++ b/src/cooling/grackle/cooling.c
@@ -256,34 +256,6 @@ void cooling_first_init_part(const struct phys_const* restrict phys_const,
 #endif
 }
 
-/**
- * @brief Returns the subgrid temperature of a particle.
- *
- * This model has no subgrid quantity. We return an error.
- *
- * @param p The particle.
- * @param xp The extended particle data.
- */
-INLINE static float cooling_get_subgrid_temperature(const struct part* p,
-                                                    const struct xpart* xp) {
-  error("This cooling model does not use subgrid quantities!");
-  return -1.f;
-}
-
-/**
- * @brief Returns the subgrid density of a particle.
- *
- * This model has no subgrid quantity. We return an error.
- *
- * @param p The particle.
- * @param xp The extended particle data.
- */
-INLINE static float cooling_get_subgrid_density(const struct part* p,
-                                                const struct xpart* xp) {
-  error("This cooling model does not use subgrid quantities!");
-  return -1.f;
-}
-
 /**
  * @brief Returns the total radiated energy by this particle.
  *
diff --git a/src/cooling/grackle/cooling.h b/src/cooling/grackle/cooling.h
index 4c58ef8fb1b2cddcfd965cf7a6c90aed2cbbe82d..22dd1d642c6d581c821cbbb8ac5bb33feacce7fa 100644
--- a/src/cooling/grackle/cooling.h
+++ b/src/cooling/grackle/cooling.h
@@ -68,10 +68,33 @@ void cooling_first_init_part(const struct phys_const* restrict phys_const,
                              const struct part* restrict p,
                              struct xpart* restrict xp);
 
-float cooling_get_subgrid_temperature(const struct part* p,
-                                      const struct xpart* xp);
+/**
+ * @brief Returns the subgrid temperature of a particle.
+ *
+ * This model has no subgrid quantity. We return an error.
+ *
+ * @param p The particle.
+ * @param xp The extended particle data.
+ */
+INLINE static float cooling_get_subgrid_temperature(const struct part* p,
+                                                    const struct xpart* xp) {
+  error("This cooling model does not use subgrid quantities!");
+  return -1.f;
+}
 
-float cooling_get_subgrid_density(const struct part* p, const struct xpart* xp);
+/**
+ * @brief Returns the subgrid density of a particle.
+ *
+ * This model has no subgrid quantity. We return an error.
+ *
+ * @param p The particle.
+ * @param xp The extended particle data.
+ */
+INLINE static float cooling_get_subgrid_density(const struct part* p,
+                                                const struct xpart* xp) {
+  error("This cooling model does not use subgrid quantities!");
+  return -1.f;
+}
 
 float cooling_get_radiated_energy(const struct xpart* restrict xp);
 void cooling_print_backend(const struct cooling_function_data* cooling);
diff --git a/src/cooling/grackle/cooling_io.h b/src/cooling/grackle/cooling_io.h
index e75ad138ee2fd71a5b490a22651fd2266f867fcc..6a1db8b68ed13c5dbfef5dfeb00e76ecd3213ae7 100644
--- a/src/cooling/grackle/cooling_io.h
+++ b/src/cooling/grackle/cooling_io.h
@@ -23,6 +23,7 @@
 #include "cooling_properties.h"
 #include "cooling_struct.h"
 #include "io_properties.h"
+#include "physical_constants.h"
 
 #ifdef HAVE_HDF5
 
diff --git a/src/distributed_io.c b/src/distributed_io.c
index f5cf8b3d3ce64a2c7a085fca544b31cb129cdccf..3d5a550a62b9afef0fdc27b3ebc1d3c3dbab803d 100644
--- a/src/distributed_io.c
+++ b/src/distributed_io.c
@@ -434,6 +434,8 @@ void write_output_distributed(struct engine* e,
                      numParticlesHighWord, swift_type_count);
   double MassTable[swift_type_count] = {0};
   io_write_attribute(h_grp, "MassTable", DOUBLE, MassTable, swift_type_count);
+  io_write_attribute(h_grp, "InitialMassTable", DOUBLE,
+                     e->s->initial_mean_mass_particles, swift_type_count);
   unsigned int flagEntropy[swift_type_count] = {0};
   flagEntropy[0] = writeEntropyFlag();
   io_write_attribute(h_grp, "Flag_Entropy_ICs", UINT, flagEntropy,
diff --git a/src/drift.h b/src/drift.h
index e994e4fde7339d6eedbe376ccc136152564106ae..7af873a542587f8ebb0313f6abff0597851c7188 100644
--- a/src/drift.h
+++ b/src/drift.h
@@ -151,9 +151,9 @@ __attribute__((always_inline)) INLINE static void drift_part(
 
   /* Predict velocities (for gravity terms) */
   if (p->gpart != NULL) {
-    p->v[0] += (p->gpart->a_grav[0] + p->gpart->a_grav_mesh[0]) * dt_kick_grav;
-    p->v[1] += (p->gpart->a_grav[1] + p->gpart->a_grav_mesh[1]) * dt_kick_grav;
-    p->v[2] += (p->gpart->a_grav[2] + p->gpart->a_grav_mesh[2]) * dt_kick_grav;
+    p->v[0] += xp->a_grav[0] * dt_kick_grav;
+    p->v[1] += xp->a_grav[1] * dt_kick_grav;
+    p->v[2] += xp->a_grav[2] * dt_kick_grav;
   }
 
   /* Predict the values of the extra fields */
diff --git a/src/engine.c b/src/engine.c
index de4ca0fe86a72729d37574629599adde34d01410..1847d0e0080f258d79dd4c59f81eee97c40fff79 100644
--- a/src/engine.c
+++ b/src/engine.c
@@ -1771,6 +1771,9 @@ void engine_init_particles(struct engine *e, int flag_entropy_ICs,
     }
   }
 
+  /* Collect initial mean mass of each particle type */
+  space_collect_mean_masses(e->s, e->verbose);
+
 #ifdef SWIFT_DEBUG_CHECKS
   /* Check that we have the correct total mass in the top-level multipoles */
   long long num_gpart_mpole = 0;
@@ -2711,6 +2714,12 @@ void engine_init(struct engine *e, struct space *s, struct swift_params *params,
   parser_get_param_string(params, "Snapshots:basename", e->snapshot_base_name);
   parser_get_opt_param_string(params, "Snapshots:subdir", e->snapshot_subdir,
                               engine_default_snapshot_subdir);
+  e->snapshot_run_on_dump =
+      parser_get_opt_param_int(params, "Snapshots:run_on_dump", 0);
+  if (e->snapshot_run_on_dump) {
+    parser_get_param_string(params, "Snapshots:dump_command",
+                            e->snapshot_dump_command);
+  }
   e->snapshot_compression =
       parser_get_opt_param_int(params, "Snapshots:compression", 0);
   e->snapshot_distributed =
diff --git a/src/engine.h b/src/engine.h
index 7d8561fb3a8d0eff7355d0d355a9d33279f660df..06bdc3be5dde1054d1590617c2154b3ceb920c9f 100644
--- a/src/engine.h
+++ b/src/engine.h
@@ -311,6 +311,8 @@ struct engine {
 
   char snapshot_base_name[PARSER_MAX_LINE_SIZE];
   char snapshot_subdir[PARSER_MAX_LINE_SIZE];
+  char snapshot_dump_command[PARSER_MAX_LINE_SIZE];
+  int snapshot_run_on_dump;
   int snapshot_distributed;
   int snapshot_compression;
   int snapshot_int_time_label_on;
@@ -566,6 +568,7 @@ void engine_check_for_dumps(struct engine *e);
 void engine_check_for_index_dump(struct engine *e);
 void engine_collect_end_of_step(struct engine *e, int apply);
 void engine_dump_snapshot(struct engine *e);
+void engine_run_on_dump(struct engine *e);
 void engine_init_output_lists(struct engine *e, struct swift_params *params);
 void engine_init(struct engine *e, struct space *s, struct swift_params *params,
                  struct output_options *output_options, long long Ngas,
diff --git a/src/engine_io.c b/src/engine_io.c
index 12e4d2f8f9be75aed7e3579bb5d97b7dead1f3f4..9307b7acecf3f7e10eec2abdb3e9e9771e3f739b 100644
--- a/src/engine_io.c
+++ b/src/engine_io.c
@@ -39,6 +39,8 @@
 #include "serial_io.h"
 #include "single_io.h"
 
+#include <stdio.h>
+
 /**
  * @brief Check whether an index file has to be written during this
  * step.
@@ -214,6 +216,34 @@ void engine_dump_snapshot(struct engine *e) {
   if (e->verbose)
     message("writing particle properties took %.3f %s.",
             (float)clocks_diff(&time1, &time2), clocks_getunit());
+
+  /* Run the post-dump command if required */
+  engine_run_on_dump(e);
+}
+
+/**
+ * @brief Runs the snapshot_dump_command if relevant. Note that we
+ *        perform no error checking on this command, and assume
+ *        it works fine.
+ *
+ * @param e The #engine.
+ */
+void engine_run_on_dump(struct engine *e) {
+  if (e->snapshot_run_on_dump) {
+    /* Generate a string containing (optionally) the snapshot number.
+     * Note that -1 is used because snapshot_output_count was just
+     * increased when the write_output_* functions are called. */
+    const int buf_size = PARSER_MAX_LINE_SIZE * 3;
+    char dump_command_buf[buf_size];
+    snprintf(dump_command_buf, buf_size, "%s %s %04d", e->snapshot_dump_command,
+             e->snapshot_base_name, e->snapshot_output_count - 1);
+
+    /* Let's trust the user's command... */
+    const int result = system(dump_command_buf);
+    if (result != 0) {
+      message("Snapshot dump command returned error code %d", result);
+    }
+  }
 }
 
 /**
diff --git a/src/equation_of_state/planetary/sesame.h b/src/equation_of_state/planetary/sesame.h
index 2af28baebf333d77b5ccacafe1d8752624c8c93a..3480de806a1c665585bd6e603448175177e4be87 100644
--- a/src/equation_of_state/planetary/sesame.h
+++ b/src/equation_of_state/planetary/sesame.h
@@ -47,9 +47,9 @@ struct SESAME_params {
   float *table_log_u_rho_T;
   float *table_P_rho_T;
   float *table_c_rho_T;
-  float *table_s_rho_T;
+  float *table_log_s_rho_T;
   int date, num_rho, num_T;
-  float P_tiny, c_tiny;
+  float u_tiny, P_tiny, c_tiny, s_tiny;
   enum eos_planetary_material_id mat_id;
 };
 
@@ -138,7 +138,7 @@ INLINE static void load_table_SESAME(struct SESAME_params *mat,
       (float *)malloc(mat->num_rho * mat->num_T * sizeof(float));
   mat->table_c_rho_T =
       (float *)malloc(mat->num_rho * mat->num_T * sizeof(float));
-  mat->table_s_rho_T =
+  mat->table_log_s_rho_T =
       (float *)malloc(mat->num_rho * mat->num_T * sizeof(float));
 
   // Densities (not log yet)
@@ -159,7 +159,8 @@ INLINE static void load_table_SESAME(struct SESAME_params *mat,
     if (c != 1) error("Failed to read the SESAME EoS table %s", table_file);
   }
 
-  // Sp. int. energies (not log yet), pressures, sound speeds, and entropies
+  // Sp. int. energies (not log yet), pressures, sound speeds, and sp.
+  // entropies (not log yet)
   for (int i_T = -1; i_T < mat->num_T; i_T++) {
     for (int i_rho = -1; i_rho < mat->num_rho; i_rho++) {
       // Ignore the first elements of rho = 0, T = 0
@@ -171,7 +172,7 @@ INLINE static void load_table_SESAME(struct SESAME_params *mat,
                    &mat->table_log_u_rho_T[i_rho * mat->num_T + i_T],
                    &mat->table_P_rho_T[i_rho * mat->num_T + i_T],
                    &mat->table_c_rho_T[i_rho * mat->num_T + i_T],
-                   &mat->table_s_rho_T[i_rho * mat->num_T + i_T]);
+                   &mat->table_log_s_rho_T[i_rho * mat->num_T + i_T]);
         if (c != 4) error("Failed to read the SESAME EoS table %s", table_file);
       }
     }
@@ -188,22 +189,11 @@ INLINE static void prepare_table_SESAME(struct SESAME_params *mat) {
     mat->table_log_rho[i_rho] = logf(mat->table_log_rho[i_rho]);
   }
 
-  // Convert sp. int. energies to log(sp. int. energy)
-  for (int i_rho = 0; i_rho < mat->num_rho; i_rho++) {
-    for (int i_T = 0; i_T < mat->num_T; i_T++) {
-      // If not positive then set very small for the log
-      if (mat->table_log_u_rho_T[i_rho * mat->num_T + i_T] <= 0) {
-        mat->table_log_u_rho_T[i_rho * mat->num_T + i_T] = 1.f;
-      }
-
-      mat->table_log_u_rho_T[i_rho * mat->num_T + i_T] =
-          logf(mat->table_log_u_rho_T[i_rho * mat->num_T + i_T]);
-    }
-  }
-
-  // Initialise tiny pressure and soundspeed
+  // Initialise tiny values
+  mat->u_tiny = FLT_MAX;
   mat->P_tiny = FLT_MAX;
   mat->c_tiny = FLT_MAX;
+  mat->s_tiny = FLT_MAX;
 
   // Enforce that the 1D arrays of u (at each rho) are monotonic
   // This is necessary because, for some high-density u slices at very low T,
@@ -223,7 +213,11 @@ INLINE static void prepare_table_SESAME(struct SESAME_params *mat) {
         break;
       }
 
-      // Smallest positive pressure and sound speed
+      // Smallest positive values
+      if ((mat->table_log_u_rho_T[i_rho * mat->num_T + i_T] < mat->u_tiny) &&
+          (mat->table_log_u_rho_T[i_rho * mat->num_T + i_T] > 0)) {
+        mat->u_tiny = mat->table_log_u_rho_T[i_rho * mat->num_T + i_T];
+      }
       if ((mat->table_P_rho_T[i_rho * mat->num_T + i_T] < mat->P_tiny) &&
           (mat->table_P_rho_T[i_rho * mat->num_T + i_T] > 0)) {
         mat->P_tiny = mat->table_P_rho_T[i_rho * mat->num_T + i_T];
@@ -232,12 +226,38 @@ INLINE static void prepare_table_SESAME(struct SESAME_params *mat) {
           (mat->table_c_rho_T[i_rho * mat->num_T + i_T] > 0)) {
         mat->c_tiny = mat->table_c_rho_T[i_rho * mat->num_T + i_T];
       }
+      if ((mat->table_log_s_rho_T[i_rho * mat->num_T + i_T] < mat->s_tiny) &&
+          (mat->table_log_s_rho_T[i_rho * mat->num_T + i_T] > 0)) {
+        mat->s_tiny = mat->table_log_s_rho_T[i_rho * mat->num_T + i_T];
+      }
     }
   }
 
-  // Tiny pressure to allow interpolation near non-positive values
+  // Tiny values to allow interpolation near non-positive values
+  mat->u_tiny *= 1e-3f;
   mat->P_tiny *= 1e-3f;
   mat->c_tiny *= 1e-3f;
+  mat->s_tiny *= 1e-3f;
+
+  // Convert sp. int. energies to log(sp. int. energy), same for sp. entropies
+  for (int i_rho = 0; i_rho < mat->num_rho; i_rho++) {
+    for (int i_T = 0; i_T < mat->num_T; i_T++) {
+      // If not positive then set very small for the log
+      if (mat->table_log_u_rho_T[i_rho * mat->num_T + i_T] <= 0) {
+        mat->table_log_u_rho_T[i_rho * mat->num_T + i_T] = mat->u_tiny;
+      }
+
+      mat->table_log_u_rho_T[i_rho * mat->num_T + i_T] =
+          logf(mat->table_log_u_rho_T[i_rho * mat->num_T + i_T]);
+
+      if (mat->table_log_s_rho_T[i_rho * mat->num_T + i_T] <= 0) {
+        mat->table_log_s_rho_T[i_rho * mat->num_T + i_T] = mat->s_tiny;
+      }
+
+      mat->table_log_s_rho_T[i_rho * mat->num_T + i_T] =
+          logf(mat->table_log_s_rho_T[i_rho * mat->num_T + i_T]);
+    }
+  }
 }
 
 // Convert to internal units
@@ -255,7 +275,7 @@ INLINE static void convert_units_SESAME(struct SESAME_params *mat,
              units_cgs_conversion_factor(us, UNIT_CONV_DENSITY));
   }
 
-  // Sp. Int. Energies (log), pressures, and sound speeds
+  // Sp. int. energies (log), pressures, sound speeds, and sp. entropies
   for (int i_rho = 0; i_rho < mat->num_rho; i_rho++) {
     for (int i_T = 0; i_T < mat->num_T; i_T++) {
       mat->table_log_u_rho_T[i_rho * mat->num_T + i_T] += logf(
@@ -267,26 +287,118 @@ INLINE static void convert_units_SESAME(struct SESAME_params *mat,
       mat->table_c_rho_T[i_rho * mat->num_T + i_T] *=
           1e3f * units_cgs_conversion_factor(&si, UNIT_CONV_SPEED) /
           units_cgs_conversion_factor(us, UNIT_CONV_SPEED);
-      mat->table_s_rho_T[i_rho * mat->num_T + i_T] *=
-          units_cgs_conversion_factor(&si, UNIT_CONV_ENERGY_PER_UNIT_MASS) /
-          units_cgs_conversion_factor(us, UNIT_CONV_ENTROPY);
+      mat->table_log_s_rho_T[i_rho * mat->num_T + i_T] +=
+          logf(units_cgs_conversion_factor(
+                   &si, UNIT_CONV_PHYSICAL_ENTROPY_PER_UNIT_MASS) /
+               units_cgs_conversion_factor(
+                   us, UNIT_CONV_PHYSICAL_ENTROPY_PER_UNIT_MASS));
     }
   }
 
-  // Tiny pressure and sound speed
+  // Tiny values
+  mat->u_tiny *=
+      units_cgs_conversion_factor(&si, UNIT_CONV_ENERGY_PER_UNIT_MASS) /
+      units_cgs_conversion_factor(us, UNIT_CONV_ENERGY_PER_UNIT_MASS);
   mat->P_tiny *= units_cgs_conversion_factor(&si, UNIT_CONV_PRESSURE) /
                  units_cgs_conversion_factor(us, UNIT_CONV_PRESSURE);
   mat->c_tiny *= 1e3f * units_cgs_conversion_factor(&si, UNIT_CONV_SPEED) /
                  units_cgs_conversion_factor(us, UNIT_CONV_SPEED);
+  mat->s_tiny *=
+      units_cgs_conversion_factor(&si,
+                                  UNIT_CONV_PHYSICAL_ENTROPY_PER_UNIT_MASS) /
+      units_cgs_conversion_factor(us, UNIT_CONV_PHYSICAL_ENTROPY_PER_UNIT_MASS);
 }
 
 // gas_internal_energy_from_entropy
 INLINE static float SESAME_internal_energy_from_entropy(
     float density, float entropy, const struct SESAME_params *mat) {
 
-  error("This EOS function is not yet implemented!");
+  float u, log_u_1, log_u_2, log_u_3, log_u_4;
 
-  return 0.f;
+  if (entropy <= 0.f) {
+    return 0.f;
+  }
+
+  int idx_rho, idx_s_1, idx_s_2;
+  float intp_rho, intp_s_1, intp_s_2;
+  const float log_rho = logf(density);
+  const float log_s = logf(entropy);
+
+  // 2D interpolation (bilinear with log(rho), log(s)) to find u(rho, s))
+  // Density index
+  idx_rho =
+      find_value_in_monot_incr_array(log_rho, mat->table_log_rho, mat->num_rho);
+
+  // Sp. entropy at this and the next density (in relevant slice of s array)
+  idx_s_1 = find_value_in_monot_incr_array(
+      log_s, mat->table_log_s_rho_T + idx_rho * mat->num_T, mat->num_T);
+  idx_s_2 = find_value_in_monot_incr_array(
+      log_s, mat->table_log_s_rho_T + (idx_rho + 1) * mat->num_T, mat->num_T);
+
+  // If outside the table then extrapolate from the edge and edge-but-one values
+  if (idx_rho <= -1) {
+    idx_rho = 0;
+  } else if (idx_rho >= mat->num_rho) {
+    idx_rho = mat->num_rho - 2;
+  }
+  if (idx_s_1 <= -1) {
+    idx_s_1 = 0;
+  } else if (idx_s_1 >= mat->num_T) {
+    idx_s_1 = mat->num_T - 2;
+  }
+  if (idx_s_2 <= -1) {
+    idx_s_2 = 0;
+  } else if (idx_s_2 >= mat->num_T) {
+    idx_s_2 = mat->num_T - 2;
+  }
+
+  // Check for duplicates in SESAME tables before interpolation
+  if (mat->table_log_rho[idx_rho + 1] != mat->table_log_rho[idx_rho]) {
+    intp_rho = (log_rho - mat->table_log_rho[idx_rho]) /
+               (mat->table_log_rho[idx_rho + 1] - mat->table_log_rho[idx_rho]);
+  } else {
+    intp_rho = 1.f;
+  }
+  if (mat->table_log_s_rho_T[idx_rho * mat->num_T + (idx_s_1 + 1)] !=
+      mat->table_log_s_rho_T[idx_rho * mat->num_T + idx_s_1]) {
+    intp_s_1 =
+        (log_s - mat->table_log_s_rho_T[idx_rho * mat->num_T + idx_s_1]) /
+        (mat->table_log_s_rho_T[idx_rho * mat->num_T + (idx_s_1 + 1)] -
+         mat->table_log_s_rho_T[idx_rho * mat->num_T + idx_s_1]);
+  } else {
+    intp_s_1 = 1.f;
+  }
+  if (mat->table_log_s_rho_T[(idx_rho + 1) * mat->num_T + (idx_s_2 + 1)] !=
+      mat->table_log_s_rho_T[(idx_rho + 1) * mat->num_T + idx_s_2]) {
+    intp_s_2 =
+        (log_s - mat->table_log_s_rho_T[(idx_rho + 1) * mat->num_T + idx_s_2]) /
+        (mat->table_log_s_rho_T[(idx_rho + 1) * mat->num_T + (idx_s_2 + 1)] -
+         mat->table_log_s_rho_T[(idx_rho + 1) * mat->num_T + idx_s_2]);
+  } else {
+    intp_s_2 = 1.f;
+  }
+
+  // Table values
+  log_u_1 = mat->table_log_u_rho_T[idx_rho * mat->num_T + idx_s_1];
+  log_u_2 = mat->table_log_u_rho_T[idx_rho * mat->num_T + idx_s_1 + 1];
+  log_u_3 = mat->table_log_u_rho_T[(idx_rho + 1) * mat->num_T + idx_s_2];
+  log_u_4 = mat->table_log_u_rho_T[(idx_rho + 1) * mat->num_T + idx_s_2 + 1];
+
+  // If below the minimum s at this rho then just use the lowest table values
+  if ((idx_rho > 0.f) && ((intp_s_1 < 0.f) || (intp_s_2 < 0.f) ||
+                          (log_u_1 > log_u_2) || (log_u_3 > log_u_4))) {
+    intp_s_1 = 0;
+    intp_s_2 = 0;
+  }
+
+  // Interpolate with the log values
+  u = (1.f - intp_rho) * ((1.f - intp_s_1) * log_u_1 + intp_s_1 * log_u_2) +
+      intp_rho * ((1.f - intp_s_2) * log_u_3 + intp_s_2 * log_u_4);
+
+  // Convert back from log
+  u = expf(u);
+
+  return u;
 }
 
 // gas_pressure_from_entropy
@@ -338,7 +450,7 @@ INLINE static float SESAME_pressure_from_internal_energy(
   const float log_rho = logf(density);
   const float log_u = logf(u);
 
-  // 2D interpolation (bilinear with log(rho), log(u)) to find P(rho, u)
+  // 2D interpolation (bilinear with log(rho), log(u)) to find P(rho, u))
   // Density index
   idx_rho =
       find_value_in_monot_incr_array(log_rho, mat->table_log_rho, mat->num_rho);
@@ -463,7 +575,7 @@ INLINE static float SESAME_soundspeed_from_internal_energy(
   const float log_rho = logf(density);
   const float log_u = logf(u);
 
-  // 2D interpolation (bilinear with log(rho), log(u)) to find c(rho, u)
+  // 2D interpolation (bilinear with log(rho), log(u)) to find c(rho, u))
   // Density index
   idx_rho =
       find_value_in_monot_incr_array(log_rho, mat->table_log_rho, mat->num_rho);
diff --git a/src/feedback/EAGLE/feedback.h b/src/feedback/EAGLE/feedback.h
index 51294394e625f88a0cf41cafbbf83a1634b537fd..0a42bf94c7581f3934519238a3edb7fe91d69522 100644
--- a/src/feedback/EAGLE/feedback.h
+++ b/src/feedback/EAGLE/feedback.h
@@ -99,10 +99,13 @@ INLINE static double feedback_get_enrichment_timestep(
     const struct cosmology* cosmo, const double time, const double dt_star) {
 
   if (with_cosmology) {
-    return cosmology_get_delta_time_from_scale_factors(
-        cosmo, (double)sp->last_enrichment_time, cosmo->a);
+    if (cosmo->a > (double)sp->last_enrichment_time)
+      return cosmology_get_delta_time_from_scale_factors(
+          cosmo, (double)sp->last_enrichment_time, cosmo->a);
+    else
+      return 0.;
   } else {
-    return time - sp->last_enrichment_time;
+    return max(time - sp->last_enrichment_time, 0.);
   }
 }
 
@@ -170,10 +173,9 @@ __attribute__((always_inline)) INLINE static void feedback_prepare_spart(
     struct spart* sp, const struct feedback_props* feedback_props) {}
 
 /**
- * @brief Evolve the stellar properties of a #spart.
+ * @brief Prepare a #spart for the feedback task.
  *
- * This function allows for example to compute the SN rate before sending
- * this information to a different MPI rank.
+ * In EAGLE, this function evolves the stellar properties of a #spart.
  *
  * @param sp The particle to act upon
  * @param feedback_props The #feedback_props structure.
@@ -187,7 +189,7 @@ __attribute__((always_inline)) INLINE static void feedback_prepare_spart(
  * @param ti_begin The integer time at the beginning of the step.
  * @param with_cosmology Are we running with cosmology on?
  */
-__attribute__((always_inline)) INLINE static void feedback_evolve_spart(
+__attribute__((always_inline)) INLINE static void feedback_prepare_feedback(
     struct spart* restrict sp, const struct feedback_props* feedback_props,
     const struct cosmology* cosmo, const struct unit_system* us,
     const struct phys_const* phys_const, const double star_age_beg_step,
@@ -225,16 +227,27 @@ __attribute__((always_inline)) INLINE static void feedback_evolve_spart(
 /**
  * @brief Will this star particle want to do feedback during the next time-step?
  *
- * @param sp The star of interest.
- * @param feedback_props The properties of the feedback model.
- * @param with_cosmology Are we running a cosmological problem?
- * @param cosmo The cosmological model.
- * @param time The current time (since the start of the run / Big Bang).
+ * This is called in the time step task and increases counters of time-steps
+ * that have been performed.
+ *
+ * @param sp The particle to act upon
+ * @param feedback_props The #feedback_props structure.
+ * @param cosmo The current cosmological model.
+ * @param us The unit system.
+ * @param phys_const The #phys_const.
+ * @param star_age_beg_step The age of the star at the star of the time-step in
+ * internal units.
+ * @param dt The time-step size of this star in internal units.
+ * @param time The physical time in internal units.
+ * @param ti_begin The integer time at the beginning of the step.
+ * @param with_cosmology Are we running with cosmology on?
  */
-__attribute__((always_inline)) INLINE static int feedback_will_do_feedback(
-    struct spart* restrict sp, const struct feedback_props* feedback_props,
-    const int with_cosmology, const struct cosmology* cosmo,
-    const double time) {
+__attribute__((always_inline)) INLINE static void feedback_will_do_feedback(
+    struct spart* sp, const struct feedback_props* feedback_props,
+    const int with_cosmology, const struct cosmology* cosmo, const double time,
+    const struct unit_system* us, const struct phys_const* phys_const,
+    const double star_age_beg_step, const double dt,
+    const integertime_t ti_begin) {
 
   /* Special case for new-born stars */
   if (with_cosmology) {
@@ -242,28 +255,25 @@ __attribute__((always_inline)) INLINE static int feedback_will_do_feedback(
 
       /* Set the counter to "let's do enrichment" */
       sp->count_since_last_enrichment = 0;
-
-      /* Say we want to do feedback */
-      return 1;
     }
   } else {
     if (sp->birth_time == (float)time) {
 
       /* Set the counter to "let's do enrichment" */
       sp->count_since_last_enrichment = 0;
-
-      /* Say we want to do feedback */
-      return 1;
     }
   }
 
   /* Calculate age of the star at current time */
   double age_of_star;
   if (with_cosmology) {
-    age_of_star = cosmology_get_delta_time_from_scale_factors(
-        cosmo, (double)sp->birth_scale_factor, cosmo->a);
+    if (cosmo->a > (double)sp->birth_scale_factor)
+      age_of_star = cosmology_get_delta_time_from_scale_factors(
+          cosmo, (double)sp->birth_scale_factor, cosmo->a);
+    else
+      age_of_star = 0.;
   } else {
-    age_of_star = time - (double)sp->birth_time;
+    age_of_star = max(time - (double)sp->birth_time, 0.);
   }
 
   /* Is the star still young? */
@@ -272,9 +282,6 @@ __attribute__((always_inline)) INLINE static int feedback_will_do_feedback(
     /* Set the counter to "let's do enrichment" */
     sp->count_since_last_enrichment = 0;
 
-    /* Say we want to do feedback */
-    return 1;
-
   } else {
 
     /* Increment counter */
@@ -285,14 +292,6 @@ __attribute__((always_inline)) INLINE static int feedback_will_do_feedback(
 
       /* Reset counter */
       sp->count_since_last_enrichment = 0;
-
-      /* Say we want to do feedback */
-      return 1;
-
-    } else {
-
-      /* Say we don't want to do feedback */
-      return 0;
     }
   }
 }
diff --git a/src/feedback/GEAR/feedback.c b/src/feedback/GEAR/feedback.c
index c418465ab5097866094adc85216c5e500ad00b21..38662719854f7a507173ed4bef4e0f9620939e48 100644
--- a/src/feedback/GEAR/feedback.c
+++ b/src/feedback/GEAR/feedback.c
@@ -87,21 +87,69 @@ void feedback_update_part(struct part* restrict p, struct xpart* restrict xp,
 }
 
 /**
- * @brief Should we do feedback for this star?
+ * @brief Will this star particle want to do feedback during the next time-step?
  *
- * @param sp The star to consider.
- * @param feedback_props The #feedback_props.
- * @param with_cosmology Is the cosmology switch on?
- * @param cosmo The #cosmology.
- * @param time The current time.
+ * This is called in the time step task.
+ *
+ * In GEAR, we compute the full stellar evolution here.
+ *
+ * @param sp The particle to act upon
+ * @param feedback_props The #feedback_props structure.
+ * @param cosmo The current cosmological model.
+ * @param us The unit system.
+ * @param phys_const The #phys_const.
+ * @param star_age_beg_step The age of the star at the star of the time-step in
+ * internal units.
+ * @param dt The time-step size of this star in internal units.
+ * @param time The physical time in internal units.
+ * @param ti_begin The integer time at the beginning of the step.
+ * @param with_cosmology Are we running with cosmology on?
  */
-int feedback_will_do_feedback(const struct spart* sp,
-                              const struct feedback_props* feedback_props,
-                              const int with_cosmology,
-                              const struct cosmology* cosmo,
-                              const double time) {
+void feedback_will_do_feedback(struct spart* sp,
+                               const struct feedback_props* feedback_props,
+                               const int with_cosmology,
+                               const struct cosmology* cosmo, const double time,
+                               const struct unit_system* us,
+                               const struct phys_const* phys_const,
+                               const double star_age_beg_step, const double dt,
+                               const integertime_t ti_begin) {
+
+  /* Zero the energy of supernovae */
+  sp->feedback_data.energy_ejected = 0;
+  sp->feedback_data.will_do_feedback = 0;
+
+#ifdef SWIFT_DEBUG_CHECKS
+  if (sp->birth_time == -1.) error("Evolving a star particle that should not!");
+
+  if (star_age_beg_step < -1e-6) {
+    error("Negative age for a star");
+  }
+#endif
+  /* Has this star been around for a while ? */
+  if (star_age_beg_step + dt <= 0.) return;
 
-  return (sp->birth_time != -1.);
+  const double star_age_beg_step_safe =
+      star_age_beg_step < 0 ? 0 : star_age_beg_step;
+
+  /* Pick the correct table. (if only one table, threshold is < 0) */
+  const float metal =
+      chemistry_get_star_total_metal_mass_fraction_for_feedback(sp);
+  const float threshold = feedback_props->metallicity_max_first_stars;
+
+  const struct stellar_model* model =
+      metal < threshold ? &feedback_props->stellar_model_first_stars
+                        : &feedback_props->stellar_model;
+
+  /* Compute the stellar evolution */
+  stellar_evolution_evolve_spart(sp, model, cosmo, us, phys_const, ti_begin,
+                                 star_age_beg_step_safe, dt);
+
+  /* Transform the number of SN to the energy */
+  sp->feedback_data.energy_ejected =
+      sp->feedback_data.number_sn * feedback_props->energy_per_supernovae;
+
+  /* Set the particle as doing some feedback */
+  sp->feedback_data.will_do_feedback = sp->feedback_data.energy_ejected != 0.;
 }
 
 /**
@@ -116,14 +164,7 @@ int feedback_is_active(const struct spart* sp, const double time,
                        const struct cosmology* cosmo,
                        const int with_cosmology) {
 
-  // TODO improve this with estimates for SNII and SNIa
-  if (sp->birth_time == -1.) return 0;
-
-  if (with_cosmology) {
-    return ((double)cosmo->a) > sp->birth_scale_factor;
-  } else {
-    return time > sp->birth_time;
-  }
+  return sp->feedback_data.will_do_feedback;
 }
 
 /**
@@ -157,16 +198,31 @@ void feedback_init_spart(struct spart* sp) {
 }
 
 /**
- * @brief Prepares a star's feedback field before computing what
- * needs to be distributed.
+ * @brief Reset the feedback field when the spart is not
+ * in a correct state for feeedback_will_do_feedback.
+ *
+ * This function is called in the timestep task.
  */
-void feedback_reset_feedback(struct spart* sp,
-                             const struct feedback_props* feedback_props) {
+void feedback_init_after_star_formation(
+    struct spart* sp, const struct feedback_props* feedback_props) {
+  feedback_init_spart(sp);
 
   /* Zero the energy of supernovae */
   sp->feedback_data.energy_ejected = 0;
+
+  /* Activate the feedback loop for the first step */
+  sp->feedback_data.will_do_feedback = 1;
 }
 
+/**
+ * @brief Prepares a star's feedback field before computing what
+ * needs to be distributed.
+ *
+ * This is called in the stars ghost.
+ */
+void feedback_reset_feedback(struct spart* sp,
+                             const struct feedback_props* feedback_props) {}
+
 /**
  * @brief Initialises the s-particles feedback props for the first time
  *
@@ -181,7 +237,11 @@ void feedback_first_init_spart(struct spart* sp,
 
   feedback_init_spart(sp);
 
-  feedback_reset_feedback(sp, feedback_props);
+  /* Zero the energy of supernovae */
+  sp->feedback_data.energy_ejected = 0;
+
+  /* Activate the feedback loop for the first step */
+  sp->feedback_data.will_do_feedback = 1;
 }
 
 /**
@@ -197,10 +257,11 @@ void feedback_prepare_spart(struct spart* sp,
                             const struct feedback_props* feedback_props) {}
 
 /**
- * @brief Evolve the stellar properties of a #spart.
+ * @brief Prepare a #spart for the feedback task.
  *
- * This function compute the SN rate and yields before sending
- * this information to a different MPI rank.
+ * This is called in the stars ghost task.
+ *
+ * In here, we only need to add the missing coefficients.
  *
  * @param sp The particle to act upon
  * @param feedback_props The #feedback_props structure.
@@ -214,50 +275,18 @@ void feedback_prepare_spart(struct spart* sp,
  * @param ti_begin The integer time at the beginning of the step.
  * @param with_cosmology Are we running with cosmology on?
  */
-void feedback_evolve_spart(struct spart* restrict sp,
-                           const struct feedback_props* feedback_props,
-                           const struct cosmology* cosmo,
-                           const struct unit_system* us,
-                           const struct phys_const* phys_const,
-                           const double star_age_beg_step, const double dt,
-                           const double time, const integertime_t ti_begin,
-                           const int with_cosmology) {
-
-#ifdef SWIFT_DEBUG_CHECKS
-  if (sp->birth_time == -1.) error("Evolving a star particle that should not!");
-
-  if (star_age_beg_step < -1e-6) {
-    error("Negative age for a star");
-  }
-#endif
-  const double star_age_beg_step_safe =
-      star_age_beg_step < 0 ? 0 : star_age_beg_step;
-
-  /* Reset the feedback */
-  feedback_reset_feedback(sp, feedback_props);
-
+void feedback_prepare_feedback(struct spart* restrict sp,
+                               const struct feedback_props* feedback_props,
+                               const struct cosmology* cosmo,
+                               const struct unit_system* us,
+                               const struct phys_const* phys_const,
+                               const double star_age_beg_step, const double dt,
+                               const double time, const integertime_t ti_begin,
+                               const int with_cosmology) {
   /* Add missing h factor */
   const float hi_inv = 1.f / sp->h;
   const float hi_inv_dim = pow_dimension(hi_inv); /* 1/h^d */
-
   sp->feedback_data.enrichment_weight *= hi_inv_dim;
-
-  /* Pick the correct table. (if only one table, threshold is < 0) */
-  const float metal =
-      chemistry_get_star_total_metal_mass_fraction_for_feedback(sp);
-  const float threshold = feedback_props->metallicity_max_first_stars;
-
-  const struct stellar_model* model =
-      metal < threshold ? &feedback_props->stellar_model_first_stars
-                        : &feedback_props->stellar_model;
-
-  /* Compute the stellar evolution */
-  stellar_evolution_evolve_spart(sp, model, cosmo, us, phys_const, ti_begin,
-                                 star_age_beg_step_safe, dt);
-
-  /* Transform the number of SN to the energy */
-  sp->feedback_data.energy_ejected =
-      sp->feedback_data.number_sn * feedback_props->energy_per_supernovae;
 }
 
 /**
diff --git a/src/feedback/GEAR/feedback.h b/src/feedback/GEAR/feedback.h
index c4b26b29e746fc192ec82e02fb077faa09ab8611..bd0306b158e41ddde1b4d825cc5a30fa81154205 100644
--- a/src/feedback/GEAR/feedback.h
+++ b/src/feedback/GEAR/feedback.h
@@ -32,10 +32,14 @@
 void feedback_update_part(struct part* restrict p, struct xpart* restrict xp,
                           const struct engine* restrict e);
 
-int feedback_will_do_feedback(const struct spart* sp,
-                              const struct feedback_props* feedback_props,
-                              const int with_cosmology,
-                              const struct cosmology* cosmo, const double time);
+void feedback_will_do_feedback(struct spart* sp,
+                               const struct feedback_props* feedback_props,
+                               const int with_cosmology,
+                               const struct cosmology* cosmo, const double time,
+                               const struct unit_system* us,
+                               const struct phys_const* phys_const,
+                               const double star_age_beg_step, const double dt,
+                               const integertime_t ti_begin);
 
 int feedback_is_active(const struct spart* sp, const double time,
                        const struct cosmology* cosmo, const int with_cosmology);
@@ -46,20 +50,22 @@ double feedback_get_enrichment_timestep(const struct spart* sp,
                                         const double dt_star);
 void feedback_init_spart(struct spart* sp);
 
+void feedback_init_after_star_formation(
+    struct spart* sp, const struct feedback_props* feedback_props);
 void feedback_reset_feedback(struct spart* sp,
                              const struct feedback_props* feedback_props);
 void feedback_first_init_spart(struct spart* sp,
                                const struct feedback_props* feedback_props);
 void feedback_prepare_spart(struct spart* sp,
                             const struct feedback_props* feedback_props);
-void feedback_evolve_spart(struct spart* restrict sp,
-                           const struct feedback_props* feedback_props,
-                           const struct cosmology* cosmo,
-                           const struct unit_system* us,
-                           const struct phys_const* phys_const,
-                           const double star_age_beg_step, const double dt,
-                           const double time, const integertime_t ti_begin,
-                           const int with_cosmology);
+void feedback_prepare_feedback(struct spart* restrict sp,
+                               const struct feedback_props* feedback_props,
+                               const struct cosmology* cosmo,
+                               const struct unit_system* us,
+                               const struct phys_const* phys_const,
+                               const double star_age_beg_step, const double dt,
+                               const double time, const integertime_t ti_begin,
+                               const int with_cosmology);
 void feedback_struct_dump(const struct feedback_props* feedback, FILE* stream);
 void feedback_struct_restore(struct feedback_props* feedback, FILE* stream);
 void feedback_clean(struct feedback_props* feedback);
diff --git a/src/feedback/GEAR/feedback_struct.h b/src/feedback/GEAR/feedback_struct.h
index bc000bea0dc0f45be9cec35b928742f9f4ae64f1..f047462555c04a4ee1a0e4bd3a45284e22cf61fa 100644
--- a/src/feedback/GEAR/feedback_struct.h
+++ b/src/feedback/GEAR/feedback_struct.h
@@ -58,6 +58,9 @@ struct feedback_spart_data {
 
   /*! Chemical composition of the mass ejected */
   double metal_mass_ejected[GEAR_CHEMISTRY_ELEMENT_COUNT];
+
+  /*! Does the particle needs the feedback loop? */
+  char will_do_feedback;
 };
 
 #endif /* SWIFT_FEEDBACK_STRUCT_GEAR_H */
diff --git a/src/feedback/none/feedback.h b/src/feedback/none/feedback.h
index e5d42f456b4b1a9246f5e24c22aeb3136c9fc799..1a0e869417fe5fc2c1d691af7dd9d811310bf95b 100644
--- a/src/feedback/none/feedback.h
+++ b/src/feedback/none/feedback.h
@@ -104,6 +104,8 @@ INLINE static double feedback_get_enrichment_timestep(
 /**
  * @brief Prepares a star's feedback field before computing what
  * needs to be distributed.
+ *
+ * This is called in the stars ghost.
  */
 __attribute__((always_inline)) INLINE static void feedback_reset_feedback(
     struct spart* sp, const struct feedback_props* feedback_props) {}
@@ -133,8 +135,9 @@ __attribute__((always_inline)) INLINE static void feedback_prepare_spart(
     struct spart* sp, const struct feedback_props* feedback_props) {}
 
 /**
- * @brief Evolve the stellar properties of a #spart.
+ * @brief Prepare a #spart for the feedback task.
  *
+ * This is called in the stars ghost task.
  * This function allows for example to compute the SN rate before sending
  * this information to a different MPI rank.
  *
@@ -150,7 +153,7 @@ __attribute__((always_inline)) INLINE static void feedback_prepare_spart(
  * @param ti_begin The integer time at the beginning of the step.
  * @param with_cosmology Are we running with cosmology on?
  */
-__attribute__((always_inline)) INLINE static void feedback_evolve_spart(
+__attribute__((always_inline)) INLINE static void feedback_prepare_feedback(
     struct spart* restrict sp, const struct feedback_props* feedback_props,
     const struct cosmology* cosmo, const struct unit_system* us,
     const struct phys_const* phys_const, const double star_age_beg_step,
@@ -160,18 +163,26 @@ __attribute__((always_inline)) INLINE static void feedback_evolve_spart(
 /**
  * @brief Will this star particle want to do feedback during the next time-step?
  *
- * @param sp The star of interest.
- * @param feedback_props The properties of the feedback model.
- * @param with_cosmology Are we running with cosmological time integration?
- * @param cosmo The #cosmology object.
- * @param time The current time (since the Big Bang).
+ * This is called in the time step task.
+ *
+ * @param sp The particle to act upon
+ * @param feedback_props The #feedback_props structure.
+ * @param cosmo The current cosmological model.
+ * @param us The unit system.
+ * @param phys_const The #phys_const.
+ * @param star_age_beg_step The age of the star at the star of the time-step in
+ * internal units.
+ * @param dt The time-step size of this star in internal units.
+ * @param time The physical time in internal units.
+ * @param ti_begin The integer time at the beginning of the step.
+ * @param with_cosmology Are we running with cosmology on?
  */
-__attribute__((always_inline)) INLINE static int feedback_will_do_feedback(
-    struct spart* restrict sp, const struct feedback_props* feedback_props,
-    const int with_cosmology, const struct cosmology* cosmo,
-    const double time) {
-  return 1;
-}
+__attribute__((always_inline)) INLINE static void feedback_will_do_feedback(
+    const struct spart* sp, const struct feedback_props* feedback_props,
+    const int with_cosmology, const struct cosmology* cosmo, const double time,
+    const struct unit_system* us, const struct phys_const* phys_const,
+    const double star_age_beg_step, const double dt,
+    const integertime_t ti_begin) {}
 
 /**
  * @brief Clean-up the memory allocated for the feedback routines
diff --git a/src/hydro/AnarchyPU/hydro_part.h b/src/hydro/AnarchyPU/hydro_part.h
index d55307bf9f73761c6e54eea096d6efbcd4fd4d56..185416f317e2a46b3f800ffbc121430652d7f930 100644
--- a/src/hydro/AnarchyPU/hydro_part.h
+++ b/src/hydro/AnarchyPU/hydro_part.h
@@ -55,6 +55,9 @@ struct xpart {
   /*! Velocity at the last full step. */
   float v_full[3];
 
+  /*! Gravitational acceleration at the end of the last step */
+  float a_grav[3];
+
   /*! Internal energy at the last full step. */
   float u_full;
 
diff --git a/src/hydro/Gadget2/hydro_part.h b/src/hydro/Gadget2/hydro_part.h
index bb614c897588a7f7ae4a00041d070c3a0c2460f1..fc34fa80260ced93741f50449c776fb3a482ccb2 100644
--- a/src/hydro/Gadget2/hydro_part.h
+++ b/src/hydro/Gadget2/hydro_part.h
@@ -54,6 +54,9 @@ struct xpart {
   /* Velocity at the last full step. */
   float v_full[3];
 
+  /*! Gravitational acceleration at the end of the last step */
+  float a_grav[3];
+
   /* Entropy at the last full step. */
   float entropy_full;
 
diff --git a/src/hydro/Gizmo/hydro_part.h b/src/hydro/Gizmo/hydro_part.h
index 3244aa1b6cc288a613db1df4a0176939b5c51dc8..17b0179916db7a65dcbba7800723fab1d049fd7d 100644
--- a/src/hydro/Gizmo/hydro_part.h
+++ b/src/hydro/Gizmo/hydro_part.h
@@ -39,6 +39,9 @@ struct xpart {
   /* Velocity at the last full step. */
   float v_full[3];
 
+  /*! Gravitational acceleration at the end of the last step */
+  float a_grav[3];
+
   /* Additional data used to record cooling information */
   struct cooling_xpart_data cooling_data;
 
diff --git a/src/hydro/Minimal/hydro.h b/src/hydro/Minimal/hydro.h
index 0ee99bf33525317128df70869ae1c5b32d87bc3b..8dab9acba8e1ef804b11e0ef9bb56c61eb9c1928 100644
--- a/src/hydro/Minimal/hydro.h
+++ b/src/hydro/Minimal/hydro.h
@@ -25,7 +25,7 @@
  * equations)
  *
  * The thermal variable is the internal energy (u). Simple constant
- * viscosity term without switches is implemented. No thermal conduction
+ * viscosity term with the Balsara (1995) switch. No thermal conduction
  * term is implemented.
  *
  * This corresponds to equations (43), (44), (45), (101), (103)  and (104) with
diff --git a/src/hydro/Minimal/hydro_debug.h b/src/hydro/Minimal/hydro_debug.h
index 514fcad2cb1ee646d0f69dabf9c79fbaf1e8427a..f66578cf85cf0c4602005905680e82fb618f2e29 100644
--- a/src/hydro/Minimal/hydro_debug.h
+++ b/src/hydro/Minimal/hydro_debug.h
@@ -24,7 +24,7 @@
  * @brief Minimal conservative implementation of SPH (Debugging routines)
  *
  * The thermal variable is the internal energy (u). Simple constant
- * viscosity term without switches is implemented. No thermal conduction
+ * viscosity term with the Balsara (1995) switch. No thermal conduction
  * term is implemented.
  *
  * This corresponds to equations (43), (44), (45), (101), (103)  and (104) with
diff --git a/src/hydro/Minimal/hydro_iact.h b/src/hydro/Minimal/hydro_iact.h
index 0c7e136d3042980f8e34698d81a3416805bb3d92..768b0358ab3457813fe503f77d69bb7828381af4 100644
--- a/src/hydro/Minimal/hydro_iact.h
+++ b/src/hydro/Minimal/hydro_iact.h
@@ -24,7 +24,7 @@
  * @brief Minimal conservative implementation of SPH (Neighbour loop equations)
  *
  * The thermal variable is the internal energy (u). Simple constant
- * viscosity term without switches is implemented. No thermal conduction
+ * viscosity term with the Balsara (1995) switch. No thermal conduction
  * term is implemented.
  *
  * This corresponds to equations (43), (44), (45), (101), (103)  and (104) with
diff --git a/src/hydro/Minimal/hydro_io.h b/src/hydro/Minimal/hydro_io.h
index dbd4d6381fadf79544d0f3536cd322d7269bc263..f5d6d51ac5d57c0789e4c140494d492e8a17bdde 100644
--- a/src/hydro/Minimal/hydro_io.h
+++ b/src/hydro/Minimal/hydro_io.h
@@ -24,7 +24,7 @@
  * @brief Minimal conservative implementation of SPH (i/o routines)
  *
  * The thermal variable is the internal energy (u). Simple constant
- * viscosity term without switches is implemented. No thermal conduction
+ * viscosity term with the Balsara (1995) switch. No thermal conduction
  * term is implemented.
  *
  * This corresponds to equations (43), (44), (45), (101), (103)  and (104) with
@@ -225,8 +225,9 @@ INLINE static void hydro_write_flavour(hid_t h_grpsph) {
   /* Viscosity and thermal conduction */
   /* Nothing in this minimal model... */
   io_write_attribute_s(h_grpsph, "Thermal Conductivity Model", "No treatment");
-  io_write_attribute_s(h_grpsph, "Viscosity Model",
-                       "Minimal treatment as in Monaghan (1992)");
+  io_write_attribute_s(
+      h_grpsph, "Viscosity Model",
+      "as in Springel (2005), i.e. Monaghan (1992) with Balsara (1995) switch");
 }
 
 /**
diff --git a/src/hydro/Minimal/hydro_part.h b/src/hydro/Minimal/hydro_part.h
index 158c42415e320d4f7893f8b26dd7c0296667eba0..94aaf51282b129ee6f3a29d4dca6a4e6983b397a 100644
--- a/src/hydro/Minimal/hydro_part.h
+++ b/src/hydro/Minimal/hydro_part.h
@@ -24,7 +24,7 @@
  * @brief Minimal conservative implementation of SPH (Particle definition)
  *
  * The thermal variable is the internal energy (u). Simple constant
- * viscosity term without switches is implemented. No thermal conduction
+ * viscosity term with the Balsara (1995) switch. No thermal conduction
  * term is implemented.
  *
  * This corresponds to equations (43), (44), (45), (101), (103)  and (104) with
@@ -58,6 +58,9 @@ struct xpart {
   /*! Velocity at the last full step. */
   float v_full[3];
 
+  /*! Gravitational acceleration at the end of the last step */
+  float a_grav[3];
+
   /*! Internal energy at the last full step. */
   float u_full;
 
diff --git a/src/hydro/Phantom/hydro_part.h b/src/hydro/Phantom/hydro_part.h
index 170c0d9e434f65e5baaf68c31ccc16797de2f40f..73a9c4bde74c10604bbdc9664c9ff391461dd4d3 100644
--- a/src/hydro/Phantom/hydro_part.h
+++ b/src/hydro/Phantom/hydro_part.h
@@ -56,6 +56,9 @@ struct xpart {
   /*! Velocity at the last full step. */
   float v_full[3];
 
+  /*! Gravitational acceleration at the end of the last step */
+  float a_grav[3];
+
   /*! Internal energy at the last full step. */
   float u_full;
 
diff --git a/src/hydro/Planetary/hydro.h b/src/hydro/Planetary/hydro.h
index 6b0357080527bf54065212a15516a4515df5d459..85c02aafc27a0656c05acda78dfcb75757e17b1b 100644
--- a/src/hydro/Planetary/hydro.h
+++ b/src/hydro/Planetary/hydro.h
@@ -392,10 +392,15 @@ hydro_set_drifted_physical_internal_energy(struct part *p,
   /* Now recompute the extra quantities */
 
   /* Compute the sound speed */
+  const float pressure =
+      gas_pressure_from_internal_energy(p->rho, p->u, p->mat_id);
   const float soundspeed = hydro_get_comoving_soundspeed(p);
 
   /* Update variables. */
+  p->force.pressure = pressure;
   p->force.soundspeed = soundspeed;
+
+  p->force.v_sig = max(p->force.v_sig, 2.f * soundspeed);
 }
 
 /**
@@ -472,6 +477,10 @@ __attribute__((always_inline)) INLINE static void hydro_init_part(
   p->density.wcount_dh = 0.f;
   p->rho = 0.f;
   p->density.rho_dh = 0.f;
+  p->density.div_v = 0.f;
+  p->density.rot_v[0] = 0.f;
+  p->density.rot_v[1] = 0.f;
+  p->density.rot_v[2] = 0.f;
 }
 
 /**
@@ -507,6 +516,17 @@ __attribute__((always_inline)) INLINE static void hydro_end_density(
   p->density.rho_dh *= h_inv_dim_plus_one;
   p->density.wcount *= h_inv_dim;
   p->density.wcount_dh *= h_inv_dim_plus_one;
+
+  const float rho_inv = 1.f / p->rho;
+  const float a_inv2 = cosmo->a2_inv;
+
+  /* Finish calculation of the (physical) velocity curl components */
+  p->density.rot_v[0] *= h_inv_dim_plus_one * a_inv2 * rho_inv;
+  p->density.rot_v[1] *= h_inv_dim_plus_one * a_inv2 * rho_inv;
+  p->density.rot_v[2] *= h_inv_dim_plus_one * a_inv2 * rho_inv;
+
+  /* Finish calculation of the (physical) velocity divergence */
+  p->density.div_v *= h_inv_dim_plus_one * a_inv2 * rho_inv;
 }
 
 /**
@@ -534,6 +554,10 @@ __attribute__((always_inline)) INLINE static void hydro_part_has_no_neighbours(
   p->density.wcount = kernel_root * h_inv_dim;
   p->density.rho_dh = 0.f;
   p->density.wcount_dh = 0.f;
+  p->density.div_v = 0.f;
+  p->density.rot_v[0] = 0.f;
+  p->density.rot_v[1] = 0.f;
+  p->density.rot_v[2] = 0.f;
 }
 
 /**
@@ -558,15 +582,22 @@ __attribute__((always_inline)) INLINE static void hydro_prepare_force(
     const struct cosmology *cosmo, const struct hydro_props *hydro_props,
     const float dt_alpha) {
 
-  const float fac_mu = cosmo->a_factor_mu;
+  const float fac_Balsara_eps = cosmo->a_factor_Balsara_eps;
 
   /* Compute the norm of the curl */
   const float curl_v = sqrtf(p->density.rot_v[0] * p->density.rot_v[0] +
                              p->density.rot_v[1] * p->density.rot_v[1] +
                              p->density.rot_v[2] * p->density.rot_v[2]);
 
-  /* Compute the norm of div v */
-  const float abs_div_v = fabsf(p->density.div_v);
+  /* Compute the norm of div v including the Hubble flow term */
+  const float div_physical_v = p->density.div_v + hydro_dimension * cosmo->H;
+  const float abs_div_physical_v = fabsf(div_physical_v);
+
+#ifdef PLANETARY_FIXED_ENTROPY
+  /* Override the internal energy to satisfy the fixed entropy */
+  p->u = gas_internal_energy_from_entropy(p->rho, p->s_fixed, p->mat_id);
+  xp->u_full = p->u;
+#endif
 
   /* Compute the pressure */
   const float pressure =
@@ -576,23 +607,30 @@ __attribute__((always_inline)) INLINE static void hydro_prepare_force(
   const float soundspeed =
       gas_soundspeed_from_internal_energy(p->rho, p->u, p->mat_id);
 
-  /* Compute the "grad h" term */
-  const float rho_inv = 1.f / p->rho;
-  float rho_dh = p->density.rho_dh;
+  /* Compute the "grad h" term  - Note here that we have \tilde{x}
+   * as 1 as we use the local number density to find neighbours. This
+   * introduces a j-component that is considered in the force loop,
+   * meaning that this cached grad_h_term gives:
+   *
+   * f_ij = 1.f - grad_h_term_i / m_j */
+  const float common_factor = p->h * hydro_dimension_inv / p->density.wcount;
+  float grad_h_term = common_factor * p->density.rho_dh /
+                      (1.f + common_factor * p->density.wcount_dh);
+
   /* Ignore changing-kernel effects when h ~= h_max */
   if (p->h > 0.9999f * hydro_props->h_max) {
-    rho_dh = 0.f;
+    grad_h_term = 0.f;
   }
-  const float grad_h_term =
-      1.f / (1.f + hydro_dimension_inv * p->h * rho_dh * rho_inv);
 
   /* Compute the Balsara switch */
 #ifdef PLANETARY_SPH_NO_BALSARA
   const float balsara = hydro_props->viscosity.alpha;
 #else
-  const float balsara =
-      hydro_props->viscosity.alpha * abs_div_v /
-      (abs_div_v + curl_v + 0.0001f * fac_mu * soundspeed / p->h);
+  /* Pre-multiply in the AV factor; hydro_props are not passed to the iact
+   * functions */
+  const float balsara = hydro_props->viscosity.alpha * abs_div_physical_v /
+                        (abs_div_physical_v + curl_v +
+                         0.0001f * fac_Balsara_eps * soundspeed / p->h);
 #endif
 
   /* Update variables. */
@@ -641,12 +679,12 @@ __attribute__((always_inline)) INLINE static void hydro_reset_predicted_values(
   p->v[1] = xp->v_full[1];
   p->v[2] = xp->v_full[2];
 
-  /* Re-set the entropy */
+  /* Re-set the internal energy */
   p->u = xp->u_full;
 
   /* Compute the pressure */
   const float pressure =
-      gas_pressure_from_internal_energy(p->rho, xp->u_full, p->mat_id);
+      gas_pressure_from_internal_energy(p->rho, p->u, p->mat_id);
 
   /* Compute the sound speed */
   const float soundspeed =
@@ -654,6 +692,8 @@ __attribute__((always_inline)) INLINE static void hydro_reset_predicted_values(
 
   p->force.pressure = pressure;
   p->force.soundspeed = soundspeed;
+
+  p->force.v_sig = max(p->force.v_sig, 2.f * soundspeed);
 }
 
 /**
@@ -683,7 +723,8 @@ __attribute__((always_inline)) INLINE static void hydro_predict_extra(
   p->u += p->u_dt * dt_therm;
 
   /* Check against absolute minimum */
-  const float min_u = hydro_props->minimal_internal_energy;
+  const float min_u =
+      hydro_props->minimal_internal_energy / cosmo->a_factor_internal_energy;
 
   p->u = max(p->u, min_u);
 
@@ -713,6 +754,8 @@ __attribute__((always_inline)) INLINE static void hydro_predict_extra(
 
   p->force.pressure = pressure;
   p->force.soundspeed = soundspeed;
+
+  p->force.v_sig = max(p->force.v_sig, 2.f * soundspeed);
 }
 
 /**
@@ -759,7 +802,8 @@ __attribute__((always_inline)) INLINE static void hydro_kick_extra(
   xp->u_full = max(xp->u_full + delta_u, 0.5f * xp->u_full);
 
   /* Check against absolute minimum */
-  const float min_u = hydro_props->minimal_internal_energy;
+  const float min_u =
+      hydro_props->minimal_internal_energy / cosmo->a_factor_internal_energy;
 
   if (xp->u_full < min_u) {
     xp->u_full = min_u;
diff --git a/src/hydro/Planetary/hydro_iact.h b/src/hydro/Planetary/hydro_iact.h
index 426a1672e66e94067be153682534f03c31bfd8ee..b9a22c8382da6c32608cdf4194ceb015fc0bc7e6 100644
--- a/src/hydro/Planetary/hydro_iact.h
+++ b/src/hydro/Planetary/hydro_iact.h
@@ -56,6 +56,13 @@ __attribute__((always_inline)) INLINE static void runner_iact_density(
 
   float wi, wj, wi_dx, wj_dx;
 
+#ifdef SWIFT_DEBUG_CHECKS
+  if (pi->time_bin >= time_bin_inhibited)
+    error("Inhibited pi in interaction function!");
+  if (pj->time_bin >= time_bin_inhibited)
+    error("Inhibited pj in interaction function!");
+#endif
+
   /* Get r. */
   const float r_inv = 1.0f / sqrtf(r2);
   const float r = r2 * r_inv;
@@ -83,6 +90,33 @@ __attribute__((always_inline)) INLINE static void runner_iact_density(
   pj->density.rho_dh -= mi * (hydro_dimension * wj + uj * wj_dx);
   pj->density.wcount += wj;
   pj->density.wcount_dh -= (hydro_dimension * wj + uj * wj_dx);
+
+  /* Compute dv dot r */
+  float dv[3], curlvr[3];
+
+  const float faci = mj * wi_dx * r_inv;
+  const float facj = mi * wj_dx * r_inv;
+
+  dv[0] = pi->v[0] - pj->v[0];
+  dv[1] = pi->v[1] - pj->v[1];
+  dv[2] = pi->v[2] - pj->v[2];
+  const float dvdr = dv[0] * dx[0] + dv[1] * dx[1] + dv[2] * dx[2];
+
+  pi->density.div_v -= faci * dvdr;
+  pj->density.div_v -= facj * dvdr;
+
+  /* Compute dv cross r */
+  curlvr[0] = dv[1] * dx[2] - dv[2] * dx[1];
+  curlvr[1] = dv[2] * dx[0] - dv[0] * dx[2];
+  curlvr[2] = dv[0] * dx[1] - dv[1] * dx[0];
+
+  pi->density.rot_v[0] += faci * curlvr[0];
+  pi->density.rot_v[1] += faci * curlvr[1];
+  pi->density.rot_v[2] += faci * curlvr[2];
+
+  pj->density.rot_v[0] += facj * curlvr[0];
+  pj->density.rot_v[1] += facj * curlvr[1];
+  pj->density.rot_v[2] += facj * curlvr[2];
 }
 
 /**
@@ -103,6 +137,13 @@ __attribute__((always_inline)) INLINE static void runner_iact_nonsym_density(
 
   float wi, wi_dx;
 
+#ifdef SWIFT_DEBUG_CHECKS
+  if (pi->time_bin >= time_bin_inhibited)
+    error("Inhibited pi in interaction function!");
+  if (pj->time_bin >= time_bin_inhibited)
+    error("Inhibited pj in interaction function!");
+#endif
+
   /* Get the masses. */
   const float mj = pj->mass;
 
@@ -118,6 +159,27 @@ __attribute__((always_inline)) INLINE static void runner_iact_nonsym_density(
   pi->density.rho_dh -= mj * (hydro_dimension * wi + ui * wi_dx);
   pi->density.wcount += wi;
   pi->density.wcount_dh -= (hydro_dimension * wi + ui * wi_dx);
+
+  /* Compute dv dot r */
+  float dv[3], curlvr[3];
+
+  const float faci = mj * wi_dx * r_inv;
+
+  dv[0] = pi->v[0] - pj->v[0];
+  dv[1] = pi->v[1] - pj->v[1];
+  dv[2] = pi->v[2] - pj->v[2];
+  const float dvdr = dv[0] * dx[0] + dv[1] * dx[1] + dv[2] * dx[2];
+
+  pi->density.div_v -= faci * dvdr;
+
+  /* Compute dv cross r */
+  curlvr[0] = dv[1] * dx[2] - dv[2] * dx[1];
+  curlvr[1] = dv[2] * dx[0] - dv[0] * dx[2];
+  curlvr[2] = dv[0] * dx[1] - dv[1] * dx[0];
+
+  pi->density.rot_v[0] += faci * curlvr[0];
+  pi->density.rot_v[1] += faci * curlvr[1];
+  pi->density.rot_v[2] += faci * curlvr[2];
 }
 
 /**
@@ -136,6 +198,13 @@ __attribute__((always_inline)) INLINE static void runner_iact_force(
     float r2, const float *dx, float hi, float hj, struct part *restrict pi,
     struct part *restrict pj, float a, float H) {
 
+#ifdef SWIFT_DEBUG_CHECKS
+  if (pi->time_bin >= time_bin_inhibited)
+    error("Inhibited pi in interaction function!");
+  if (pj->time_bin >= time_bin_inhibited)
+    error("Inhibited pj in interaction function!");
+#endif
+
   /* Cosmological factors entering the EoMs */
   const float fac_mu = pow_three_gamma_minus_five_over_two(a);
   const float a2_Hubble = a * a * H;
@@ -168,9 +237,13 @@ __attribute__((always_inline)) INLINE static void runner_iact_force(
   kernel_deval(xj, &wj, &wj_dx);
   const float wj_dr = hjd_inv * wj_dx;
 
+  /* Variable smoothing length term */
+  const float f_ij = 1.f - pi->force.f / mj;
+  const float f_ji = 1.f - pj->force.f / mi;
+
   /* Compute gradient terms */
-  const float P_over_rho2_i = pressurei / (rhoi * rhoi) * pi->force.f;
-  const float P_over_rho2_j = pressurej / (rhoj * rhoj) * pj->force.f;
+  const float P_over_rho2_i = pressurei / (rhoi * rhoi) * f_ij;
+  const float P_over_rho2_j = pressurej / (rhoj * rhoj) * f_ji;
 
   /* Compute dv dot r. */
   const float dvdr = (pi->v[0] - pj->v[0]) * dx[0] +
@@ -195,7 +268,8 @@ __attribute__((always_inline)) INLINE static void runner_iact_force(
   const float visc = -0.25f * v_sig * mu_ij * (balsara_i + balsara_j) / rho_ij;
 
   /* Convolve with the kernel */
-  const float visc_acc_term = 0.5f * visc * (wi_dr + wj_dr) * r_inv;
+  const float visc_acc_term =
+      0.5f * visc * (wi_dr * f_ij + wj_dr * f_ji) * r_inv;
 
   /* SPH acceleration term */
   const float sph_acc_term =
@@ -253,6 +327,13 @@ __attribute__((always_inline)) INLINE static void runner_iact_nonsym_force(
     float r2, const float *dx, float hi, float hj, struct part *restrict pi,
     const struct part *restrict pj, float a, float H) {
 
+#ifdef SWIFT_DEBUG_CHECKS
+  if (pi->time_bin >= time_bin_inhibited)
+    error("Inhibited pi in interaction function!");
+  if (pj->time_bin >= time_bin_inhibited)
+    error("Inhibited pj in interaction function!");
+#endif
+
   /* Cosmological factors entering the EoMs */
   const float fac_mu = pow_three_gamma_minus_five_over_two(a);
   const float a2_Hubble = a * a * H;
@@ -262,7 +343,7 @@ __attribute__((always_inline)) INLINE static void runner_iact_nonsym_force(
   const float r = r2 * r_inv;
 
   /* Recover some data */
-  // const float mi = pi->mass;
+  const float mi = pi->mass;
   const float mj = pj->mass;
   const float rhoi = pi->rho;
   const float rhoj = pj->rho;
@@ -285,9 +366,13 @@ __attribute__((always_inline)) INLINE static void runner_iact_nonsym_force(
   kernel_deval(xj, &wj, &wj_dx);
   const float wj_dr = hjd_inv * wj_dx;
 
+  /* Variable smoothing length term */
+  const float f_ij = 1.f - pi->force.f / mj;
+  const float f_ji = 1.f - pj->force.f / mi;
+
   /* Compute gradient terms */
-  const float P_over_rho2_i = pressurei / (rhoi * rhoi) * pi->force.f;
-  const float P_over_rho2_j = pressurej / (rhoj * rhoj) * pj->force.f;
+  const float P_over_rho2_i = pressurei / (rhoi * rhoi) * f_ij;
+  const float P_over_rho2_j = pressurej / (rhoj * rhoj) * f_ji;
 
   /* Compute dv dot r. */
   const float dvdr = (pi->v[0] - pj->v[0]) * dx[0] +
@@ -314,7 +399,8 @@ __attribute__((always_inline)) INLINE static void runner_iact_nonsym_force(
   const float visc = -0.25f * v_sig * mu_ij * (balsara_i + balsara_j) / rho_ij;
 
   /* Convolve with the kernel */
-  const float visc_acc_term = 0.5f * visc * (wi_dr + wj_dr) * r_inv;
+  const float visc_acc_term =
+      0.5f * visc * (wi_dr * f_ij + wj_dr * f_ji) * r_inv;
 
   /* SPH acceleration term */
   const float sph_acc_term =
diff --git a/src/hydro/Planetary/hydro_io.h b/src/hydro/Planetary/hydro_io.h
index 9872441e2d5a641069e589154ed4b52600db3f45..eabf77ebb83ed974208e559e0910e387a94ab555 100644
--- a/src/hydro/Planetary/hydro_io.h
+++ b/src/hydro/Planetary/hydro_io.h
@@ -51,7 +51,11 @@ INLINE static void hydro_read_particles(struct part* parts,
                                         struct io_props* list,
                                         int* num_fields) {
 
+#ifdef PLANETARY_FIXED_ENTROPY
+  *num_fields = 10;
+#else
   *num_fields = 9;
+#endif
 
   /* List what we want to read */
   list[0] = io_make_input_field("Coordinates", DOUBLE, 3, COMPULSORY,
@@ -72,6 +76,11 @@ INLINE static void hydro_read_particles(struct part* parts,
                                 UNIT_CONV_DENSITY, parts, rho);
   list[8] = io_make_input_field("MaterialIDs", INT, 1, COMPULSORY,
                                 UNIT_CONV_NO_UNITS, parts, mat_id);
+#ifdef PLANETARY_FIXED_ENTROPY
+  list[9] = io_make_input_field("Entropies", FLOAT, 1, COMPULSORY,
+                                UNIT_CONV_PHYSICAL_ENTROPY_PER_UNIT_MASS, parts,
+                                s_fixed);
+#endif
 }
 
 INLINE static void convert_S(const struct engine* e, const struct part* p,
diff --git a/src/hydro/Planetary/hydro_part.h b/src/hydro/Planetary/hydro_part.h
index 482f44288c9218fb55bf6d199978b423b84c6caa..40d2f778356bdc296e344b9821f44ada213ce229 100644
--- a/src/hydro/Planetary/hydro_part.h
+++ b/src/hydro/Planetary/hydro_part.h
@@ -61,6 +61,9 @@ struct xpart {
   /*! Velocity at the last full step. */
   float v_full[3];
 
+  /*! Gravitational acceleration at the end of the last step */
+  float a_grav[3];
+
   /*! Internal energy at the last full step. */
   float u_full;
 
@@ -208,6 +211,11 @@ struct part {
 
 #endif
 
+#ifdef PLANETARY_FIXED_ENTROPY
+  /* Fixed specific entropy */
+  float s_fixed;
+#endif
+
 } SWIFT_STRUCT_ALIGN;
 
 #endif /* SWIFT_PLANETARY_HYDRO_PART_H */
diff --git a/src/hydro/PressureEnergy/hydro_part.h b/src/hydro/PressureEnergy/hydro_part.h
index 66b44d56c8818969ca10c199fe0f26bf73a38b70..8be543aadc8ad494bec4e7f1127e865f6f9f9d41 100644
--- a/src/hydro/PressureEnergy/hydro_part.h
+++ b/src/hydro/PressureEnergy/hydro_part.h
@@ -58,6 +58,9 @@ struct xpart {
   /*! Velocity at the last full step. */
   float v_full[3];
 
+  /*! Gravitational acceleration at the end of the last step */
+  float a_grav[3];
+
   /*! Internal energy at the last full step. */
   float u_full;
 
diff --git a/src/hydro/PressureEnergyMorrisMonaghanAV/hydro_part.h b/src/hydro/PressureEnergyMorrisMonaghanAV/hydro_part.h
index 30db47c9d2e94046c312f31a096033bee1c6a667..605bdfb8fb36932842256f54b754feca08f018a2 100644
--- a/src/hydro/PressureEnergyMorrisMonaghanAV/hydro_part.h
+++ b/src/hydro/PressureEnergyMorrisMonaghanAV/hydro_part.h
@@ -58,6 +58,9 @@ struct xpart {
   /*! Velocity at the last full step. */
   float v_full[3];
 
+  /*! Gravitational acceleration at the end of the last step */
+  float a_grav[3];
+
   /*! Internal energy at the last full step. */
   float u_full;
 
diff --git a/src/hydro/PressureEntropy/hydro_part.h b/src/hydro/PressureEntropy/hydro_part.h
index 9a03c8862a491165376d16a03fe27429572ec9ab..3177475b1c71c69026707052273841288679b98a 100644
--- a/src/hydro/PressureEntropy/hydro_part.h
+++ b/src/hydro/PressureEntropy/hydro_part.h
@@ -50,6 +50,9 @@ struct xpart {
   /*! Velocity at the last full step. */
   float v_full[3];
 
+  /*! Gravitational acceleration at the end of the last step */
+  float a_grav[3];
+
   /*! Entropy at the last full step. */
   float entropy_full;
 
diff --git a/src/hydro/SPHENIX/hydro_part.h b/src/hydro/SPHENIX/hydro_part.h
index 4a3c2877909bcbf5a11769e6f6d30ce71a1d4010..f13e29c0e3e4f89ea6bd09e594b7e6dcfb4bfafd 100644
--- a/src/hydro/SPHENIX/hydro_part.h
+++ b/src/hydro/SPHENIX/hydro_part.h
@@ -53,6 +53,9 @@ struct xpart {
   /*! Velocity at the last full step. */
   float v_full[3];
 
+  /*! Gravitational acceleration at the end of the last step */
+  float a_grav[3];
+
   /*! Internal energy at the last full step. */
   float u_full;
 
diff --git a/src/hydro/Shadowswift/hydro_part.h b/src/hydro/Shadowswift/hydro_part.h
index bc374e5500c23fc3e9c6210adb1079f6c55accde..2ed26b8568dce0aa6b43f8bdeced26d6e2caf3ba 100644
--- a/src/hydro/Shadowswift/hydro_part.h
+++ b/src/hydro/Shadowswift/hydro_part.h
@@ -41,6 +41,9 @@ struct xpart {
   /* Velocity at the last full step. */
   float v_full[3];
 
+  /*! Gravitational acceleration at the end of the last step */
+  float a_grav[3];
+
   /* Additional data used to record cooling information */
   struct cooling_xpart_data cooling_data;
 
diff --git a/src/line_of_sight.c b/src/line_of_sight.c
index 3c686f99280f951a5d466d9b172c49e7c7c21a89..dec2e7c947b335c2a741a55a37be01d6bf1695be 100644
--- a/src/line_of_sight.c
+++ b/src/line_of_sight.c
@@ -531,6 +531,14 @@ void write_hdf5_header(hid_t h_file, const struct engine *e,
                      swift_type_count);
   io_write_attribute(h_grp, "NumPart_Total_HighWord", UINT,
                      numParticlesHighWord, swift_type_count);
+  double MassTable[swift_type_count] = {0};
+  io_write_attribute(h_grp, "MassTable", DOUBLE, MassTable, swift_type_count);
+  io_write_attribute(h_grp, "InitialMassTable", DOUBLE,
+                     e->s->initial_mean_mass_particles, swift_type_count);
+  unsigned int flagEntropy[swift_type_count] = {0};
+  flagEntropy[0] = writeEntropyFlag();
+  io_write_attribute(h_grp, "Flag_Entropy_ICs", UINT, flagEntropy,
+                     swift_type_count);
   io_write_attribute_i(h_grp, "NumFilesPerSnapshot", 1);
   io_write_attribute_i(h_grp, "ThisFile", 0);
   io_write_attribute_s(h_grp, "OutputType", "LineOfSight");
diff --git a/src/parallel_io.c b/src/parallel_io.c
index 55b345912da6b57060bf587179089afcf3c88f9a..de756d3e2fa35f4fb7e65260a3f2afaf049c723d 100644
--- a/src/parallel_io.c
+++ b/src/parallel_io.c
@@ -1208,6 +1208,8 @@ void prepare_file(struct engine* e, const char* fileName,
                      numParticlesHighWord, swift_type_count);
   double MassTable[6] = {0., 0., 0., 0., 0., 0.};
   io_write_attribute(h_grp, "MassTable", DOUBLE, MassTable, swift_type_count);
+  io_write_attribute(h_grp, "InitialMassTable", DOUBLE,
+                     e->s->initial_mean_mass_particles, swift_type_count);
   unsigned int flagEntropy[swift_type_count] = {0};
   flagEntropy[0] = writeEntropyFlag();
   io_write_attribute(h_grp, "Flag_Entropy_ICs", UINT, flagEntropy,
diff --git a/src/restart.c b/src/restart.c
index 63d1bee0345cb378f098218fdcd0e061bdc33bf9..9e913240ba5af3d3f415c5d449b362574391f510 100644
--- a/src/restart.c
+++ b/src/restart.c
@@ -358,5 +358,7 @@ void restart_resubmit(const char *command) {
 
   /* Let's trust the user's command... */
   const int result = system(command);
-  if (result != 0) message("Command returned error code %d", result);
+  if (result != 0) {
+    message("Restart resubmit command returned error code %d", result);
+  }
 }
diff --git a/src/runner_doiact_functions_hydro.h b/src/runner_doiact_functions_hydro.h
index b0dbbc1b97ba91975fc2669e547ee9ddf6fd10f4..4bdbcec05cf8fd9e5d5b06b13737d02886c94164 100644
--- a/src/runner_doiact_functions_hydro.h
+++ b/src/runner_doiact_functions_hydro.h
@@ -40,6 +40,11 @@ void DOPAIR1_NAIVE(struct runner *r, struct cell *restrict ci,
 
   const struct engine *e = r->e;
   const struct cosmology *cosmo = e->cosmology;
+#if (FUNCTION_TASK_LOOP == TASK_LOOP_FORCE)
+  const double time_base = e->time_base;
+  const integertime_t t_current = e->ti_current;
+  const int with_cosmology = (e->policy & engine_policy_cosmology);
+#endif
 
   TIMER_TIC;
 
@@ -119,6 +124,8 @@ void DOPAIR1_NAIVE(struct runner *r, struct cell *restrict ci,
 #endif
 #if (FUNCTION_TASK_LOOP == TASK_LOOP_FORCE)
         runner_iact_nonsym_timebin(r2, dx, hi, hj, pi, pj, a, H);
+        runner_iact_nonsym_diffusion(r2, dx, hi, hj, pi, pj, a, H, time_base,
+                                     t_current, cosmo, with_cosmology);
 #endif
       }
       if (r2 < hjg2 && pj_active) {
@@ -135,6 +142,8 @@ void DOPAIR1_NAIVE(struct runner *r, struct cell *restrict ci,
 #endif
 #if (FUNCTION_TASK_LOOP == TASK_LOOP_FORCE)
         runner_iact_nonsym_timebin(r2, dx, hj, hi, pj, pi, a, H);
+        runner_iact_nonsym_diffusion(r2, dx, hj, hi, pj, pi, a, H, time_base,
+                                     t_current, cosmo, with_cosmology);
 #endif
       }
     } /* loop over the parts in cj. */
@@ -157,6 +166,11 @@ void DOPAIR2_NAIVE(struct runner *r, struct cell *restrict ci,
 
   const struct engine *e = r->e;
   const struct cosmology *cosmo = e->cosmology;
+#if (FUNCTION_TASK_LOOP == TASK_LOOP_FORCE)
+  const double time_base = e->time_base;
+  const integertime_t t_current = e->ti_current;
+  const int with_cosmology = (e->policy & engine_policy_cosmology);
+#endif
 
   TIMER_TIC;
 
@@ -238,6 +252,8 @@ void DOPAIR2_NAIVE(struct runner *r, struct cell *restrict ci,
 #endif
 #if (FUNCTION_TASK_LOOP == TASK_LOOP_FORCE)
           runner_iact_timebin(r2, dx, hi, hj, pi, pj, a, H);
+          runner_iact_diffusion(r2, dx, hi, hj, pi, pj, a, H, time_base,
+                                t_current, cosmo, with_cosmology);
 #endif
         } else if (pi_active) {
 
@@ -249,6 +265,8 @@ void DOPAIR2_NAIVE(struct runner *r, struct cell *restrict ci,
 #endif
 #if (FUNCTION_TASK_LOOP == TASK_LOOP_FORCE)
           runner_iact_nonsym_timebin(r2, dx, hi, hj, pi, pj, a, H);
+          runner_iact_nonsym_diffusion(r2, dx, hi, hj, pi, pj, a, H, time_base,
+                                       t_current, cosmo, with_cosmology);
 #endif
         } else if (pj_active) {
 
@@ -264,6 +282,8 @@ void DOPAIR2_NAIVE(struct runner *r, struct cell *restrict ci,
 #endif
 #if (FUNCTION_TASK_LOOP == TASK_LOOP_FORCE)
           runner_iact_nonsym_timebin(r2, dx, hj, hi, pj, pi, a, H);
+          runner_iact_nonsym_diffusion(r2, dx, hj, hi, pj, pi, a, H, time_base,
+                                       t_current, cosmo, with_cosmology);
 #endif
         }
       }
@@ -285,6 +305,11 @@ void DOSELF1_NAIVE(struct runner *r, struct cell *restrict c) {
 
   const struct engine *e = r->e;
   const struct cosmology *cosmo = e->cosmology;
+#if (FUNCTION_TASK_LOOP == TASK_LOOP_FORCE)
+  const double time_base = e->time_base;
+  const integertime_t t_current = e->ti_current;
+  const int with_cosmology = (e->policy & engine_policy_cosmology);
+#endif
 
   TIMER_TIC;
 
@@ -356,6 +381,8 @@ void DOSELF1_NAIVE(struct runner *r, struct cell *restrict c) {
 #endif
 #if (FUNCTION_TASK_LOOP == TASK_LOOP_FORCE)
         runner_iact_timebin(r2, dx, hi, hj, pi, pj, a, H);
+        runner_iact_diffusion(r2, dx, hi, hj, pi, pj, a, H, time_base,
+                              t_current, cosmo, with_cosmology);
 #endif
       } else if (doi) {
 
@@ -367,6 +394,8 @@ void DOSELF1_NAIVE(struct runner *r, struct cell *restrict c) {
 #endif
 #if (FUNCTION_TASK_LOOP == TASK_LOOP_FORCE)
         runner_iact_nonsym_timebin(r2, dx, hi, hj, pi, pj, a, H);
+        runner_iact_nonsym_diffusion(r2, dx, hi, hj, pi, pj, a, H, time_base,
+                                     t_current, cosmo, with_cosmology);
 #endif
       } else if (doj) {
 
@@ -382,6 +411,8 @@ void DOSELF1_NAIVE(struct runner *r, struct cell *restrict c) {
 #endif
 #if (FUNCTION_TASK_LOOP == TASK_LOOP_FORCE)
         runner_iact_nonsym_timebin(r2, dx, hj, hi, pj, pi, a, H);
+        runner_iact_nonsym_diffusion(r2, dx, hj, hi, pj, pi, a, H, time_base,
+                                     t_current, cosmo, with_cosmology);
 #endif
       }
     } /* loop over the parts in cj. */
@@ -402,6 +433,11 @@ void DOSELF2_NAIVE(struct runner *r, struct cell *restrict c) {
 
   const struct engine *e = r->e;
   const struct cosmology *cosmo = e->cosmology;
+#if (FUNCTION_TASK_LOOP == TASK_LOOP_FORCE)
+  const double time_base = e->time_base;
+  const integertime_t t_current = e->ti_current;
+  const int with_cosmology = (e->policy & engine_policy_cosmology);
+#endif
 
   TIMER_TIC;
 
@@ -473,6 +509,8 @@ void DOSELF2_NAIVE(struct runner *r, struct cell *restrict c) {
 #endif
 #if (FUNCTION_TASK_LOOP == TASK_LOOP_FORCE)
         runner_iact_timebin(r2, dx, hi, hj, pi, pj, a, H);
+        runner_iact_diffusion(r2, dx, hi, hj, pi, pj, a, H, time_base,
+                              t_current, cosmo, with_cosmology);
 #endif
       } else if (doi) {
 
@@ -484,6 +522,8 @@ void DOSELF2_NAIVE(struct runner *r, struct cell *restrict c) {
 #endif
 #if (FUNCTION_TASK_LOOP == TASK_LOOP_FORCE)
         runner_iact_nonsym_timebin(r2, dx, hi, hj, pi, pj, a, H);
+        runner_iact_nonsym_diffusion(r2, dx, hi, hj, pi, pj, a, H, time_base,
+                                     t_current, cosmo, with_cosmology);
 #endif
       } else if (doj) {
 
@@ -499,6 +539,8 @@ void DOSELF2_NAIVE(struct runner *r, struct cell *restrict c) {
 #endif
 #if (FUNCTION_TASK_LOOP == TASK_LOOP_FORCE)
         runner_iact_nonsym_timebin(r2, dx, hj, hi, pj, pi, a, H);
+        runner_iact_nonsym_diffusion(r2, dx, hj, hi, pj, pi, a, H, time_base,
+                                     t_current, cosmo, with_cosmology);
 #endif
       }
     } /* loop over the parts in cj. */
@@ -528,6 +570,11 @@ void DOPAIR_SUBSET_NAIVE(struct runner *r, struct cell *restrict ci,
 
   const struct engine *e = r->e;
   const struct cosmology *cosmo = e->cosmology;
+#if (FUNCTION_TASK_LOOP == TASK_LOOP_FORCE)
+  const double time_base = e->time_base;
+  const integertime_t t_current = e->ti_current;
+  const int with_cosmology = (e->policy & engine_policy_cosmology);
+#endif
 
   TIMER_TIC;
 
@@ -589,6 +636,8 @@ void DOPAIR_SUBSET_NAIVE(struct runner *r, struct cell *restrict ci,
 #endif
 #if (FUNCTION_TASK_LOOP == TASK_LOOP_FORCE)
         runner_iact_nonsym_timebin(r2, dx, hi, pj->h, pi, pj, a, H);
+        runner_iact_nonsym_diffusion(r2, dx, hi, pj->h, pi, pj, a, H, time_base,
+                                     t_current, cosmo, with_cosmology);
 #endif
       }
     } /* loop over the parts in cj. */
@@ -618,6 +667,11 @@ void DOPAIR_SUBSET(struct runner *r, struct cell *restrict ci,
 
   const struct engine *e = r->e;
   const struct cosmology *cosmo = e->cosmology;
+#if (FUNCTION_TASK_LOOP == TASK_LOOP_FORCE)
+  const double time_base = e->time_base;
+  const integertime_t t_current = e->ti_current;
+  const int with_cosmology = (e->policy & engine_policy_cosmology);
+#endif
 
   TIMER_TIC;
 
@@ -686,6 +740,8 @@ void DOPAIR_SUBSET(struct runner *r, struct cell *restrict ci,
 #endif
 #if (FUNCTION_TASK_LOOP == TASK_LOOP_FORCE)
           runner_iact_nonsym_timebin(r2, dx, hi, hj, pi, pj, a, H);
+          runner_iact_nonsym_diffusion(r2, dx, hi, hj, pi, pj, a, H, time_base,
+                                       t_current, cosmo, with_cosmology);
 #endif
         }
       } /* loop over the parts in cj. */
@@ -746,6 +802,8 @@ void DOPAIR_SUBSET(struct runner *r, struct cell *restrict ci,
 #endif
 #if (FUNCTION_TASK_LOOP == TASK_LOOP_FORCE)
           runner_iact_nonsym_timebin(r2, dx, hi, hj, pi, pj, a, H);
+          runner_iact_nonsym_diffusion(r2, dx, hi, hj, pi, pj, a, H, time_base,
+                                       t_current, cosmo, with_cosmology);
 #endif
         }
       } /* loop over the parts in cj. */
@@ -831,6 +889,11 @@ void DOSELF_SUBSET(struct runner *r, struct cell *restrict ci,
 
   const struct engine *e = r->e;
   const struct cosmology *cosmo = e->cosmology;
+#if (FUNCTION_TASK_LOOP == TASK_LOOP_FORCE)
+  const double time_base = e->time_base;
+  const integertime_t t_current = e->ti_current;
+  const int with_cosmology = (e->policy & engine_policy_cosmology);
+#endif
 
   TIMER_TIC;
 
@@ -892,6 +955,8 @@ void DOSELF_SUBSET(struct runner *r, struct cell *restrict ci,
 #endif
 #if (FUNCTION_TASK_LOOP == TASK_LOOP_FORCE)
         runner_iact_nonsym_timebin(r2, dx, hi, hj, pi, pj, a, H);
+        runner_iact_nonsym_diffusion(r2, dx, hi, hj, pi, pj, a, H, time_base,
+                                     t_current, cosmo, with_cosmology);
 #endif
       }
     } /* loop over the parts in cj. */
@@ -935,6 +1000,11 @@ void DOPAIR1(struct runner *r, struct cell *ci, struct cell *cj, const int sid,
 
   const struct engine *restrict e = r->e;
   const struct cosmology *restrict cosmo = e->cosmology;
+#if (FUNCTION_TASK_LOOP == TASK_LOOP_FORCE)
+  const double time_base = e->time_base;
+  const integertime_t t_current = e->ti_current;
+  const int with_cosmology = (e->policy & engine_policy_cosmology);
+#endif
 
   TIMER_TIC;
 
@@ -1060,6 +1130,8 @@ void DOPAIR1(struct runner *r, struct cell *ci, struct cell *cj, const int sid,
 #endif
 #if (FUNCTION_TASK_LOOP == TASK_LOOP_FORCE)
           runner_iact_nonsym_timebin(r2, dx, hi, hj, pi, pj, a, H);
+          runner_iact_nonsym_diffusion(r2, dx, hi, hj, pi, pj, a, H, time_base,
+                                       t_current, cosmo, with_cosmology);
 #endif
         }
       } /* loop over the parts in cj. */
@@ -1152,6 +1224,8 @@ void DOPAIR1(struct runner *r, struct cell *ci, struct cell *cj, const int sid,
 #endif
 #if (FUNCTION_TASK_LOOP == TASK_LOOP_FORCE)
           runner_iact_nonsym_timebin(r2, dx, hj, hi, pj, pi, a, H);
+          runner_iact_nonsym_diffusion(r2, dx, hj, hi, pj, pi, a, H, time_base,
+                                       t_current, cosmo, with_cosmology);
 #endif
         }
       } /* loop over the parts in ci. */
@@ -1268,6 +1342,11 @@ void DOPAIR2(struct runner *r, struct cell *ci, struct cell *cj, const int sid,
 
   const struct engine *restrict e = r->e;
   const struct cosmology *restrict cosmo = e->cosmology;
+#if (FUNCTION_TASK_LOOP == TASK_LOOP_FORCE)
+  const double time_base = e->time_base;
+  const integertime_t t_current = e->ti_current;
+  const int with_cosmology = (e->policy & engine_policy_cosmology);
+#endif
 
   TIMER_TIC;
 
@@ -1458,6 +1537,8 @@ void DOPAIR2(struct runner *r, struct cell *ci, struct cell *cj, const int sid,
 #endif
 #if (FUNCTION_TASK_LOOP == TASK_LOOP_FORCE)
           runner_iact_nonsym_timebin(r2, dx, hj, hi, pj, pi, a, H);
+          runner_iact_nonsym_diffusion(r2, dx, hj, hi, pj, pi, a, H, time_base,
+                                       t_current, cosmo, with_cosmology);
 #endif
         }
       } /* loop over the active parts in cj. */
@@ -1533,6 +1614,8 @@ void DOPAIR2(struct runner *r, struct cell *ci, struct cell *cj, const int sid,
 #endif
 #if (FUNCTION_TASK_LOOP == TASK_LOOP_FORCE)
             runner_iact_timebin(r2, dx, hi, hj, pi, pj, a, H);
+            runner_iact_diffusion(r2, dx, hi, hj, pi, pj, a, H, time_base,
+                                  t_current, cosmo, with_cosmology);
 #endif
           } else {
             IACT_NONSYM(r2, dx, hi, hj, pi, pj, a, H);
@@ -1543,6 +1626,9 @@ void DOPAIR2(struct runner *r, struct cell *ci, struct cell *cj, const int sid,
 #endif
 #if (FUNCTION_TASK_LOOP == TASK_LOOP_FORCE)
             runner_iact_nonsym_timebin(r2, dx, hi, hj, pi, pj, a, H);
+            runner_iact_nonsym_diffusion(r2, dx, hi, hj, pi, pj, a, H,
+                                         time_base, t_current, cosmo,
+                                         with_cosmology);
 #endif
           }
         }
@@ -1650,6 +1736,8 @@ void DOPAIR2(struct runner *r, struct cell *ci, struct cell *cj, const int sid,
 #endif
 #if (FUNCTION_TASK_LOOP == TASK_LOOP_FORCE)
           runner_iact_nonsym_timebin(r2, dx, hi, hj, pi, pj, a, H);
+          runner_iact_nonsym_diffusion(r2, dx, hi, hj, pi, pj, a, H, time_base,
+                                       t_current, cosmo, with_cosmology);
 #endif
         }
       } /* loop over the active parts in ci. */
@@ -1727,6 +1815,8 @@ void DOPAIR2(struct runner *r, struct cell *ci, struct cell *cj, const int sid,
 #endif
 #if (FUNCTION_TASK_LOOP == TASK_LOOP_FORCE)
             runner_iact_timebin(r2, dx, hj, hi, pj, pi, a, H);
+            runner_iact_diffusion(r2, dx, hj, hi, pj, pi, a, H, time_base,
+                                  t_current, cosmo, with_cosmology);
 #endif
           } else {
             IACT_NONSYM(r2, dx, hj, hi, pj, pi, a, H);
@@ -1737,6 +1827,9 @@ void DOPAIR2(struct runner *r, struct cell *ci, struct cell *cj, const int sid,
 #endif
 #if (FUNCTION_TASK_LOOP == TASK_LOOP_FORCE)
             runner_iact_nonsym_timebin(r2, dx, hj, hi, pj, pi, a, H);
+            runner_iact_nonsym_diffusion(r2, dx, hj, hi, pj, pi, a, H,
+                                         time_base, t_current, cosmo,
+                                         with_cosmology);
 #endif
           }
         }
@@ -1856,6 +1949,11 @@ void DOSELF1(struct runner *r, struct cell *restrict c) {
 
   const struct engine *e = r->e;
   const struct cosmology *cosmo = e->cosmology;
+#if (FUNCTION_TASK_LOOP == TASK_LOOP_FORCE)
+  const double time_base = e->time_base;
+  const integertime_t t_current = e->ti_current;
+  const int with_cosmology = (e->policy & engine_policy_cosmology);
+#endif
 
   TIMER_TIC;
 
@@ -1930,6 +2028,8 @@ void DOSELF1(struct runner *r, struct cell *restrict c) {
 #endif
 #if (FUNCTION_TASK_LOOP == TASK_LOOP_FORCE)
           runner_iact_nonsym_timebin(r2, dx, hj, hi, pj, pi, a, H);
+          runner_iact_nonsym_diffusion(r2, dx, hj, hi, pj, pi, a, H, time_base,
+                                       t_current, cosmo, with_cosmology);
 #endif
         }
       } /* loop over all other particles. */
@@ -1986,6 +2086,8 @@ void DOSELF1(struct runner *r, struct cell *restrict c) {
 #endif
 #if (FUNCTION_TASK_LOOP == TASK_LOOP_FORCE)
             runner_iact_timebin(r2, dx, hi, hj, pi, pj, a, H);
+            runner_iact_diffusion(r2, dx, hi, hj, pi, pj, a, H, time_base,
+                                  t_current, cosmo, with_cosmology);
 #endif
           } else if (doi) {
 
@@ -1997,6 +2099,9 @@ void DOSELF1(struct runner *r, struct cell *restrict c) {
 #endif
 #if (FUNCTION_TASK_LOOP == TASK_LOOP_FORCE)
             runner_iact_nonsym_timebin(r2, dx, hi, hj, pi, pj, a, H);
+            runner_iact_nonsym_diffusion(r2, dx, hi, hj, pi, pj, a, H,
+                                         time_base, t_current, cosmo,
+                                         with_cosmology);
 #endif
           } else if (doj) {
 
@@ -2011,6 +2116,9 @@ void DOSELF1(struct runner *r, struct cell *restrict c) {
 #endif
 #if (FUNCTION_TASK_LOOP == TASK_LOOP_FORCE)
             runner_iact_nonsym_timebin(r2, dx, hj, hi, pj, pi, a, H);
+            runner_iact_nonsym_diffusion(r2, dx, hj, hi, pj, pi, a, H,
+                                         time_base, t_current, cosmo,
+                                         with_cosmology);
 #endif
           }
         }
@@ -2068,6 +2176,11 @@ void DOSELF2(struct runner *r, struct cell *restrict c) {
 
   const struct engine *e = r->e;
   const struct cosmology *cosmo = e->cosmology;
+#if (FUNCTION_TASK_LOOP == TASK_LOOP_FORCE)
+  const double time_base = e->time_base;
+  const integertime_t t_current = e->ti_current;
+  const int with_cosmology = (e->policy & engine_policy_cosmology);
+#endif
 
   TIMER_TIC;
 
@@ -2142,6 +2255,8 @@ void DOSELF2(struct runner *r, struct cell *restrict c) {
 #endif
 #if (FUNCTION_TASK_LOOP == TASK_LOOP_FORCE)
           runner_iact_nonsym_timebin(r2, dx, hj, hi, pj, pi, a, H);
+          runner_iact_nonsym_diffusion(r2, dx, hj, hi, pj, pi, a, H, time_base,
+                                       t_current, cosmo, with_cosmology);
 #endif
         }
       } /* loop over all other particles. */
@@ -2193,6 +2308,8 @@ void DOSELF2(struct runner *r, struct cell *restrict c) {
 #endif
 #if (FUNCTION_TASK_LOOP == TASK_LOOP_FORCE)
             runner_iact_timebin(r2, dx, hi, hj, pi, pj, a, H);
+            runner_iact_diffusion(r2, dx, hi, hj, pi, pj, a, H, time_base,
+                                  t_current, cosmo, with_cosmology);
 #endif
           } else {
             IACT_NONSYM(r2, dx, hi, hj, pi, pj, a, H);
@@ -2203,6 +2320,9 @@ void DOSELF2(struct runner *r, struct cell *restrict c) {
 #endif
 #if (FUNCTION_TASK_LOOP == TASK_LOOP_FORCE)
             runner_iact_nonsym_timebin(r2, dx, hi, hj, pi, pj, a, H);
+            runner_iact_nonsym_diffusion(r2, dx, hi, hj, pi, pj, a, H,
+                                         time_base, t_current, cosmo,
+                                         with_cosmology);
 #endif
           }
         }
diff --git a/src/runner_doiact_functions_stars.h b/src/runner_doiact_functions_stars.h
index 7e802d6f4afd885697c3915d475d38a4e5da7a73..b8bea6fa2125bbddb0deb013c46876e97ec61927 100644
--- a/src/runner_doiact_functions_stars.h
+++ b/src/runner_doiact_functions_stars.h
@@ -1098,7 +1098,7 @@ void DOSELF1_BRANCH_STARS(struct runner *r, struct cell *c) {
             "ci->nodeID=%d d=%e sort_j[pjd].d=%e cj->" #TYPE                \
             ".dx_max_sort=%e "                                              \
             "cj->" #TYPE                                                    \
-            ".dx_max_sort_old=%e, cellID=%i super->cellID=%i"               \
+            ".dx_max_sort_old=%e, cellID=%lld super->cellID=%lld"           \
             "cj->depth=%d cj->maxdepth=%d",                                 \
             cj->nodeID, ci->nodeID, d, sort_j[pjd].d, cj->TYPE.dx_max_sort, \
             cj->TYPE.dx_max_sort_old, cj->cellID, cj->hydro.super->cellID,  \
diff --git a/src/runner_doiact_grav.c b/src/runner_doiact_grav.c
index e0e040e38eb0d7551769f87df01eb646af21ac1e..f3b044f31bcb5e21f4c6afc397f843d32228dc8d 100644
--- a/src/runner_doiact_grav.c
+++ b/src/runner_doiact_grav.c
@@ -243,9 +243,9 @@ static INLINE void runner_dopair_grav_pp_full_no_cache(
       /* Update the M2P interaction counter and forces. */
       accumulate_add_ll(&gparts_i[i].num_interacted_m2p,
                         multi_j->m_pole.num_gpart);
-      gparts_i[i].a_grav_m2p[0] += f_x;
-      gparts_i[i].a_grav_m2p[1] += f_y;
-      gparts_i[i].a_grav_m2p[2] += f_z;
+      accumulate_add_f(&gparts_i[i].a_grav_m2p[0], f_x);
+      accumulate_add_f(&gparts_i[i].a_grav_m2p[1], f_y);
+      accumulate_add_f(&gparts_i[i].a_grav_m2p[2], f_z);
 #endif
 
     } else {
@@ -307,9 +307,9 @@ static INLINE void runner_dopair_grav_pp_full_no_cache(
 #ifdef SWIFT_GRAVITY_FORCE_CHECKS
         /* Update the p2p interaction counter */
         accumulate_inc_ll(&gparts_i[i].num_interacted_p2p);
-        gparts_i[i].a_grav_p2p[0] += a_x;
-        gparts_i[i].a_grav_p2p[1] += a_y;
-        gparts_i[i].a_grav_p2p[2] += a_z;
+        accumulate_add_f(&gparts_i[i].a_grav_p2p[0], a_x);
+        accumulate_add_f(&gparts_i[i].a_grav_p2p[1], a_y);
+        accumulate_add_f(&gparts_i[i].a_grav_p2p[2], a_z);
 #endif
       }
     }
@@ -436,9 +436,9 @@ static INLINE void runner_dopair_grav_pp_truncated_no_cache(
       /* Update the M2P interaction counter and forces. */
       accumulate_add_ll(&gparts_i[i].num_interacted_m2p,
                         multi_j->m_pole.num_gpart);
-      gparts_i[i].a_grav_m2p[0] += f_x;
-      gparts_i[i].a_grav_m2p[1] += f_y;
-      gparts_i[i].a_grav_m2p[2] += f_z;
+      accumulate_add_f(&gparts_i[i].a_grav_m2p[0], f_x);
+      accumulate_add_f(&gparts_i[i].a_grav_m2p[1], f_y);
+      accumulate_add_f(&gparts_i[i].a_grav_m2p[2], f_z);
 #endif
 
     } else {
@@ -505,9 +505,9 @@ static INLINE void runner_dopair_grav_pp_truncated_no_cache(
 #ifdef SWIFT_GRAVITY_FORCE_CHECKS
         /* Update the p2p interaction counter */
         accumulate_inc_ll(&gparts_i[i].num_interacted_p2p);
-        gparts_i[i].a_grav_p2p[0] += a_x;
-        gparts_i[i].a_grav_p2p[1] += a_y;
-        gparts_i[i].a_grav_p2p[2] += a_z;
+        accumulate_add_f(&gparts_i[i].a_grav_p2p[0], a_x);
+        accumulate_add_f(&gparts_i[i].a_grav_p2p[1], a_y);
+        accumulate_add_f(&gparts_i[i].a_grav_p2p[2], a_z);
 #endif
       }
     }
@@ -672,9 +672,9 @@ static INLINE void runner_dopair_grav_pp_full(
     ci_cache->pot[pid] += pot;
 
 #ifdef SWIFT_GRAVITY_FORCE_CHECKS
-    gparts_i[pid].a_grav_p2p[0] += a_x;
-    gparts_i[pid].a_grav_p2p[1] += a_y;
-    gparts_i[pid].a_grav_p2p[2] += a_z;
+    accumulate_add_f(&gparts_i[pid].a_grav_p2p[0], a_x);
+    accumulate_add_f(&gparts_i[pid].a_grav_p2p[1], a_y);
+    accumulate_add_f(&gparts_i[pid].a_grav_p2p[2], a_z);
 #endif
   }
 }
@@ -828,9 +828,9 @@ static INLINE void runner_dopair_grav_pp_truncated(
     ci_cache->pot[pid] += pot;
 
 #ifdef SWIFT_GRAVITY_FORCE_CHECKS
-    gparts_i[pid].a_grav_p2p[0] += a_x;
-    gparts_i[pid].a_grav_p2p[1] += a_y;
-    gparts_i[pid].a_grav_p2p[2] += a_z;
+    accumulate_add_f(&gparts_i[pid].a_grav_p2p[0], a_x);
+    accumulate_add_f(&gparts_i[pid].a_grav_p2p[1], a_y);
+    accumulate_add_f(&gparts_i[pid].a_grav_p2p[2], a_z);
 #endif
   }
 }
@@ -882,7 +882,8 @@ static INLINE void runner_dopair_grav_pm_full(
   const float multi_epsilon = multi_j->max_softening;
 
   /* Loop over all particles in ci... */
-#if !defined(SWIFT_DEBUG_CHECKS) && _OPENMP >= 201307
+#if !defined(SWIFT_DEBUG_CHECKS) && !defined(SWIFT_GRAVITY_FORCE_CHECKS) && \
+    _OPENMP >= 201307
 #pragma omp simd
 #endif
   for (int pid = 0; pid < gcount_padded_i; pid++) {
@@ -963,9 +964,9 @@ static INLINE void runner_dopair_grav_pm_full(
     if (pid < gcount_i) {
       accumulate_add_ll(&gparts_i[pid].num_interacted_m2p,
                         cj->grav.multipole->m_pole.num_gpart);
-      gparts_i[pid].a_grav_m2p[0] += f_x;
-      gparts_i[pid].a_grav_m2p[1] += f_y;
-      gparts_i[pid].a_grav_m2p[2] += f_z;
+      accumulate_add_f(&gparts_i[pid].a_grav_m2p[0], f_x);
+      accumulate_add_f(&gparts_i[pid].a_grav_m2p[1], f_y);
+      accumulate_add_f(&gparts_i[pid].a_grav_m2p[2], f_z);
     }
 #endif
   }
@@ -1025,7 +1026,8 @@ static INLINE void runner_dopair_grav_pm_truncated(
   const float multi_epsilon = multi_j->max_softening;
 
   /* Loop over all particles in ci... */
-#if !defined(SWIFT_DEBUG_CHECKS) && _OPENMP >= 201307
+#if !defined(SWIFT_DEBUG_CHECKS) && !defined(SWIFT_GRAVITY_FORCE_CHECKS) && \
+    _OPENMP >= 201307
 #pragma omp simd
 #endif
   for (int pid = 0; pid < gcount_padded_i; pid++) {
@@ -1104,9 +1106,9 @@ static INLINE void runner_dopair_grav_pm_truncated(
     if (pid < gcount_i) {
       accumulate_add_ll(&gparts_i[pid].num_interacted_m2p,
                         cj->grav.multipole->m_pole.num_gpart);
-      gparts_i[pid].a_grav_m2p[0] += f_x;
-      gparts_i[pid].a_grav_m2p[1] += f_y;
-      gparts_i[pid].a_grav_m2p[2] += f_z;
+      accumulate_add_f(&gparts_i[pid].a_grav_m2p[0], f_x);
+      accumulate_add_f(&gparts_i[pid].a_grav_m2p[1], f_y);
+      accumulate_add_f(&gparts_i[pid].a_grav_m2p[2], f_z);
     }
 #endif
   }
@@ -1535,9 +1537,9 @@ static INLINE void runner_doself_grav_pp_full(
     ci_cache->pot[pid] += pot;
 
 #ifdef SWIFT_GRAVITY_FORCE_CHECKS
-    gparts[pid].a_grav_p2p[0] += a_x;
-    gparts[pid].a_grav_p2p[1] += a_y;
-    gparts[pid].a_grav_p2p[2] += a_z;
+    accumulate_add_f(&gparts[pid].a_grav_p2p[0], a_x);
+    accumulate_add_f(&gparts[pid].a_grav_p2p[1], a_y);
+    accumulate_add_f(&gparts[pid].a_grav_p2p[2], a_z);
 #endif
   }
 }
@@ -1674,9 +1676,9 @@ static INLINE void runner_doself_grav_pp_truncated(
     ci_cache->pot[pid] += pot;
 
 #ifdef SWIFT_GRAVITY_FORCE_CHECKS
-    gparts[pid].a_grav_p2p[0] += a_x;
-    gparts[pid].a_grav_p2p[1] += a_y;
-    gparts[pid].a_grav_p2p[2] += a_z;
+    accumulate_add_f(&gparts[pid].a_grav_p2p[0], a_x);
+    accumulate_add_f(&gparts[pid].a_grav_p2p[1], a_y);
+    accumulate_add_f(&gparts[pid].a_grav_p2p[2], a_z);
 #endif
   }
 }
diff --git a/src/runner_ghost.c b/src/runner_ghost.c
index 8ef7214d80f375cfdd6ace1fcd521c085231e256..b98e47a0f31c4c5f235f6f20589597015934fb7e 100644
--- a/src/runner_ghost.c
+++ b/src/runner_ghost.c
@@ -242,9 +242,10 @@ void runner_do_stars_ghost(struct runner *r, struct cell *c, int timer) {
                     star_age_end_of_step - dt_enrichment;
 
                 /* Compute the stellar evolution  */
-                feedback_evolve_spart(sp, feedback_props, cosmo, us, phys_const,
-                                      star_age_beg_of_step, dt_enrichment,
-                                      e->time, ti_begin, with_cosmology);
+                feedback_prepare_feedback(sp, feedback_props, cosmo, us,
+                                          phys_const, star_age_beg_of_step,
+                                          dt_enrichment, e->time, ti_begin,
+                                          with_cosmology);
               } else {
 
                 /* Reset the feedback fields of the star particle */
@@ -386,9 +387,9 @@ void runner_do_stars_ghost(struct runner *r, struct cell *c, int timer) {
                 star_age_end_of_step - dt_enrichment;
 
             /* Compute the stellar evolution  */
-            feedback_evolve_spart(sp, feedback_props, cosmo, us, phys_const,
-                                  star_age_beg_of_step, dt_enrichment, e->time,
-                                  ti_begin, with_cosmology);
+            feedback_prepare_feedback(sp, feedback_props, cosmo, us, phys_const,
+                                      star_age_beg_of_step, dt_enrichment,
+                                      e->time, ti_begin, with_cosmology);
           } else {
 
             /* Reset the feedback fields of the star particle */
diff --git a/src/runner_others.c b/src/runner_others.c
index 7efb403213f15c7ad8a539cc30047d4b053e8912..479da517a7e435d91161ba4397d7acda9dbd76fe 100644
--- a/src/runner_others.c
+++ b/src/runner_others.c
@@ -313,7 +313,7 @@ void runner_do_star_formation(struct runner *r, struct cell *c, int timer) {
             /* Did we get a star? (Or did we run out of spare ones?) */
             if (sp != NULL) {
 
-              /* message("We formed a star id=%lld cellID=%d", sp->id,
+              /* message("We formed a star id=%lld cellID=%lld", sp->id,
                * c->cellID); */
 
               /* Copy the properties of the gas particle to the star particle */
@@ -496,6 +496,7 @@ void runner_do_sink_formation(struct runner *r, struct cell *c) {
 void runner_do_end_hydro_force(struct runner *r, struct cell *c, int timer) {
 
   const struct engine *e = r->e;
+  const int with_cosmology = e->policy & engine_policy_cosmology;
 
   TIMER_TIC;
 
@@ -518,12 +519,24 @@ void runner_do_end_hydro_force(struct runner *r, struct cell *c, int timer) {
       /* Get a handle on the part. */
       struct part *restrict p = &parts[k];
 
+      double dt = 0;
       if (part_is_active(p, e)) {
 
+        if (with_cosmology) {
+          /* Compute the time step. */
+          const integertime_t ti_step = get_integer_timestep(p->time_bin);
+          const integertime_t ti_begin =
+              get_integer_time_begin(e->ti_current - 1, p->time_bin);
+
+          dt = cosmology_get_delta_time(cosmo, ti_begin, ti_begin + ti_step);
+        } else {
+          dt = get_timestep(p->time_bin, e->time_base);
+        }
+
         /* Finish the force loop */
         hydro_end_force(p, cosmo);
         timestep_limiter_end_force(p);
-        chemistry_end_force(p, cosmo);
+        chemistry_end_force(p, cosmo, with_cosmology, e->time, dt);
 
 #ifdef SWIFT_BOUNDARY_PARTICLES
 
diff --git a/src/runner_time_integration.c b/src/runner_time_integration.c
index 656490bd3491a5166936af682623e6de668b28fc..bfb2611e5f4c2f9545465aca23792dbd839954e4 100644
--- a/src/runner_time_integration.c
+++ b/src/runner_time_integration.c
@@ -177,6 +177,14 @@ void runner_do_kick1(struct runner *r, struct cell *c, const int timer) {
         kick_part(p, xp, dt_kick_hydro, dt_kick_grav, dt_kick_mesh_grav,
                   dt_kick_therm, dt_kick_corr, cosmo, hydro_props,
                   entropy_floor, ti_begin, ti_end, ti_begin_mesh, ti_end_mesh);
+
+        /* Update the accelerations to be used in the drift for hydro */
+        if (p->gpart != NULL) {
+
+          xp->a_grav[0] = p->gpart->a_grav[0] + p->gpart->a_grav_mesh[0];
+          xp->a_grav[1] = p->gpart->a_grav[1] + p->gpart->a_grav_mesh[1];
+          xp->a_grav[2] = p->gpart->a_grav[2] + p->gpart->a_grav_mesh[2];
+        }
       }
     }
 
@@ -624,6 +632,10 @@ void runner_do_timestep(struct runner *r, struct cell *c, const int timer) {
 
   const struct engine *e = r->e;
   const integertime_t ti_current = e->ti_current;
+  const struct cosmology *cosmo = e->cosmology;
+  const struct feedback_props *feedback_props = e->feedback_props;
+  const struct unit_system *us = e->internal_units;
+  const struct phys_const *phys_const = e->physical_constants;
   const int with_cosmology = (e->policy & engine_policy_cosmology);
   const int with_feedback = (e->policy & engine_policy_feedback);
   const int count = c->hydro.count;
@@ -839,9 +851,44 @@ void runner_do_timestep(struct runner *r, struct cell *c, const int timer) {
         sp->gpart->time_bin = get_time_bin(ti_new_step);
 
         /* Update feedback related counters */
-        if (with_feedback)
-          feedback_will_do_feedback(sp, e->feedback_props, with_cosmology,
-                                    e->cosmology, e->time);
+        if (with_feedback) {
+          const integertime_t ti_step = get_integer_timestep(sp->time_bin);
+          const integertime_t ti_begin_star =
+              get_integer_time_begin(e->ti_current, sp->time_bin);
+
+          /* Get particle time-step */
+          double dt_star;
+          if (with_cosmology) {
+            dt_star = cosmology_get_delta_time(e->cosmology, ti_begin_star,
+                                               ti_begin_star + ti_step);
+          } else {
+            dt_star = get_timestep(sp->time_bin, e->time_base);
+          }
+
+          /* Calculate age of the star at current time */
+          double star_age_end_of_step;
+          if (with_cosmology) {
+            if (cosmo->a > (double)sp->birth_scale_factor)
+              star_age_end_of_step =
+                  cosmology_get_delta_time_from_scale_factors(
+                      cosmo, (double)sp->birth_scale_factor, cosmo->a);
+            else
+              star_age_end_of_step = 0.;
+          } else {
+            star_age_end_of_step = max(e->time - (double)sp->birth_time, 0.);
+          }
+
+          /* Get the length of the enrichment time-step */
+          const double dt_enrichment = feedback_get_enrichment_timestep(
+              sp, with_cosmology, cosmo, e->time, dt_star);
+          const double star_age_beg_of_step =
+              star_age_end_of_step - dt_enrichment;
+
+          /* Compute the stellar evolution  */
+          feedback_will_do_feedback(
+              sp, feedback_props, with_cosmology, cosmo, e->time, us,
+              phys_const, star_age_beg_of_step, dt_enrichment, ti_begin_star);
+        }
 
         /* Number of updated s-particles */
         s_updated++;
@@ -1178,7 +1225,7 @@ void runner_do_limiter(struct runner *r, struct cell *c, int force,
       /* Bip, bip, bip... wake-up time */
       if (p->limiter_data.wakeup != time_bin_not_awake) {
 
-        // message("Limiting particle %lld in cell %d", p->id, c->cellID);
+        // message("Limiting particle %lld in cell %lld", p->id, c->cellID);
 
         /* Apply the limiter and get the new end of time-step */
         const integertime_t ti_end_new = timestep_limit_part(p, xp, e);
diff --git a/src/serial_io.c b/src/serial_io.c
index eb801c8f10a21f1336e6dc13897f57925a0905bb..c76b0c23ea777beba7500ec523838a87ba522426 100644
--- a/src/serial_io.c
+++ b/src/serial_io.c
@@ -1094,6 +1094,8 @@ void write_output_serial(struct engine* e,
                        numParticlesHighWord, swift_type_count);
     double MassTable[6] = {0., 0., 0., 0., 0., 0.};
     io_write_attribute(h_grp, "MassTable", DOUBLE, MassTable, swift_type_count);
+    io_write_attribute(h_grp, "InitialMassTable", DOUBLE,
+                       e->s->initial_mean_mass_particles, swift_type_count);
     unsigned int flagEntropy[swift_type_count] = {0};
     flagEntropy[0] = writeEntropyFlag();
     io_write_attribute(h_grp, "Flag_Entropy_ICs", UINT, flagEntropy,
diff --git a/src/single_io.c b/src/single_io.c
index e0912ba45e53945a3ee98d9d945343df77292b53..b46e331e6cc506d5a3faa895f2518f2d3a94f9fc 100644
--- a/src/single_io.c
+++ b/src/single_io.c
@@ -933,6 +933,8 @@ void write_output_single(struct engine* e,
                      numParticlesHighWord, swift_type_count);
   double MassTable[swift_type_count] = {0};
   io_write_attribute(h_grp, "MassTable", DOUBLE, MassTable, swift_type_count);
+  io_write_attribute(h_grp, "InitialMassTable", DOUBLE,
+                     e->s->initial_mean_mass_particles, swift_type_count);
   unsigned int flagEntropy[swift_type_count] = {0};
   flagEntropy[0] = writeEntropyFlag();
   io_write_attribute(h_grp, "Flag_Entropy_ICs", UINT, flagEntropy,
diff --git a/src/space.c b/src/space.c
index 9c8afa141aed1d638c94da3fea4ac5d390074cfd..52fbdeb96b9b861c1de23431397258bdddfa95c9 100644
--- a/src/space.c
+++ b/src/space.c
@@ -94,9 +94,9 @@ int engine_star_resort_task_depth = engine_star_resort_task_depth_default;
 /*! Expected maximal number of strays received at a rebuild */
 int space_expected_max_nr_strays = space_expected_max_nr_strays_default;
 
-/*! Counter for cell IDs (when debugging) */
+/*! Counter for cell IDs (when debugging + max vals for unique IDs exceeded) */
 #if defined(SWIFT_DEBUG_CHECKS) || defined(SWIFT_CELL_GRAPH)
-int last_cell_id;
+long long last_cell_id;
 #endif
 
 /**
@@ -842,6 +842,140 @@ void space_convert_quantities(struct space *s, int verbose) {
             clocks_getunit());
 }
 
+void space_collect_sum_part_mass(void *restrict map_data, int count,
+                                 void *restrict extra_data) {
+
+  struct space *s = (struct space *)extra_data;
+  const struct part *parts = (const struct part *)map_data;
+
+  /* Local collection */
+  double sum = 0.;
+  for (int i = 0; i < count; ++i) sum += hydro_get_mass(&parts[i]);
+
+  /* Store back */
+  atomic_add_d(&s->initial_mean_mass_particles[0], sum);
+  atomic_add(&s->initial_count_particles[0], count);
+}
+
+void space_collect_sum_gpart_mass(void *restrict map_data, int count,
+                                  void *restrict extra_data) {
+
+  struct space *s = (struct space *)extra_data;
+  const struct gpart *gparts = (const struct gpart *)map_data;
+
+  /* Local collection */
+  double sum_DM = 0., sum_DM_background = 0.;
+  long long count_DM = 0, count_DM_background = 0;
+  for (int i = 0; i < count; ++i) {
+    if (gparts[i].type == swift_type_dark_matter) {
+      sum_DM += gparts[i].mass;
+      count_DM++;
+    }
+    if (gparts[i].type == swift_type_dark_matter_background) {
+      sum_DM_background += gparts[i].mass;
+      count_DM_background++;
+    }
+  }
+
+  /* Store back */
+  atomic_add_d(&s->initial_mean_mass_particles[1], sum_DM);
+  atomic_add(&s->initial_count_particles[1], count_DM);
+  atomic_add_d(&s->initial_mean_mass_particles[2], sum_DM_background);
+  atomic_add(&s->initial_count_particles[2], count_DM_background);
+}
+
+void space_collect_sum_sink_mass(void *restrict map_data, int count,
+                                 void *restrict extra_data) {
+
+  struct space *s = (struct space *)extra_data;
+  const struct sink *sinks = (const struct sink *)map_data;
+
+  /* Local collection */
+  double sum = 0.;
+  for (int i = 0; i < count; ++i) sum += sinks[i].mass;
+
+  /* Store back */
+  atomic_add_d(&s->initial_mean_mass_particles[3], sum);
+  atomic_add(&s->initial_count_particles[3], count);
+}
+
+void space_collect_sum_spart_mass(void *restrict map_data, int count,
+                                  void *restrict extra_data) {
+
+  struct space *s = (struct space *)extra_data;
+  const struct spart *sparts = (const struct spart *)map_data;
+
+  /* Local collection */
+  double sum = 0.;
+  for (int i = 0; i < count; ++i) sum += sparts[i].mass;
+
+  /* Store back */
+  atomic_add_d(&s->initial_mean_mass_particles[4], sum);
+  atomic_add(&s->initial_count_particles[4], count);
+}
+
+void space_collect_sum_bpart_mass(void *restrict map_data, int count,
+                                  void *restrict extra_data) {
+
+  struct space *s = (struct space *)extra_data;
+  const struct bpart *bparts = (const struct bpart *)map_data;
+
+  /* Local collection */
+  double sum = 0.;
+  for (int i = 0; i < count; ++i) sum += bparts[i].mass;
+
+  /* Store back */
+  atomic_add_d(&s->initial_mean_mass_particles[5], sum);
+  atomic_add(&s->initial_count_particles[5], count);
+}
+
+/**
+ * @breif Collect the mean mass of each particle type in the #space.
+ */
+void space_collect_mean_masses(struct space *s, int verbose) {
+
+  /* Init counters */
+  for (int i = 0; i < swift_type_count; ++i)
+    s->initial_mean_mass_particles[i] = 0.;
+  for (int i = 0; i < swift_type_count; ++i) s->initial_count_particles[i] = 0;
+
+  /* Collect each particle type */
+  threadpool_map(&s->e->threadpool, space_collect_sum_part_mass, s->parts,
+                 s->nr_parts, sizeof(struct part), threadpool_auto_chunk_size,
+                 s);
+  threadpool_map(&s->e->threadpool, space_collect_sum_gpart_mass, s->gparts,
+                 s->nr_gparts, sizeof(struct gpart), threadpool_auto_chunk_size,
+                 s);
+  threadpool_map(&s->e->threadpool, space_collect_sum_spart_mass, s->sparts,
+                 s->nr_sparts, sizeof(struct spart), threadpool_auto_chunk_size,
+                 s);
+  threadpool_map(&s->e->threadpool, space_collect_sum_sink_mass, s->sinks,
+                 s->nr_sinks, sizeof(struct sink), threadpool_auto_chunk_size,
+                 s);
+  threadpool_map(&s->e->threadpool, space_collect_sum_bpart_mass, s->bparts,
+                 s->nr_bparts, sizeof(struct bpart), threadpool_auto_chunk_size,
+                 s);
+
+#ifdef WITH_MPI
+  MPI_Allreduce(MPI_IN_PLACE, s->initial_mean_mass_particles, swift_type_count,
+                MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);
+  MPI_Allreduce(MPI_IN_PLACE, s->initial_count_particles, swift_type_count,
+                MPI_LONG_LONG, MPI_SUM, MPI_COMM_WORLD);
+#endif
+
+  /* Get means
+   *
+   * Note: the Intel compiler vectorizes this loop and creates FPEs from
+   * the masked bit of the vector... Silly ICC... */
+#if defined(__ICC)
+#pragma novector
+#endif
+  for (int i = 0; i < swift_type_count; ++i)
+    if (s->initial_count_particles[i] > 0)
+      s->initial_mean_mass_particles[i] /=
+          (double)s->initial_count_particles[i];
+}
+
 /**
  * @brief Split the space into cells given the array of particles.
  *
@@ -2284,24 +2418,24 @@ void space_write_cell(const struct space *s, FILE *f, const struct cell *c) {
   if (c == NULL) return;
 
   /* Get parent ID */
-  int parent = root_cell_id;
+  long long parent = root_cell_id;
   if (c->parent != NULL) parent = c->parent->cellID;
 
   /* Get super ID */
   char superID[100] = "";
-  if (c->super != NULL) sprintf(superID, "%i", c->super->cellID);
+  if (c->super != NULL) sprintf(superID, "%lld", c->super->cellID);
 
   /* Get hydro super ID */
   char hydro_superID[100] = "";
   if (c->hydro.super != NULL)
-    sprintf(hydro_superID, "%i", c->hydro.super->cellID);
+    sprintf(hydro_superID, "%lld", c->hydro.super->cellID);
 
   /* Write line for current cell */
-  fprintf(f, "%i,%i,%i,", c->cellID, parent, c->nodeID);
+  fprintf(f, "%lld,%lld,%i,", c->cellID, parent, c->nodeID);
   fprintf(f, "%i,%i,%i,%s,%s,%g,%g,%g,%g,%g,%g, ", c->hydro.count,
           c->stars.count, c->grav.count, superID, hydro_superID, c->loc[0],
           c->loc[1], c->loc[2], c->width[0], c->width[1], c->width[2]);
-  fprintf(f, "%g, %g %i %i\n", c->hydro.h_max, c->stars.h_max, c->depth,
+  fprintf(f, "%g, %g, %i, %i\n", c->hydro.h_max, c->stars.h_max, c->depth,
           c->maxdepth);
 
   /* Write children */
@@ -2338,9 +2472,9 @@ void space_write_cell_hierarchy(const struct space *s, int j) {
 
     /* Write root data */
     fprintf(f, "%i, ,-1,", root_id);
-    fprintf(f, "%li,%li,%li, , , , , , , , , ", s->nr_parts, s->nr_sparts,
+    fprintf(f, "%li,%li,%li, , , , , , , , ,", s->nr_parts, s->nr_sparts,
             s->nr_gparts);
-    fprintf(f, ",\n");
+    fprintf(f, " , , ,\n");
   }
 
   /* Write all the top level cells (and their children) */
diff --git a/src/space.h b/src/space.h
index 17949392f46f63ff8d980a5fcc8da1236dad1c45..c726a68d9873f7473c191f649fb3d7fa53b10400 100644
--- a/src/space.h
+++ b/src/space.h
@@ -276,6 +276,12 @@ struct space {
   /*! Sum of the norm of the velocity of all the #bpart */
   float sum_bpart_vel_norm;
 
+  /* Initial mean mass of each particle type in the system. */
+  double initial_mean_mass_particles[swift_type_count];
+
+  /* Initial count of each particle type in the system. */
+  long long initial_count_particles[swift_type_count];
+
   /*! Initial value of the smoothing length read from the parameter file */
   float initial_spart_h;
 
@@ -385,6 +391,7 @@ void space_first_init_gparts(struct space *s, int verbose);
 void space_first_init_sparts(struct space *s, int verbose);
 void space_first_init_bparts(struct space *s, int verbose);
 void space_first_init_sinks(struct space *s, int verbose);
+void space_collect_mean_masses(struct space *s, int verbose);
 void space_init_parts(struct space *s, int verbose);
 void space_init_gparts(struct space *s, int verbose);
 void space_init_sparts(struct space *s, int verbose);
diff --git a/src/space_extras.c b/src/space_extras.c
index 5eda872340a14534b1f561ee9b81337574aac384..8890251eb5d4af1909a2455e95cf357dec56b4b2 100644
--- a/src/space_extras.c
+++ b/src/space_extras.c
@@ -117,15 +117,25 @@ void space_allocate_extras(struct space *s, int verbose) {
   }
 
   if (expected_num_extra_parts < s->nr_extra_parts)
-    error("Reduction in top-level cells number not handled.");
+    error(
+        "Reduction in top-level cells number not handled. Please set a lower "
+        "h_max or reduce the number of top level cells.");
   if (expected_num_extra_gparts < s->nr_extra_gparts)
-    error("Reduction in top-level cells number not handled.");
+    error(
+        "Reduction in top-level cells number not handled. Please set a lower "
+        "h_max or reduce the number of top level cells.");
   if (expected_num_extra_sparts < s->nr_extra_sparts)
-    error("Reduction in top-level cells number not handled.");
+    error(
+        "Reduction in top-level cells number not handled. Please set a lower "
+        "h_max or reduce the number of top level cells.");
   if (expected_num_extra_bparts < s->nr_extra_bparts)
-    error("Reduction in top-level cells number not handled.");
+    error(
+        "Reduction in top-level cells number not handled. Please set a lower "
+        "h_max or reduce the number of top level cells.");
   if (expected_num_extra_sinks < s->nr_extra_sinks)
-    error("Reduction in top-level cells number not handled.");
+    error(
+        "Reduction in top-level cells number not handled. Please set a lower "
+        "h_max or reduce the number of top level cells.");
 
   /* Do we have enough space for the extra gparts (i.e. we haven't used up any)
    * ? */
diff --git a/src/space_rebuild.c b/src/space_rebuild.c
index f456be124731830a7c91f65caca4d1517644d701..0b68ca2b0cf9cfaea16ba95908fa38cacff71a25 100644
--- a/src/space_rebuild.c
+++ b/src/space_rebuild.c
@@ -35,7 +35,7 @@ extern int space_expected_max_nr_strays;
 
 /*! Counter for cell IDs (when debugging) */
 #if defined(SWIFT_DEBUG_CHECKS) || defined(SWIFT_CELL_GRAPH)
-extern int last_cell_id;
+extern long long last_cell_id;
 #endif
 
 /**
@@ -911,8 +911,7 @@ void space_rebuild(struct space *s, int repartitioned, int verbose) {
     c->black_holes.ti_old_part = ti_current;
 
 #if defined(SWIFT_DEBUG_CHECKS) || defined(SWIFT_CELL_GRAPH)
-    c->cellID = -last_cell_id;
-    last_cell_id++;
+    cell_assign_top_level_cell_index(c, s->cdim, s->dim, s->width);
 #endif
 
     const int is_local = (c->nodeID == engine_rank);
diff --git a/src/space_regrid.c b/src/space_regrid.c
index b20a3b4894bf3ce6418cd288f74fe91cedc36ab7..e9ca3e238e9d8e3cdaea5a7ac14d9b359e958965 100644
--- a/src/space_regrid.c
+++ b/src/space_regrid.c
@@ -30,11 +30,6 @@
 #include "engine.h"
 #include "scheduler.h"
 
-/*! Counter for cell IDs (when debugging) */
-#if defined(SWIFT_DEBUG_CHECKS) || defined(SWIFT_CELL_GRAPH)
-extern int last_cell_id;
-#endif
-
 /**
  * @brief Re-build the top-level cell grid.
  *
@@ -149,8 +144,8 @@ void space_regrid(struct space *s, int verbose) {
  * global partition is recomputed and the particles redistributed.
  * Be prepared to do that. */
 #ifdef WITH_MPI
-  double oldwidth[3];
-  double oldcdim[3];
+  double oldwidth[3] = {0., 0., 0.};
+  double oldcdim[3] = {0., 0., 0.};
   int *oldnodeIDs = NULL;
   if (cdim[0] < s->cdim[0] || cdim[1] < s->cdim[1] || cdim[2] < s->cdim[2]) {
 
@@ -321,8 +316,7 @@ void space_regrid(struct space *s, int verbose) {
 #endif  // WITH_MPI
           if (s->with_self_gravity) c->grav.multipole = &s->multipoles_top[cid];
 #if defined(SWIFT_DEBUG_CHECKS) || defined(SWIFT_CELL_GRAPH)
-          c->cellID = -last_cell_id;
-          last_cell_id++;
+          cell_assign_top_level_cell_index(c, s->cdim, s->dim, s->width);
 #endif
         }
 
diff --git a/src/space_split.c b/src/space_split.c
index 728c6f8984b3476bfc0729546059477c3d1f38cf..38d75ad48b3e10c21d7047bb6471f41ee16bd881 100644
--- a/src/space_split.c
+++ b/src/space_split.c
@@ -33,11 +33,6 @@
 #include "star_formation_logger.h"
 #include "threadpool.h"
 
-/*! Counter for cell IDs (when debugging) */
-#if defined(SWIFT_DEBUG_CHECKS) || defined(SWIFT_CELL_GRAPH)
-extern int last_cell_id;
-#endif
-
 /**
  * @brief Recursively split a cell.
  *
@@ -250,7 +245,7 @@ void space_split_recursive(struct space *s, struct cell *c,
       cp->mpi.tag = -1;
 #endif  // WITH_MPI
 #if defined(SWIFT_DEBUG_CHECKS) || defined(SWIFT_CELL_GRAPH)
-      cp->cellID = last_cell_id++;
+      cell_assign_cell_index(cp, c);
 #endif
     }
 
diff --git a/src/star_formation/EAGLE/star_formation.h b/src/star_formation/EAGLE/star_formation.h
index 318c2610e12c8d4c6561b8f81de22aa5e27f4e9a..9794d71efe70987346791a470ae3280268344534 100644
--- a/src/star_formation/EAGLE/star_formation.h
+++ b/src/star_formation/EAGLE/star_formation.h
@@ -905,6 +905,20 @@ INLINE static void starformation_print_backend(
           starform->gas_density_direct_HpCM3);
 }
 
+/**
+ * @brief Return the star formation rate of a particle.
+ *
+ * @param p The particle.
+ * @param xp The extended data of the particle.
+ */
+INLINE static float star_formation_get_SFR(const struct part* p,
+                                           const struct xpart* xp) {
+  if (xp->sf_data.SFR <= 0.)
+    return 0.f;
+  else
+    return xp->sf_data.SFR;
+}
+
 /**
  * @brief Finishes the density calculation.
  *
diff --git a/src/star_formation/GEAR/star_formation.h b/src/star_formation/GEAR/star_formation.h
index 341ed815b1699fef9ad87d30d44e2c3b0a08f459..c69c7a962dd246289cb8ba26df5858d10d4497b3 100644
--- a/src/star_formation/GEAR/star_formation.h
+++ b/src/star_formation/GEAR/star_formation.h
@@ -27,6 +27,7 @@
 #include "engine.h"
 #include "entropy_floor.h"
 #include "error.h"
+#include "feedback.h"
 #include "hydro_properties.h"
 #include "parser.h"
 #include "part.h"
@@ -307,6 +308,9 @@ INLINE static void star_formation_copy_properties(
     const struct cooling_function_data* restrict cooling,
     const int convert_part) {
 
+  /* Initialize the feedback */
+  feedback_init_after_star_formation(sp, e->feedback_props);
+
   /* Store the current mass */
   const float mass_gas = hydro_get_mass(p);
   if (!convert_part) {
@@ -361,6 +365,19 @@ INLINE static void starformation_print_backend(
   message("Star formation law is 'GEAR'");
 }
 
+/**
+ * @brief Return the star formation rate of a particle.
+ *
+ * This scheme does not store the SFR in the particles. We return 0.
+ *
+ * @param p The particle.
+ * @param xp The extended data of the particle.
+ */
+INLINE static float star_formation_get_SFR(const struct part* p,
+                                           const struct xpart* xp) {
+  return 0.f;
+}
+
 /**
  * @brief Finishes the density calculation.
  *
diff --git a/src/star_formation/QLA/star_formation.h b/src/star_formation/QLA/star_formation.h
index dd0244bf7731c7809a3003502e8b6e9c75f6b370..91f5c0592d9ea8aefa111c533234b9e072bdf852 100644
--- a/src/star_formation/QLA/star_formation.h
+++ b/src/star_formation/QLA/star_formation.h
@@ -210,6 +210,20 @@ INLINE static void starformation_print_backend(
   message("Over-density for star formation: %f", starform->over_density);
 }
 
+/**
+ * @brief Return the star formation rate of a particle.
+ *
+ * In this simple model, particles are directly turned into stars. There is no
+ * point at which SFR > 0.
+ *
+ * @param p The particle.
+ * @param xp The extended data of the particle.
+ */
+INLINE static float star_formation_get_SFR(const struct part* p,
+                                           const struct xpart* xp) {
+  return 0.f;
+}
+
 /**
  * @brief Finishes the density calculation.
  *
diff --git a/src/star_formation/none/star_formation.h b/src/star_formation/none/star_formation.h
index 57544c837368cb0c4290f2dbdbaedebd369d971d..db6ed8658afac741ca26c86eefd4b24572976a6b 100644
--- a/src/star_formation/none/star_formation.h
+++ b/src/star_formation/none/star_formation.h
@@ -189,6 +189,19 @@ INLINE static void starformation_print_backend(
   message("Star formation law is 'No Star Formation'");
 }
 
+/**
+ * @brief Return the star formation rate of a particle.
+ *
+ * No star formation in this model --> we return 0
+ *
+ * @param p The particle.
+ * @param xp The extended data of the particle.
+ */
+INLINE static float star_formation_get_SFR(const struct part* p,
+                                           const struct xpart* xp) {
+  return 0.f;
+}
+
 /**
  * @brief Finishes the density calculation.
  *
diff --git a/src/stars/Default/stars_part.h b/src/stars/Default/stars_part.h
index 049795a5753113e839c98164d9a7c2fea5f97ee4..6dfc5960d2745b19136199ddf2cd51c7369a0d2d 100644
--- a/src/stars/Default/stars_part.h
+++ b/src/stars/Default/stars_part.h
@@ -89,7 +89,7 @@ struct spart {
   struct tracers_xpart_data tracers_data;
 
   /*! Chemistry structure */
-  struct chemistry_part_data chemistry_data;
+  struct chemistry_spart_data chemistry_data;
 
 #ifdef WITH_LOGGER
   /* Additional data for the particle logger */
diff --git a/src/stars/EAGLE/stars_part.h b/src/stars/EAGLE/stars_part.h
index 07072add21eb5d8dae41afcad4da6478a0554fde..63a995dcfccbc30e04669e61df42ab375fb74d93 100644
--- a/src/stars/EAGLE/stars_part.h
+++ b/src/stars/EAGLE/stars_part.h
@@ -111,7 +111,7 @@ struct spart {
   struct tracers_xpart_data tracers_data;
 
   /*! Chemistry structure */
-  struct chemistry_part_data chemistry_data;
+  struct chemistry_spart_data chemistry_data;
 
   /*! Particle time bin */
   timebin_t time_bin;
diff --git a/src/threadpool.c b/src/threadpool.c
index d147533960172f7f219eb2a592464f547ad76c32..932c74a271f8c27d7ba64271017c0adb304c2ab1 100644
--- a/src/threadpool.c
+++ b/src/threadpool.c
@@ -43,8 +43,8 @@
 /**
  * @brief Store a log entry of the given chunk.
  */
-void threadpool_log(struct threadpool *tp, int tid, size_t chunk_size,
-                    ticks tic, ticks toc) {
+static void threadpool_log(struct threadpool *tp, int tid, size_t chunk_size,
+                           ticks tic, ticks toc) {
   struct mapper_log *log = &tp->logs[tid > 0 ? tid : 0];
 
   /* Check if we need to re-allocate the log buffer. */
@@ -132,15 +132,15 @@ void threadpool_dump_log(struct threadpool *tp, const char *filename,
 /**
  * @brief Runner main loop, get a chunk and call the mapper function.
  */
-void threadpool_chomp(struct threadpool *tp, int tid) {
+static void threadpool_chomp(struct threadpool *tp, int tid) {
 
   /* Loop until we can't get a chunk. */
   while (1) {
     /* Compute the desired chunk size. */
     ptrdiff_t chunk_size;
     if (tp->map_data_chunk == threadpool_uniform_chunk_size) {
-      chunk_size = (int)((tid + 1) * tp->map_data_size / tp->num_threads) -
-                   (int)(tid * tp->map_data_size / tp->num_threads);
+      chunk_size = ((tid + 1) * tp->map_data_size / tp->num_threads) -
+                   (tid * tp->map_data_size / tp->num_threads);
     } else {
       chunk_size =
           (tp->map_data_size - tp->map_data_count) / (2 * tp->num_threads);
@@ -148,6 +148,9 @@ void threadpool_chomp(struct threadpool *tp, int tid) {
     }
     if (chunk_size < 1) chunk_size = 1;
 
+    /* A chunk cannot exceed INT_MAX, as we use int elements in map_function. */
+    if (chunk_size > INT_MAX) chunk_size = INT_MAX;
+
     /* Get a chunk and check its size. */
     size_t task_ind = atomic_add(&tp->map_data_count, chunk_size);
     if (task_ind >= tp->map_data_size) break;
@@ -166,7 +169,7 @@ void threadpool_chomp(struct threadpool *tp, int tid) {
   }
 }
 
-void *threadpool_runner(void *data) {
+static void *threadpool_runner(void *data) {
 
   /* Our threadpool. */
   struct threadpool *tp = (struct threadpool *)data;
@@ -270,16 +273,45 @@ void threadpool_map(struct threadpool *tp, threadpool_map_function map_function,
                     void *extra_data) {
 
 #ifdef SWIFT_DEBUG_THREADPOOL
-  ticks tic = getticks();
+  ticks tic_total = getticks();
 #endif
 
   /* If we just have a single thread, call the map function directly. */
   if (tp->num_threads == 1) {
-    map_function(map_data, N, extra_data);
+
+    if (N <= INT_MAX) {
+      map_function(map_data, N, extra_data);
+
+#ifdef SWIFT_DEBUG_THREADPOOL
+      tp->map_function = map_function;
+      threadpool_log(tp, 0, N, tic_total, getticks());
+#endif
+    } else {
+
+      /* N > INT_MAX, we need to do this in chunks as map_function only takes
+       * an int. */
+      size_t chunk_size = INT_MAX;
+      size_t data_size = N;
+      size_t data_count = 0;
+      while (1) {
+
+/* Call the mapper function. */
+#ifdef SWIFT_DEBUG_THREADPOOL
+        ticks tic = getticks();
+#endif
+        map_function((char *)map_data + (stride * data_count), chunk_size,
+                     extra_data);
 #ifdef SWIFT_DEBUG_THREADPOOL
-    tp->map_function = map_function;
-    threadpool_log(tp, 0, N, tic, getticks());
+        threadpool_log(tp, 0, chunk_size, tic, getticks());
 #endif
+        /* Get the next chunk and check its size. */
+        data_count += chunk_size;
+        if (data_count >= data_size) break;
+        if (data_count + chunk_size > data_size)
+          chunk_size = data_size - data_count;
+      }
+    }
+
     return;
   }
 
@@ -289,7 +321,7 @@ void threadpool_map(struct threadpool *tp, threadpool_map_function map_function,
   tp->map_data_count = 0;
   if (chunk == threadpool_auto_chunk_size) {
     tp->map_data_chunk =
-        max((int)(N / (tp->num_threads * threadpool_default_chunk_ratio)), 1);
+        max((N / (tp->num_threads * threadpool_default_chunk_ratio)), 1U);
   } else if (chunk == threadpool_uniform_chunk_size) {
     tp->map_data_chunk = threadpool_uniform_chunk_size;
   } else {
@@ -311,7 +343,7 @@ void threadpool_map(struct threadpool *tp, threadpool_map_function map_function,
 
 #ifdef SWIFT_DEBUG_THREADPOOL
   /* Log the total call time to thread id -1. */
-  threadpool_log(tp, -1, N, tic, getticks());
+  threadpool_log(tp, -1, N, tic_total, getticks());
 #endif
 }
 
diff --git a/src/units.c b/src/units.c
index 00c53d89953b38d09de876aea6cccf167b79f8cb..074b553c3da6f0c5209dcc2a98a528c9f70ec838 100644
--- a/src/units.c
+++ b/src/units.c
@@ -329,6 +329,12 @@ void units_get_base_unit_exponents_array(float baseUnitsExp[5],
       baseUnitsExp[UNIT_TIME] = -2.f;
       break;
 
+    case UNIT_CONV_PHYSICAL_ENTROPY_PER_UNIT_MASS:
+      baseUnitsExp[UNIT_LENGTH] = 2.f;
+      baseUnitsExp[UNIT_TIME] = -2.f;
+      baseUnitsExp[UNIT_TEMPERATURE] = -1.f;
+      break;
+
     case UNIT_CONV_POWER:
       baseUnitsExp[UNIT_MASS] = 1.f;
       baseUnitsExp[UNIT_LENGTH] = 2.f;
diff --git a/src/units.h b/src/units.h
index 87ab963e0ac640521501475121fa77ad1faaf2bc..933ab6a66c69511d43a901739a407ff7e14d3a43 100644
--- a/src/units.h
+++ b/src/units.h
@@ -84,6 +84,7 @@ enum unit_conversion_factor {
   UNIT_CONV_ENERGY_PER_UNIT_MASS,
   UNIT_CONV_ENTROPY,
   UNIT_CONV_ENTROPY_PER_UNIT_MASS,
+  UNIT_CONV_PHYSICAL_ENTROPY_PER_UNIT_MASS,
   UNIT_CONV_POWER,
   UNIT_CONV_PRESSURE,
   UNIT_CONV_FREQUENCY,
diff --git a/tests/testSymmetry.c b/tests/testSymmetry.c
index a5a1be2ce21d42af6e86513fb7293a6376a351de..3a3359d9cc33c872758b8ddf9285859b0ac7e8da 100644
--- a/tests/testSymmetry.c
+++ b/tests/testSymmetry.c
@@ -46,6 +46,8 @@ void test(void) {
   /* Start with some values for the cosmological parameters */
   const float a = (float)random_uniform(0.8, 1.);
   const float H = 1.f;
+  const integertime_t ti_current = 1;
+  const double time_base = 1e-5;
 
   /* Create two random particles (don't do this at home !) */
   struct part pi, pj;
@@ -219,15 +221,23 @@ void test(void) {
 
   /* Call the symmetric version */
   runner_iact_force(r2, dx, pi.h, pj.h, &pi, &pj, a, H);
+  runner_iact_diffusion(r2, dx, pi.h, pj.h, &pi, &pj, a, H, time_base,
+                        ti_current, NULL, /*with_cosmology=*/0);
   runner_iact_timebin(r2, dx, pi.h, pj.h, &pi, &pj, a, H);
 
   /* Call the non-symmetric version */
   runner_iact_nonsym_force(r2, dx, pi2.h, pj2.h, &pi2, &pj2, a, H);
+  runner_iact_nonsym_diffusion(r2, dx, pi2.h, pj2.h, &pi2, &pj2, a, H,
+                               time_base, ti_current, NULL,
+                               /*with_cosmology=*/0);
   runner_iact_nonsym_timebin(r2, dx, pi2.h, pj2.h, &pi2, &pj2, a, H);
   dx[0] = -dx[0];
   dx[1] = -dx[1];
   dx[2] = -dx[2];
   runner_iact_nonsym_force(r2, dx, pj2.h, pi2.h, &pj2, &pi2, a, H);
+  runner_iact_nonsym_diffusion(r2, dx, pj2.h, pi2.h, &pj2, &pi2, a, H,
+                               time_base, ti_current, NULL,
+                               /*with_cosmology=*/0);
   runner_iact_nonsym_timebin(r2, dx, pj2.h, pi2.h, &pj2, &pi2, a, H);
 
 /* Check that the particles are the same */