diff --git a/src/adiabatic_index.h b/src/adiabatic_index.h
index a0c9ce09e3e004af07e8b208ef9f1af5f46c9e81..763510fdd7b7fa59cae383aad8aca53f7d8c38db 100644
--- a/src/adiabatic_index.h
+++ b/src/adiabatic_index.h
@@ -432,4 +432,40 @@ __attribute__((always_inline)) INLINE static float pow_one_over_gamma(float x) {
 #endif
 }
 
+/**
+ * @brief Return the argument to the power one minus two over the adiabatic
+ * index
+ *
+ * Computes \f$x^{1 - \frac{2}{\gamma}}\f$.
+ *
+ * @param x Argument
+ * @return Argument to the power one minus two over the adiabatic index
+ */
+__attribute__((always_inline)) INLINE static float pow_one_minus_two_over_gamma(
+    float x) {
+
+#if defined(HYDRO_GAMMA_5_3)
+
+  return powf(x, -0.2f);
+
+#elif defined(HYDRO_GAMMA_7_5)
+
+  return powf(x, -0.428571429f);
+
+#elif defined(HYDRO_GAMMA_4_3)
+
+  return 1.f / sqrtf(x);
+
+#elif defined(HYDRO_GAMMA_2_1)
+
+  return 1.f;
+
+#else
+
+  error("The adiabatic index is not defined !");
+  return 0.f;
+
+#endif
+}
+
 #endif /* SWIFT_ADIABATIC_INDEX_H */
diff --git a/src/hydro/PressureEntropy/hydro.h b/src/hydro/PressureEntropy/hydro.h
index c97f6699b88b6c9643f47c3a58dc460a34ced161..ab41ab8263a2f73aa5947da0a63851b44c5b4ad7 100644
--- a/src/hydro/PressureEntropy/hydro.h
+++ b/src/hydro/PressureEntropy/hydro.h
@@ -20,8 +20,9 @@
 #define SWIFT_PRESSURE_ENTROPY_HYDRO_H
 
 /**
- * @file PressureEntropy/hydro_part.h
- * @brief Pressure-Entropy implementation of SPH (Particle definition)
+ * @file PressureEntropy/hydro.h
+ * @brief Pressure-Entropy implementation of SPH (Non-neighbour loop
+ * equations)
  *
  * The thermal variable is the entropy (S) and the entropy is smoothed over
  * contact discontinuities to prevent spurious surface tension.
@@ -88,6 +89,28 @@ __attribute__((always_inline)) INLINE static float hydro_get_soundspeed(
   return p->force.soundspeed;
 }
 
+/**
+ * @brief Returns the density of a particle
+ *
+ * @param p The particle of interest
+ */
+__attribute__((always_inline)) INLINE static float hydro_get_density(
+    const struct part *restrict p) {
+
+  return p->rho;
+}
+
+/**
+ * @brief Returns the mass of a particle
+ *
+ * @param p The particle of interest
+ */
+__attribute__((always_inline)) INLINE static float hydro_get_mass(
+    const struct part *restrict p) {
+
+  return p->mass;
+}
+
 /**
  * @brief Modifies the thermal state of a particle to the imposed internal
  * energy
@@ -301,8 +324,8 @@ __attribute__((always_inline)) INLINE static void hydro_reset_acceleration(
  * @param timeBase The minimal time-step size
  */
 __attribute__((always_inline)) INLINE static void hydro_predict_extra(
-    struct part *restrict p, const struct xpart *restrict xp, int t0, int t1,
-    double timeBase) {
+    struct part *restrict p, const struct xpart *restrict xp, float dt, int t0,
+    int t1, double timeBase) {
 
   /* Drift the pressure */
   const float dt_entr = (t1 - (p->ti_begin + p->ti_end) / 2) * timeBase;
diff --git a/src/hydro/PressureEntropy/hydro_debug.h b/src/hydro/PressureEntropy/hydro_debug.h
index eda1fa2e923638d700ae5aa673c4f0621e688466..daa0d2a7fd386617314b449b81d086944e20f932 100644
--- a/src/hydro/PressureEntropy/hydro_debug.h
+++ b/src/hydro/PressureEntropy/hydro_debug.h
@@ -21,7 +21,7 @@
 
 /**
  * @file PressureEntropy/hydro_debug.h
- * @brief Pressure-Entropy implementation of SPH (Particle definition)
+ * @brief Pressure-Entropy implementation of SPH (Debugging routines)
  *
  * The thermal variable is the entropy (S) and the entropy is smoothed over
  * contact discontinuities to prevent spurious surface tension.
diff --git a/src/hydro/PressureEntropy/hydro_iact.h b/src/hydro/PressureEntropy/hydro_iact.h
index da38f417fc71c3cbc155061da9dfea2103c474b3..c1c0641b296207a62382aeecf23f5c3ff33dfe1d 100644
--- a/src/hydro/PressureEntropy/hydro_iact.h
+++ b/src/hydro/PressureEntropy/hydro_iact.h
@@ -20,8 +20,8 @@
 #define SWIFT_PRESSURE_ENTROPY_HYDRO_IACT_H
 
 /**
- * @file PressureEntropy/hydro_part.h
- * @brief Pressure-Entropy implementation of SPH (Particle definition)
+ * @file PressureEntropy/hydro_iact.h
+ * @brief Pressure-Entropy implementation of SPH (Neighbour loop equations)
  *
  * The thermal variable is the entropy (S) and the entropy is smoothed over
  * contact discontinuities to prevent spurious surface tension.
diff --git a/src/hydro/PressureEntropy/hydro_io.h b/src/hydro/PressureEntropy/hydro_io.h
index 8a3dc60210e1cd40678154253615882a63647d65..746c38d0e3359fa07b32603c43e73308a3897db1 100644
--- a/src/hydro/PressureEntropy/hydro_io.h
+++ b/src/hydro/PressureEntropy/hydro_io.h
@@ -21,7 +21,7 @@
 
 /**
  * @file PressureEntropy/hydro_io.h
- * @brief Pressure-Entropy implementation of SPH (Particle definition)
+ * @brief Pressure-Entropy implementation of SPH (i/o routines)
  *
  * The thermal variable is the entropy (S) and the entropy is smoothed over
  * contact discontinuities to prevent spurious surface tension.
diff --git a/src/hydro/PressureEntropy/hydro_part.h b/src/hydro/PressureEntropy/hydro_part.h
index 711ddd4bbea356b82e0482498c6944d59fd6ec53..cc71f9b994c493238ff1bb8bc678ea14c4e541e2 100644
--- a/src/hydro/PressureEntropy/hydro_part.h
+++ b/src/hydro/PressureEntropy/hydro_part.h
@@ -29,100 +29,105 @@
  * Follows Hopkins, P., MNRAS, 2013, Volume 428, Issue 4, pp. 2840-2856
  */
 
+#include "cooling_struct.h"
+
 /* Extra particle data not needed during the SPH loops over neighbours. */
 struct xpart {
 
-  /* Offset between current position and position at last tree rebuild. */
+  /*! Offset between current position and position at last tree rebuild. */
   float x_diff[3];
 
-  /* Velocity at the last full step. */
+  /*! Velocity at the last full step. */
   float v_full[3];
 
-} __attribute__((aligned(xpart_align)));
+  /*! Additional data used to record cooling information */
+  struct cooling_xpart_data cooling_data;
+
+} SWIFT_STRUCT_ALIGN;
 
 /* Data of a single particle. */
 struct part {
 
-  /* Particle position. */
+  /*! Particle position. */
   double x[3];
 
-  /* Particle predicted velocity. */
+  /*! Particle predicted velocity. */
   float v[3];
 
-  /* Particle acceleration. */
+  /*! Particle acceleration. */
   float a_hydro[3];
 
-  /* Particle cutoff radius. */
+  /*! Particle cutoff radius. */
   float h;
 
-  /* Particle mass. */
+  /*! Particle mass. */
   float mass;
 
-  /* Particle time of beginning of time-step. */
+  /*! Particle time of beginning of time-step. */
   int ti_begin;
 
-  /* Particle time of end of time-step. */
+  /*! Particle time of end of time-step. */
   int ti_end;
 
-  /* Particle density. */
+  /*! Particle density. */
   float rho;
 
-  /* Particle weighted pressure. */
+  /*! Particle weighted pressure. */
   float weightedPressure;
 
-  /* Particle entropy. */
+  /*! Particle entropy. */
   float entropy;
 
-  /* Entropy time derivative */
+  /*! Entropy time derivative */
   float entropy_dt;
 
   union {
 
     struct {
 
-      /* Number of neighbours. */
+      /*! Number of neighbours. */
       float wcount;
 
-      /* Number of neighbours spatial derivative. */
+      /*! Number of neighbours spatial derivative. */
       float wcount_dh;
 
-      /* Derivative of particle weighted pressure with h. */
+      /*! Derivative of particle weighted pressure with h. */
       float weightedPressure_dh;
 
-      /* Particle velocity curl. */
+      /*! Particle velocity curl. */
       // float rot_v[3];
 
-      /* Particle velocity divergence. */
+      /*! Particle velocity divergence. */
       // float div_v;
 
     } density;
 
     struct {
 
-      /* "Grad h" term */
+      /*! "Grad h" term */
       float f_ij;
 
-      /* Pressure term */
+      /*! Pressure term */
       float pressure_term;
 
-      /* Particle sound speed. */
+      /*! Particle sound speed. */
       float soundspeed;
 
-      /* Signal velocity. */
+      /*! Signal velocity. */
       float v_sig;
 
-      /* Time derivative of the smoothing length */
+      /*! Time derivative of the smoothing length */
       float h_dt;
 
     } force;
   };
 
-  /* Particle ID. */
+  /*! Particle ID. */
   long long id;
 
-  /* Pointer to corresponding gravity part. */
+  /*! Pointer to corresponding gravity part. */
   struct gpart* gpart;
 
-} __attribute__((aligned(part_align)));
+} SWIFT_STRUCT_ALIGN;
 
 #endif /* SWIFT_PRESSURE_ENTROPY_HYDRO_PART_H */