diff --git a/configure.ac b/configure.ac
index de2cbc2cf7e690440be6e46ac2216235af9ab5fa..4c547d3466dac2afc45087396977af46c9604551 100644
--- a/configure.ac
+++ b/configure.ac
@@ -1489,7 +1489,7 @@ AC_SUBST([NUMA_INCS])
 # As an example for this, see the call to AC_ARG_WITH for cooling.
 AC_ARG_WITH([subgrid],
 	[AS_HELP_STRING([--with-subgrid=<subgrid>],
-		[Master switch for subgrid methods. Inexperienced user should start here. Options are: @<:@none, GEAR, QLA, EAGLE default: none@:>@]
+		[Master switch for subgrid methods. Inexperienced user should start here. Options are: @<:@none, GEAR, QLA, EAGLE, EAGLE-XL default: none@:>@]
 	)],
 	[with_subgrid="$withval"],
 	[with_subgrid=none]
diff --git a/examples/QuickLymanAlpha/L050N0752/qla_50.yml b/examples/QuickLymanAlpha/L050N0752/qla_50.yml
index 0748fa112b10c5d8e93502f69984bbbda54edca1..e62bccc526fa6448a7d3f0286da46d2951aa2333 100644
--- a/examples/QuickLymanAlpha/L050N0752/qla_50.yml
+++ b/examples/QuickLymanAlpha/L050N0752/qla_50.yml
@@ -86,12 +86,13 @@ LineOfSight:
 
 # Quick Lyman-alpha cooling (EAGLE with fixed primoridal Z)
 QLACooling:
-  dir_name:                ./coolingtables/
-  H_reion_z:               7.5                 # Planck 2018
+  dir_name:                ./UV_dust1_CR1_G1_shield1.hdf5 # Location of the cooling tables
+  H_reion_z:               7.5               # Redshift of Hydrogen re-ionization (Planck 2018)
   H_reion_eV_p_H:          2.0
-  He_reion_z_centre:       3.5
-  He_reion_z_sigma:        0.5
-  He_reion_eV_p_H:         2.0
+  He_reion_z_centre:       3.5               # Redshift of the centre of the Helium re-ionization Gaussian
+  He_reion_z_sigma:        0.5               # Spread in redshift of the  Helium re-ionization Gaussian
+  He_reion_eV_p_H:         2.0               # Energy inject by Helium re-ionization in electron-volt per Hydrogen atom
+  rapid_cooling_threshold: 0.333333          # Switch to rapid cooling regime for dt / t_cool above this threshold.
 
 # Quick Lyman-alpha star formation parameters
 QLAStarFormation:
diff --git a/examples/QuickLymanAlpha/getQLACoolingTable.sh b/examples/QuickLymanAlpha/getQLACoolingTable.sh
index 5cfd93ef0f4603e40b7675f3f2c254b2250f699f..0945dc565a4c8b1723b21e685b03369e782cd3be 100755
--- a/examples/QuickLymanAlpha/getQLACoolingTable.sh
+++ b/examples/QuickLymanAlpha/getQLACoolingTable.sh
@@ -1,3 +1,2 @@
 #!/bin/bash
-wget http://virgodb.cosma.dur.ac.uk/swift-webstorage/CoolingTables/EAGLE/coolingtables.tar.gz
-tar -xf coolingtables.tar.gz 
+wget http://virgodb.cosma.dur.ac.uk/swift-webstorage/CoolingTables/COLIBRE/UV_dust1_CR1_G1_shield1.hdf5
diff --git a/examples/parameter_example.yml b/examples/parameter_example.yml
index b5179809aeb3a7b802ce7647d415796c2077e56a..e34c280c8862c29f41ffffa130927342bdbfb3b2 100644
--- a/examples/parameter_example.yml
+++ b/examples/parameter_example.yml
@@ -389,14 +389,26 @@ EAGLECooling:
   Ca_over_Si_in_solar:       1.                # (Optional) Ratio of Ca/Si to use in units of solar. If set to 1, the code uses [Ca/Si] = 0, i.e. Ca/Si = 0.0941736.
   S_over_Si_in_solar:        1.                # (Optional) Ratio of S/Si to use in units of solar. If set to 1, the code uses [S/Si] = 0, i.e. S/Si = 0.6054160.
 
-# Quick Lyman-alpha cooling (EAGLE with fixed primoridal Z)
+# Quick Lyman-alpha cooling (EAGLE-XL with fixed primoridal Z)
 QLACooling:
-  dir_name:                ./coolingtables/   # Location of the Wiersma+08 cooling tables
-  H_reion_z:               7.5                # Redshift of Hydrogen re-ionization
-  H_reion_eV_p_H:          2.0                # Energy inject by Hydrogen re-ionization in electron-volt per Hydrogen atom
-  He_reion_z_centre:       3.5                # Redshift of the centre of the Helium re-ionization Gaussian
-  He_reion_z_sigma:        0.5                # Spread in redshift of the  Helium re-ionization Gaussian
-  He_reion_eV_p_H:         2.0                # Energy inject by Helium re-ionization in electron-volt per Hydrogen atom
+  dir_name:                ./UV_dust1_CR1_G1_shield1.hdf5 # Location of the Ploeckinger+20 cooling tables
+  H_reion_z:               7.5               # Redshift of Hydrogen re-ionization (Planck 2018)
+  H_reion_eV_p_H:          2.0               # Energy inject by Hydrogen re-ionization in electron-volt per Hydrogen atom
+  He_reion_z_centre:       3.5               # Redshift of the centre of the Helium re-ionization Gaussian
+  He_reion_z_sigma:        0.5               # Spread in redshift of the  Helium re-ionization Gaussian
+  He_reion_eV_p_H:         2.0               # Energy inject by Helium re-ionization in electron-volt per Hydrogen atom
+  rapid_cooling_threshold: 0.333333          # Switch to rapid cooling regime for dt / t_cool above this threshold.
+
+# COLIBRE cooling parameters (EAGLE-XL)
+COLIBRECooling:
+  dir_name:                ./UV_dust1_CR1_G1_shield1.hdf5 # Location of the cooling tables
+  H_reion_z:               7.5               # Redshift of Hydrogen re-ionization (Planck 2018)
+  H_reion_eV_p_H:          2.0               # Energy inject by Hydrogen re-ionization in electron-volt per Hydrogen atom
+  He_reion_z_centre:       3.5               # Redshift of the centre of the Helium re-ionization Gaussian
+  He_reion_z_sigma:        0.5               # Spread in redshift of the  Helium re-ionization Gaussian
+  He_reion_eV_p_H:         2.0               # Energy inject by Helium re-ionization in electron-volt per Hydrogen atom
+  rapid_cooling_threshold: 0.333333          # Switch to rapid cooling regime for dt / t_cool above this threshold.
+  delta_logTEOS_subgrid_properties: 0.3      # delta log T above the EOS below which the subgrid properties use Teq assumption
   
 # Cooling with Grackle 3.0
 GrackleCooling:
diff --git a/src/cooling/QLA/cooling.c b/src/cooling/QLA/cooling.c
index 8c2d1313cb0a5db7fbfdb6247d5888e0e204a03f..acce27ea7e4dc7c85a4a96f915587f26ea2fba91 100644
--- a/src/cooling/QLA/cooling.c
+++ b/src/cooling/QLA/cooling.c
@@ -18,7 +18,7 @@
  ******************************************************************************/
 /**
  * @file src/cooling/QLA/cooling.c
- * @brief Quick Lyman-alpha cooling functions
+ * @brief QLA cooling functions
  */
 
 /* Config parameters. */
@@ -31,6 +31,7 @@
 #include <time.h>
 
 /* Local includes. */
+#include "adiabatic_index.h"
 #include "chemistry.h"
 #include "cooling.h"
 #include "cooling_rates.h"
@@ -48,7 +49,8 @@
 #include "space.h"
 #include "units.h"
 
-/* Maximum number of iterations for bisection scheme */
+/* Maximum number of iterations for
+ * bisection integration schemes */
 static const int bisection_max_iterations = 150;
 
 /* Tolerances for termination criteria. */
@@ -56,64 +58,6 @@ static const float explicit_tolerance = 0.05;
 static const float bisection_tolerance = 1.0e-6;
 static const double bracket_factor = 1.5;
 
-/**
- * @brief Find the index of the current redshift along the redshift dimension
- * of the cooling tables.
- *
- * Since the redshift table is not evenly spaced, compare z with each
- * table value in decreasing order starting with the previous redshift index
- *
- * The returned difference is expressed in units of the table separation. This
- * means dx = (x - table[i]) / (table[i+1] - table[i]). It is always between
- * 0 and 1.
- *
- * @param z Redshift we are searching for.
- * @param z_index (return) Index of the redshift in the table.
- * @param dz (return) Difference in redshift between z and table[z_index].
- * @param cooling #cooling_function_data structure containing redshift table.
- */
-__attribute__((always_inline)) INLINE void get_redshift_index(
-    const float z, int *z_index, float *dz,
-    struct cooling_function_data *restrict cooling) {
-
-  /* Before the earliest redshift or before hydrogen reionization, flag for
-   * collisional cooling */
-  if (z > cooling->H_reion_z) {
-    *z_index = qla_cooling_N_redshifts;
-    *dz = 0.0;
-  }
-
-  /* From reionization use the cooling tables */
-  else if (z > cooling->Redshifts[qla_cooling_N_redshifts - 1] &&
-           z <= cooling->H_reion_z) {
-    *z_index = qla_cooling_N_redshifts + 1;
-    *dz = 0.0;
-  }
-
-  /* At the end, just use the last value */
-  else if (z <= cooling->Redshifts[0]) {
-    *z_index = 0;
-    *dz = 0.0;
-  }
-
-  /* Normal case: search... */
-  else {
-
-    /* start at the previous index and search */
-    for (int i = cooling->previous_z_index; i >= 0; i--) {
-      if (z > cooling->Redshifts[i]) {
-
-        *z_index = i;
-        cooling->previous_z_index = i;
-
-        *dz = (z - cooling->Redshifts[i]) /
-              (cooling->Redshifts[i + 1] - cooling->Redshifts[i]);
-        break;
-      }
-    }
-  }
-}
-
 /**
  * @brief Common operations performed on the cooling function at a
  * given time-step or redshift. Predominantly used to read cooling tables
@@ -128,15 +72,6 @@ __attribute__((always_inline)) INLINE void get_redshift_index(
 void cooling_update(const struct cosmology *cosmo,
                     struct cooling_function_data *cooling, struct space *s) {
 
-  /* Current redshift */
-  const float redshift = cosmo->z;
-
-  /* What is the current table index along the redshift axis? */
-  int z_index = -1;
-  float dz = 0.f;
-  get_redshift_index(redshift, &z_index, &dz, cooling);
-  cooling->dz = dz;
-
   /* Extra energy for reionization? */
   if (!cooling->H_reion_done) {
 
@@ -153,35 +88,164 @@ void cooling_update(const struct cosmology *cosmo,
       cooling->H_reion_done = 1;
     }
   }
+}
 
-  /* Do we already have the correct tables loaded? */
-  if (cooling->z_index == z_index) return;
+/**
+ * @brief Compute the internal energy of a #part based on the cooling function
+ * but for a given temperature.
+ *
+ * @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.
+ * @param T temperature of the gas (internal units).
+ */
+float cooling_get_internalenergy_for_temperature(
+    const struct phys_const *phys_const, const struct hydro_props *hydro_props,
+    const struct unit_system *us, const struct cosmology *cosmo,
+    const struct cooling_function_data *cooling, const struct part *p,
+    const struct xpart *xp, const float T) {
 
-  /* Which table should we load ? */
-  if (z_index >= qla_cooling_N_redshifts) {
+#ifdef SWIFT_DEBUG_CHECKS
+  if (cooling->Redshifts == NULL)
+    error(
+        "Cooling function has not been initialised. Did you forget the "
+        "--temperature runtime flag?");
+#endif
 
-    if (z_index == qla_cooling_N_redshifts + 1) {
+  /* Get the Hydrogen mass fraction */
+  const float XH = 1. - phys_const->const_primordial_He_fraction;
 
-      /* Bewtween re-ionization and first table */
-      get_redshift_invariant_table(cooling, /* photodis=*/0);
+  /* Convert Hydrogen mass fraction into Hydrogen number density */
+  const float rho = hydro_get_physical_density(p, cosmo);
+  const double n_H = rho * XH / phys_const->const_proton_mass;
+  const double n_H_cgs = n_H * cooling->number_density_to_cgs;
 
-    } else {
+  /* Get this particle's metallicity ratio to solar.
+   *
+   * Note that we do not need the individual element's ratios that
+   * the function also computes. */
+  float dummy[qla_cooling_N_elementtypes];
+  const float logZZsol =
+      abundance_ratio_to_solar(p, cooling, phys_const, dummy);
 
-      /* Above re-ionization */
-      get_redshift_invariant_table(cooling, /* photodis=*/1);
-    }
+  /* compute hydrogen number density, metallicity and redshift indices and
+   * offsets  */
 
-  } else {
+  float d_red, d_met, d_n_H;
+  int red_index, met_index, n_H_index;
 
-    /* Normal case: two tables bracketing the current z */
-    const int low_z_index = z_index;
-    const int high_z_index = z_index + 1;
+  get_index_1d(cooling->Redshifts, qla_cooling_N_redshifts, cosmo->z,
+               &red_index, &d_red);
+  get_index_1d(cooling->Metallicity, qla_cooling_N_metallicity, logZZsol,
+               &met_index, &d_met);
+  get_index_1d(cooling->nH, qla_cooling_N_density, log10(n_H_cgs), &n_H_index,
+               &d_n_H);
 
-    get_cooling_table(cooling, low_z_index, high_z_index);
-  }
+  /* Compute the log10 of the temperature by interpolating the table */
+  const double log_10_U =
+      qla_convert_temp_to_u(log10(T), cosmo->z, n_H_index, d_n_H, met_index,
+                            d_met, red_index, d_red, cooling);
+
+  /* Undo the log! */
+  return exp10(log_10_U);
+}
+
+/**
+ * @brief Compute the temperature based on gas properties.
+ *
+ * The temperature returned is consistent with the cooling rates.
+ *
+ * @param phys_const #phys_const data structure.
+ * @param cosmo #cosmology data structure.
+ * @param cooling #cooling_function_data struct.
+ * @param rho_phys Density of the gas in internal physical units.
+ * @param logZZsol Logarithm base 10 of the gas' metallicity in units of solar
+ * metallicity.
+ * @param XH The Hydrogen abundance of the gas.
+ * @param u_phys Internal energy of the gas in internal physical units.
+ */
+float cooling_get_temperature_from_gas(
+    const struct phys_const *phys_const, const struct cosmology *cosmo,
+    const struct cooling_function_data *cooling, const float rho_phys,
+    const float logZZsol, const float XH, const float u_phys) {
+
+  /* Convert to CGS */
+  const double u_cgs = u_phys * cooling->internal_energy_to_cgs;
+  const double n_H = rho_phys * XH / phys_const->const_proton_mass;
+  const double n_H_cgs = n_H * cooling->number_density_to_cgs;
+
+  /* compute hydrogen number density, metallicity and redshift indices and
+   * offsets  */
+
+  float d_red, d_met, d_n_H;
+  int red_index, met_index, n_H_index;
+
+  get_index_1d(cooling->Redshifts, qla_cooling_N_redshifts, cosmo->z,
+               &red_index, &d_red);
+  get_index_1d(cooling->Metallicity, qla_cooling_N_metallicity, logZZsol,
+               &met_index, &d_met);
+  get_index_1d(cooling->nH, qla_cooling_N_density, log10(n_H_cgs), &n_H_index,
+               &d_n_H);
+
+  /* Compute the log10 of the temperature by interpolating the table */
+  const double log_10_T =
+      qla_convert_u_to_temp(log10(u_cgs), cosmo->z, n_H_index, d_n_H, met_index,
+                            d_met, red_index, d_red, cooling);
+
+  /* Undo the log! */
+  return exp10(log_10_T);
+}
+
+/**
+ * @brief Compute the temperature of a #part based on the cooling function.
+ *
+ * The temperature returned is consistent with the cooling rates.
+ *
+ * @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.
+ */
+float cooling_get_temperature(const struct phys_const *phys_const,
+                              const struct hydro_props *hydro_props,
+                              const struct unit_system *us,
+                              const struct cosmology *cosmo,
+                              const struct cooling_function_data *cooling,
+                              const struct part *p, const struct xpart *xp) {
+
+#ifdef SWIFT_DEBUG_CHECKS
+  if (cooling->Redshifts == NULL)
+    error(
+        "Cooling function has not been initialised. Did you forget the "
+        "--temperature runtime flag?");
+#endif
+
+  /* Get physical internal energy */
+  const float u_phys = hydro_get_physical_internal_energy(p, xp, cosmo);
+
+  /* Get the Hydrogen mass fraction */
+  const float XH = 1. - phys_const->const_primordial_He_fraction;
 
-  /* Store the currently loaded index */
-  cooling->z_index = z_index;
+  /* Convert Hydrogen mass fraction into Hydrogen number density */
+  const float rho_phys = hydro_get_physical_density(p, cosmo);
+
+  /* Get this particle's metallicity ratio to solar.
+   *
+   * Note that we do not need the individual element's ratios that
+   * the function also computes. */
+  float dummy[qla_cooling_N_elementtypes];
+  const float logZZsol =
+      abundance_ratio_to_solar(p, cooling, phys_const, dummy);
+
+  return cooling_get_temperature_from_gas(phys_const, cosmo, cooling, rho_phys,
+                                          logZZsol, XH, u_phys);
 }
 
 /**
@@ -192,8 +256,10 @@ void cooling_update(const struct cosmology *cosmo,
  * @param redshift Current redshift.
  * @param n_H_index Particle hydrogen number density index.
  * @param d_n_H Particle hydrogen number density offset.
- * @param He_index Particle helium fraction index.
- * @param d_He Particle helium fraction offset.
+ * @param met_index Particle metallicity index.
+ * @param d_met Particle metallicity offset.
+ * @param red_index Redshift index.
+ * @param d_red Redshift offset.
  * @param Lambda_He_reion_cgs Cooling rate coming from He reionization.
  * @param ratefact_cgs Multiplication factor to get a cooling rate.
  * @param cooling #cooling_function_data structure.
@@ -201,18 +267,17 @@ void cooling_update(const struct cosmology *cosmo,
  * @param dt_cgs timestep in CGS.
  * @param ID ID of the particle (for debugging).
  */
-INLINE static double bisection_iter(
+static INLINE double bisection_iter(
     const double u_ini_cgs, const double n_H_cgs, const double redshift,
-    const int n_H_index, const float d_n_H, const int He_index,
-    const float d_He, const double Lambda_He_reion_cgs,
-    const double ratefact_cgs,
-    const struct cooling_function_data *restrict cooling,
-    const float abundance_ratio[qla_cooling_N_abundances], const double dt_cgs,
-    const long long ID) {
+    int n_H_index, float d_n_H, int met_index, float d_met, int red_index,
+    float d_red, double Lambda_He_reion_cgs, double ratefact_cgs,
+    const struct cooling_function_data *cooling,
+    const float abundance_ratio[qla_cooling_N_elementtypes], double dt_cgs,
+    long long ID) {
 
   /* Bracketing */
-  double u_lower_cgs = u_ini_cgs;
-  double u_upper_cgs = u_ini_cgs;
+  double u_lower_cgs = max(u_ini_cgs, cooling->umin_cgs);
+  double u_upper_cgs = max(u_ini_cgs, cooling->umin_cgs);
 
   /*************************************/
   /* Let's get a first guess           */
@@ -221,7 +286,8 @@ INLINE static double bisection_iter(
   double LambdaNet_cgs =
       Lambda_He_reion_cgs +
       qla_cooling_rate(log10(u_ini_cgs), redshift, n_H_cgs, abundance_ratio,
-                       n_H_index, d_n_H, He_index, d_He, cooling);
+                       n_H_index, d_n_H, met_index, d_met, red_index, d_red,
+                       cooling, 0, 0, 0, 0);
 
   /*************************************/
   /* Let's try to bracket the solution */
@@ -230,36 +296,52 @@ INLINE static double bisection_iter(
   if (LambdaNet_cgs < 0) {
 
     /* we're cooling! */
-    u_lower_cgs /= bracket_factor;
-    u_upper_cgs *= bracket_factor;
+    u_lower_cgs = max(u_lower_cgs / bracket_factor, cooling->umin_cgs);
+    u_upper_cgs = max(u_upper_cgs * bracket_factor, cooling->umin_cgs);
 
     /* Compute a new rate */
     LambdaNet_cgs =
         Lambda_He_reion_cgs +
         qla_cooling_rate(log10(u_lower_cgs), redshift, n_H_cgs, abundance_ratio,
-                         n_H_index, d_n_H, He_index, d_He, cooling);
+                         n_H_index, d_n_H, met_index, d_met, red_index, d_red,
+                         cooling, 0, 0, 0, 0);
 
     int i = 0;
     while (u_lower_cgs - u_ini_cgs - LambdaNet_cgs * ratefact_cgs * dt_cgs >
                0 &&
            i < bisection_max_iterations) {
 
-      u_lower_cgs /= bracket_factor;
-      u_upper_cgs /= bracket_factor;
+      u_lower_cgs = max(u_lower_cgs / bracket_factor, cooling->umin_cgs);
+      u_upper_cgs = max(u_upper_cgs / bracket_factor, cooling->umin_cgs);
 
       /* Compute a new rate */
-      LambdaNet_cgs = Lambda_He_reion_cgs +
-                      qla_cooling_rate(log10(u_lower_cgs), redshift, n_H_cgs,
-                                       abundance_ratio, n_H_index, d_n_H,
-                                       He_index, d_He, cooling);
+      LambdaNet_cgs =
+          Lambda_He_reion_cgs +
+          qla_cooling_rate(log10(u_lower_cgs), redshift, n_H_cgs,
+                           abundance_ratio, n_H_index, d_n_H, met_index, d_met,
+                           red_index, d_red, cooling, 0, 0, 0, 0);
+
+      /* If the energy is below or equal the minimum energy and we are still
+       * cooling, return the minimum energy */
+      if ((u_lower_cgs <= cooling->umin_cgs) && (LambdaNet_cgs < 0.))
+        return cooling->umin_cgs;
+
       i++;
     }
 
     if (i >= bisection_max_iterations) {
       error(
           "particle %llu exceeded max iterations searching for bounds when "
-          "cooling, u_ini_cgs %.5e n_H_cgs %.5e",
-          ID, u_ini_cgs, n_H_cgs);
+          "cooling \n more info: n_H_cgs = %.4e, u_ini_cgs = %.4e, redshift = "
+          "%.4f\n"
+          "n_H_index = %i, d_n_H = %.4f\n"
+          "met_index = %i, d_met = %.4f, red_index = %i, d_red = %.4f, initial "
+          "Lambda = %.4e",
+          ID, n_H_cgs, u_ini_cgs, redshift, n_H_index, d_n_H, met_index, d_met,
+          red_index, d_red,
+          qla_cooling_rate(log10(u_ini_cgs), redshift, n_H_cgs, abundance_ratio,
+                           n_H_index, d_n_H, met_index, d_met, red_index, d_red,
+                           cooling, 0, 0, 0, 0));
     }
   } else {
 
@@ -271,7 +353,8 @@ INLINE static double bisection_iter(
     LambdaNet_cgs =
         Lambda_He_reion_cgs +
         qla_cooling_rate(log10(u_upper_cgs), redshift, n_H_cgs, abundance_ratio,
-                         n_H_index, d_n_H, He_index, d_He, cooling);
+                         n_H_index, d_n_H, met_index, d_met, red_index, d_red,
+                         cooling, 0, 0, 0, 0);
 
     int i = 0;
     while (u_upper_cgs - u_ini_cgs - LambdaNet_cgs * ratefact_cgs * dt_cgs <
@@ -282,18 +365,33 @@ INLINE static double bisection_iter(
       u_upper_cgs *= bracket_factor;
 
       /* Compute a new rate */
-      LambdaNet_cgs = Lambda_He_reion_cgs +
-                      qla_cooling_rate(log10(u_upper_cgs), redshift, n_H_cgs,
-                                       abundance_ratio, n_H_index, d_n_H,
-                                       He_index, d_He, cooling);
+      LambdaNet_cgs =
+          Lambda_He_reion_cgs +
+          qla_cooling_rate(log10(u_upper_cgs), redshift, n_H_cgs,
+                           abundance_ratio, n_H_index, d_n_H, met_index, d_met,
+                           red_index, d_red, cooling, 0, 0, 0, 0);
       i++;
     }
 
     if (i >= bisection_max_iterations) {
+      message("Aborting...");
+      message("particle %llu", ID);
+      message("n_H_cgs = %.4e", n_H_cgs);
+      message("u_ini_cgs = %.4e", u_ini_cgs);
+      message("redshift = %.4f", redshift);
+      message("indices nH, met, red = %i, %i, %i", n_H_index, met_index,
+              red_index);
+      message("index weights nH, met, red = %.4e, %.4e, %.4e", d_n_H, d_met,
+              d_red);
+      fflush(stdout);
+      message("cooling rate = %.4e",
+              qla_cooling_rate(log10(u_ini_cgs), redshift, n_H_cgs,
+                               abundance_ratio, n_H_index, d_n_H, met_index,
+                               d_met, red_index, d_red, cooling, 0, 0, 0, 0));
       error(
           "particle %llu exceeded max iterations searching for bounds when "
-          "heating, u_ini_cgs %.5e n_H_cgs %.5e",
-          ID, u_ini_cgs, n_H_cgs);
+          "cooling",
+          ID);
     }
   }
 
@@ -315,14 +413,8 @@ INLINE static double bisection_iter(
     LambdaNet_cgs =
         Lambda_He_reion_cgs +
         qla_cooling_rate(log10(u_next_cgs), redshift, n_H_cgs, abundance_ratio,
-                         n_H_index, d_n_H, He_index, d_He, cooling);
-#ifdef SWIFT_DEBUG_CHECKS
-    if (u_next_cgs <= 0)
-      error(
-          "Got negative energy! u_next_cgs=%.5e u_upper=%.5e u_lower=%.5e "
-          "Lambda=%.5e",
-          u_next_cgs, u_upper_cgs, u_lower_cgs, LambdaNet_cgs);
-#endif
+                         n_H_index, d_n_H, met_index, d_met, red_index, d_red,
+                         cooling, 0, 0, 0, 0);
 
     /* Where do we go next? */
     if (u_next_cgs - u_ini_cgs - LambdaNet_cgs * ratefact_cgs * dt_cgs > 0.0) {
@@ -356,9 +448,7 @@ INLINE static double bisection_iter(
  *
  * f(u_new) = u_new - u_old - dt * du/dt(u_new, X) = 0
  *
- * We first try a few Newton-Raphson iteration if it does not converge, we
- * revert to a bisection scheme.
- *
+ * A bisection scheme is used.
  * This is done by first bracketing the solution and then iterating
  * towards the solution by reducing the window down to a certain tolerance.
  * Note there is always at least one solution since
@@ -374,8 +464,7 @@ INLINE static double bisection_iter(
  * @param xp Pointer to the extended particle data.
  * @param dt The cooling time-step of this particle.
  * @param dt_therm The hydro time-step of this particle.
- * @param time The current time (since the Big Bang or start of the run) in
- * internal units.
+ * @param time Time since Big Bang
  */
 void cooling_cool_part(const struct phys_const *phys_const,
                        const struct unit_system *us,
@@ -383,12 +472,13 @@ void cooling_cool_part(const struct phys_const *phys_const,
                        const struct hydro_props *hydro_properties,
                        const struct entropy_floor_properties *floor_props,
                        const struct cooling_function_data *cooling,
-                       struct part *restrict p, struct xpart *restrict xp,
-                       const float dt, const float dt_therm,
-                       const double time) {
+                       struct part *p, struct xpart *xp, const float dt,
+                       const float dt_therm, const double time) {
 
   /* No cooling happens over zero time */
-  if (dt == 0.) return;
+  if (dt == 0.) {
+    return;
+  }
 
 #ifdef SWIFT_DEBUG_CHECKS
   if (cooling->Redshifts == NULL)
@@ -403,7 +493,7 @@ void cooling_cool_part(const struct phys_const *phys_const,
   /* Get the change in internal energy due to hydro forces */
   const float hydro_du_dt = hydro_get_physical_internal_energy_dt(p, cosmo);
 
-  /* Get internal energy at the end of the step (assuming dt does not
+  /* Get internal energy at the end of the next kick step (assuming dt does not
    * increase) */
   double u_0 = (u_start + hydro_du_dt * dt_therm);
 
@@ -421,19 +511,15 @@ void cooling_cool_part(const struct phys_const *phys_const,
   /* Get this particle's abundance ratios compared to solar
    * Note that we need to add S and Ca that are in the tables but not tracked
    * by the particles themselves.
-   * The order is [H, He, C, N, O, Ne, Mg, Si, S, Ca, Fe] */
-  float abundance_ratio[qla_cooling_N_abundances];
-  abundance_ratio_to_solar(p, cooling, phys_const, abundance_ratio);
+   * The order is [H, He, C, N, O, Ne, Mg, Si, S, Ca, Fe, OA] */
+  float abundance_ratio[qla_cooling_N_elementtypes];
+  float logZZsol =
+      abundance_ratio_to_solar(p, cooling, phys_const, abundance_ratio);
 
   /* Get the Hydrogen and Helium mass fractions */
   const float XH = 1. - phys_const->const_primordial_He_fraction;
-  const float XHe = phys_const->const_primordial_He_fraction;
-
-  /* Get the Helium mass fraction. Note that this is He / (H + He), i.e. a
-   * metal-free Helium mass fraction as per the Wiersma+08 definition */
-  const float HeFrac = XHe / (XH + XHe);
 
-  /* convert Hydrogen mass fraction into physical Hydrogen number density */
+  /* convert Hydrogen mass fraction into Hydrogen number density */
   const double n_H =
       hydro_get_physical_density(p, cosmo) * XH / phys_const->const_proton_mass;
   const double n_H_cgs = n_H * cooling->number_density_to_cgs;
@@ -442,13 +528,17 @@ void cooling_cool_part(const struct phys_const *phys_const,
    * equivalent expression  below */
   const double ratefact_cgs = n_H_cgs * (XH * cooling->inv_proton_mass_cgs);
 
-  /* compute hydrogen number density and helium fraction table indices and
+  /* compute hydrogen number density, metallicity and redshift indices and
    * offsets (These are fixed for any value of u, so no need to recompute them)
    */
-  int He_index, n_H_index;
-  float d_He, d_n_H;
-  get_index_1d(cooling->HeFrac, qla_cooling_N_He_frac, HeFrac, &He_index,
-               &d_He);
+
+  float d_red, d_met, d_n_H;
+  int red_index, met_index, n_H_index;
+
+  get_index_1d(cooling->Redshifts, qla_cooling_N_redshifts, cosmo->z,
+               &red_index, &d_red);
+  get_index_1d(cooling->Metallicity, qla_cooling_N_metallicity, logZZsol,
+               &met_index, &d_met);
   get_index_1d(cooling->nH, qla_cooling_N_density, log10(n_H_cgs), &n_H_index,
                &d_n_H);
 
@@ -457,22 +547,21 @@ void cooling_cool_part(const struct phys_const *phys_const,
 
   /* Get helium and hydrogen reheating term */
   const double Helium_reion_heat_cgs =
-      qla_helium_reionization_extraheat(cosmo->z, delta_redshift, cooling);
+      eagle_helium_reionization_extraheat(cosmo->z, delta_redshift, cooling);
 
   /* Convert this into a rate */
   const double Lambda_He_reion_cgs =
       Helium_reion_heat_cgs / (dt_cgs * ratefact_cgs);
 
   /* Let's compute the internal energy at the end of the step */
-  /* Initialise to the initial energy to appease compiler; this will never not
-     be overwritten. */
-  double u_final_cgs = u_0_cgs;
+  double u_final_cgs;
 
   /* First try an explicit integration (note we ignore the derivative) */
   const double LambdaNet_cgs =
       Lambda_He_reion_cgs + qla_cooling_rate(log10(u_0_cgs), cosmo->z, n_H_cgs,
                                              abundance_ratio, n_H_index, d_n_H,
-                                             He_index, d_He, cooling);
+                                             met_index, d_met, red_index, d_red,
+                                             cooling, 0, 0, 0, 0);
 
   /* if cooling rate is small, take the explicit solution */
   if (fabs(ratefact_cgs * LambdaNet_cgs * dt_cgs) <
@@ -482,11 +571,10 @@ void cooling_cool_part(const struct phys_const *phys_const,
 
   } else {
 
-    /* Otherwise, go the bisection route. */
     u_final_cgs =
-        bisection_iter(u_0_cgs, n_H_cgs, cosmo->z, n_H_index, d_n_H, He_index,
-                       d_He, Lambda_He_reion_cgs, ratefact_cgs, cooling,
-                       abundance_ratio, dt_cgs, p->id);
+        bisection_iter(u_0_cgs, n_H_cgs, cosmo->z, n_H_index, d_n_H, met_index,
+                       d_met, red_index, d_red, Lambda_He_reion_cgs,
+                       ratefact_cgs, cooling, abundance_ratio, dt_cgs, p->id);
   }
 
   /* Convert back to internal units */
@@ -509,11 +597,34 @@ void cooling_cool_part(const struct phys_const *phys_const,
      (assuming no change in dt) */
   const double delta_u = u_final - max(u_start, u_floor);
 
-  /* Turn this into a rate of change (including cosmology term) */
-  const float cooling_du_dt = delta_u / dt_therm;
+  /* Determine if we are in the slow- or rapid-cooling regime,
+   * by comparing dt / t_cool to the rapid_cooling_threshold.
+   *
+   * Note that dt / t_cool = fabs(delta_u) / u_start. */
+  const double dt_over_t_cool = fabs(delta_u) / max(u_start, u_floor);
+
+  /* If rapid_cooling_threshold < 0, always use the slow-cooling
+   * regime. */
+  if ((cooling->rapid_cooling_threshold >= 0.0) &&
+      (dt_over_t_cool >= cooling->rapid_cooling_threshold)) {
 
-  /* Update the internal energy time derivative */
-  hydro_set_physical_internal_energy_dt(p, cosmo, cooling_du_dt);
+    /* Rapid-cooling regime. */
+
+    /* Update the particle's u and du/dt */
+    hydro_set_physical_internal_energy(p, xp, cosmo, u_final);
+    hydro_set_drifted_physical_internal_energy(p, cosmo, u_final);
+    hydro_set_physical_internal_energy_dt(p, cosmo, 0.);
+
+  } else {
+
+    /* Slow-cooling regime. */
+
+    /* Update du/dt so that we can subsequently drift internal energy. */
+    const float cooling_du_dt = delta_u / dt_therm;
+
+    /* Update the internal energy time derivative */
+    hydro_set_physical_internal_energy_dt(p, cosmo, cooling_du_dt);
+  }
 }
 
 /**
@@ -530,12 +641,10 @@ void cooling_cool_part(const struct phys_const *phys_const,
  * @param xp extended particle data.
  */
 __attribute__((always_inline)) INLINE 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) {
+    const struct cooling_function_data *cooling,
+    const struct phys_const *phys_const, const struct cosmology *cosmo,
+    const struct unit_system *us, const struct hydro_props *hydro_props,
+    const struct part *p, const struct xpart *xp) {
 
   return FLT_MAX;
 }
@@ -553,78 +662,10 @@ __attribute__((always_inline)) INLINE float cooling_timestep(
  * @param xp Pointer to the #xpart data.
  */
 __attribute__((always_inline)) INLINE 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) {}
-
-/**
- * @brief Compute the temperature of a #part based on the cooling function.
- *
- * We use the Temperature table of the Wiersma+08 set. This computes the
- * equilibirum temperature of a gas for a given redshift, Hydrogen density,
- * internal energy per unit mass and Helium fraction.
- *
- * The temperature returned is consistent with the cooling rates.
- *
- * @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.
- */
-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) {
-
-#ifdef SWIFT_DEBUG_CHECKS
-  if (cooling->Redshifts == NULL)
-    error(
-        "Cooling function has not been initialised. Did you forget the "
-        "--temperature runtime flag?");
-#endif
-
-  /* Get physical internal energy */
-  const float u = hydro_get_physical_internal_energy(p, xp, cosmo);
-  const double u_cgs = u * cooling->internal_energy_to_cgs;
-
-  /* Get the Hydrogen and Helium mass fractions */
-  const float XH = 1. - phys_const->const_primordial_He_fraction;
-  const float XHe = phys_const->const_primordial_He_fraction;
-
-  /* Get the Helium mass fraction. Note that this is He / (H + He), i.e. a
-   * metal-free Helium mass fraction as per the Wiersma+08 definition */
-  const float HeFrac = XHe / (XH + XHe);
-
-  /* Convert Hydrogen mass fraction into Hydrogen number density */
-  const float rho = hydro_get_physical_density(p, cosmo);
-  const double n_H = rho * XH / phys_const->const_proton_mass;
-  const double n_H_cgs = n_H * cooling->number_density_to_cgs;
-
-  /* compute hydrogen number density and helium fraction table indices and
-   * offsets */
-  int He_index, n_H_index;
-  float d_He, d_n_H;
-  get_index_1d(cooling->HeFrac, qla_cooling_N_He_frac, HeFrac, &He_index,
-               &d_He);
-  get_index_1d(cooling->nH, qla_cooling_N_density, log10(n_H_cgs), &n_H_index,
-               &d_n_H);
-
-  /* Compute the log10 of the temperature by interpolating the table */
-  const double log_10_T = qla_convert_u_to_temp(
-      log10(u_cgs), cosmo->z, n_H_index, He_index, d_n_H, d_He, cooling);
-
-  /* Undo the log! */
-  return exp10(log_10_T);
-}
+    const struct phys_const *phys_const, const struct unit_system *us,
+    const struct hydro_props *hydro_props, const struct cosmology *cosmo,
+    const struct cooling_function_data *cooling, struct part *p,
+    struct xpart *xp) {}
 
 /**
  * @brief Returns the total radiated energy by this particle.
@@ -632,9 +673,8 @@ float cooling_get_temperature(
  * @param xp #xpart data struct
  */
 __attribute__((always_inline)) INLINE float cooling_get_radiated_energy(
-    const struct xpart *restrict xp) {
-
-  return 0.f;
+    const struct xpart *xp) {
+  return -1.f;
 }
 
 /**
@@ -665,7 +705,7 @@ void cooling_Hydrogen_reionization(const struct cooling_function_data *cooling,
   const float extra_heat =
       cooling->H_reion_heat_cgs * cooling->internal_energy_from_cgs;
 
-  message("Applying extra energy for H reionization!");
+  message("Applying extra energy for H reionization to non-star-forming gas!");
 
   /* Loop through particles and set new heat */
   for (size_t i = 0; i < s->nr_parts; i++) {
@@ -684,11 +724,11 @@ void cooling_Hydrogen_reionization(const struct cooling_function_data *cooling,
 /**
  * @brief Initialises properties stored in the cooling_function_data struct
  *
- * @param parameter_file The parsed parameter file
- * @param us Internal system of units data structure
- * @param phys_const #phys_const data structure
- * @param hydro_props The properties of the hydro scheme.
- * @param cooling #cooling_function_data struct to initialize
+ * @param parameter_file The parsed parameter file.
+ * @param us Internal system of units data structure.
+ * @param hydro_props the properties of the hydro scheme.
+ * @param phys_const #phys_const data structure.
+ * @param cooling #cooling_function_data struct to initialize.
  */
 void cooling_init_backend(struct swift_params *parameter_file,
                           const struct unit_system *us,
@@ -696,9 +736,8 @@ void cooling_init_backend(struct swift_params *parameter_file,
                           const struct hydro_props *hydro_props,
                           struct cooling_function_data *cooling) {
 
-  /* Read model parameters */
+  /* read some parameters */
 
-  /* Directory for cooling tables */
   parser_get_param_string(parameter_file, "QLACooling:dir_name",
                           cooling->cooling_table_path);
 
@@ -732,32 +771,46 @@ void cooling_init_backend(struct swift_params *parameter_file,
       phys_const->const_electron_volt / phys_const->const_proton_mass *
       units_cgs_conversion_factor(us, UNIT_CONV_ENERGY_PER_UNIT_MASS);
 
-  /* Read in the list of redshifts */
-  get_cooling_redshifts(cooling);
-
-  /* Read in cooling table header */
-  char fname[qla_table_path_name_length + 12];
-  sprintf(fname, "%sz_0.000.hdf5", cooling->cooling_table_path);
-  read_cooling_header(fname, cooling);
-
-  /* Allocate space for cooling tables */
-  allocate_cooling_tables(cooling);
-
   /* Compute conversion factors */
+  cooling->pressure_to_cgs =
+      units_cgs_conversion_factor(us, UNIT_CONV_PRESSURE);
   cooling->internal_energy_to_cgs =
       units_cgs_conversion_factor(us, UNIT_CONV_ENERGY_PER_UNIT_MASS);
   cooling->internal_energy_from_cgs = 1. / cooling->internal_energy_to_cgs;
   cooling->number_density_to_cgs =
       units_cgs_conversion_factor(us, UNIT_CONV_NUMBER_DENSITY);
+  cooling->number_density_from_cgs = 1. / cooling->number_density_to_cgs;
 
   /* Store some constants in CGS units */
+  const float units_kB[5] = {1, 2, -2, 0, -1};
+  const double kB_cgs = phys_const->const_boltzmann_k *
+                        units_general_cgs_conversion_factor(us, units_kB);
   const double proton_mass_cgs =
       phys_const->const_proton_mass *
       units_cgs_conversion_factor(us, UNIT_CONV_MASS);
+
+  cooling->log10_kB_cgs = log10(kB_cgs);
   cooling->inv_proton_mass_cgs = 1. / proton_mass_cgs;
   cooling->T_CMB_0 = phys_const->const_T_CMB_0 *
                      units_cgs_conversion_factor(us, UNIT_CONV_TEMPERATURE);
 
+  /* Get the minimal temperature allowed */
+  cooling->Tmin = hydro_props->minimal_temperature;
+  if (cooling->Tmin < 10.)
+    error("QLA cooling cannot handle a minimal temperature below 10 K");
+
+  /* Recover the minimal energy allowed (in internal units) */
+  const double u_min = hydro_props->minimal_internal_energy;
+
+  /* Convert to CGS units */
+  cooling->umin_cgs = u_min * cooling->internal_energy_to_cgs;
+
+#ifdef SWIFT_DEBUG_CHECKS
+  /* Basic cross-check... */
+  if (kB_cgs > 1.381e-16 || kB_cgs < 1.380e-16)
+    error("Boltzmann's constant not initialised properly!");
+#endif
+
   /* 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;
@@ -785,11 +838,16 @@ void cooling_init_backend(struct swift_params *parameter_file,
                               cooling->T_CMB_0 * cooling->T_CMB_0 *
                               cooling->T_CMB_0;
 
-  /* Set the redshift indices to invalid values */
-  cooling->z_index = -10;
+  /* Threshold in dt / t_cool above which we
+   * are in the rapid cooling regime. If negative,
+   * we never use this scheme (i.e. always drift
+   * the internal energies). */
+  cooling->rapid_cooling_threshold = parser_get_param_double(
+      parameter_file, "QLACooling:rapid_cooling_threshold");
 
-  /* set previous_z_index and to last value of redshift table*/
-  cooling->previous_z_index = qla_cooling_N_redshifts - 2;
+  /* Finally, read the tables */
+  read_cooling_header(cooling);
+  read_cooling_tables(cooling);
 }
 
 /**
@@ -802,20 +860,9 @@ void cooling_init_backend(struct swift_params *parameter_file,
 void cooling_restore_tables(struct cooling_function_data *cooling,
                             const struct cosmology *cosmo) {
 
-  /* Read redshifts */
-  get_cooling_redshifts(cooling);
-
-  /* Read cooling header */
-  char fname[qla_table_path_name_length + 12];
-  sprintf(fname, "%sz_0.000.hdf5", cooling->cooling_table_path);
-  read_cooling_header(fname, cooling);
-
-  /* Allocate memory for the tables */
-  allocate_cooling_tables(cooling);
+  read_cooling_header(cooling);
+  read_cooling_tables(cooling);
 
-  /* Force a re-read of the cooling tables */
-  cooling->z_index = -10;
-  cooling->previous_z_index = qla_cooling_N_redshifts - 2;
   cooling_update(cosmo, cooling, /*space=*/NULL);
 }
 
@@ -827,7 +874,7 @@ void cooling_restore_tables(struct cooling_function_data *cooling,
 void cooling_print_backend(const struct cooling_function_data *cooling) {
 
   message(
-      "Cooling function is 'Quick Lyman-alpha (EAGLE with primordial Z "
+      "Cooling function is 'Quick Lyman-alpha (COLIBRE with primordial Z "
       "only)'.");
 }
 
@@ -841,20 +888,37 @@ void cooling_print_backend(const struct cooling_function_data *cooling) {
 void cooling_clean(struct cooling_function_data *cooling) {
 
   /* Free the side arrays */
-  swift_free("cooling", cooling->Redshifts);
-  swift_free("cooling", cooling->nH);
-  swift_free("cooling", cooling->Temp);
-  swift_free("cooling", cooling->HeFrac);
-  swift_free("cooling", cooling->Therm);
-  swift_free("cooling", cooling->SolarAbundances);
-  swift_free("cooling", cooling->SolarAbundances_inv);
+  free(cooling->Redshifts);
+  free(cooling->nH);
+  free(cooling->Temp);
+  free(cooling->Metallicity);
+  free(cooling->Therm);
+  free(cooling->LogAbundances);
+  free(cooling->Abundances);
+  free(cooling->Abundances_inv);
+  free(cooling->atomicmass);
+  free(cooling->atomicmass_inv);
+  free(cooling->Zsol);
+  free(cooling->Zsol_inv);
+  free(cooling->LogMassFractions);
+  free(cooling->MassFractions);
 
   /* Free the tables */
-  swift_free("cooling-tables", cooling->table.metal_heating);
-  swift_free("cooling-tables", cooling->table.electron_abundance);
-  swift_free("cooling-tables", cooling->table.temperature);
-  swift_free("cooling-tables", cooling->table.H_plus_He_heating);
-  swift_free("cooling-tables", cooling->table.H_plus_He_electron_abundance);
+  free(cooling->table.Tcooling);
+  free(cooling->table.Ucooling);
+  free(cooling->table.Theating);
+  free(cooling->table.Uheating);
+  free(cooling->table.Telectron_fraction);
+  free(cooling->table.Uelectron_fraction);
+  free(cooling->table.T_from_U);
+  free(cooling->table.U_from_T);
+  free(cooling->table.Umu);
+  free(cooling->table.Tmu);
+  free(cooling->table.meanpartmass_Teq);
+  free(cooling->table.logHfracs_Teq);
+  free(cooling->table.logHfracs_all);
+  free(cooling->table.logTeq);
+  free(cooling->table.logPeq);
 }
 
 /**
@@ -873,14 +937,23 @@ void cooling_struct_dump(const struct cooling_function_data *cooling,
   cooling_copy.Redshifts = NULL;
   cooling_copy.nH = NULL;
   cooling_copy.Temp = NULL;
+  cooling_copy.Metallicity = NULL;
   cooling_copy.Therm = NULL;
-  cooling_copy.SolarAbundances = NULL;
-  cooling_copy.SolarAbundances_inv = NULL;
-  cooling_copy.table.metal_heating = NULL;
-  cooling_copy.table.H_plus_He_heating = NULL;
-  cooling_copy.table.H_plus_He_electron_abundance = NULL;
-  cooling_copy.table.temperature = NULL;
-  cooling_copy.table.electron_abundance = NULL;
+  cooling_copy.LogAbundances = NULL;
+  cooling_copy.Abundances = NULL;
+  cooling_copy.Abundances_inv = NULL;
+  cooling_copy.atomicmass = NULL;
+  cooling_copy.LogMassFractions = NULL;
+  cooling_copy.MassFractions = NULL;
+
+  cooling_copy.table.Tcooling = NULL;
+  cooling_copy.table.Theating = NULL;
+  cooling_copy.table.Telectron_fraction = NULL;
+  cooling_copy.table.Ucooling = NULL;
+  cooling_copy.table.Uheating = NULL;
+  cooling_copy.table.Uelectron_fraction = NULL;
+  cooling_copy.table.T_from_U = NULL;
+  cooling_copy.table.U_from_T = NULL;
 
   restart_write_blocks((void *)&cooling_copy,
                        sizeof(struct cooling_function_data), 1, stream,
diff --git a/src/cooling/QLA/cooling.h b/src/cooling/QLA/cooling.h
index 5a0ef4d55a007b932b5b892db32d39e8baec17b5..a0e881d9779b7211d6943323dc3276b069ccf761 100644
--- a/src/cooling/QLA/cooling.h
+++ b/src/cooling/QLA/cooling.h
@@ -20,8 +20,9 @@
 #define SWIFT_COOLING_QLA_H
 
 /**
- * @file src/cooling/EAGLE/cooling.h
- * @brief EAGLE cooling function declarations
+ * @file src/cooling/QLA/cooling.h
+ * @brief Quick Lyman alpha cooling (COLIBRE with primoridal Z) function
+ * declarations
  */
 
 /* Local includes. */
@@ -32,10 +33,8 @@ struct xpart;
 struct cosmology;
 struct hydro_props;
 struct entropy_floor_properties;
-struct space;
-struct unit_system;
 struct phys_const;
-struct swift_params;
+struct space;
 
 void cooling_update(const struct cosmology *cosmo,
                     struct cooling_function_data *cooling, struct space *s);
@@ -46,34 +45,36 @@ void cooling_cool_part(const struct phys_const *phys_const,
                        const struct hydro_props *hydro_properties,
                        const struct entropy_floor_properties *floor_props,
                        const struct cooling_function_data *cooling,
-                       struct part *restrict p, struct xpart *restrict xp,
-                       const float dt, const float dt_therm, const double time);
+                       struct part *p, struct xpart *xp, const float dt,
+                       const float dt_therm, const double time);
 
-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,
+float cooling_timestep(const struct cooling_function_data *cooling,
+                       const struct phys_const *phys_const,
+                       const struct cosmology *cosmo,
+                       const struct unit_system *us,
                        const struct hydro_props *hydro_props,
-                       const struct part *restrict p,
-                       const struct xpart *restrict xp);
-
-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);
-
-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);
-
-float cooling_get_radiated_energy(const struct xpart *restrict xp);
+                       const struct part *p, const struct xpart *xp);
+
+void cooling_first_init_part(const struct phys_const *phys_const,
+                             const struct unit_system *us,
+                             const struct hydro_props *hydro_props,
+                             const struct cosmology *cosmo,
+                             const struct cooling_function_data *cooling,
+                             struct part *p, struct xpart *xp);
+
+float cooling_get_temperature_from_gas(
+    const struct phys_const *phys_const, const struct cosmology *cosmo,
+    const struct cooling_function_data *cooling, const float rho_phys,
+    const float XH, const float logZZsol, const float u_phys);
+
+float cooling_get_temperature(const struct phys_const *phys_const,
+                              const struct hydro_props *hydro_props,
+                              const struct unit_system *us,
+                              const struct cosmology *cosmo,
+                              const struct cooling_function_data *cooling,
+                              const struct part *p, const struct xpart *xp);
+
+float cooling_get_radiated_energy(const struct xpart *xp);
 
 void cooling_split_part(struct part *p, struct xpart *xp, double n);
 
diff --git a/src/cooling/QLA/cooling_io.h b/src/cooling/QLA/cooling_io.h
index 779ef844f20159bf8d08196d78e85496345f8590..5e2ebc9e3727df0f0173cfa6bd53a9bc188b93ae 100644
--- a/src/cooling/QLA/cooling_io.h
+++ b/src/cooling/QLA/cooling_io.h
@@ -41,7 +41,7 @@ __attribute__((always_inline)) INLINE static void cooling_write_flavour(
     const struct cooling_function_data* cooling) {
 
   io_write_attribute_s(h_grp, "Cooling Model",
-                       "Quick Lyman-alpha (EAGLE with primordial Z only)");
+                       "Quick Lyman-alpha (COLIBRE with primordial Z only)");
 }
 #endif
 
@@ -70,6 +70,7 @@ __attribute__((always_inline)) INLINE static int cooling_write_particles(
   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;
 }
 
diff --git a/src/cooling/QLA/cooling_rates.h b/src/cooling/QLA/cooling_rates.h
index db4ba670f7f96a1da2e51a1d9bc89e75a4099b61..bc422a1a8cc8ccb62a7465d53521a7b80d432a0f 100644
--- a/src/cooling/QLA/cooling_rates.h
+++ b/src/cooling/QLA/cooling_rates.h
@@ -16,15 +16,16 @@
  * along with this program.  If not, see <http://www.gnu.org/licenses/>.
  *
  ******************************************************************************/
-
-/* Config parameters. */
 #ifndef SWIFT_QLA_COOLING_RATES_H
 #define SWIFT_QLA_COOLING_RATES_H
 
+/* Config parameters. */
 #include "../config.h"
 
 /* Local includes. */
+#include "chemistry_struct.h"
 #include "cooling_tables.h"
+#include "error.h"
 #include "exp10.h"
 #include "interpolate.h"
 
@@ -34,44 +35,142 @@
  *
  * The solar abundances are taken from the tables themselves.
  *
- * The quick Lyman-alpha chemistry model does not track any elements.
- * We use the primordial H and He to fill in the element abundance
- * array.
+ * The COLIBRE chemistry model does not track S and Ca. We assume
+ * that their abundance with respect to solar is the same as
+ * the ratio for Si.
+ *
+ * The other un-tracked elements are scaled with the total metallicity.
+ *
+ * We optionally apply a correction if the user asked for a different
+ * ratio.
  *
  * We also re-order the elements such that they match the order of the
- * tables. This is [H, He, C, N, O, Ne, Mg, Si, S, Ca, Fe].
+ * tables. This is [H, He, C, N, O, Ne, Mg, Si, S, Ca, Fe, OA].
  *
  * The solar abundances table (from the cooling struct) is arranged as
  * [H, He, C, N, O, Ne, Mg, Si, S, Ca, Fe].
  *
  * @param p Pointer to #part struct.
  * @param cooling #cooling_function_data struct.
- * @param phys_const the physical constants in internal units
  * @param ratio_solar (return) Array of ratios to solar abundances.
+ *
+ * @return The log10 of the total metallicity with respect to solar, i.e.
+ * log10(Z / Z_sun).
  */
-__attribute__((always_inline)) INLINE static void abundance_ratio_to_solar(
+__attribute__((always_inline)) INLINE static float abundance_ratio_to_solar(
     const struct part *p, const struct cooling_function_data *cooling,
     const struct phys_const *phys_const,
-    float ratio_solar[qla_cooling_N_abundances]) {
-
-  /* We just construct an array of primoridal abundances as fractions
-     of solar */
-
-  ratio_solar[0] = (1. - phys_const->const_primordial_He_fraction) *
-                   cooling->SolarAbundances_inv[0 /* H */];
-
-  ratio_solar[1] = phys_const->const_primordial_He_fraction *
-                   cooling->SolarAbundances_inv[1 /* He */];
-
-  ratio_solar[2] = 0.f;
-  ratio_solar[3] = 0.f;
-  ratio_solar[4] = 0.f;
-  ratio_solar[5] = 0.f;
-  ratio_solar[6] = 0.f;
-  ratio_solar[7] = 0.f;
-  ratio_solar[8] = 0.f;
-  ratio_solar[9] = 0.f;
-  ratio_solar[10] = 0.f;
+    float ratio_solar[qla_cooling_N_elementtypes]) {
+
+  /* Convert mass fractions to abundances (nx/nH) and compute metal mass */
+  float totmass = 0., metalmass = 0.;
+  for (enum qla_cooling_element elem = element_H; elem < element_OA; elem++) {
+
+    double Z_mass_frac = -1.f;
+    double Z_mass_frac_H = 1. - phys_const->const_primordial_He_fraction;
+
+    /* Assume primordial abundances */
+    if (elem == element_H) {
+      Z_mass_frac = 1. - phys_const->const_primordial_He_fraction;
+    } else if (elem == element_He) {
+      Z_mass_frac = phys_const->const_primordial_He_fraction;
+    } else {
+      Z_mass_frac = 0.f;
+    }
+
+    const int indx1d =
+        row_major_index_2d(cooling->indxZsol, elem, qla_cooling_N_metallicity,
+                           qla_cooling_N_elementtypes);
+
+    const float Mfrac = Z_mass_frac;
+
+    /* ratio_X = ((M_x/M) / (M_H/M)) * (m_H / m_X) * (1 / Z_sun_X) */
+    ratio_solar[elem] =
+        (Mfrac / Z_mass_frac_H) * cooling->atomicmass[element_H] *
+        cooling->atomicmass_inv[elem] * cooling->Abundances_inv[indx1d];
+
+    totmass += Mfrac;
+    if (elem > element_He) metalmass += Mfrac;
+  }
+
+  /* Compute metallicity (Z / Z_sun) of the elements we explicitly track. */
+  /* float ZZsol = (metalmass / totmass) * cooling->Zsol_inv[0];
+  if (ZZsol <= 0.0f) ZZsol = FLT_MIN;
+  const float logZZsol = log10f(ZZsol); */
+
+  /* Get total metallicity [Z/Z_sun] from the particle data */
+  const float Z_total =
+      (float)chemistry_get_total_metal_mass_fraction_for_cooling(p);
+  float ZZsol = Z_total * cooling->Zsol_inv[0];
+  if (ZZsol <= 0.0f) ZZsol = FLT_MIN;
+  const float logZZsol = log10f(ZZsol);
+
+  /* All other elements (element_OA): scale with metallicity */
+  ratio_solar[element_OA] = ZZsol;
+
+  /* Get index and offset from the metallicity table conresponding to this
+   * metallicity */
+  int met_index;
+  float d_met;
+
+  get_index_1d(cooling->Metallicity, qla_cooling_N_metallicity, logZZsol,
+               &met_index, &d_met);
+
+  /* At this point ratio_solar is (nx/nH) / (nx/nH)_sol.
+   * To multiply with the tables, we want the individual abundance ratio
+   * relative to what is used in the tables for each metallicity
+   */
+
+  /* For example: for a metallicity of 1 per cent solar, the metallicity bin
+   * for logZZsol = -2 has already the reduced cooling rates for each element;
+   * it should therefore NOT be multiplied by 0.01 again.
+   *
+   * BUT: if e.g. Carbon is twice as abundant as the solar abundance ratio,
+   * i.e. nC / nH = 0.02 * (nC/nH)_sol for the overall metallicity of 0.01,
+   * the Carbon cooling rate is multiplied by 2
+   *
+   * We only do this if we are not in the primodial metallicity bin in which
+   * case the ratio to solar should be 0.
+   */
+
+  for (int i = 0; i < qla_cooling_N_elementtypes; i++) {
+
+    /* Are we considering a metal and are *not* in the primodial metallicity
+     * bin? Or are we looking at H or He? */
+    if ((met_index > 0) || (i == element_H) || (i == element_He)) {
+
+      /* Get min/max abundances */
+      const float log_nx_nH_min = cooling->LogAbundances[row_major_index_2d(
+          met_index, i, qla_cooling_N_metallicity, qla_cooling_N_elementtypes)];
+
+      const float log_nx_nH_max = cooling->LogAbundances[row_major_index_2d(
+          met_index + 1, i, qla_cooling_N_metallicity,
+          qla_cooling_N_elementtypes)];
+
+      /* Get solar abundances */
+      const float log_nx_nH_sol = cooling->LogAbundances[row_major_index_2d(
+          cooling->indxZsol, i, qla_cooling_N_metallicity,
+          qla_cooling_N_elementtypes)];
+
+      /* Interpolate ! (linearly in log-space) */
+      const float log_nx_nH =
+          (log_nx_nH_min * (1.f - d_met) + log_nx_nH_max * d_met) -
+          log_nx_nH_sol;
+
+      ratio_solar[i] *= exp10f(-log_nx_nH);
+
+    } else {
+
+      /* Primordial bin --> Z/Z_sun is 0 for that element */
+      ratio_solar[i] = 0.;
+    }
+  }
+
+  /* at this point ratio_solar is (nx/nH) / (nx/nH)_table [Z],
+   * the metallicity dependent abundance ratio for solar abundances.
+   * We also return the total metallicity */
+
+  return logZZsol;
 }
 
 /**
@@ -87,9 +186,9 @@ __attribute__((always_inline)) INLINE static void abundance_ratio_to_solar(
  * @param cooling The #cooling_function_data used in the run.
  * @return Helium reionization energy in CGS units.
  */
-__attribute__((always_inline)) INLINE double qla_helium_reionization_extraheat(
-    const double z, const double delta_z,
-    const struct cooling_function_data *cooling) {
+__attribute__((always_inline)) INLINE static double
+eagle_helium_reionization_extraheat(
+    double z, double delta_z, const struct cooling_function_data *cooling) {
 
 #ifdef SWIFT_DEBUG_CHECKS
   if (delta_z > 0.f) error("Invalid value for delta_z. Should be negative.");
@@ -114,68 +213,43 @@ __attribute__((always_inline)) INLINE double qla_helium_reionization_extraheat(
 
 /**
  * @brief Computes the log_10 of the temperature corresponding to a given
- * internal energy, hydrogen number density, Helium fraction and redshift.
- *
- * Note that the redshift is implicitly passed in via the currently loaded
- * tables in the #cooling_function_data.
- *
- * For the low-z case, we interpolate the flattened 4D table 'u_to_temp' that
- * is arranged in the following way:
- * - 1st dim: redshift, length = eagle_cooling_N_loaded_redshifts
- * - 2nd dim: Hydrogen density, length = eagle_cooling_N_density
- * - 3rd dim: Helium fraction, length = eagle_cooling_N_He_frac
- * - 4th dim: Internal energy, length = eagle_cooling_N_temperature
- *
- * For the high-z case, we interpolate the flattened 3D table 'u_to_temp' that
- * is arranged in the following way:
- * - 1st dim: Hydrogen density, length = eagle_cooling_N_density
- * - 2nd dim: Helium fraction, length = eagle_cooling_N_He_frac
- * - 3rd dim: Internal energy, length = eagle_cooling_N_temperature
+ * internal energy, hydrogen number density, metallicity and redshift
  *
  * @param log_10_u_cgs Log base 10 of internal energy in cgs.
  * @param redshift Current redshift.
  * @param n_H_index Index along the Hydrogen density dimension.
- * @param He_index Index along the Helium fraction dimension.
  * @param d_n_H Offset between Hydrogen density and table[n_H_index].
- * @param d_He Offset between helium fraction and table[He_index].
+ * @param met_index Index along the metallicity dimension.
+ * @param d_met Offset between metallicity and table[met_index].
+ * @param red_index Index along the redshift dimension.
+ * @param d_red Offset between redshift and table[red_index].
  * @param cooling #cooling_function_data structure.
  *
  * @return log_10 of the temperature.
+ *
+ * TO DO: outside table ranges, it uses at the moment the minimum, maximu value
  */
-__attribute__((always_inline)) INLINE double qla_convert_u_to_temp(
-    const double log_10_u_cgs, const float redshift, const int n_H_index,
-    const int He_index, const float d_n_H, const float d_He,
+__attribute__((always_inline)) INLINE static float qla_convert_u_to_temp(
+    const double log_10_u_cgs, const float redshift, int n_H_index, float d_n_H,
+    int met_index, float d_met, int red_index, float d_red,
     const struct cooling_function_data *cooling) {
 
   /* Get index of u along the internal energy axis */
   int u_index;
   float d_u;
-  get_index_1d(cooling->Therm, qla_cooling_N_temperature, log_10_u_cgs,
+
+  get_index_1d(cooling->Therm, qla_cooling_N_internalenergy, log_10_u_cgs,
                &u_index, &d_u);
 
   /* Interpolate temperature table to return temperature for current
-   * internal energy (use 3D interpolation for high redshift table,
-   * otherwise 4D) */
+   * internal energy */
   float log_10_T;
-  if (redshift > cooling->Redshifts[qla_cooling_N_redshifts - 1]) {
-
-    log_10_T = interpolation_3d(cooling->table.temperature,   /* */
-                                n_H_index, He_index, u_index, /* */
-                                d_n_H, d_He, d_u,             /* */
-                                qla_cooling_N_density,        /* */
-                                qla_cooling_N_He_frac,        /* */
-                                qla_cooling_N_temperature);   /* */
-  } else {
-
-    log_10_T =
-        interpolation_4d(cooling->table.temperature,                  /* */
-                         /*z_index=*/0, n_H_index, He_index, u_index, /* */
-                         cooling->dz, d_n_H, d_He, d_u,               /* */
-                         qla_cooling_N_loaded_redshifts,              /* */
-                         qla_cooling_N_density,                       /* */
-                         qla_cooling_N_He_frac,                       /* */
-                         qla_cooling_N_temperature);                  /* */
-  }
+
+  /* Temperature from internal energy */
+  log_10_T = interpolation_4d(
+      cooling->table.T_from_U, red_index, u_index, met_index, n_H_index, d_red,
+      d_u, d_met, d_n_H, qla_cooling_N_redshifts, qla_cooling_N_internalenergy,
+      qla_cooling_N_metallicity, qla_cooling_N_density);
 
   /* Special case for temperatures below the start of the table */
   if (u_index == 0 && d_u == 0.f) {
@@ -189,331 +263,455 @@ __attribute__((always_inline)) INLINE double qla_convert_u_to_temp(
 }
 
 /**
- * @brief Compute the Compton cooling rate from the CMB at a given
- * redshift, electron abundance, temperature and Hydrogen density.
+ * @brief Computes the log_10 of the internal energy corresponding to a given
+ * temperature, hydrogen number density, metallicity and redshift
+ *
+ * @param log_10_T Log base 10 of temperature in K
+ * @param redshift Current redshift.
+ * @param n_H_index Index along the Hydrogen density dimension.
+ * @param d_n_H Offset between Hydrogen density and table[n_H_index].
+ * @param met_index Index along the metallicity dimension.
+ * @param d_met Offset between metallicity and table[met_index].
+ * @param red_index Index along the redshift dimension.
+ * @param d_red Offset between redshift and table[red_index].
+ * @param cooling #cooling_function_data structure.
  *
- * Uses an analytic formula.
+ * @return log_10 of the internal energy in cgs
  *
- * @param cooling The #cooling_function_data used in the run.
- * @param redshift The current redshift.
- * @param n_H_cgs The Hydrogen number density in CGS units.
- * @param temperature The temperature.
- * @param electron_abundance The electron abundance.
+ * TO DO: outside table ranges, it uses at the moment the minimum, maximu value
  */
-__attribute__((always_inline)) INLINE double qla_Compton_cooling_rate(
-    const struct cooling_function_data *cooling, const double redshift,
-    const double n_H_cgs, const double temperature,
-    const double electron_abundance) {
+__attribute__((always_inline)) INLINE static float qla_convert_temp_to_u(
+    const double log_10_T, const float redshift, int n_H_index, float d_n_H,
+    int met_index, float d_met, int red_index, float d_red,
+    const struct cooling_function_data *cooling) {
+
+  /* Get index of u along the internal energy axis */
+  int T_index;
+  float d_T;
 
-  const double zp1 = 1. + redshift;
-  const double zp1p2 = zp1 * zp1;
-  const double zp1p4 = zp1p2 * zp1p2;
+  get_index_1d(cooling->Temp, qla_cooling_N_temperature, log_10_T, &T_index,
+               &d_T);
 
-  /* CMB temperature at this redshift */
-  const double T_CMB = cooling->T_CMB_0 * zp1;
+  /* Interpolate internal energy table to return internal energy for current
+   * temperature */
+  float log_10_U;
 
-  /* Compton cooling rate */
-  return cooling->compton_rate_cgs * (temperature - T_CMB) * zp1p4 *
-         electron_abundance / n_H_cgs;
+  /* Internal energy from temperature*/
+  log_10_U = interpolation_4d(
+      cooling->table.U_from_T, red_index, T_index, met_index, n_H_index, d_red,
+      d_T, d_met, d_n_H, qla_cooling_N_redshifts, qla_cooling_N_temperature,
+      qla_cooling_N_metallicity, qla_cooling_N_density);
+
+  return log_10_U;
 }
 
 /**
- * @brief Computes the cooling rate corresponding to a given internal energy,
- * hydrogen number density, Helium fraction, redshift and metallicity from
- * all the possible channels.
- *
- * 1) Metal-free cooling:
- * We interpolate the flattened 4D table 'H_and_He_net_heating' that is
- * arranged in the following way:
- * - 1st dim: redshift, length = eagle_cooling_N_loaded_redshifts
- * - 2nd dim: Hydrogen density, length = eagle_cooling_N_density
- * - 3rd dim: Helium fraction, length = eagle_cooling_N_He_frac
- * - 4th dim: Internal energy, length = eagle_cooling_N_temperature
- *
- * 2) Electron abundance
- * We compute the electron abundance by interpolating the flattened 4d table
- * 'H_and_He_electron_abundance' that is arranged in the following way:
- * - 1st dim: redshift, length = eagle_cooling_N_loaded_redshifts
- * - 2nd dim: Hydrogen density, length = eagle_cooling_N_density
- * - 3rd dim: Helium fraction, length = eagle_cooling_N_He_frac
- * - 4th dim: Internal energy, length = eagle_cooling_N_temperature
- *
- * 3) Compton cooling is applied via the analytic formula.
- *
- * 4) Solar electron abudance
- * We compute the solar electron abundance by interpolating the flattened 3d
- * table 'solar_electron_abundance' that is arranged in the following way:
- * - 1st dim: redshift, length = eagle_cooling_N_loaded_redshifts
- * - 2nd dim: Hydrogen density, length = eagle_cooling_N_density
- * - 3rd dim: Internal energy, length = eagle_cooling_N_temperature
- *
- * 5) Metal-line cooling
- * For each tracked element we interpolate the flattened 4D table
- * 'table_metals_net_heating' that is arrange in the following way:
- * - 1st dim: element, length = eagle_cooling_N_metal
- * - 2nd dim: redshift, length = eagle_cooling_N_loaded_redshifts
- * - 3rd dim: Hydrogen density, length = eagle_cooling_N_density
- * - 4th dim: Internal energy, length = eagle_cooling_N_temperature
- *
- * Note that this is a fake 4D interpolation as we do not interpolate
- * along the 1st dimension. We just do this once per element.
- *
- * Since only the temperature changes when cooling a given particle,
- * the redshift, hydrogen number density and helium fraction indices
- * and offsets passed in.
- *
- * If the argument element_lambda is non-NULL, the routine also
- * returns the cooling rate per element in the array.
- *
- * @param log10_u_cgs Log base 10 of internal energy per unit mass in CGS units.
- * @param redshift The current redshift
- * @param n_H_cgs The Hydrogen number density in CGS units.
- * @param solar_ratio Array of ratios of particle metal abundances
- * to solar metal abundances
- *
- * @param n_H_index Particle hydrogen number density index
- * @param d_n_H Particle hydrogen number density offset
- * @param He_index Particle helium fraction index
- * @param d_He Particle helium fraction offset
- * @param cooling Cooling data structure
- *
- * @param element_lambda (return) Cooling rate from each element
- *
- * @return The cooling rate
+ * @brief Computes the mean particle mass for a given
+ * metallicity, temperature, redshift, and density.
+ *
+ * @param log_T_cgs Log base 10 of temperature in K
+ * @param redshift Current redshift
+ * @param n_H_cgs Hydrogen number density in cgs
+ * @param ZZsol Metallicity relative to the solar value from the tables
+ * @param n_H_index Index along the Hydrogen number density dimension
+ * @param d_n_H Offset between Hydrogen density and table[n_H_index]
+ * @param met_index Index along the metallicity dimension
+ * @param d_met Offset between metallicity and table[met_index]
+ * @param red_index Index along redshift dimension
+ * @param d_red Offset between redshift and table[red_index]
+ * @param cooling #cooling_function_data structure
+ *
+ * @return linear electron density in cm-3 (NOT the electron fraction)
  */
-INLINE static double qla_metal_cooling_rate(
-    const double log10_u_cgs, const double redshift, const double n_H_cgs,
-    const float solar_ratio[qla_cooling_N_abundances], const int n_H_index,
-    const float d_n_H, const int He_index, const float d_He,
-    const struct cooling_function_data *cooling, double *element_lambda) {
+INLINE static float qla_meanparticlemass_temperature(
+    double log_T_cgs, double redshift, double n_H_cgs, float ZZsol,
+    int n_H_index, float d_n_H, int met_index, float d_met, int red_index,
+    float d_red, const struct cooling_function_data *cooling) {
 
-  /* Temperature */
-  const double log_10_T = qla_convert_u_to_temp(
-      log10_u_cgs, redshift, n_H_index, He_index, d_n_H, d_He, cooling);
-
-  /* Get index along temperature dimension of the tables */
+  /* Get index of T along the temperature axis */
   int T_index;
   float d_T;
-  get_index_1d(cooling->Temp, qla_cooling_N_temperature, log_10_T, &T_index,
+
+  get_index_1d(cooling->Temp, qla_cooling_N_temperature, log_T_cgs, &T_index,
                &d_T);
 
-  /**********************/
-  /* Metal-free cooling */
-  /**********************/
-
-  double Lambda_free;
-
-  if (redshift > cooling->Redshifts[qla_cooling_N_redshifts - 1]) {
-
-    /* If we're using the high redshift tables then we don't interpolate
-     * in redshift */
-    Lambda_free = interpolation_3d(cooling->table.H_plus_He_heating, /* */
-                                   n_H_index, He_index, T_index,     /* */
-                                   d_n_H, d_He, d_T,                 /* */
-                                   qla_cooling_N_density,            /* */
-                                   qla_cooling_N_He_frac,            /* */
-                                   qla_cooling_N_temperature);       /* */
-  } else {
-
-    /* Using normal tables, have to interpolate in redshift */
-    Lambda_free =
-        interpolation_4d(cooling->table.H_plus_He_heating,            /* */
-                         /*z_index=*/0, n_H_index, He_index, T_index, /* */
-                         cooling->dz, d_n_H, d_He, d_T,               /* */
-                         qla_cooling_N_loaded_redshifts,              /* */
-                         qla_cooling_N_density,                       /* */
-                         qla_cooling_N_He_frac,                       /* */
-                         qla_cooling_N_temperature);                  /* */
-  }
+  const float mu = interpolation_4d(
+      cooling->table.Tmu, red_index, T_index, met_index, n_H_index, d_red, d_T,
+      d_met, d_n_H, qla_cooling_N_redshifts, qla_cooling_N_temperature,
+      qla_cooling_N_metallicity, qla_cooling_N_density);
 
-  /* If we're testing cooling rate contributions write to array */
-  if (element_lambda != NULL) {
-    element_lambda[0] = Lambda_free;
-  }
+  return mu;
+}
 
-  /**********************/
-  /* Electron abundance */
-  /**********************/
+/**
+ * @brief Computes the electron density for a given element
+ * abundance ratio, internal energy, redshift, and density.
+ *
+ * @param log_u_cgs Log base 10 of internal energy in cgs [erg g-1]
+ * @param redshift Current redshift
+ * @param n_H_cgs Hydrogen number density in cgs
+ * @param ZZsol Metallicity relative to the solar value from the tables
+ * @param abundance_ratio Abundance ratio for each element x relative to solar
+ * @param n_H_index Index along the Hydrogen number density dimension
+ * @param d_n_H Offset between Hydrogen density and table[n_H_index]
+ * @param met_index Index along the metallicity dimension
+ * @param d_met Offset between metallicity and table[met_index]
+ * @param red_index Index along redshift dimension
+ * @param d_red Offset between redshift and table[red_index]
+ * @param cooling #cooling_function_data structure
+ *
+ * @return linear electron density in cm-3 (NOT the electron fraction)
+ */
+INLINE static float qla_electron_density(
+    double log_u_cgs, double redshift, double n_H_cgs, float ZZsol,
+    const float abundance_ratio[qla_cooling_N_elementtypes], int n_H_index,
+    float d_n_H, int met_index, float d_met, int red_index, float d_red,
+    const struct cooling_function_data *cooling) {
 
-  double H_plus_He_electron_abundance;
+  /* Get index of u along the internal energy axis */
+  int U_index;
+  float d_U;
 
-  if (redshift > cooling->Redshifts[qla_cooling_N_redshifts - 1]) {
+  get_index_1d(cooling->Therm, qla_cooling_N_internalenergy, log_u_cgs,
+               &U_index, &d_U);
 
-    H_plus_He_electron_abundance =
-        interpolation_3d(cooling->table.H_plus_He_electron_abundance, /* */
-                         n_H_index, He_index, T_index,                /* */
-                         d_n_H, d_He, d_T,                            /* */
-                         qla_cooling_N_density,                       /* */
-                         qla_cooling_N_He_frac,                       /* */
-                         qla_cooling_N_temperature);                  /* */
+  /* n_e / n_H */
+  const float electron_fraction = interpolation4d_plus_summation(
+      cooling->table.Uelectron_fraction, abundance_ratio, element_H,
+      qla_cooling_N_electrontypes - 4, red_index, U_index, met_index, n_H_index,
+      d_red, d_U, d_met, d_n_H, qla_cooling_N_redshifts,
+      qla_cooling_N_internalenergy, qla_cooling_N_metallicity,
+      qla_cooling_N_density, qla_cooling_N_electrontypes);
 
-  } else {
+  return electron_fraction * n_H_cgs;
+}
 
-    H_plus_He_electron_abundance =
-        interpolation_4d(cooling->table.H_plus_He_electron_abundance, /* */
-                         /*z_index=*/0, n_H_index, He_index, T_index, /* */
-                         cooling->dz, d_n_H, d_He, d_T,               /* */
-                         qla_cooling_N_loaded_redshifts,              /* */
-                         qla_cooling_N_density,                       /* */
-                         qla_cooling_N_He_frac,                       /* */
-                         qla_cooling_N_temperature);                  /* */
-  }
+/**
+ * @brief Computes the electron density for a given element
+ * abundance ratio, temperature, redshift, and density.
+ *
+ * @param log_T_cgs Log base 10 of temperature
+ * @param redshift Current redshift
+ * @param n_H_cgs Hydrogen number density in cgs
+ * @param ZZsol Metallicity relative to the solar value from the tables
+ * @param abundance_ratio Abundance ratio for each element x relative to solar
+ * @param n_H_index Index along the Hydrogen number density dimension
+ * @param d_n_H Offset between Hydrogen density and table[n_H_index]
+ * @param met_index Index along the metallicity dimension
+ * @param d_met Offset between metallicity and table[met_index]
+ * @param red_index Index along redshift dimension
+ * @param d_red Offset between redshift and table[red_index]
+ * @param cooling #cooling_function_data structure
+ *
+ * @return linear electron density in cm-3 (NOT the electron fraction)
+ */
+INLINE static float qla_electron_density_temperature(
+    double log_T_cgs, double redshift, double n_H_cgs, float ZZsol,
+    const float abundance_ratio[qla_cooling_N_elementtypes], int n_H_index,
+    float d_n_H, int met_index, float d_met, int red_index, float d_red,
+    const struct cooling_function_data *cooling) {
 
-  /**********************/
-  /* Compton cooling    */
-  /**********************/
+  /* Get index of u along the internal energy axis */
+  int T_index;
+  float d_T;
 
-  double Lambda_Compton = 0.;
+  get_index_1d(cooling->Temp, qla_cooling_N_temperature, log_T_cgs, &T_index,
+               &d_T);
 
-  /* Do we need to add the inverse Compton cooling? */
-  /* It is *not* stored in the tables before re-ionisation */
-  if ((redshift > cooling->Redshifts[qla_cooling_N_redshifts - 1]) ||
-      (redshift > cooling->H_reion_z)) {
+  /* n_e / n_H */
+  const float electron_fraction = interpolation4d_plus_summation(
+      cooling->table.Telectron_fraction, abundance_ratio, element_H,
+      qla_cooling_N_electrontypes - 4, red_index, T_index, met_index, n_H_index,
+      d_red, d_T, d_met, d_n_H, qla_cooling_N_redshifts,
+      qla_cooling_N_internalenergy, qla_cooling_N_metallicity,
+      qla_cooling_N_density, qla_cooling_N_electrontypes);
 
-    const double T = exp10(log_10_T);
+  return electron_fraction * n_H_cgs;
+}
 
-    /* Note the minus sign */
-    Lambda_Compton -= qla_Compton_cooling_rate(cooling, redshift, n_H_cgs, T,
-                                               H_plus_He_electron_abundance);
+/**
+ * @brief Computes the net cooling rate (cooling - heating) for a given element
+ * abundance ratio, internal energy, redshift, and density. The unit of the net
+ * cooling rate is Lambda / nH**2 [erg cm^3 s-1] and all input values are in
+ * cgs. The Compton cooling is not taken from the tables but calculated
+ * analytically and added separately
+ *
+ * @param log_u_cgs Log base 10 of internal energy in cgs [erg g-1]
+ * @param redshift Current redshift
+ * @param n_H_cgs Hydrogen number density in cgs
+ * @param abundance_ratio Abundance ratio for each element x relative to solar
+ * @param n_H_index Index along the Hydrogen number density dimension
+ * @param d_n_H Offset between Hydrogen density and table[n_H_index]
+ * @param met_index Index along the metallicity dimension
+ * @param d_met Offset between metallicity and table[met_index]
+ * @param red_index Index along redshift dimension
+ * @param d_red Offset between redshift and table[red_index]
+ * @param cooling #cooling_function_data structure
+ *
+ * @param onlyicool if true / 1 only plot cooling channel icool
+ * @param onlyiheat if true / 1 only plot cooling channel iheat
+ * @param icool cooling channel to be used
+ * @param iheat heating channel to be used
+ *
+ * Throughout the code: onlyicool = onlyiheat = icool = iheat = 0
+ * These are only used for testing.
+ */
+INLINE static double qla_cooling_rate(
+    double log_u_cgs, double redshift, double n_H_cgs,
+    const float abundance_ratio[qla_cooling_N_elementtypes], int n_H_index,
+    float d_n_H, int met_index, float d_met, int red_index, float d_red,
+    const struct cooling_function_data *cooling, int onlyicool, int onlyiheat,
+    int icool, int iheat) {
+
+  /* Set weights for cooling rates */
+  float weights_cooling[qla_cooling_N_cooltypes - 2];
+  for (int i = 0; i < qla_cooling_N_cooltypes - 2; i++) {
+
+    if (i < qla_cooling_N_elementtypes) {
+      /* Use abundance ratios */
+      weights_cooling[i] = abundance_ratio[i];
+    } else if (i == cooltype_Compton) {
+      /* added analytically later, do not use value from table*/
+      weights_cooling[i] = 0.f;
+    } else {
+      /* use same abundances as in the tables */
+      weights_cooling[i] = 1.f;
+    }
   }
 
-  /* If we're testing cooling rate contributions write to array */
-  if (element_lambda != NULL) {
-    element_lambda[1] = Lambda_Compton;
+  /* If we care only about one channel, cancel all the other ones */
+  if (onlyicool != 0) {
+    for (int i = 0; i < qla_cooling_N_cooltypes - 2; i++) {
+      if (i != icool) weights_cooling[i] = 0.f;
+    }
   }
 
-  /*******************************/
-  /* Solar electron abundance    */
-  /*******************************/
-
-  double solar_electron_abundance;
-
-  if (redshift > cooling->Redshifts[qla_cooling_N_redshifts - 1]) {
-
-    /* If we're using the high redshift tables then we don't interpolate
-     * in redshift */
-    solar_electron_abundance =
-        interpolation_2d(cooling->table.electron_abundance, /* */
-                         n_H_index, T_index,                /* */
-                         d_n_H, d_T,                        /* */
-                         qla_cooling_N_density,             /* */
-                         qla_cooling_N_temperature);        /* */
-
-  } else {
-
-    /* Using normal tables, have to interpolate in redshift */
-    solar_electron_abundance =
-        interpolation_3d(cooling->table.electron_abundance, /* */
-                         /*z_index=*/0, n_H_index, T_index, /* */
-                         cooling->dz, d_n_H, d_T,           /* */
-                         qla_cooling_N_loaded_redshifts,    /* */
-                         qla_cooling_N_density,             /* */
-                         qla_cooling_N_temperature);        /* */
+  /* Set weights for heating rates */
+  float weights_heating[qla_cooling_N_heattypes - 2];
+  for (int i = 0; i < qla_cooling_N_heattypes - 2; i++) {
+    if (i < qla_cooling_N_elementtypes) {
+      weights_heating[i] = abundance_ratio[i];
+    } else {
+      weights_heating[i] = 1.f; /* use same abundances as in the tables */
+    }
   }
 
-  const double electron_abundance_ratio =
-      H_plus_He_electron_abundance / solar_electron_abundance;
-
-  /**********************/
-  /* Metal-line cooling */
-  /**********************/
-
-  /* for each element the cooling rate is multiplied by the ratio of H, He
-   * electron abundance to solar electron abundance then by the ratio of the
-   * particle metal abundance to solar metal abundance. */
-
-  double lambda_metal[qla_cooling_N_metal + 2] = {0.};
-
-  if (redshift > cooling->Redshifts[qla_cooling_N_redshifts - 1]) {
-
-    /* Loop over the metals (ignore H and He) */
-    for (int elem = 2; elem < qla_cooling_N_metal + 2; elem++) {
-
-      if (solar_ratio[elem] > 0.) {
-
-        /* Note that we do not interpolate along the x-axis
-         * (element dimension) */
-        lambda_metal[elem] =
-            interpolation_3d_no_x(cooling->table.metal_heating,   /* */
-                                  elem - 2, n_H_index, T_index,   /* */
-                                  /*delta_elem=*/0.f, d_n_H, d_T, /* */
-                                  qla_cooling_N_metal,            /* */
-                                  qla_cooling_N_density,          /* */
-                                  qla_cooling_N_temperature);     /* */
-
-        lambda_metal[elem] *= electron_abundance_ratio;
-        lambda_metal[elem] *= solar_ratio[elem];
-      }
+  /* If we care only about one channel, cancel all the other ones */
+  if (onlyiheat != 0) {
+    for (int i = 0; i < qla_cooling_N_heattypes - 2; i++) {
+      if (i != iheat) weights_heating[i] = 0.f;
     }
+  }
 
-  } else {
-
-    /* Loop over the metals (ignore H and He) */
-    for (int elem = 2; elem < qla_cooling_N_metal + 2; elem++) {
-
-      if (solar_ratio[elem] > 0.) {
+  /* Get index of u along the internal energy axis */
+  int U_index;
+  float d_U;
+  get_index_1d(cooling->Therm, qla_cooling_N_internalenergy, log_u_cgs,
+               &U_index, &d_U);
+
+  /* n_e / n_H */
+  const double electron_fraction = interpolation4d_plus_summation(
+      cooling->table.Uelectron_fraction, abundance_ratio, /* */
+      element_H, qla_cooling_N_electrontypes - 4,         /* */
+      red_index, U_index, met_index, n_H_index,           /* */
+      d_red, d_U, d_met, d_n_H,                           /* */
+      qla_cooling_N_redshifts,                            /* */
+      qla_cooling_N_internalenergy,                       /* */
+      qla_cooling_N_metallicity,                          /* */
+      qla_cooling_N_density,                              /* */
+      qla_cooling_N_electrontypes);                       /* */
+
+  /* Lambda / n_H**2 */
+  const double cooling_rate = interpolation4d_plus_summation(
+      cooling->table.Ucooling, weights_cooling, /* */
+      element_H, qla_cooling_N_cooltypes - 3,   /* */
+      red_index, U_index, met_index, n_H_index, /* */
+      d_red, d_U, d_met, d_n_H,                 /* */
+      qla_cooling_N_redshifts,                  /* */
+      qla_cooling_N_internalenergy,             /* */
+      qla_cooling_N_metallicity,                /* */
+      qla_cooling_N_density,                    /* */
+      qla_cooling_N_cooltypes);                 /* */
+
+  /* Gamma / n_H**2 */
+  const double heating_rate = interpolation4d_plus_summation(
+      cooling->table.Uheating, weights_heating, /* */
+      element_H, qla_cooling_N_heattypes - 3,   /* */
+      red_index, U_index, met_index, n_H_index, /* */
+      d_red, d_U, d_met, d_n_H,                 /* */
+      qla_cooling_N_redshifts,                  /* */
+      qla_cooling_N_internalenergy,             /* */
+      qla_cooling_N_metallicity,                /* */
+      qla_cooling_N_density,                    /* */
+      qla_cooling_N_heattypes);                 /* */
+
+  /* Temperature from internal energy */
+  const double logtemp =
+      interpolation_4d(cooling->table.T_from_U,                  /* */
+                       red_index, U_index, met_index, n_H_index, /* */
+                       d_red, d_U, d_met, d_n_H,                 /* */
+                       qla_cooling_N_redshifts,                  /* */
+                       qla_cooling_N_internalenergy,             /* */
+                       qla_cooling_N_metallicity,                /* */
+                       qla_cooling_N_density);                   /* */
+                                                                 /* */
+  const double temp = exp10(logtemp);
+
+  /* Compton cooling/heating */
+  double Compton_cooling_rate = 0.;
+  if (onlyicool == 0 || (onlyicool == 1 && icool == cooltype_Compton)) {
+
+    const double zp1 = 1. + redshift;
+    const double zp1p2 = zp1 * zp1;
+    const double zp1p4 = zp1p2 * zp1p2;
+
+    /* CMB temperature at this redshift */
+    const double T_CMB = cooling->T_CMB_0 * zp1;
+
+    /* Analytic Compton cooling rate: Lambda_Compton / n_H**2 */
+    Compton_cooling_rate = cooling->compton_rate_cgs * (temp - T_CMB) * zp1p4 *
+                           electron_fraction / n_H_cgs;
+  }
 
-        /* Note that we do not interpolate along the x-axis
-         * (element dimension) */
-        lambda_metal[elem] = interpolation_4d_no_x(
-            cooling->table.metal_heating,                /* */
-            elem - 2, /*z_index=*/0, n_H_index, T_index, /* */
-            /*delta_elem=*/0.f, cooling->dz, d_n_H, d_T, /* */
-            qla_cooling_N_metal,                         /* */
-            qla_cooling_N_loaded_redshifts,              /* */
-            qla_cooling_N_density,                       /* */
-            qla_cooling_N_temperature);                  /* */
+  /* Return the net heating rate (Lambda_heat - Lambda_cool) */
+  return heating_rate - cooling_rate - Compton_cooling_rate;
+}
 
-        lambda_metal[elem] *= electron_abundance_ratio;
-        lambda_metal[elem] *= solar_ratio[elem];
+/**
+ * @brief Computes the net cooling rate (cooling - heating) for a given element
+ * abundance ratio, temperature, redshift, and density. The unit of the net
+ * cooling rate is Lambda / nH**2 [erg cm^3 s-1] and all input values are in
+ * cgs. The Compton cooling is not taken from the tables but calculated
+ * analytically and added separately
+ *
+ * @param log_T_cgs Log base 10 of temperature in K
+ * @param redshift Current redshift
+ * @param n_H_cgs Hydrogen number density in cgs
+ * @param abundance_ratio Abundance ratio for each element x relative to solar
+ * @param n_H_index Index along the Hydrogen number density dimension
+ * @param d_n_H Offset between Hydrogen density and table[n_H_index]
+ * @param met_index Index along the metallicity dimension
+ * @param d_met Offset between metallicity and table[met_index]
+ * @param red_index Index along redshift dimension
+ * @param d_red Offset between redshift and table[red_index]
+ * @param cooling #cooling_function_data structure
+ *
+ * @param onlyicool if true / 1 only plot cooling channel icool
+ * @param onlyiheat if true / 1 only plot cooling channel iheat
+ * @param icool cooling channel to be used
+ * @param iheat heating channel to be used
+ *
+ * Throughout the code: onlyicool = onlyiheat = icool = iheat = 0
+ * These are only used for testing.
+ */
+INLINE static double qla_cooling_rate_temperature(
+    double log_T_cgs, double redshift, double n_H_cgs,
+    const float abundance_ratio[qla_cooling_N_elementtypes], int n_H_index,
+    float d_n_H, int met_index, float d_met, int red_index, float d_red,
+    const struct cooling_function_data *cooling, int onlyicool, int onlyiheat,
+    int icool, int iheat) {
+
+  /* Set weights for cooling rates */
+  float weights_cooling[qla_cooling_N_cooltypes - 2];
+  for (int i = 0; i < qla_cooling_N_cooltypes - 2; i++) {
+
+    if (i < qla_cooling_N_elementtypes) {
+      /* Use abundance ratios */
+      weights_cooling[i] = abundance_ratio[i];
+    } else if (i == cooltype_Compton) {
+      /* added analytically later, do not use value from table*/
+      weights_cooling[i] = 0.f;
+    } else {
+      /* use same abundances as in the tables */
+      weights_cooling[i] = 1.f;
+    }
+  }
 
-        // message("lambda[%d]=%e", elem, lambda_metal[elem]);
-      }
+  /* If we care only about one channel, cancel all the other ones */
+  if (onlyicool != 0) {
+    for (int i = 0; i < qla_cooling_N_cooltypes - 2; i++) {
+      if (i != icool) weights_cooling[i] = 0.f;
     }
   }
 
-  if (element_lambda != NULL) {
-    for (int elem = 2; elem < qla_cooling_N_metal + 2; ++elem) {
-      element_lambda[elem] = lambda_metal[elem];
+  /* Set weights for heating rates */
+  float weights_heating[qla_cooling_N_heattypes - 2];
+  for (int i = 0; i < qla_cooling_N_heattypes - 2; i++) {
+    if (i < qla_cooling_N_elementtypes) {
+      weights_heating[i] = abundance_ratio[i];
+    } else {
+      weights_heating[i] = 1.f; /* use same abundances as in the tables */
     }
   }
 
-  /* Sum up all the contributions */
-  double Lambda_net = Lambda_free + Lambda_Compton;
-  for (int elem = 2; elem < qla_cooling_N_metal + 2; ++elem) {
-    Lambda_net += lambda_metal[elem];
+  /* If we care only about one channel, cancel all the other ones */
+  if (onlyiheat != 0) {
+    for (int i = 0; i < qla_cooling_N_heattypes - 2; i++) {
+      if (i != iheat) weights_heating[i] = 0.f;
+    }
   }
 
-  return Lambda_net;
-}
+  /* Get index of T along the internal energy axis */
+  int T_index;
+  float d_T;
+  get_index_1d(cooling->Temp, qla_cooling_N_temperature, log_T_cgs, &T_index,
+               &d_T);
 
-/**
- * @brief Wrapper function used to calculate cooling rate.
- * Table indices and offsets for redshift, hydrogen number density and
- * helium fraction are passed it so as to compute them only once per particle.
- *
- * @param log10_u_cgs Log base 10 of internal energy per unit mass in CGS units.
- * @param redshift The current redshift.
- * @param n_H_cgs Hydrogen number density in CGS units.
- * @param abundance_ratio Ratio of element abundance to solar.
- *
- * @param n_H_index Particle hydrogen number density index
- * @param d_n_H Particle hydrogen number density offset
- * @param He_index Particle helium fraction index
- * @param d_He Particle helium fraction offset
- * @param cooling #cooling_function_data structure
- *
- * @return The cooling rate
- */
-INLINE static double qla_cooling_rate(
-    const double log10_u_cgs, const double redshift, const double n_H_cgs,
-    const float abundance_ratio[qla_cooling_N_abundances], const int n_H_index,
-    const float d_n_H, const int He_index, const float d_He,
-    const struct cooling_function_data *cooling) {
+  /* n_e / n_H */
+  const double electron_fraction = interpolation4d_plus_summation(
+      cooling->table.Telectron_fraction, abundance_ratio, /* */
+      element_H, qla_cooling_N_electrontypes - 4,         /* */
+      red_index, T_index, met_index, n_H_index,           /* */
+      d_red, d_T, d_met, d_n_H,                           /* */
+      qla_cooling_N_redshifts,                            /* */
+      qla_cooling_N_temperature,                          /* */
+      qla_cooling_N_metallicity,                          /* */
+      qla_cooling_N_density,                              /* */
+      qla_cooling_N_electrontypes);                       /* */
+
+  /* Lambda / n_H**2 */
+  const double cooling_rate = interpolation4d_plus_summation(
+      cooling->table.Tcooling, weights_cooling, /* */
+      element_H, qla_cooling_N_cooltypes - 3,   /* */
+      red_index, T_index, met_index, n_H_index, /* */
+      d_red, d_T, d_met, d_n_H,                 /* */
+      qla_cooling_N_redshifts,                  /* */
+      qla_cooling_N_temperature,                /* */
+      qla_cooling_N_metallicity,                /* */
+      qla_cooling_N_density,                    /* */
+      qla_cooling_N_cooltypes);                 /* */
+
+  /* Gamma / n_H**2 */
+  const double heating_rate = interpolation4d_plus_summation(
+      cooling->table.Theating, weights_heating, /* */
+      element_H, qla_cooling_N_heattypes - 3,   /* */
+      red_index, T_index, met_index, n_H_index, /* */
+      d_red, d_T, d_met, d_n_H,                 /* */
+      qla_cooling_N_redshifts,                  /* */
+      qla_cooling_N_temperature,                /* */
+      qla_cooling_N_metallicity,                /* */
+      qla_cooling_N_density,                    /* */
+      qla_cooling_N_heattypes);                 /* */
+
+  const double temp = exp10(log_T_cgs);
+
+  double Compton_cooling_rate = 0.;
+  if (onlyicool == 0 || (onlyicool == 1 && icool == cooltype_Compton)) {
+
+    const double zp1 = 1. + redshift;
+    const double zp1p2 = zp1 * zp1;
+    const double zp1p4 = zp1p2 * zp1p2;
+
+    /* CMB temperature at this redshift */
+    const double T_CMB = cooling->T_CMB_0 * zp1;
+
+    /* Analytic Compton cooling rate: Lambda_Compton / n_H**2 */
+    Compton_cooling_rate = cooling->compton_rate_cgs * (temp - T_CMB) * zp1p4 *
+                           electron_fraction / n_H_cgs;
+  }
 
-  return qla_metal_cooling_rate(log10_u_cgs, redshift, n_H_cgs, abundance_ratio,
-                                n_H_index, d_n_H, He_index, d_He, cooling,
-                                /* element_lambda=*/NULL);
+  /* Return the net heating rate (Lambda_heat - Lambda_cool) */
+  return heating_rate - cooling_rate - Compton_cooling_rate;
 }
 
 #endif /* SWIFT_QLA_COOLING_RATES_H */
diff --git a/src/cooling/QLA/cooling_struct.h b/src/cooling/QLA/cooling_struct.h
index 0e9405cc1d401b41d0d03e1b53bf11ddd0e8c606..ce0887af4e293a0368b53c7e900ee9a539ac20e8 100644
--- a/src/cooling/QLA/cooling_struct.h
+++ b/src/cooling/QLA/cooling_struct.h
@@ -26,20 +26,50 @@
  */
 struct cooling_tables {
 
-  /* array of heating rates due to metals */
-  float *metal_heating;
+  /* array of all mean particle masses mu (temperature) */
+  float *Tmu;
 
-  /* array of heating rates due to hydrogen and helium */
-  float *H_plus_He_heating;
+  /* array of all mean particle masses mu (internal energy) */
+  float *Umu;
 
-  /* array of electron abundances due to hydrogen and helium */
-  float *H_plus_He_electron_abundance;
+  /* array of all cooling processes (temperature) */
+  float *Tcooling;
 
-  /* array of temperatures */
-  float *temperature;
+  /* array of all cooling processes (internal energy) */
+  float *Ucooling;
 
-  /* array of electron abundances due to metals */
-  float *electron_abundance;
+  /* array of all heating processes (temperature) */
+  float *Theating;
+
+  /* array of all heating processes (internal energy) */
+  float *Uheating;
+
+  /* array of all electron abundances (temperature) */
+  float *Telectron_fraction;
+
+  /* array of all electron abundances (internal energy) */
+  float *Uelectron_fraction;
+
+  /* array to get T from U */
+  float *T_from_U;
+
+  /* array to get U from T */
+  float *U_from_T;
+
+  /* array of equilibrium temperatures */
+  float *logTeq;
+
+  /* array of mean particle masses at equilibrium temperatures */
+  float *meanpartmass_Teq;
+
+  /* array of pressures at equilibrium temperatures */
+  float *logPeq;
+
+  /* array of hydrogen fractions at equilibrium temperature */
+  float *logHfracs_Teq;
+
+  /* array of all hydrogen fractions */
+  float *logHfracs_all;
 };
 
 /**
@@ -59,17 +89,33 @@ struct cooling_function_data {
   /*! Temperature bins */
   float *Temp;
 
-  /*! Helium fraction bins */
-  float *HeFrac;
+  /*! Metallicity bins */
+  float *Metallicity;
 
   /*! Internal energy bins */
   float *Therm;
 
-  /*! Mass fractions of elements for solar abundances (from the tables) */
-  float *SolarAbundances;
+  /*! Abundance ratios for each metallicity bin and for each included element */
+  float *LogAbundances;
+  float *Abundances;
+  float *Abundances_inv;
+
+  /*! Atomic masses for all included elements */
+  float *atomicmass;
+  float *atomicmass_inv;
+
+  /*! Mass fractions of all included elements */
+  float *LogMassFractions;
+  float *MassFractions;
+
+  /*! Index for solar metallicity in the metallicity dimension */
+  int indxZsol;
 
-  /*! Inverse of the solar mass fractions */
-  float *SolarAbundances_inv;
+  /*! Solar metallicity (metal mass fraction) */
+  float *Zsol;
+
+  /*! Inverse of solar metallicity (metal mass fraction) */
+  float *Zsol_inv;
 
   /*! Filepath to the directory containing the HDF5 cooling tables */
   char cooling_table_path[qla_table_path_name_length];
@@ -80,7 +126,7 @@ struct cooling_function_data {
   /*! H reionization energy in CGS units */
   float H_reion_heat_cgs;
 
-  /*! Have we already done H reioisation? */
+  /*! Have we already done H reionization? */
   int H_reion_done;
 
   /*! Redshift of He reionization */
@@ -100,26 +146,35 @@ struct cooling_function_data {
    */
   double internal_energy_from_cgs;
 
+  /*! Pressure conversion from internal units to CGS (for quick access) */
+  double pressure_to_cgs;
+
   /*! Number density conversion from internal units to CGS (for quick access) */
   double number_density_to_cgs;
 
+  /*! Number density conversion from CGS to internal units (for quick access) */
+  double number_density_from_cgs;
+
   /*! Inverse of proton mass in cgs (for quick access) */
   double inv_proton_mass_cgs;
 
-  /*! Temperature of the CMB at present day (for quick access) */
+  /*! Logarithm base 10 of the Boltzmann constant in CGS (for quick access) */
+  double log10_kB_cgs;
+
+  /*! Temperatur of the CMB at present day (for quick access) */
   double T_CMB_0;
 
   /*! Compton rate in cgs units */
   double compton_rate_cgs;
 
-  /*! Index of the current redshift along the redshift index of the tables */
-  int z_index;
+  /*! Minimal temperature allowed for the gas particles */
+  double Tmin;
 
-  /*! Distance between the current redshift and table[z_index] */
-  float dz;
+  /*! Minimal internal energy in cgs allowed for the gas particles */
+  double umin_cgs;
 
-  /*! Index of the previous tables along the redshift index of the tables */
-  int previous_z_index;
+  /*! Threshold to switch between rapid and slow cooling regimes. */
+  double rapid_cooling_threshold;
 };
 
 /**
diff --git a/src/cooling/QLA/cooling_tables.c b/src/cooling/QLA/cooling_tables.c
index 22fbc6da0d321937c124c581300f535ed182e4e1..20d07767970e45dc2fa52e11cd1f4f8d585874c5 100644
--- a/src/cooling/QLA/cooling_tables.c
+++ b/src/cooling/QLA/cooling_tables.c
@@ -25,6 +25,10 @@
 /* Config parameters. */
 #include "../config.h"
 
+/* This file's header */
+#include "cooling_tables.h"
+
+/* Standard includes */
 #include <hdf5.h>
 #include <math.h>
 #include <stdlib.h>
@@ -33,268 +37,191 @@
 /* Local includes. */
 #include "chemistry_struct.h"
 #include "cooling_struct.h"
-#include "cooling_tables.h"
 #include "error.h"
+#include "exp10.h"
 #include "interpolate.h"
 
 /**
- * @brief Names of the elements in the order they are stored in the files
- */
-static const char *qla_tables_element_names[qla_cooling_N_metal] = {
-    "Carbon",  "Nitrogen", "Oxygen",  "Neon", "Magnesium",
-    "Silicon", "Sulphur",  "Calcium", "Iron"};
-
-/*! Number of elements in a z-slice of the H+He cooling rate tables */
-static const size_t num_elements_cooling_rate =
-    qla_cooling_N_temperature * qla_cooling_N_density;
-
-/*! Number of elements in a z-slice of the metal cooling rate tables */
-static const size_t num_elements_metal_heating =
-    qla_cooling_N_metal * qla_cooling_N_temperature * qla_cooling_N_density;
-
-/*! Number of elements in a z-slice of the metal electron abundance tables */
-static const size_t num_elements_electron_abundance =
-    qla_cooling_N_temperature * qla_cooling_N_density;
-
-/*! Number of elements in a z-slice of the temperature tables */
-static const size_t num_elements_temperature =
-    qla_cooling_N_He_frac * qla_cooling_N_temperature * qla_cooling_N_density;
-
-/*! Number of elements in a z-slice of the H+He cooling rate tables */
-static const size_t num_elements_HpHe_heating =
-    qla_cooling_N_He_frac * qla_cooling_N_temperature * qla_cooling_N_density;
-
-/*! Number of elements in a z-slice of the H+He electron abundance tables */
-static const size_t num_elements_HpHe_electron_abundance =
-    qla_cooling_N_He_frac * qla_cooling_N_temperature * qla_cooling_N_density;
-
-/**
- * @brief Reads in EAGLE table of redshift values
+ * @brief Reads in COLIBRE cooling table header. Consists of tables
+ * of values for temperature, hydrogen number density, metallicity,
+ * abundance ratios, and elements used to index the cooling tables.
  *
- * @param cooling #cooling_function_data structure
+ * @param cooling Cooling data structure
  */
-void get_cooling_redshifts(struct cooling_function_data *cooling) {
-
-  /* Read the list of table redshifts */
-  char redshift_filename[qla_table_path_name_length + 16];
-  sprintf(redshift_filename, "%s/redshifts.dat", cooling->cooling_table_path);
-
-  FILE *infile = fopen(redshift_filename, "r");
-  if (infile == NULL) {
-    error("Cannot open the list of cooling table redshifts (%s)",
-          redshift_filename);
-  }
-
-  int N_Redshifts = -1;
-
-  /* Read the file */
-  if (!feof(infile)) {
-
-    char buffer[50];
+void read_cooling_header(struct cooling_function_data *cooling) {
 
-    /* Read the number of redshifts (1st line in the file) */
-    if (fgets(buffer, 50, infile) != NULL)
-      N_Redshifts = atoi(buffer);
-    else
-      error("Impossible to read the number of redshifts");
-
-    /* Be verbose about it */
-    message("Found cooling tables at %d redhsifts", N_Redshifts);
-
-    /* Check value */
-    if (N_Redshifts != qla_cooling_N_redshifts)
-      error("Invalid redshift length array.");
+#ifdef HAVE_HDF5
 
-    /* Allocate the list of redshifts */
-    if (swift_memalign("cooling", (void **)&cooling->Redshifts,
-                       SWIFT_STRUCT_ALIGNMENT,
-                       qla_cooling_N_redshifts * sizeof(float)) != 0)
-      error("Failed to allocate redshift table");
+  hid_t dataset;
+  herr_t status;
 
-    /* Read all the redshift values */
-    int count = 0;
-    while (!feof(infile)) {
-      if (fgets(buffer, 50, infile) != NULL) {
-        cooling->Redshifts[count] = atof(buffer);
-        count++;
-      }
-    }
+  /* read sizes of array dimensions */
+  hid_t tempfile_id =
+      H5Fopen(cooling->cooling_table_path, H5F_ACC_RDONLY, H5P_DEFAULT);
+  if (tempfile_id < 0)
+    error("unable to open file %s\n", cooling->cooling_table_path);
+
+  /* allocate arrays of bins */
+  if (posix_memalign((void **)&cooling->Temp, SWIFT_STRUCT_ALIGNMENT,
+                     qla_cooling_N_temperature * sizeof(float)) != 0)
+    error("Failed to allocate temperature table\n");
+
+  if (posix_memalign((void **)&cooling->Redshifts, SWIFT_STRUCT_ALIGNMENT,
+                     qla_cooling_N_redshifts * sizeof(float)) != 0)
+    error("Failed to allocate redshift table\n");
+
+  if (posix_memalign((void **)&cooling->nH, SWIFT_STRUCT_ALIGNMENT,
+                     qla_cooling_N_density * sizeof(float)) != 0)
+    error("Failed to allocate density table\n");
+
+  if (posix_memalign((void **)&cooling->Metallicity, SWIFT_STRUCT_ALIGNMENT,
+                     qla_cooling_N_metallicity * sizeof(float)) != 0)
+    error("Failed to allocate metallicity table\n");
+
+  if (posix_memalign((void **)&cooling->Therm, SWIFT_STRUCT_ALIGNMENT,
+                     qla_cooling_N_internalenergy * sizeof(float)) != 0)
+    error("Failed to allocate internal energy table\n");
+
+  if (posix_memalign((void **)&cooling->LogAbundances, SWIFT_STRUCT_ALIGNMENT,
+                     qla_cooling_N_metallicity * qla_cooling_N_elementtypes *
+                         sizeof(float)) != 0)
+    error("Failed to allocate abundance array\n");
 
-    /* Verify that the file was self-consistent */
-    if (count != N_Redshifts) {
-      error(
-          "Redshift file (%s) does not contain the correct number of redshifts "
-          "(%d vs. %d)",
-          redshift_filename, count, N_Redshifts);
-    }
-  } else {
-    error("Redshift file (%s) is empty!", redshift_filename);
-  }
+  if (posix_memalign((void **)&cooling->Abundances, SWIFT_STRUCT_ALIGNMENT,
+                     qla_cooling_N_metallicity * qla_cooling_N_elementtypes *
+                         sizeof(float)) != 0)
+    error("Failed to allocate abundance array\n");
 
-  /* We are done with this file */
-  fclose(infile);
+  if (posix_memalign((void **)&cooling->Abundances_inv, SWIFT_STRUCT_ALIGNMENT,
+                     qla_cooling_N_metallicity * qla_cooling_N_elementtypes *
+                         sizeof(float)) != 0)
+    error("Failed to allocate abundance array\n");
 
-  /* QLA cooling assumes cooling->Redshifts table is in increasing order. Test
-   * this. */
-  for (int i = 0; i < N_Redshifts - 1; i++) {
-    if (cooling->Redshifts[i + 1] < cooling->Redshifts[i]) {
-      error("table should be in increasing order\n");
-    }
-  }
-}
+  if (posix_memalign((void **)&cooling->atomicmass, SWIFT_STRUCT_ALIGNMENT,
+                     qla_cooling_N_elementtypes * sizeof(float)) != 0)
+    error("Failed to allocate atomic masses array\n");
 
-/**
- * @brief Reads in QLA cooling table header. Consists of tables
- * of values for temperature, hydrogen number density, helium fraction
- * solar element abundances, and elements used to index the cooling tables.
- *
- * @param fname Filepath for cooling table from which to read header
- * @param cooling Cooling data structure
- */
-void read_cooling_header(const char *fname,
-                         struct cooling_function_data *cooling) {
+  if (posix_memalign((void **)&cooling->atomicmass_inv, SWIFT_STRUCT_ALIGNMENT,
+                     qla_cooling_N_elementtypes * sizeof(float)) != 0)
+    error("Failed to allocate inverse atomic masses array\n");
 
-#ifdef HAVE_HDF5
+  if (posix_memalign((void **)&cooling->Zsol, SWIFT_STRUCT_ALIGNMENT,
+                     1 * sizeof(float)) != 0)
+    error("Failed to allocate solar metallicity array\n");
 
-  int N_Temp, N_nH, N_He, N_SolarAbundances, N_Elements;
+  if (posix_memalign((void **)&cooling->Zsol_inv, SWIFT_STRUCT_ALIGNMENT,
+                     1 * sizeof(float)) != 0)
+    error("Failed to allocate inverse solar metallicity array\n");
 
-  /* read sizes of array dimensions */
-  hid_t tempfile_id = H5Fopen(fname, H5F_ACC_RDONLY, H5P_DEFAULT);
-  if (tempfile_id < 0) error("unable to open file %s\n", fname);
-
-  /* read size of each table of values */
-  hid_t dataset =
-      H5Dopen(tempfile_id, "/Header/Number_of_temperature_bins", H5P_DEFAULT);
-  herr_t status =
-      H5Dread(dataset, H5T_NATIVE_INT, H5S_ALL, H5S_ALL, H5P_DEFAULT, &N_Temp);
-  if (status < 0) error("error reading number of temperature bins");
-  status = H5Dclose(dataset);
-  if (status < 0) error("error closing cooling dataset");
+  if (posix_memalign((void **)&cooling->LogMassFractions,
+                     SWIFT_STRUCT_ALIGNMENT,
+                     qla_cooling_N_metallicity * qla_cooling_N_elementtypes *
+                         sizeof(float)) != 0)
+    error("Failed to allocate log mass fraction array\n");
 
-  /* Check value */
-  if (N_Temp != qla_cooling_N_temperature)
-    error("Invalid temperature array length.");
+  if (posix_memalign((void **)&cooling->MassFractions, SWIFT_STRUCT_ALIGNMENT,
+                     qla_cooling_N_metallicity * qla_cooling_N_elementtypes *
+                         sizeof(float)) != 0)
+    error("Failed to allocate mass fraction array\n");
 
-  dataset = H5Dopen(tempfile_id, "/Header/Number_of_density_bins", H5P_DEFAULT);
-  status =
-      H5Dread(dataset, H5T_NATIVE_INT, H5S_ALL, H5S_ALL, H5P_DEFAULT, &N_nH);
-  if (status < 0) error("error reading number of density bins");
+  /* read in bins and misc information */
+  dataset = H5Dopen(tempfile_id, "/TableBins/TemperatureBins", H5P_DEFAULT);
+  status = H5Dread(dataset, H5T_NATIVE_FLOAT, H5S_ALL, H5S_ALL, H5P_DEFAULT,
+                   cooling->Temp);
+  if (status < 0) error("error reading temperature bins\n");
   status = H5Dclose(dataset);
   if (status < 0) error("error closing cooling dataset");
 
-  /* Check value */
-  if (N_nH != qla_cooling_N_density) error("Invalid density array length.");
-
-  dataset =
-      H5Dopen(tempfile_id, "/Header/Number_of_helium_fractions", H5P_DEFAULT);
-  status =
-      H5Dread(dataset, H5T_NATIVE_INT, H5S_ALL, H5S_ALL, H5P_DEFAULT, &N_He);
-  if (status < 0) error("error reading number of He fraction bins");
+  dataset = H5Dopen(tempfile_id, "/TableBins/RedshiftBins", H5P_DEFAULT);
+  status = H5Dread(dataset, H5T_NATIVE_FLOAT, H5S_ALL, H5S_ALL, H5P_DEFAULT,
+                   cooling->Redshifts);
+  if (status < 0) error("error reading redshift bins\n");
   status = H5Dclose(dataset);
   if (status < 0) error("error closing cooling dataset");
 
-  /* Check value */
-  if (N_He != qla_cooling_N_He_frac)
-    error("Invalid Helium fraction array length.");
-
-  dataset = H5Dopen(tempfile_id, "/Header/Abundances/Number_of_abundances",
-                    H5P_DEFAULT);
-  status = H5Dread(dataset, H5T_NATIVE_INT, H5S_ALL, H5S_ALL, H5P_DEFAULT,
-                   &N_SolarAbundances);
-  if (status < 0) error("error reading number of solar abundance bins");
+  dataset = H5Dopen(tempfile_id, "/TableBins/DensityBins", H5P_DEFAULT);
+  status = H5Dread(dataset, H5T_NATIVE_FLOAT, H5S_ALL, H5S_ALL, H5P_DEFAULT,
+                   cooling->nH);
+  if (status < 0) error("error reading density bins\n");
   status = H5Dclose(dataset);
   if (status < 0) error("error closing cooling dataset");
 
-  /* Check value */
-  if (N_SolarAbundances != qla_cooling_N_abundances)
-    error("Invalid solar abundances array length.");
-
-  dataset = H5Dopen(tempfile_id, "/Header/Number_of_metals", H5P_DEFAULT);
-  status = H5Dread(dataset, H5T_NATIVE_INT, H5S_ALL, H5S_ALL, H5P_DEFAULT,
-                   &N_Elements);
-  if (status < 0) error("error reading number of metal bins");
+  dataset = H5Dopen(tempfile_id, "/TableBins/MetallicityBins", H5P_DEFAULT);
+  status = H5Dread(dataset, H5T_NATIVE_FLOAT, H5S_ALL, H5S_ALL, H5P_DEFAULT,
+                   cooling->Metallicity);
+  if (status < 0) error("error reading metallicity bins\n");
   status = H5Dclose(dataset);
   if (status < 0) error("error closing cooling dataset");
 
-  /* Check value */
-  if (N_Elements != qla_cooling_N_metal) error("Invalid metal array length.");
-
-  /* allocate arrays of values for each of the above quantities */
-  if (swift_memalign("cooling", (void **)&cooling->Temp, SWIFT_STRUCT_ALIGNMENT,
-                     N_Temp * sizeof(float)) != 0)
-    error("Failed to allocate temperature table");
-  if (swift_memalign("cooling", (void **)&cooling->Therm,
-                     SWIFT_STRUCT_ALIGNMENT, N_Temp * sizeof(float)) != 0)
-    error("Failed to allocate internal energy table");
-  if (swift_memalign("cooling", (void **)&cooling->nH, SWIFT_STRUCT_ALIGNMENT,
-                     N_nH * sizeof(float)) != 0)
-    error("Failed to allocate nH table");
-  if (swift_memalign("cooling", (void **)&cooling->HeFrac,
-                     SWIFT_STRUCT_ALIGNMENT, N_He * sizeof(float)) != 0)
-    error("Failed to allocate HeFrac table");
-  if (swift_memalign("cooling", (void **)&cooling->SolarAbundances,
-                     SWIFT_STRUCT_ALIGNMENT,
-                     N_SolarAbundances * sizeof(float)) != 0)
-    error("Failed to allocate Solar abundances table");
-  if (swift_memalign("cooling", (void **)&cooling->SolarAbundances_inv,
-                     SWIFT_STRUCT_ALIGNMENT,
-                     N_SolarAbundances * sizeof(float)) != 0)
-    error("Failed to allocate Solar abundances inverses table");
-
-  /* read in values for each of the arrays */
-  dataset = H5Dopen(tempfile_id, "/Solar/Temperature_bins", H5P_DEFAULT);
+  dataset = H5Dopen(tempfile_id, "/TableBins/InternalEnergyBins", H5P_DEFAULT);
   status = H5Dread(dataset, H5T_NATIVE_FLOAT, H5S_ALL, H5S_ALL, H5P_DEFAULT,
-                   cooling->Temp);
-  if (status < 0) error("error reading temperature bins");
+                   cooling->Therm);
+  if (status < 0) error("error reading internal energy bins\n");
   status = H5Dclose(dataset);
   if (status < 0) error("error closing cooling dataset");
 
-  dataset = H5Dopen(tempfile_id, "/Solar/Hydrogen_density_bins", H5P_DEFAULT);
+  dataset = H5Dopen(tempfile_id, "/TotalAbundances", H5P_DEFAULT);
   status = H5Dread(dataset, H5T_NATIVE_FLOAT, H5S_ALL, H5S_ALL, H5P_DEFAULT,
-                   cooling->nH);
-  if (status < 0) error("error reading H density bins");
+                   cooling->LogAbundances);
+  if (status < 0) error("error reading total abundances\n");
   status = H5Dclose(dataset);
   if (status < 0) error("error closing cooling dataset");
 
-  dataset = H5Dopen(tempfile_id, "/Metal_free/Helium_mass_fraction_bins",
-                    H5P_DEFAULT);
+  dataset = H5Dopen(tempfile_id, "/TotalMassFractions", H5P_DEFAULT);
   status = H5Dread(dataset, H5T_NATIVE_FLOAT, H5S_ALL, H5S_ALL, H5P_DEFAULT,
-                   cooling->HeFrac);
-  if (status < 0) error("error reading He fraction bins");
+                   cooling->LogMassFractions);
+  if (status < 0) error("error reading total mass fractions\n");
   status = H5Dclose(dataset);
   if (status < 0) error("error closing cooling dataset");
 
-  dataset = H5Dopen(tempfile_id, "/Header/Abundances/Solar_mass_fractions",
-                    H5P_DEFAULT);
+  dataset = H5Dopen(tempfile_id, "/ElementMasses", H5P_DEFAULT);
   status = H5Dread(dataset, H5T_NATIVE_FLOAT, H5S_ALL, H5S_ALL, H5P_DEFAULT,
-                   cooling->SolarAbundances);
-  if (status < 0) error("error reading solar mass fraction bins");
+                   cooling->atomicmass);
+  if (status < 0) error("error reading element masses\n");
   status = H5Dclose(dataset);
   if (status < 0) error("error closing cooling dataset");
 
-  dataset = H5Dopen(tempfile_id, "/Metal_free/Temperature/Energy_density_bins",
-                    H5P_DEFAULT);
+  dataset = H5Dopen(tempfile_id, "/SolarMetallicity", H5P_DEFAULT);
   status = H5Dread(dataset, H5T_NATIVE_FLOAT, H5S_ALL, H5S_ALL, H5P_DEFAULT,
-                   cooling->Therm);
-  if (status < 0) error("error reading internal energy bins");
+                   cooling->Zsol);
+  if (status < 0) error("error reading solar metallicity \n");
   status = H5Dclose(dataset);
   if (status < 0) error("error closing cooling dataset");
 
-  /* Convert to temperature, density and internal energy arrays to log10 */
-  for (int i = 0; i < N_Temp; i++) {
-    cooling->Temp[i] = log10(cooling->Temp[i]);
-    cooling->Therm[i] = log10(cooling->Therm[i]);
+  /* Close the file */
+  H5Fclose(tempfile_id);
+
+  cooling->Zsol_inv[0] = 1.f / cooling->Zsol[0];
+
+  /* find the metallicity bin that refers to solar metallicity */
+  const float tol = 1.e-3;
+  for (int i = 0; i < qla_cooling_N_metallicity; i++) {
+    if (fabsf(cooling->Metallicity[i]) < tol) {
+      cooling->indxZsol = i;
+    }
   }
-  for (int i = 0; i < N_nH; i++) {
-    cooling->nH[i] = log10(cooling->nH[i]);
+
+#if defined(__ICC)
+#pragma novector
+#endif
+  for (int i = 0; i < qla_cooling_N_elementtypes; i++) {
+    cooling->atomicmass_inv[i] = 1.f / cooling->atomicmass[i];
   }
 
-    /* Compute inverse of solar mass fractions */
+  /* set some additional useful abundance arrays */
+  for (int i = 0; i < qla_cooling_N_metallicity; i++) {
+
 #if defined(__ICC)
 #pragma novector
 #endif
-  for (int i = 0; i < N_SolarAbundances; ++i) {
-    cooling->SolarAbundances_inv[i] = 1.f / cooling->SolarAbundances[i];
+    for (int j = 0; j < qla_cooling_N_elementtypes; j++) {
+      const int indx1d = row_major_index_2d(i, j, qla_cooling_N_metallicity,
+                                            qla_cooling_N_elementtypes);
+      cooling->Abundances[indx1d] = exp10f(cooling->LogAbundances[indx1d]);
+      cooling->Abundances_inv[indx1d] = 1.f / cooling->Abundances[indx1d];
+      cooling->MassFractions[indx1d] =
+          exp10f(cooling->LogMassFractions[indx1d]);
+    }
   }
 
 #else
@@ -303,461 +230,255 @@ void read_cooling_header(const char *fname,
 }
 
 /**
- * @brief Allocate space for cooling tables.
+ * @brief Allocate space for cooling tables and read them
  *
  * @param cooling #cooling_function_data structure
  */
-void allocate_cooling_tables(struct cooling_function_data *restrict cooling) {
+void read_cooling_tables(struct cooling_function_data *restrict cooling) {
 
-  /* Allocate arrays to store cooling tables. Arrays contain two tables of
-   * cooling rates with one table being for the redshift above current redshift
-   * and one below. */
+#ifdef HAVE_HDF5
+  hid_t dataset;
+  herr_t status;
 
-  if (swift_memalign("cooling-tables", (void **)&cooling->table.metal_heating,
-                     SWIFT_STRUCT_ALIGNMENT,
-                     qla_cooling_N_loaded_redshifts *
-                         num_elements_metal_heating * sizeof(float)) != 0)
-    error("Failed to allocate metal_heating array");
+  /* open hdf5 file */
+  hid_t tempfile_id =
+      H5Fopen(cooling->cooling_table_path, H5F_ACC_RDONLY, H5P_DEFAULT);
+  if (tempfile_id < 0)
+    error("unable to open file %s\n", cooling->cooling_table_path);
 
-  if (swift_memalign("cooling-tables",
-                     (void **)&cooling->table.electron_abundance,
-                     SWIFT_STRUCT_ALIGNMENT,
-                     qla_cooling_N_loaded_redshifts *
-                         num_elements_electron_abundance * sizeof(float)) != 0)
-    error("Failed to allocate electron_abundance array");
+  /* Allocate and read arrays to store cooling tables. */
 
-  if (swift_memalign("cooling-tables", (void **)&cooling->table.temperature,
-                     SWIFT_STRUCT_ALIGNMENT,
-                     qla_cooling_N_loaded_redshifts * num_elements_temperature *
+  /* Mean particle mass (temperature) */
+  if (posix_memalign((void **)&cooling->table.Tmu, SWIFT_STRUCT_ALIGNMENT,
+                     qla_cooling_N_redshifts * qla_cooling_N_temperature *
+                         qla_cooling_N_metallicity * qla_cooling_N_density *
                          sizeof(float)) != 0)
-    error("Failed to allocate temperature array");
+    error("Failed to allocate Tmu array\n");
 
-  if (swift_memalign("cooling-tables",
-                     (void **)&cooling->table.H_plus_He_heating,
-                     SWIFT_STRUCT_ALIGNMENT,
-                     qla_cooling_N_loaded_redshifts *
-                         num_elements_HpHe_heating * sizeof(float)) != 0)
-    error("Failed to allocate H_plus_He_heating array");
+  dataset = H5Dopen(tempfile_id, "/Tdep/MeanParticleMass", H5P_DEFAULT);
+  status = H5Dread(dataset, H5T_NATIVE_FLOAT, H5S_ALL, H5S_ALL, H5P_DEFAULT,
+                   cooling->table.Tmu);
+  if (status < 0) error("error reading Tmu\n");
+  status = H5Dclose(dataset);
+  if (status < 0) error("error closing mean particle mass dataset");
 
-  if (swift_memalign("cooling-tables",
-                     (void **)&cooling->table.H_plus_He_electron_abundance,
-                     SWIFT_STRUCT_ALIGNMENT,
-                     qla_cooling_N_loaded_redshifts *
-                         num_elements_HpHe_electron_abundance *
+  /* Mean particle mass (internal energy) */
+  if (posix_memalign((void **)&cooling->table.Umu, SWIFT_STRUCT_ALIGNMENT,
+                     qla_cooling_N_redshifts * qla_cooling_N_internalenergy *
+                         qla_cooling_N_metallicity * qla_cooling_N_density *
                          sizeof(float)) != 0)
-    error("Failed to allocate H_plus_He_electron_abundance array");
-}
+    error("Failed to allocate Umu array\n");
 
-/**
- * @brief Get the redshift invariant table of cooling rates (before reionization
- * at redshift ~9) Reads in table of cooling rates and electron abundances due
- * to metals (depending on temperature, hydrogen number density), cooling rates
- * and electron abundances due to hydrogen and helium (depending on temperature,
- * hydrogen number density and helium fraction), and temperatures (depending on
- * internal energy, hydrogen number density and helium fraction; note: this is
- * distinct from table of temperatures read in ReadCoolingHeader, as that table
- * is used to index the cooling, electron abundance tables, whereas this one is
- * used to obtain temperature of particle)
- *
- * @param cooling #cooling_function_data structure
- * @param photodis Are we loading the photo-dissociation table?
- */
-void get_redshift_invariant_table(
-    struct cooling_function_data *restrict cooling, const int photodis) {
-#ifdef HAVE_HDF5
+  dataset = H5Dopen(tempfile_id, "/Udep/MeanParticleMass", H5P_DEFAULT);
+  status = H5Dread(dataset, H5T_NATIVE_FLOAT, H5S_ALL, H5S_ALL, H5P_DEFAULT,
+                   cooling->table.Umu);
+  if (status < 0) error("error reading Umu\n");
+  status = H5Dclose(dataset);
+  if (status < 0) error("error closing mean particle mass dataset");
 
-  /* Temporary tables */
-  float *net_cooling_rate = NULL;
-  float *electron_abundance = NULL;
-  float *temperature = NULL;
-  float *he_net_cooling_rate = NULL;
-  float *he_electron_abundance = NULL;
+  /* Cooling (temperature) */
+  if (posix_memalign((void **)&cooling->table.Tcooling, SWIFT_STRUCT_ALIGNMENT,
+                     qla_cooling_N_redshifts * qla_cooling_N_temperature *
+                         qla_cooling_N_metallicity * qla_cooling_N_density *
+                         qla_cooling_N_cooltypes * sizeof(float)) != 0)
+    error("Failed to allocate Tcooling array\n");
 
-  /* Allocate arrays for reading in cooling tables.  */
-  if (swift_memalign("cooling-temp", (void **)&net_cooling_rate,
-                     SWIFT_STRUCT_ALIGNMENT,
-                     num_elements_cooling_rate * sizeof(float)) != 0)
-    error("Failed to allocate net_cooling_rate array");
-  if (swift_memalign("cooling-temp", (void **)&electron_abundance,
-                     SWIFT_STRUCT_ALIGNMENT,
-                     num_elements_electron_abundance * sizeof(float)) != 0)
-    error("Failed to allocate electron_abundance array");
-  if (swift_memalign("cooling-temp", (void **)&temperature,
-                     SWIFT_STRUCT_ALIGNMENT,
-                     num_elements_temperature * sizeof(float)) != 0)
-    error("Failed to allocate temperature array");
-  if (swift_memalign("cooling-temp", (void **)&he_net_cooling_rate,
-                     SWIFT_STRUCT_ALIGNMENT,
-                     num_elements_HpHe_heating * sizeof(float)) != 0)
-    error("Failed to allocate he_net_cooling_rate array");
-  if (swift_memalign("cooling-temp", (void **)&he_electron_abundance,
-                     SWIFT_STRUCT_ALIGNMENT,
-                     num_elements_HpHe_electron_abundance * sizeof(float)) != 0)
-    error("Failed to allocate he_electron_abundance array");
-
-  /* Decide which high redshift table to read. Indices set in cooling_update */
-  char filename[qla_table_path_name_length + 21];
-  if (photodis) {
-    sprintf(filename, "%sz_photodis.hdf5", cooling->cooling_table_path);
-    message("Reading cooling table 'z_photodis.hdf5'");
-  } else {
-    sprintf(filename, "%sz_8.989nocompton.hdf5", cooling->cooling_table_path);
-    message("Reading cooling table 'z_8.989nocompton.hdf5'");
-  }
+  dataset = H5Dopen(tempfile_id, "/Tdep/Cooling", H5P_DEFAULT);
+  status = H5Dread(dataset, H5T_NATIVE_FLOAT, H5S_ALL, H5S_ALL, H5P_DEFAULT,
+                   cooling->table.Tcooling);
+  if (status < 0) error("error reading Tcooling\n");
+  status = H5Dclose(dataset);
+  if (status < 0) error("error closing cooling dataset");
 
-  hid_t file_id = H5Fopen(filename, H5F_ACC_RDONLY, H5P_DEFAULT);
-  if (file_id < 0) error("unable to open file %s\n", filename);
-
-  char set_name[64];
-
-  /* read in cooling rates due to metals */
-  for (int specs = 0; specs < qla_cooling_N_metal; specs++) {
-
-    /* Read in the cooling rate for this metal */
-    sprintf(set_name, "/%s/Net_Cooling", qla_tables_element_names[specs]);
-    hid_t dataset = H5Dopen(file_id, set_name, H5P_DEFAULT);
-    herr_t status = H5Dread(dataset, H5T_NATIVE_FLOAT, H5S_ALL, H5S_ALL,
-                            H5P_DEFAULT, net_cooling_rate);
-    if (status < 0) error("error reading metal cooling rate table");
-    status = H5Dclose(dataset);
-    if (status < 0) error("error closing cooling dataset");
-
-    /* Transpose from order tables are stored in (temperature, nH)
-     * to (metal species, nH, temperature) where fastest
-     * varying index is on right. Tables contain cooling rates but we
-     * want rate of change of internal energy, hence minus sign. */
-    for (int j = 0; j < qla_cooling_N_temperature; j++) {
-      for (int k = 0; k < qla_cooling_N_density; k++) {
-
-        /* Index in the HDF5 table */
-        const int hdf5_index = row_major_index_2d(
-            j, k, qla_cooling_N_temperature, qla_cooling_N_density);
-
-        /* Index in the internal table */
-        const int internal_index = row_major_index_3d(
-            specs, k, j, qla_cooling_N_metal, qla_cooling_N_density,
-            qla_cooling_N_temperature);
-
-        /* Change the sign and transpose */
-        cooling->table.metal_heating[internal_index] =
-            -net_cooling_rate[hdf5_index];
-      }
-    }
-  }
+  /* Cooling (internal energy) */
+  if (posix_memalign((void **)&cooling->table.Ucooling, SWIFT_STRUCT_ALIGNMENT,
+                     qla_cooling_N_redshifts * qla_cooling_N_internalenergy *
+                         qla_cooling_N_metallicity * qla_cooling_N_density *
+                         qla_cooling_N_cooltypes * sizeof(float)) != 0)
+    error("Failed to allocate Ucooling array\n");
 
-  /* read in cooling rates due to H + He */
-  strcpy(set_name, "/Metal_free/Net_Cooling");
-  hid_t dataset = H5Dopen(file_id, set_name, H5P_DEFAULT);
-  herr_t status = H5Dread(dataset, H5T_NATIVE_FLOAT, H5S_ALL, H5S_ALL,
-                          H5P_DEFAULT, he_net_cooling_rate);
-  if (status < 0) error("error reading metal free cooling rate table");
+  dataset = H5Dopen(tempfile_id, "/Udep/Cooling", H5P_DEFAULT);
+  status = H5Dread(dataset, H5T_NATIVE_FLOAT, H5S_ALL, H5S_ALL, H5P_DEFAULT,
+                   cooling->table.Ucooling);
+  if (status < 0) error("error reading Ucooling\n");
   status = H5Dclose(dataset);
   if (status < 0) error("error closing cooling dataset");
 
-  /* read in Temperatures */
-  strcpy(set_name, "/Metal_free/Temperature/Temperature");
-  dataset = H5Dopen(file_id, set_name, H5P_DEFAULT);
+  /* Heating (temperature) */
+  if (posix_memalign((void **)&cooling->table.Theating, SWIFT_STRUCT_ALIGNMENT,
+                     qla_cooling_N_redshifts * qla_cooling_N_temperature *
+                         qla_cooling_N_metallicity * qla_cooling_N_density *
+                         qla_cooling_N_heattypes * sizeof(float)) != 0)
+    error("Failed to allocate Theating array\n");
+
+  dataset = H5Dopen(tempfile_id, "/Tdep/Heating", H5P_DEFAULT);
   status = H5Dread(dataset, H5T_NATIVE_FLOAT, H5S_ALL, H5S_ALL, H5P_DEFAULT,
-                   temperature);
-  if (status < 0) error("error reading temperature table");
+                   cooling->table.Theating);
+  if (status < 0) error("error reading Theating\n");
   status = H5Dclose(dataset);
   if (status < 0) error("error closing cooling dataset");
 
-  /* read in H + He electron abundances */
-  strcpy(set_name, "/Metal_free/Electron_density_over_n_h");
-  dataset = H5Dopen(file_id, set_name, H5P_DEFAULT);
+  /* Heating (internal energy) */
+  if (posix_memalign((void **)&cooling->table.Uheating, SWIFT_STRUCT_ALIGNMENT,
+                     qla_cooling_N_redshifts * qla_cooling_N_internalenergy *
+                         qla_cooling_N_metallicity * qla_cooling_N_density *
+                         qla_cooling_N_heattypes * sizeof(float)) != 0)
+    error("Failed to allocate Uheating array\n");
+
+  dataset = H5Dopen(tempfile_id, "/Udep/Heating", H5P_DEFAULT);
   status = H5Dread(dataset, H5T_NATIVE_FLOAT, H5S_ALL, H5S_ALL, H5P_DEFAULT,
-                   he_electron_abundance);
-  if (status < 0) error("error reading electron density table");
+                   cooling->table.Uheating);
+  if (status < 0) error("error reading Uheating\n");
   status = H5Dclose(dataset);
   if (status < 0) error("error closing cooling dataset");
 
-  /* Transpose from order tables are stored in (helium fraction, temperature,
-   * nH) to (nH, helium fraction, temperature) where fastest
-   * varying index is on right. Tables contain cooling rates but we
-   * want rate of change of internal energy, hence minus sign. */
-  for (int i = 0; i < qla_cooling_N_He_frac; i++) {
-    for (int j = 0; j < qla_cooling_N_temperature; j++) {
-      for (int k = 0; k < qla_cooling_N_density; k++) {
-
-        /* Index in the HDF5 table */
-        const int hdf5_index = row_major_index_3d(
-            i, j, k, qla_cooling_N_He_frac, qla_cooling_N_temperature,
-            qla_cooling_N_density);
-
-        /* Index in the internal table */
-        const int internal_index = row_major_index_3d(
-            k, i, j, qla_cooling_N_density, qla_cooling_N_He_frac,
-            qla_cooling_N_temperature);
-
-        /* Change the sign and transpose */
-        cooling->table.H_plus_He_heating[internal_index] =
-            -he_net_cooling_rate[hdf5_index];
-
-        /* Convert to log T and transpose */
-        cooling->table.temperature[internal_index] =
-            log10(temperature[hdf5_index]);
-
-        /* Just transpose */
-        cooling->table.H_plus_He_electron_abundance[internal_index] =
-            he_electron_abundance[hdf5_index];
-      }
-    }
-  }
+  /* Electron fraction (temperature) */
+  if (posix_memalign((void **)&cooling->table.Telectron_fraction,
+                     SWIFT_STRUCT_ALIGNMENT,
+                     qla_cooling_N_redshifts * qla_cooling_N_temperature *
+                         qla_cooling_N_metallicity * qla_cooling_N_density *
+                         qla_cooling_N_electrontypes * sizeof(float)) != 0)
+    error("Failed to allocate Telectron_fraction array\n");
 
-  /* read in electron densities due to metals */
-  strcpy(set_name, "/Solar/Electron_density_over_n_h");
-  dataset = H5Dopen(file_id, set_name, H5P_DEFAULT);
+  dataset = H5Dopen(tempfile_id, "/Tdep/ElectronFractionsVol", H5P_DEFAULT);
   status = H5Dread(dataset, H5T_NATIVE_FLOAT, H5S_ALL, H5S_ALL, H5P_DEFAULT,
-                   electron_abundance);
-  if (status < 0) error("error reading solar electron density table");
+                   cooling->table.Telectron_fraction);
+  if (status < 0) error("error reading electron_fraction (temperature)\n");
   status = H5Dclose(dataset);
   if (status < 0) error("error closing cooling dataset");
 
-  /* Transpose from order tables are stored in (temperature, nH) to
-   * (nH, temperature) where fastest varying index is on right. */
-  for (int i = 0; i < qla_cooling_N_temperature; i++) {
-    for (int j = 0; j < qla_cooling_N_density; j++) {
-
-      /* Index in the HDF5 table */
-      const int hdf5_index = row_major_index_2d(i, j, qla_cooling_N_temperature,
-                                                qla_cooling_N_density);
+  /* Electron fraction (internal energy) */
+  if (posix_memalign((void **)&cooling->table.Uelectron_fraction,
+                     SWIFT_STRUCT_ALIGNMENT,
+                     qla_cooling_N_redshifts * qla_cooling_N_internalenergy *
+                         qla_cooling_N_metallicity * qla_cooling_N_density *
+                         qla_cooling_N_electrontypes * sizeof(float)) != 0)
+    error("Failed to allocate Uelectron_fraction array\n");
 
-      /* Index in the internal table */
-      const int internal_index = row_major_index_2d(j, i, qla_cooling_N_density,
-                                                    qla_cooling_N_temperature);
+  dataset = H5Dopen(tempfile_id, "/Udep/ElectronFractionsVol", H5P_DEFAULT);
+  status = H5Dread(dataset, H5T_NATIVE_FLOAT, H5S_ALL, H5S_ALL, H5P_DEFAULT,
+                   cooling->table.Uelectron_fraction);
+  if (status < 0) error("error reading electron_fraction (internal energy)\n");
+  status = H5Dclose(dataset);
+  if (status < 0) error("error closing cooling dataset");
 
-      /* Just transpose */
-      cooling->table.electron_abundance[internal_index] =
-          electron_abundance[hdf5_index];
-    }
-  }
+  /* Internal energy from temperature */
+  if (posix_memalign((void **)&cooling->table.U_from_T, SWIFT_STRUCT_ALIGNMENT,
+                     qla_cooling_N_redshifts * qla_cooling_N_temperature *
+                         qla_cooling_N_metallicity * qla_cooling_N_density *
+                         sizeof(float)) != 0)
+    error("Failed to allocate U_from_T array\n");
 
-  status = H5Fclose(file_id);
-  if (status < 0) error("error closing file");
+  dataset = H5Dopen(tempfile_id, "/Tdep/U_from_T", H5P_DEFAULT);
+  status = H5Dread(dataset, H5T_NATIVE_FLOAT, H5S_ALL, H5S_ALL, H5P_DEFAULT,
+                   cooling->table.U_from_T);
+  if (status < 0) error("error reading U_from_T array\n");
+  status = H5Dclose(dataset);
+  if (status < 0) error("error closing cooling dataset");
 
-  swift_free("cooling-temp", net_cooling_rate);
-  swift_free("cooling-temp", electron_abundance);
-  swift_free("cooling-temp", temperature);
-  swift_free("cooling-temp", he_net_cooling_rate);
-  swift_free("cooling-temp", he_electron_abundance);
+  /* Temperature from interal energy */
+  if (posix_memalign((void **)&cooling->table.T_from_U, SWIFT_STRUCT_ALIGNMENT,
+                     qla_cooling_N_redshifts * qla_cooling_N_internalenergy *
+                         qla_cooling_N_metallicity * qla_cooling_N_density *
+                         sizeof(float)) != 0)
+    error("Failed to allocate T_from_U array\n");
 
-#ifdef SWIFT_DEBUG_CHECKS
-  message("done reading in redshift invariant table");
-#endif
+  dataset = H5Dopen(tempfile_id, "/Udep/T_from_U", H5P_DEFAULT);
+  status = H5Dread(dataset, H5T_NATIVE_FLOAT, H5S_ALL, H5S_ALL, H5P_DEFAULT,
+                   cooling->table.T_from_U);
+  if (status < 0) error("error reading T_from_U array\n");
+  status = H5Dclose(dataset);
+  if (status < 0) error("error closing cooling dataset");
 
-#else
-  error("Need HDF5 to read cooling tables");
-#endif
-}
+  /* Thermal equilibrium temperature */
+  if (posix_memalign((void **)&cooling->table.logTeq, SWIFT_STRUCT_ALIGNMENT,
+                     qla_cooling_N_redshifts * qla_cooling_N_metallicity *
+                         qla_cooling_N_density * sizeof(float)) != 0)
+    error("Failed to allocate logTeq array\n");
 
-/**
- * @brief Get redshift dependent table of cooling rates.
- * Reads in table of cooling rates and electron abundances due to
- * metals (depending on temperature, hydrogen number density), cooling rates and
- * electron abundances due to hydrogen and helium (depending on temperature,
- * hydrogen number density and helium fraction), and temperatures (depending on
- * internal energy, hydrogen number density and helium fraction; note: this is
- * distinct from table of temperatures read in ReadCoolingHeader, as that table
- * is used to index the cooling, electron abundance tables, whereas this one is
- * used to obtain temperature of particle)
- *
- * @param cooling #cooling_function_data structure
- * @param low_z_index Index of the lowest redshift table to load.
- * @param high_z_index Index of the highest redshift table to load.
- */
-void get_cooling_table(struct cooling_function_data *restrict cooling,
-                       const int low_z_index, const int high_z_index) {
+  dataset = H5Dopen(tempfile_id, "/ThermEq/Temperature", H5P_DEFAULT);
+  status = H5Dread(dataset, H5T_NATIVE_FLOAT, H5S_ALL, H5S_ALL, H5P_DEFAULT,
+                   cooling->table.logTeq);
+  if (status < 0) error("error reading Teq array\n");
+  status = H5Dclose(dataset);
+  if (status < 0) error("error closing logTeq dataset");
 
-#ifdef HAVE_HDF5
+  /* Mean particle mass at thermal equilibrium temperature */
+  if (posix_memalign((void **)&cooling->table.meanpartmass_Teq,
+                     SWIFT_STRUCT_ALIGNMENT,
+                     qla_cooling_N_redshifts * qla_cooling_N_metallicity *
+                         qla_cooling_N_density * sizeof(float)) != 0)
+    error("Failed to allocate mu array\n");
 
-  /* Temporary tables */
-  float *net_cooling_rate = NULL;
-  float *electron_abundance = NULL;
-  float *temperature = NULL;
-  float *he_net_cooling_rate = NULL;
-  float *he_electron_abundance = NULL;
+  dataset = H5Dopen(tempfile_id, "/ThermEq/MeanParticleMass", H5P_DEFAULT);
+  status = H5Dread(dataset, H5T_NATIVE_FLOAT, H5S_ALL, H5S_ALL, H5P_DEFAULT,
+                   cooling->table.meanpartmass_Teq);
+  if (status < 0) error("error reading mu array\n");
+  status = H5Dclose(dataset);
+  if (status < 0) error("error closing mu dataset");
 
-  /* Allocate arrays for reading in cooling tables.  */
-  if (swift_memalign("cooling-temp", (void **)&net_cooling_rate,
-                     SWIFT_STRUCT_ALIGNMENT,
-                     num_elements_cooling_rate * sizeof(float)) != 0)
-    error("Failed to allocate net_cooling_rate array");
-  if (swift_memalign("cooling-temp", (void **)&electron_abundance,
-                     SWIFT_STRUCT_ALIGNMENT,
-                     num_elements_electron_abundance * sizeof(float)) != 0)
-    error("Failed to allocate electron_abundance array");
-  if (swift_memalign("cooling-temp", (void **)&temperature,
-                     SWIFT_STRUCT_ALIGNMENT,
-                     num_elements_temperature * sizeof(float)) != 0)
-    error("Failed to allocate temperature array");
-  if (swift_memalign("cooling-temp", (void **)&he_net_cooling_rate,
+  /* Hydrogen fractions at thermal equilibirum temperature */
+  if (posix_memalign((void **)&cooling->table.logHfracs_Teq,
                      SWIFT_STRUCT_ALIGNMENT,
-                     num_elements_HpHe_heating * sizeof(float)) != 0)
-    error("Failed to allocate he_net_cooling_rate array");
-  if (swift_memalign("cooling-temp", (void **)&he_electron_abundance,
+                     qla_cooling_N_redshifts * qla_cooling_N_metallicity *
+                         qla_cooling_N_density * 3 * sizeof(float)) != 0)
+    error("Failed to allocate hydrogen fractions array\n");
+
+  dataset = H5Dopen(tempfile_id, "/ThermEq/HydrogenFractionsVol", H5P_DEFAULT);
+  status = H5Dread(dataset, H5T_NATIVE_FLOAT, H5S_ALL, H5S_ALL, H5P_DEFAULT,
+                   cooling->table.logHfracs_Teq);
+  if (status < 0) error("error reading hydrogen fractions array\n");
+  status = H5Dclose(dataset);
+  if (status < 0) error("error closing hydrogen fractions dataset");
+
+  /* All hydrogen fractions */
+  if (posix_memalign((void **)&cooling->table.logHfracs_all,
                      SWIFT_STRUCT_ALIGNMENT,
-                     num_elements_HpHe_electron_abundance * sizeof(float)) != 0)
-    error("Failed to allocate he_electron_abundance array");
+                     qla_cooling_N_redshifts * qla_cooling_N_temperature *
+                         qla_cooling_N_metallicity * qla_cooling_N_density * 3 *
+                         sizeof(float)) != 0)
+    error("Failed to allocate big hydrogen fractions array\n");
 
-  /* Read in tables, transpose so that values for indices which vary most are
-   * adjacent. Repeat for redshift above and redshift below current value.  */
-  for (int z_index = low_z_index; z_index <= high_z_index; z_index++) {
+  dataset = H5Dopen(tempfile_id, "/Tdep/HydrogenFractionsVol", H5P_DEFAULT);
+  status = H5Dread(dataset, H5T_NATIVE_FLOAT, H5S_ALL, H5S_ALL, H5P_DEFAULT,
+                   cooling->table.logHfracs_all);
+  if (status < 0) error("error reading big hydrogen fractions array\n");
+  status = H5Dclose(dataset);
+  if (status < 0) error("error closing big hydrogen fractions dataset");
 
-    /* Index along redhsift dimension for the subset of tables we read */
-    const int local_z_index = z_index - low_z_index;
+  /* Close the file */
+  H5Fclose(tempfile_id);
 
-#ifdef SWIFT_DEBUG_CHECKS
-    if (local_z_index >= qla_cooling_N_loaded_redshifts)
-      error("Reading invalid number of tables along z axis.");
-#endif
+  /* Pressure at thermal equilibrium temperature */
+  if (posix_memalign((void **)&cooling->table.logPeq, SWIFT_STRUCT_ALIGNMENT,
+                     qla_cooling_N_redshifts * qla_cooling_N_metallicity *
+                         qla_cooling_N_density * sizeof(float)) != 0)
+    error("Failed to allocate logPeq array\n");
 
-    /* Open table for this redshift index */
-    char fname[qla_table_path_name_length + 12];
-    sprintf(fname, "%sz_%1.3f.hdf5", cooling->cooling_table_path,
-            cooling->Redshifts[z_index]);
-    message("Reading cooling table 'z_%1.3f.hdf5'",
-            cooling->Redshifts[z_index]);
-
-    hid_t file_id = H5Fopen(fname, H5F_ACC_RDONLY, H5P_DEFAULT);
-    if (file_id < 0) error("unable to open file %s", fname);
-
-    char set_name[64];
-
-    /* read in cooling rates due to metals */
-    for (int specs = 0; specs < qla_cooling_N_metal; specs++) {
-
-      sprintf(set_name, "/%s/Net_Cooling", qla_tables_element_names[specs]);
-      hid_t dataset = H5Dopen(file_id, set_name, H5P_DEFAULT);
-      herr_t status = H5Dread(dataset, H5T_NATIVE_FLOAT, H5S_ALL, H5S_ALL,
-                              H5P_DEFAULT, net_cooling_rate);
-      if (status < 0) error("error reading metal cooling rate table");
-      status = H5Dclose(dataset);
-      if (status < 0) error("error closing cooling dataset");
-
-      /* Transpose from order tables are stored in (temperature, nH)
-       * to (metal species, redshift, nH, temperature) where fastest
-       * varying index is on right. Tables contain cooling rates but we
-       * want rate of change of internal energy, hence minus sign. */
-      for (int i = 0; i < qla_cooling_N_density; i++) {
-        for (int j = 0; j < qla_cooling_N_temperature; j++) {
-
-          /* Index in the HDF5 table */
-          const int hdf5_index = row_major_index_2d(
-              j, i, qla_cooling_N_temperature, qla_cooling_N_density);
-
-          /* Index in the internal table */
-          const int internal_index = row_major_index_4d(
-              specs, local_z_index, i, j, qla_cooling_N_metal,
-              qla_cooling_N_loaded_redshifts, qla_cooling_N_density,
-              qla_cooling_N_temperature);
-
-          /* Change the sign and transpose */
-          cooling->table.metal_heating[internal_index] =
-              -net_cooling_rate[hdf5_index];
-        }
-      }
-    }
+  const float log10_kB_cgs = cooling->log10_kB_cgs;
 
-    /* read in cooling rates due to H + He */
-    strcpy(set_name, "/Metal_free/Net_Cooling");
-    hid_t dataset = H5Dopen(file_id, set_name, H5P_DEFAULT);
-    herr_t status = H5Dread(dataset, H5T_NATIVE_FLOAT, H5S_ALL, H5S_ALL,
-                            H5P_DEFAULT, he_net_cooling_rate);
-    if (status < 0) error("error reading metal free cooling rate table");
-    status = H5Dclose(dataset);
-    if (status < 0) error("error closing cooling dataset");
-
-    /* read in Temperature */
-    strcpy(set_name, "/Metal_free/Temperature/Temperature");
-    dataset = H5Dopen(file_id, set_name, H5P_DEFAULT);
-    status = H5Dread(dataset, H5T_NATIVE_FLOAT, H5S_ALL, H5S_ALL, H5P_DEFAULT,
-                     temperature);
-    if (status < 0) error("error reading temperature table");
-    status = H5Dclose(dataset);
-    if (status < 0) error("error closing cooling dataset");
-
-    /* Read in H + He electron abundance */
-    strcpy(set_name, "/Metal_free/Electron_density_over_n_h");
-    dataset = H5Dopen(file_id, set_name, H5P_DEFAULT);
-    status = H5Dread(dataset, H5T_NATIVE_FLOAT, H5S_ALL, H5S_ALL, H5P_DEFAULT,
-                     he_electron_abundance);
-    if (status < 0) error("error reading electron density table");
-    status = H5Dclose(dataset);
-    if (status < 0) error("error closing cooling dataset");
-
-    /* Transpose from order tables are stored in (helium fraction, temperature,
-     * nH) to (redshift, nH, helium fraction, temperature) where fastest
-     * varying index is on right. */
-    for (int i = 0; i < qla_cooling_N_He_frac; i++) {
-      for (int j = 0; j < qla_cooling_N_temperature; j++) {
-        for (int k = 0; k < qla_cooling_N_density; k++) {
-
-          /* Index in the HDF5 table */
-          const int hdf5_index = row_major_index_3d(
-              i, j, k, qla_cooling_N_He_frac, qla_cooling_N_temperature,
-              qla_cooling_N_density);
-
-          /* Index in the internal table */
-          const int internal_index = row_major_index_4d(
-              local_z_index, k, i, j, qla_cooling_N_loaded_redshifts,
-              qla_cooling_N_density, qla_cooling_N_He_frac,
-              qla_cooling_N_temperature);
-
-          /* Change the sign and transpose */
-          cooling->table.H_plus_He_heating[internal_index] =
-              -he_net_cooling_rate[hdf5_index];
-
-          /* Convert to log T and transpose */
-          cooling->table.temperature[internal_index] =
-              log10(temperature[hdf5_index]);
-
-          /* Just transpose */
-          cooling->table.H_plus_He_electron_abundance[internal_index] =
-              he_electron_abundance[hdf5_index];
-        }
-      }
-    }
+  /* Compute the pressures at thermal eq. */
+  for (int ired = 0; ired < qla_cooling_N_redshifts; ired++) {
+    for (int imet = 0; imet < qla_cooling_N_metallicity; imet++) {
+
+      const int index_XH = row_major_index_2d(
+          imet, 0, qla_cooling_N_metallicity, qla_cooling_N_elementtypes);
+
+      const float log10_XH = cooling->LogMassFractions[index_XH];
 
-    /* read in electron densities due to metals */
-    strcpy(set_name, "/Solar/Electron_density_over_n_h");
-    dataset = H5Dopen(file_id, set_name, H5P_DEFAULT);
-    status = H5Dread(dataset, H5T_NATIVE_FLOAT, H5S_ALL, H5S_ALL, H5P_DEFAULT,
-                     electron_abundance);
-    if (status < 0) error("error reading solar electron density table");
-    status = H5Dclose(dataset);
-    if (status < 0) error("error closing cooling dataset");
-
-    /* Transpose from order tables are stored in (temperature, nH) to
-     * (redshift, nH, temperature) where fastest varying index is on right. */
-    for (int i = 0; i < qla_cooling_N_temperature; i++) {
-      for (int j = 0; j < qla_cooling_N_density; j++) {
-
-        /* Index in the HDF5 table */
-        const int hdf5_index = row_major_index_2d(
-            i, j, qla_cooling_N_temperature, qla_cooling_N_density);
-
-        /* Index in the internal table */
-        const int internal_index = row_major_index_3d(
-            local_z_index, j, i, qla_cooling_N_loaded_redshifts,
-            qla_cooling_N_density, qla_cooling_N_temperature);
-
-        /* Just transpose */
-        cooling->table.electron_abundance[internal_index] =
-            electron_abundance[hdf5_index];
+      for (int iden = 0; iden < qla_cooling_N_density; iden++) {
+
+        const int index_Peq = row_major_index_3d(
+            ired, imet, iden, qla_cooling_N_redshifts,
+            qla_cooling_N_metallicity, qla_cooling_N_density);
+
+        cooling->table.logPeq[index_Peq] =
+            cooling->nH[iden] + cooling->table.logTeq[index_Peq] - log10_XH -
+            log10(cooling->table.meanpartmass_Teq[index_Peq]) + log10_kB_cgs;
       }
     }
-
-    status = H5Fclose(file_id);
-    if (status < 0) error("error closing file");
   }
 
-  swift_free("cooling-temp", net_cooling_rate);
-  swift_free("cooling-temp", electron_abundance);
-  swift_free("cooling-temp", temperature);
-  swift_free("cooling-temp", he_net_cooling_rate);
-  swift_free("cooling-temp", he_electron_abundance);
-
 #ifdef SWIFT_DEBUG_CHECKS
   message("Done reading in general cooling table");
 #endif
diff --git a/src/cooling/QLA/cooling_tables.h b/src/cooling/QLA/cooling_tables.h
index 67aa437fd9823f07a8543a0ef3cfc7ef2bab1b4e..06610abcefbf8ee15621a142bbaec39af59e57f1 100644
--- a/src/cooling/QLA/cooling_tables.h
+++ b/src/cooling/QLA/cooling_tables.h
@@ -20,46 +20,99 @@
 #define SWIFT_QLA_COOL_TABLES_H
 
 /**
- * @file src/cooling/QLA/cooling.h
- * @brief Quick Lyman-alpha cooling function (EAGLE with primordial Z only)
+ * @file src/cooling/QLA/cooling_tables.h
+ * @brief QLA cooling function
  */
 
 /* Config parameters. */
-#include "../config.h"
+#include "config.h"
 
-#include "cooling_struct.h"
+/*! Number of different bins along the temperature axis of the tables */
+#define qla_cooling_N_temperature 86
 
-/*! Number of different bins along the redhsift axis of the tables */
-#define qla_cooling_N_redshifts 49
+/*! Number of different bins along the redshift axis of the tables */
+#define qla_cooling_N_redshifts 46
 
-/*! Number of redshift bins loaded at any given point int time */
-#define qla_cooling_N_loaded_redshifts 2
+/*! Number of different bins along the density axis of the tables */
+#define qla_cooling_N_density 71
 
-/*! Number of different bins along the temperature axis of the tables */
-#define qla_cooling_N_temperature 176
+/*! Number of different bins along the metallicity axis of the tables */
+#define qla_cooling_N_metallicity 11
 
-/*! Number of different bins along the density axis of the tables */
-#define qla_cooling_N_density 41
+/*! Number of different bins along the internal energy axis of the tables */
+#define qla_cooling_N_internalenergy 191
 
-/*! Number of different bins along the metal axis of the tables */
-#define qla_cooling_N_metal 9
+/*! Number of different cooling channels in the tables */
+#define qla_cooling_N_cooltypes 22
 
-/*! Number of different bins along the metal axis of the tables */
-#define qla_cooling_N_He_frac 7
+/*! Number of different heating channels in the tables */
+#define qla_cooling_N_heattypes 24
 
-/*! Number of different bins along the abundances axis of the tables */
-#define qla_cooling_N_abundances 11
+/*! Number of different electron fractions (each element - other atoms
+ *  + tot prim + tot metal + tot)  in the tables */
+#define qla_cooling_N_electrontypes 14
 
-void get_cooling_redshifts(struct cooling_function_data *cooling);
+/*! Number of different elements in the tables */
+#define qla_cooling_N_elementtypes 12
 
-void read_cooling_header(const char *fname,
-                         struct cooling_function_data *cooling);
+/**
+ * @brief Elements present in the tables
+ */
+enum qla_cooling_element {
+  element_H,
+  element_He,
+  element_C,
+  element_N,
+  element_O,
+  element_Ne,
+  element_Mg,
+  element_Si,
+  element_S,
+  element_Ca,
+  element_Fe,
+  element_OA
+};
 
-void allocate_cooling_tables(struct cooling_function_data *restrict cooling);
+/**
+ * @brief Hydrogen species
+ */
+enum qla_hydrogen_species { neutral = 0, ionized = 1, molecular = 2 };
 
-void get_redshift_invariant_table(
-    struct cooling_function_data *restrict cooling, const int photodis);
-void get_cooling_table(struct cooling_function_data *restrict cooling,
-                       const int low_z_index, const int high_z_index);
+/**
+ * @brief Cooling channels beyond the metal lines
+ */
+enum qla_cooling_channels {
+  cooltype_H2 = element_OA + 1,
+  cooltype_molecules,
+  cooltype_HD,
+  cooltype_NetFFH,
+  cooltype_NetFFM,
+  cooltype_eeBrems,
+  cooltype_Compton,
+  cooltype_Dust
+};
+
+/**
+ * @brief Heating channels beyond the metal lines
+ */
+enum qla_heating_channels {
+  heattype_H2 = element_OA + 1,
+  heattype_COdiss,
+  heattype_CosmicRay,
+  heattype_UTA,
+  heattype_line,
+  heattype_Hlin,
+  heattype_ChaT,
+  heattype_HFF,
+  heattype_Compton,
+  heattype_Dust
+};
+
+/* Pre-declaration */
+struct cooling_function_data;
+
+void get_cooling_redshifts(struct cooling_function_data *cooling);
+void read_cooling_header(struct cooling_function_data *cooling);
+void read_cooling_tables(struct cooling_function_data *cooling);
 
-#endif /* SWIFT_QLA_COOL_TABLES_H */
+#endif
diff --git a/src/cooling/QLA/interpolate.h b/src/cooling/QLA/interpolate.h
index d49a482b1e0d031128052a0a79a905e5dcc56a48..b1cacae696393cdfc84e72a7a811362ab4fcc9e4 100644
--- a/src/cooling/QLA/interpolate.h
+++ b/src/cooling/QLA/interpolate.h
@@ -21,7 +21,7 @@
 
 /**
  * @file src/cooling/QLA/interpolate.h
- * @brief Interpolation functions for Quick Lyman-alpha cooling tables
+ * @brief Interpolation functions for Quick Lyman-alpha cooling tables (COLIBRE)
  */
 
 /* Config parameters. */
@@ -30,6 +30,7 @@
 /* Local includes. */
 #include "align.h"
 #include "error.h"
+#include "exp10.h"
 #include "inline.h"
 
 /**
@@ -87,6 +88,28 @@ __attribute__((always_inline)) INLINE int row_major_index_4d(
   return x * Ny * Nz * Nw + y * Nz * Nw + z * Nw + w;
 }
 
+/**
+ * @brief Returns the 1d index of element with 5d indices x,y,z,w
+ * from a flattened 5d array in row major order
+ *
+ * @param x, y, z, v, w Indices of element of interest
+ * @param Nx, Ny, Nz, Nv, Nw Sizes of array dimensions
+ */
+__attribute__((always_inline)) INLINE int row_major_index_5d(
+    const int x, const int y, const int z, const int w, const int v,
+    const int Nx, const int Ny, const int Nz, const int Nw, const int Nv) {
+
+#ifdef SWIFT_DEBUG_CHECKS
+  assert(x < Nx);
+  assert(y < Ny);
+  assert(z < Nz);
+  assert(w < Nw);
+  assert(v < Nv);
+#endif
+
+  return x * Ny * Nz * Nw * Nv + y * Nz * Nw * Nv + z * Nw * Nv + w * Nv + v;
+}
+
 /**
  * @brief Finds the index of a value in a table and compute delta to nearest
  * element.
@@ -119,29 +142,32 @@ __attribute__((always_inline)) INLINE void get_index_1d(
   swift_align_information(float, table, SWIFT_STRUCT_ALIGNMENT);
 
   /* Distance between elements in the array */
-  const float delta = (size - 1) / (table[size - 1] - table[0]);
+  /* Do not use first or last entry, might be an extra bin with uneven spacing
+   */
+  const float delta = (size - 3) / (table[size - 2] - table[1]);
 
-  if (x < table[0] + epsilon) {
-    /* We are below the first element */
-    *i = 0;
-    *dx = 0.f;
-  } else if (x < table[size - 1] - epsilon) {
-    /* Normal case */
-    *i = (x - table[0]) * delta;
+  /* Check for an extra entry at the beginning (e.g. metallicity) */
+  int istart = 0;
+  int iend = size - 1;
 
-#ifdef SWIFT_DEBUG_CHECKS
-    if (*i > size || *i < 0) {
-      error(
-          "trying to get index for value outside table range. Table size: %d, "
-          "calculated index: %d, value: %.5e, table[0]: %.5e, grid size: %.5e",
-          size, *i, x, table[0], delta);
-    }
-#endif
+  if (fabsf(table[1] - table[0]) > delta + epsilon) {
+    istart = 1;
+  }
+  if (fabsf(table[size - 1] - table[size - 2]) > delta + epsilon) {
+    iend = size - 2;
+  }
 
+  /*extra array at the beginning */
+  if (x < table[istart] + epsilon) {
+    /* We are before the first element */
+    *i = 0;
+    *dx = 0.f;
+  } else if (x < table[iend] - epsilon) {
+    *i = (x - table[1]) * delta + 1;
     *dx = (x - table[*i]) * delta;
   } else {
     /* We are after the last element */
-    *i = size - 2;
+    *i = iend - 1;
     *dx = 1.f;
   }
 
@@ -242,6 +268,67 @@ __attribute__((always_inline)) INLINE float interpolation_3d(
   return result;
 }
 
+/**
+ * @brief Interpolate a flattened 3D table at a given position but avoid the
+ * z-dimension.
+ *
+ * This function uses linear interpolation along each axis.
+ * We look at the zi coordoniate but do not interpolate around it. We just
+ * interpolate the remaining 2 dimensions.
+ * The function also assumes that the table is aligned on
+ * SWIFT_STRUCT_ALIGNMENT.
+ *
+ * @param table The 3D table to interpolate.
+ * @param xi, yi, zi Indices of element of interest.
+ * @param Nx, Ny, Nz Sizes of array dimensions.
+ * @param dx, dy, dz Distance between the point and the index in units of
+ * the grid spacing.
+ */
+__attribute__((always_inline)) INLINE float interpolation_3d_no_z(
+    const float *table, const int xi, const int yi, const int zi,
+    const float dx, const float dy, const float dz, const int Nx, const int Ny,
+    const int Nz) {
+
+#ifdef SWIFT_DEBUG_CHECKS
+  if (dx < -0.001f || dx > 1.001f) error("Invalid dx=%e", dx);
+  if (dy < -0.001f || dy > 1.001f) error("Invalid dy=%e", dy);
+  if (dz != 0.f) error("Attempting to interpolate along z!");
+#endif
+
+  const float tx = 1.f - dx;
+  const float ty = 1.f - dy;
+  const float tz = 1.f;
+
+  /* Indicate that the whole array is aligned on page boundaries */
+  swift_align_information(float, table, SWIFT_STRUCT_ALIGNMENT);
+
+  /* Linear interpolation along each axis. We read the table 2^2=4 times */
+  /* Note that we intentionally kept the table access along the axis where */
+  /* we do not interpolate as comments in the code to allow readers to */
+  /* understand what is going on. */
+  float result = tx * ty * tz *
+                 table[row_major_index_3d(xi + 0, yi + 0, zi + 0, Nx, Ny, Nz)];
+
+  /* result += tx * ty * dz *
+            table[row_major_index_3d(xi + 0, yi + 0, zi + 1, Nx, Ny, Nz)]; */
+  result += tx * dy * tz *
+            table[row_major_index_3d(xi + 0, yi + 1, zi + 0, Nx, Ny, Nz)];
+  result += dx * ty * tz *
+            table[row_major_index_3d(xi + 1, yi + 0, zi + 0, Nx, Ny, Nz)];
+
+  /* result += tx * dy * dz *
+            table[row_major_index_3d(xi + 0, yi + 1, zi + 1, Nx, Ny, Nz)]; */
+  /* result += dx * ty * dz *
+            table[row_major_index_3d(xi + 1, yi + 0, zi + 1, Nx, Ny, Nz)]; */
+  result += dx * dy * tz *
+            table[row_major_index_3d(xi + 1, yi + 1, zi + 0, Nx, Ny, Nz)];
+
+  /* result += dx * dy * dz *
+             table[row_major_index_3d(xi + 1, yi + 1, zi + 1, Nx, Ny, Nz)]; */
+
+  return result;
+}
+
 /**
  * @brief Interpolate a flattened 3D table at a given position but avoid the
  * x-dimension.
@@ -392,6 +479,96 @@ __attribute__((always_inline)) INLINE float interpolation_4d(
   return result;
 }
 
+/**
+ * @brief Interpolates a 5 dimensional array in the first 4 dimensions and
+ * adds the individual contributions from the 5th dimension according to their
+ * weights
+ *
+ * @param table The table to interpolate
+ * @param weights The weights for summing up the individual contributions
+ * @param istart, iend Start and stop index for 5th dimension
+ * @param xi, yi, zi, wi Indices of table element
+ * @param dx, dy, dz, dw Distance between the point and the index in units of
+ * the grid spacing.
+ * @param Nx, Ny, Nz, Nw, Nv Sizes of array dimensions
+ */
+__attribute__((always_inline)) INLINE double interpolation4d_plus_summation(
+    const float *table, const float *weights, const int istart, const int iend,
+    const int xi, const int yi, const int zi, const int wi, const float dx,
+    const float dy, const float dz, const float dw, const int Nx, const int Ny,
+    const int Nz, const int Nw, const int Nv) {
+
+  const float tx = 1.f - dx;
+  const float ty = 1.f - dy;
+  const float tz = 1.f - dz;
+  const float tw = 1.f - dw;
+
+  float result;
+  double result_global = 0.;
+
+  for (int i = istart; i <= iend; i++) {
+
+    /* Linear interpolation along each axis. We read the table 2^4=16 times */
+    result = tx * ty * tz * tw *
+             table[row_major_index_5d(xi + 0, yi + 0, zi + 0, wi + 0, i, Nx, Ny,
+                                      Nz, Nw, Nv)];
+
+    result += tx * ty * tz * dw *
+              table[row_major_index_5d(xi + 0, yi + 0, zi + 0, wi + 1, i, Nx,
+                                       Ny, Nz, Nw, Nv)];
+
+    result += tx * ty * dz * tw *
+              table[row_major_index_5d(xi + 0, yi + 0, zi + 1, wi + 0, i, Nx,
+                                       Ny, Nz, Nw, Nv)];
+    result += tx * dy * tz * tw *
+              table[row_major_index_5d(xi + 0, yi + 1, zi + 0, wi + 0, i, Nx,
+                                       Ny, Nz, Nw, Nv)];
+    result += dx * ty * tz * tw *
+              table[row_major_index_5d(xi + 1, yi + 0, zi + 0, wi + 0, i, Nx,
+                                       Ny, Nz, Nw, Nv)];
+
+    result += tx * ty * dz * dw *
+              table[row_major_index_5d(xi + 0, yi + 0, zi + 1, wi + 1, i, Nx,
+                                       Ny, Nz, Nw, Nv)];
+    result += tx * dy * tz * dw *
+              table[row_major_index_5d(xi + 0, yi + 1, zi + 0, wi + 1, i, Nx,
+                                       Ny, Nz, Nw, Nv)];
+    result += dx * ty * tz * dw *
+              table[row_major_index_5d(xi + 1, yi + 0, zi + 0, wi + 1, i, Nx,
+                                       Ny, Nz, Nw, Nv)];
+    result += tx * dy * dz * tw *
+              table[row_major_index_5d(xi + 0, yi + 1, zi + 1, wi + 0, i, Nx,
+                                       Ny, Nz, Nw, Nv)];
+    result += dx * ty * dz * tw *
+              table[row_major_index_5d(xi + 1, yi + 0, zi + 1, wi + 0, i, Nx,
+                                       Ny, Nz, Nw, Nv)];
+    result += dx * dy * tz * tw *
+              table[row_major_index_5d(xi + 1, yi + 1, zi + 0, wi + 0, i, Nx,
+                                       Ny, Nz, Nw, Nv)];
+
+    result += dx * dy * dz * tw *
+              table[row_major_index_5d(xi + 1, yi + 1, zi + 1, wi + 0, i, Nx,
+                                       Ny, Nz, Nw, Nv)];
+    result += dx * dy * tz * dw *
+              table[row_major_index_5d(xi + 1, yi + 1, zi + 0, wi + 1, i, Nx,
+                                       Ny, Nz, Nw, Nv)];
+    result += dx * ty * dz * dw *
+              table[row_major_index_5d(xi + 1, yi + 0, zi + 1, wi + 1, i, Nx,
+                                       Ny, Nz, Nw, Nv)];
+    result += tx * dy * dz * dw *
+              table[row_major_index_5d(xi + 0, yi + 1, zi + 1, wi + 1, i, Nx,
+                                       Ny, Nz, Nw, Nv)];
+
+    result += dx * dy * dz * dw *
+              table[row_major_index_5d(xi + 1, yi + 1, zi + 1, wi + 1, i, Nx,
+                                       Ny, Nz, Nw, Nv)];
+
+    result_global += weights[i] * exp10f(result);
+  }
+
+  return result_global;
+}
+
 /**
  * @brief Interpolate a flattened 4D table at a given position but avoid the
  * x-dimension.
@@ -496,4 +673,205 @@ __attribute__((always_inline)) INLINE float interpolation_4d_no_x(
   return result;
 }
 
+/**
+ * @brief Interpolate a flattened 4D table at a given position but avoid the
+ * w-dimension (ie. final dimension).
+ *
+ * This function uses linear interpolation along each axis.
+ * We look at the wi coordoniate but do not interpolate around it. We just
+ * interpolate the remaining 3 dimensions.
+ * The function also assumes that the table is aligned on
+ * SWIFT_STRUCT_ALIGNMENT.
+ *
+ * @param table The 4D table to interpolate.
+ * @param xi, yi, zi, wi Indices of element of interest.
+ * @param Nx, Ny, Nz, Nw Sizes of array dimensions.
+ * @param dx, dy, dz, dw Distance between the point and the index in units of
+ * the grid spacing.
+ */
+__attribute__((always_inline)) INLINE float interpolation_4d_no_w(
+    const float *table, const int xi, const int yi, const int zi, const int wi,
+    const float dx, const float dy, const float dz, const float dw,
+    const int Nx, const int Ny, const int Nz, const int Nw) {
+
+#ifdef SWIFT_DEBUG_CHECKS
+  if (dx < -0.001f || dx > 1.001f) error("Invalid dx=%e", dx);
+  if (dy < -0.001f || dy > 1.001f) error("Invalid dy=%e", dy);
+  if (dz < -0.001f || dz > 1.001f) error("Invalid dz=%e", dz);
+  if (dw != 0.0f) error("Attempting to interpolate along w!");
+#endif
+
+  const float tx = 1.f - dx;
+  const float ty = 1.f - dy;
+  const float tz = 1.f - dz;
+  const float tw = 1.f;
+
+  /* Indicate that the whole array is aligned on boundaries */
+  swift_align_information(float, table, SWIFT_STRUCT_ALIGNMENT);
+
+  /* Linear interpolation along each axis. We read the table 2^3=8 times */
+  /* Note that we intentionally kept the table access along the axis where */
+  /* we do not interpolate as comments in the code to allow readers to */
+  /* understand what is going on. */
+  float result =
+      tx * ty * tz * tw *
+      table[row_major_index_4d(xi + 0, yi + 0, zi + 0, wi + 0, Nx, Ny, Nz, Nw)];
+
+  /* result +=
+      tx * ty * tz * dw *
+      table[row_major_index_4d(xi + 0, yi + 0, zi + 0, wi + 1, Nx, Ny, Nz,
+     Nw)];*/
+  result +=
+      tx * ty * dz * tw *
+      table[row_major_index_4d(xi + 0, yi + 0, zi + 1, wi + 0, Nx, Ny, Nz, Nw)];
+  result +=
+      tx * dy * tz * tw *
+      table[row_major_index_4d(xi + 0, yi + 1, zi + 0, wi + 0, Nx, Ny, Nz, Nw)];
+  result +=
+      dx * ty * tz * tw *
+      table[row_major_index_4d(xi + 1, yi + 0, zi + 0, wi + 0, Nx, Ny, Nz, Nw)];
+
+  /* result +=
+      tx * ty * dz * dw *
+      table[row_major_index_4d(xi + 0, yi + 0, zi + 1, wi + 1, Nx, Ny, Nz, Nw)];
+  result +=
+      tx * dy * tz * dw *
+      table[row_major_index_4d(xi + 0, yi + 1, zi + 0, wi + 1, Nx, Ny, Nz, Nw)];
+  result +=
+      dx * ty * tz * dw *
+      table[row_major_index_4d(xi + 1, yi + 0, zi + 0, wi + 1, Nx, Ny, Nz, Nw)];
+*/
+  result +=
+      tx * dy * dz * tw *
+      table[row_major_index_4d(xi + 0, yi + 1, zi + 1, wi + 0, Nx, Ny, Nz, Nw)];
+  result +=
+      dx * ty * dz * tw *
+      table[row_major_index_4d(xi + 1, yi + 0, zi + 1, wi + 0, Nx, Ny, Nz, Nw)];
+  result +=
+      dx * dy * tz * tw *
+      table[row_major_index_4d(xi + 1, yi + 1, zi + 0, wi + 0, Nx, Ny, Nz, Nw)];
+
+  result +=
+      dx * dy * dz * tw *
+      table[row_major_index_4d(xi + 1, yi + 1, zi + 1, wi + 0, Nx, Ny, Nz, Nw)];
+  /* result +=
+     dx * dy * tz * dw *
+     table[row_major_index_4d(xi + 1, yi + 1, zi + 0, wi + 1, Nx, Ny, Nz, Nw)];
+  result +=
+      dx * ty * dz * dw *
+      table[row_major_index_4d(xi + 1, yi + 0, zi + 1, wi + 1, Nx, Ny, Nz, Nw)];
+  result +=
+      tx * dy * dz * dw *
+      table[row_major_index_4d(xi + 0, yi + 1, zi + 1, wi + 1, Nx, Ny, Nz, Nw)];
+
+  result +=
+      dx * dy * dz * dw *
+      table[row_major_index_4d(xi + 1, yi + 1, zi + 1, wi + 1, Nx, Ny, Nz, Nw)];
+*/
+
+  return result;
+}
+
+/**
+ * @brief Interpolate a flattened 4D table at a given position but avoid the
+ * z and w-dimension (ie. the two final dimensions).
+ *
+ * This function uses linear interpolation along each axis.
+ * We look at the wi coordoniate but do not interpolate around it. We just
+ * interpolate the remaining 3 dimensions.
+ * The function also assumes that the table is aligned on
+ * SWIFT_STRUCT_ALIGNMENT.
+ *
+ * @param table The 4D table to interpolate.
+ * @param xi, yi, zi, wi Indices of element of interest.
+ * @param Nx, Ny, Nz, Nw Sizes of array dimensions.
+ * @param dx, dy, dz, dw Distance between the point and the index in units of
+ * the grid spacing.
+ */
+__attribute__((always_inline)) INLINE float interpolation_4d_no_z_no_w(
+    const float *table, const int xi, const int yi, const int zi, const int wi,
+    const float dx, const float dy, const float dz, const float dw,
+    const int Nx, const int Ny, const int Nz, const int Nw) {
+
+#ifdef SWIFT_DEBUG_CHECKS
+  if (dx < -0.001f || dx > 1.001f) error("Invalid dx=%e", dx);
+  if (dy < -0.001f || dy > 1.001f) error("Invalid dy=%e", dy);
+  if (dz != 0.0f) error("Attempting to interpolate along z!");
+  if (dw != 0.0f) error("Attempting to interpolate along w!");
+#endif
+
+  const float tx = 1.f - dx;
+  const float ty = 1.f - dy;
+  const float tz = 1.f;
+  const float tw = 1.f;
+
+  /* Indicate that the whole array is aligned on boundaries */
+  swift_align_information(float, table, SWIFT_STRUCT_ALIGNMENT);
+
+  /* Linear interpolation along each axis. We read the table 2^3=8 times */
+  /* Note that we intentionally kept the table access along the axis where */
+  /* we do not interpolate as comments in the code to allow readers to */
+  /* understand what is going on. */
+  float result =
+      tx * ty * tz * tw *
+      table[row_major_index_4d(xi + 0, yi + 0, zi + 0, wi + 0, Nx, Ny, Nz, Nw)];
+
+  /* result +=
+      tx * ty * tz * dw *
+      table[row_major_index_4d(xi + 0, yi + 0, zi + 0, wi + 1, Nx, Ny, Nz,
+     Nw)];*/
+  /* result +=
+      tx * ty * dz * tw *
+      table[row_major_index_4d(xi + 0, yi + 0, zi + 1, wi + 0, Nx, Ny, Nz,
+     Nw)];*/
+  result +=
+      tx * dy * tz * tw *
+      table[row_major_index_4d(xi + 0, yi + 1, zi + 0, wi + 0, Nx, Ny, Nz, Nw)];
+  result +=
+      dx * ty * tz * tw *
+      table[row_major_index_4d(xi + 1, yi + 0, zi + 0, wi + 0, Nx, Ny, Nz, Nw)];
+
+  /* result +=
+      tx * ty * dz * dw *
+      table[row_major_index_4d(xi + 0, yi + 0, zi + 1, wi + 1, Nx, Ny, Nz, Nw)];
+  result +=
+      tx * dy * tz * dw *
+      table[row_major_index_4d(xi + 0, yi + 1, zi + 0, wi + 1, Nx, Ny, Nz, Nw)];
+  result +=
+      dx * ty * tz * dw *
+      table[row_major_index_4d(xi + 1, yi + 0, zi + 0, wi + 1, Nx, Ny, Nz, Nw)];
+*/
+  /* result +=
+      tx * dy * dz * tw *
+      table[row_major_index_4d(xi + 0, yi + 1, zi + 1, wi + 0, Nx, Ny, Nz, Nw)];
+  result +=
+      dx * ty * dz * tw *
+      table[row_major_index_4d(xi + 1, yi + 0, zi + 1, wi + 0, Nx, Ny, Nz, Nw)];
+*/
+  result +=
+      dx * dy * tz * tw *
+      table[row_major_index_4d(xi + 1, yi + 1, zi + 0, wi + 0, Nx, Ny, Nz, Nw)];
+
+  /* result +=
+      dx * dy * dz * tw *
+      table[row_major_index_4d(xi + 1, yi + 1, zi + 1, wi + 0, Nx, Ny, Nz, Nw)];
+   */
+  /* result +=
+     dx * dy * tz * dw *
+     table[row_major_index_4d(xi + 1, yi + 1, zi + 0, wi + 1, Nx, Ny, Nz, Nw)];
+  result +=
+      dx * ty * dz * dw *
+      table[row_major_index_4d(xi + 1, yi + 0, zi + 1, wi + 1, Nx, Ny, Nz, Nw)];
+  result +=
+      tx * dy * dz * dw *
+      table[row_major_index_4d(xi + 0, yi + 1, zi + 1, wi + 1, Nx, Ny, Nz, Nw)];
+
+  result +=
+      dx * dy * dz * dw *
+      table[row_major_index_4d(xi + 1, yi + 1, zi + 1, wi + 1, Nx, Ny, Nz, Nw)];
+*/
+
+  return result;
+}
+
 #endif