diff --git a/src/cooling/EAGLE/cooling.h b/src/cooling/EAGLE/cooling.h
index aca68ec39b677333a0e8a2b65604f2c12c96678f..99d0c7d7cc1a99ba21fa21fc0f46b26f4a3203d3 100644
--- a/src/cooling/EAGLE/cooling.h
+++ b/src/cooling/EAGLE/cooling.h
@@ -52,6 +52,7 @@ __attribute__((always_inline)) INLINE static void cooling_write_flavour(
  *
  * @param phys_const The physical constants in internal units.
  * @param us The internal system of units.
+ * @param cosmo The current cosmological model.
  * @param cooling The #cooling_function_data used in the run.
  * @param p Pointer to the particle data.
  * @param xp Pointer to the extended particle data.
@@ -60,6 +61,7 @@ __attribute__((always_inline)) INLINE static void cooling_write_flavour(
 __attribute__((always_inline)) INLINE static void cooling_cool_part(
     const struct phys_const* restrict phys_const,
     const struct unit_system* restrict us,
+    const struct cosmology* restrict cosmo,
     const struct cooling_function_data* restrict cooling,
     struct part* restrict p, struct xpart* restrict xp, float dt) {}
 
@@ -69,11 +71,13 @@ __attribute__((always_inline)) INLINE static void cooling_cool_part(
  * @param cooling The #cooling_function_data used in the run.
  * @param phys_const The physical constants in internal units.
  * @param us The internal system of units.
+ * @param cosmo The current cosmological model.
  * @param p Pointer to the particle data.
  */
 __attribute__((always_inline)) INLINE static float cooling_timestep(
     const struct cooling_function_data* restrict cooling,
     const struct phys_const* restrict phys_const,
+    const struct cosmology* restrict cosmo,
     const struct unit_system* restrict us, const struct part* restrict p) {
 
   return FLT_MAX;
diff --git a/src/cooling/const_du/cooling.h b/src/cooling/const_du/cooling.h
index 6881e7d0093583b8ba60b3d9723a2214ef677f8a..83698c06ffe8f5dfdf14b140d3befe5e72e0f6d0 100644
--- a/src/cooling/const_du/cooling.h
+++ b/src/cooling/const_du/cooling.h
@@ -62,6 +62,7 @@ __attribute__((always_inline)) INLINE static void cooling_write_flavour(
  *
  * @param phys_const The physical constants in internal units.
  * @param us The internal system of units.
+ * @param cosmo The current cosmological model.
  * @param cooling The #cooling_function_data used in the run.
  * @param p Pointer to the particle data.
  * @param xp Pointer to the extended particle data.
@@ -70,6 +71,7 @@ __attribute__((always_inline)) INLINE static void cooling_write_flavour(
 __attribute__((always_inline)) INLINE static void cooling_cool_part(
     const struct phys_const* restrict phys_const,
     const struct unit_system* restrict us,
+    const struct cosmology* restrict cosmo,
     const struct cooling_function_data* restrict cooling,
     struct part* restrict p, struct xpart* restrict xp, float dt) {
 
@@ -77,7 +79,7 @@ __attribute__((always_inline)) INLINE static void cooling_cool_part(
   const float u_floor = cooling->min_energy;
 
   /* Get current internal energy */
-  const float u_old = hydro_get_internal_energy(p);
+  const float u_old = hydro_get_physical_internal_energy(p, cosmo);
 
   /* Current du_dt */
   const float hydro_du_dt = hydro_get_internal_energy_dt(p);
@@ -108,16 +110,18 @@ __attribute__((always_inline)) INLINE static void cooling_cool_part(
  *
  * @param cooling The #cooling_function_data used in the run.
  * @param phys_const The physical constants in internal units.
+ * @param cosmo The current cosmological model.
  * @param us The internal system of units.
  * @param p Pointer to the particle data.
  */
 __attribute__((always_inline)) INLINE static float cooling_timestep(
     const struct cooling_function_data* restrict cooling,
     const struct phys_const* restrict phys_const,
+    const struct cosmology* restrict cosmo,
     const struct unit_system* restrict us, const struct part* restrict p) {
 
   const float cooling_rate = cooling->cooling_rate;
-  const float internal_energy = hydro_get_internal_energy(p);
+  const float internal_energy = hydro_get_physical_internal_energy(p, cosmo);
   return cooling->cooling_tstep_mult * internal_energy / fabsf(cooling_rate);
 }
 
diff --git a/src/cooling/const_lambda/cooling.h b/src/cooling/const_lambda/cooling.h
index f4fe4ac4df16af53c88a901edc97e5d70f4ae806..af648b904356dc7ceea68a7474426347cdfce4d0 100644
--- a/src/cooling/const_lambda/cooling.h
+++ b/src/cooling/const_lambda/cooling.h
@@ -52,15 +52,17 @@ __attribute__((always_inline)) INLINE static void cooling_write_flavour(
  *
  * @param phys_const The physical constants in internal units.
  * @param us The internal system of units.
+ * @param cosmo The current cosmological model.
  * @param cooling The #cooling_function_data used in the run.
  * @param p Pointer to the particle data..
  */
 __attribute__((always_inline)) INLINE static float cooling_rate(
     const struct phys_const* const phys_const, const struct unit_system* us,
+    const struct cosmology* restrict cosmo,
     const struct cooling_function_data* cooling, const struct part* p) {
 
   /* Get particle density */
-  const float rho = hydro_get_density(p);
+  const float rho = hydro_get_physical_density(p, cosmo);
 
   /* Get cooling function properties */
   const float X_H = cooling->hydrogen_mass_abundance;
@@ -77,6 +79,7 @@ __attribute__((always_inline)) INLINE static float cooling_rate(
  *
  * @param phys_const The physical constants in internal units.
  * @param us The internal system of units.
+ * @param cosmo The current cosmological model.
  * @param cooling The #cooling_function_data used in the run.
  * @param p Pointer to the particle data.
  * @param dt The time-step of this particle.
@@ -84,6 +87,7 @@ __attribute__((always_inline)) INLINE static float cooling_rate(
 __attribute__((always_inline)) INLINE static void cooling_cool_part(
     const struct phys_const* restrict phys_const,
     const struct unit_system* restrict us,
+    const struct cosmology* restrict cosmo,
     const struct cooling_function_data* restrict cooling,
     struct part* restrict p, struct xpart* restrict xp, float dt) {
 
@@ -91,13 +95,13 @@ __attribute__((always_inline)) INLINE static void cooling_cool_part(
   const float u_floor = cooling->min_energy;
 
   /* Current energy */
-  const float u_old = hydro_get_internal_energy(p);
+  const float u_old = hydro_get_physical_internal_energy(p, cosmo);
 
   /* Current du_dt */
   const float hydro_du_dt = hydro_get_internal_energy_dt(p);
 
   /* Calculate cooling du_dt */
-  float cooling_du_dt = cooling_rate(phys_const, us, cooling, p);
+  float cooling_du_dt = cooling_rate(phys_const, us, cosmo, cooling, p);
 
   /* Integrate cooling equation to enforce energy floor */
   /* Factor of 1.5 included since timestep could potentially double */
@@ -117,17 +121,19 @@ __attribute__((always_inline)) INLINE static void cooling_cool_part(
  *
  * @param cooling The #cooling_function_data used in the run.
  * @param phys_const The physical constants in internal units.
+ * @param cosmo The current cosmological model.
  * @param us The internal system of units.
  * @param p Pointer to the particle data.
  */
 __attribute__((always_inline)) INLINE static float cooling_timestep(
     const struct cooling_function_data* restrict cooling,
     const struct phys_const* restrict phys_const,
+    const struct cosmology* restrict cosmo,
     const struct unit_system* restrict us, const struct part* restrict p) {
 
   /* Get current internal energy */
-  const float u = hydro_get_internal_energy(p);
-  const float du_dt = cooling_rate(phys_const, us, cooling, p);
+  const float u = hydro_get_physical_internal_energy(p, cosmo);
+  const float du_dt = cooling_rate(phys_const, us, cosmo, cooling, p);
 
   /* If we are close to (or below) the energy floor, we ignore the condition */
   if (u < 1.01f * cooling->min_energy)
diff --git a/src/cooling/grackle/cooling.h b/src/cooling/grackle/cooling.h
index 29e00435e16751f01bdb58fe26af151b422bd5a1..b7a09f1907604f06d0d5534022228f9ab159a0ff 100644
--- a/src/cooling/grackle/cooling.h
+++ b/src/cooling/grackle/cooling.h
@@ -120,6 +120,7 @@ __attribute__((always_inline)) INLINE static void cooling_print_backend(
 __attribute__((always_inline)) INLINE static double cooling_rate(
     const struct phys_const* restrict phys_const,
     const struct unit_system* restrict us,
+    const struct cosmology* restrict cosmo,
     const struct cooling_function_data* restrict cooling,
     struct part* restrict p, float dt) {
 
@@ -147,8 +148,8 @@ __attribute__((always_inline)) INLINE static double cooling_rate(
   data.grid_end = grid_end;
 
   /* general particle data */
-  gr_float density = hydro_get_density(p);
-  const double energy_before = hydro_get_internal_energy(p);
+  gr_float density = hydro_get_physical_density(p, cosmo);
+  const double energy_before = hydro_get_physical_internal_energy(p, cosmo);
   gr_float energy = energy_before;
   gr_float vx = 0;
   gr_float vy = 0;
@@ -221,6 +222,7 @@ __attribute__((always_inline)) INLINE static double cooling_rate(
  *
  * @param phys_const The physical constants in internal units.
  * @param us The internal system of units.
+ * @param cosmo The current cosmological model.
  * @param cooling The #cooling_function_data used in the run.
  * @param p Pointer to the particle data.
  * @param dt The time-step of this particle.
@@ -228,6 +230,7 @@ __attribute__((always_inline)) INLINE static double cooling_rate(
 __attribute__((always_inline)) INLINE static void cooling_cool_part(
     const struct phys_const* restrict phys_const,
     const struct unit_system* restrict us,
+    const struct cosmology* restrict cosmo,
     const struct cooling_function_data* restrict cooling,
     struct part* restrict p, struct xpart* restrict xp, float dt) {
 
@@ -237,7 +240,7 @@ __attribute__((always_inline)) INLINE static void cooling_cool_part(
   const float hydro_du_dt = hydro_get_internal_energy_dt(p);
 
   /* compute cooling rate */
-  const float du_dt = cooling_rate(phys_const, us, cooling, p, dt);
+  const float du_dt = cooling_rate(phys_const, us, cosmo, cooling, p, dt);
 
   /* record energy lost */
   xp->cooling_data.radiated_energy += -du_dt * dt * hydro_get_mass(p);
@@ -253,12 +256,14 @@ __attribute__((always_inline)) INLINE static void cooling_cool_part(
  *
  * @param cooling The #cooling_function_data used in the run.
  * @param phys_const The physical constants in internal units.
+ * @param cosmo The current cosmological model.
  * @param us The internal system of units.
  * @param p Pointer to the particle data.
  */
 __attribute__((always_inline)) INLINE static float cooling_timestep(
     const struct cooling_function_data* restrict cooling,
     const struct phys_const* restrict phys_const,
+    const struct cosmology* restrict cosmo,
     const struct unit_system* restrict us, const struct part* restrict p) {
 
   return FLT_MAX;
diff --git a/src/cooling/none/cooling.h b/src/cooling/none/cooling.h
index 861159ec1022fcf6f7617fc6ca0c9b470aa6e7fa..4715c06f9fe5db0af78f95d43e37a064688c6ee3 100644
--- a/src/cooling/none/cooling.h
+++ b/src/cooling/none/cooling.h
@@ -54,6 +54,7 @@ __attribute__((always_inline)) INLINE static void cooling_write_flavour(
  *
  * @param phys_const The physical constants in internal units.
  * @param us The internal system of units.
+ * @param cosmo The current cosmological model.
  * @param cooling The #cooling_function_data used in the run.
  * @param p Pointer to the particle data.
  * @param xp Pointer to the extended particle data.
@@ -62,6 +63,7 @@ __attribute__((always_inline)) INLINE static void cooling_write_flavour(
 __attribute__((always_inline)) INLINE static void cooling_cool_part(
     const struct phys_const* restrict phys_const,
     const struct unit_system* restrict us,
+    const struct cosmology* restrict cosmo,
     const struct cooling_function_data* restrict cooling,
     struct part* restrict p, struct xpart* restrict xp, float dt) {}
 
@@ -72,12 +74,14 @@ __attribute__((always_inline)) INLINE static void cooling_cool_part(
  *
  * @param cooling The #cooling_function_data used in the run.
  * @param phys_const The physical constants in internal units.
+ * @param cosmo The current cosmological model.
  * @param us The internal system of units.
  * @param p Pointer to the particle data.
  */
 __attribute__((always_inline)) INLINE static float cooling_timestep(
     const struct cooling_function_data* restrict cooling,
     const struct phys_const* restrict phys_const,
+    const struct cosmology* restrict cosmo,
     const struct unit_system* restrict us, const struct part* restrict p) {
 
   return FLT_MAX;
diff --git a/src/runner.c b/src/runner.c
index fb74b20564482cbf88ddcd07fd99f325f93d3414..1b8f9a2fc84b08c590396f1251a6bf3c5cc3189e 100644
--- a/src/runner.c
+++ b/src/runner.c
@@ -230,7 +230,7 @@ void runner_do_cooling(struct runner *r, struct cell *c, int timer) {
         }
 
         /* Let's cool ! */
-        cooling_cool_part(constants, us, cooling_func, p, xp, dt_cool);
+        cooling_cool_part(constants, us, cosmo, cooling_func, p, xp, dt_cool);
       }
     }
   }
diff --git a/src/timestep.h b/src/timestep.h
index 395f980ff27d0b5c9f0b4fa7be9ee9eaf94528aa..4e715e1f15dd525d89b44ae3e716278cebe28a8d 100644
--- a/src/timestep.h
+++ b/src/timestep.h
@@ -119,7 +119,7 @@ __attribute__((always_inline)) INLINE static integertime_t get_part_timestep(
   float new_dt_cooling = FLT_MAX;
   if (e->policy & engine_policy_cooling)
     new_dt_cooling = cooling_timestep(e->cooling_func, e->physical_constants,
-                                      e->internal_units, p);
+                                      e->cosmology, e->internal_units, p);
 
   /* Compute the next timestep (gravity condition) */
   float new_dt_grav = FLT_MAX, new_dt_self_grav = FLT_MAX,