Extracted the physics from the tasks into a new set of files. The correct file...

Extracted the physics from the tasks into a new set of files. The correct file is included when the hydro or gravity models are selected in const.h

src/gravity.h

+/*******************************************************************************
+ * This file is part of SWIFT.
+ * Coypright (c) 2015 Matthieu Schaller (matthieu.schaller@durham.ac.uk)
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Lesser General Public License as published
+ * by the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ *
+ ******************************************************************************/
+#ifndef SWIFT_GRAVITY_H
+#define SWIFT_GRAVITY_H
+
+#include "./const.h"
+
+/* So far only one model here */
+/* Straight-forward import */
+#include "./gravity/Default/gravity.h"
+
+#endif

src/timestep.h → src/gravity/Default/gravity.h

@@ -16,33 +16,6 @@
  * along with this program.  If not, see <http://www.gnu.org/licenses/>.
  *
  ******************************************************************************/
-
-/**
- * @brief Computes the hydro time-step of a given particle
- *
- * @param p Pointer to the particle data
- * @param xp Pointer to the extended particle data
- *
- */
-__attribute__((always_inline)) INLINE static float compute_timestep_hydro(
-    struct part* p, struct xpart* xp) {
-
-  /* CFL condition */
-  float dt_cfl = const_cfl * p->h / p->force.v_sig;
-
-  /* Limit change in h */
-  float dt_h_change = (p->force.h_dt != 0.0f)
-                          ? fabsf(const_ln_max_h_change * p->h / p->force.h_dt)
-                          : FLT_MAX;
-
-  /* Limit change in u */
-  float dt_u_change = (p->force.u_dt != 0.0f)
-                          ? fabsf(const_max_u_change * p->u / p->force.u_dt)
-                          : FLT_MAX;
-
-  return fminf(dt_cfl, fminf(dt_h_change, dt_u_change));
-}
-
 /**
  * @brief Computes the gravity time-step of a given particle
  *
@@ -51,9 +24,10 @@ __attribute__((always_inline)) INLINE static float compute_timestep_hydro(
  *
  */
 
-__attribute__((always_inline)) INLINE static float compute_timestep_grav(
+__attribute__((always_inline)) INLINE static float gravity_compute_timestep(
     struct part* p, struct xpart* xp) {
 
   /* Currently no limit is imposed */
   return FLT_MAX;
 }
+

src/hydro.h

+/*******************************************************************************
+ * This file is part of SWIFT.
+ * Coypright (c) 2015 Matthieu Schaller (matthieu.schaller@durham.ac.uk)
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Lesser General Public License as published
+ * by the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ *
+ ******************************************************************************/
+#ifndef SWIFT_HYDRO_H
+#define SWIFT_HYDRO_H
+
+#include "./const.h"
+
+/* Import the right functions */
+#ifdef LEGACY_GADGET2_SPH
+#include "./hydro/Gadget2/hydro.h"
+#else
+#include "./hydro/Default/hydro.h"
+#endif
+
+#endif

src/hydro/Default/hydro.h

+/*******************************************************************************
+ * This file is part of SWIFT.
+ * Coypright (c) 2015 Matthieu Schaller (matthieu.schaller@durham.ac.uk)
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Lesser General Public License as published
+ * by the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ *
+ ******************************************************************************/
+
+/**
+ * @brief Computes the hydro time-step of a given particle
+ *
+ * @param p Pointer to the particle data
+ * @param xp Pointer to the extended particle data
+ *
+ */
+__attribute__((always_inline)) INLINE static float hydro_compute_timestep(
+    struct part* p, struct xpart* xp) {
+
+  /* CFL condition */
+  float dt_cfl = const_cfl * p->h / p->force.v_sig;
+
+  /* Limit change in h */
+  float dt_h_change = (p->force.h_dt != 0.0f)
+                          ? fabsf(const_ln_max_h_change * p->h / p->force.h_dt)
+                          : FLT_MAX;
+
+  /* Limit change in u */
+  float dt_u_change = (p->force.u_dt != 0.0f)
+                          ? fabsf(const_max_u_change * p->u / p->force.u_dt)
+                          : FLT_MAX;
+
+  return fminf(dt_cfl, fminf(dt_h_change, dt_u_change));
+}
+
+
+/**
+ * @brief Prepares a particle for the density calculation.
+ *
+ * Zeroes all the relevant arrays in preparation for the sums taking place in 
+ * the variaous density tasks
+ *
+ * @param p The particle to act upon
+ */
+__attribute__((always_inline)) INLINE static void hydro_init_part(struct part* p) {
+  p->density.wcount = 0.f;
+  p->density.wcount_dh = 0.f;
+  p->rho = 0.f;
+  p->rho_dh = 0.f;
+  p->density.div_v = 0.f;
+  p->density.curl_v[0] = 0.f;
+  p->density.curl_v[1] = 0.f;
+  p->density.curl_v[2] = 0.f;
+}
+
+
+/**
+ * @brief Finishes the density calculation. 
+ *
+ * Multiplies the density and number of neighbours by the appropiate constants 
+ * and add the self-contribution term.
+ *
+ * @param p The particle to act upon
+ */
+__attribute__((always_inline)) INLINE static void hydro_end_density(struct part* p) {
+
+  /* Some smoothing length multiples. */
+  const float h = p->h;
+  const float ih = 1.0f / h;
+  const float ih2 = ih * ih;
+  const float ih4 = ih2 * ih2;
+  
+  /* Final operation on the density. */
+  p->rho = ih * ih2 * (p->rho + p->mass * kernel_root);
+  p->rho_dh = (p->rho_dh - 3.0f * p->mass * kernel_root) * ih4;
+  p->density.wcount = (p->density.wcount + kernel_root) *
+    (4.0f / 3.0 * M_PI * kernel_gamma3);
+  p->density.wcount_dh = p->density.wcount_dh * ih * (4.0f / 3.0 * M_PI * kernel_gamma3);
+
+}
+
+
+/**
+ * @brief Prepare a particle for the force calculation.
+ *
+ * Computes viscosity term, conduction term and smoothing length gradient terms.
+ *
+ * @param p The particle to act upon
+ * @param xp The extended particle data to act upon
+ */
+__attribute__((always_inline)) INLINE static void hydro_prepare_force(struct part* p,
+								      struct xpart* xp) {
+
+  /* Some smoothing length multiples. */
+  const float h = p->h;
+  const float ih = 1.0f / h;
+  const float ih2 = ih * ih;
+  const float ih4 = ih2 * ih2;
+
+  /* Pre-compute some stuff for the balsara switch. */
+  const float normDiv_v = fabs(p->density.div_v / p->rho * ih4);
+  const float normCurl_v = sqrtf(p->density.curl_v[0] * p->density.curl_v[0] +
+		     p->density.curl_v[1] * p->density.curl_v[1] +
+		     p->density.curl_v[2] * p->density.curl_v[2]) /
+    p->rho * ih4;
+
+  /* Compute this particle's sound speed. */
+  const float u = p->u;
+  const float fc = p->force.c  =
+    sqrtf(const_hydro_gamma * (const_hydro_gamma - 1.0f) * u);
+
+  /* Compute the P/Omega/rho2. */
+  xp->omega = 1.0f + 0.3333333333f * h * p->rho_dh / p->rho;
+  p->force.POrho2 = u * (const_hydro_gamma - 1.0f) / (p->rho * xp->omega);
+
+  /* Balsara switch */
+  p->force.balsara =
+    normDiv_v / (normDiv_v + normCurl_v + 0.0001f * fc * ih);
+
+  /* Viscosity parameter decay time */
+  const float tau = h / (2.f * const_viscosity_length * p->force.c);
+
+  /* Viscosity source term */
+  const float S = fmaxf(-normDiv_v, 0.f);
+
+  /* Compute the particle's viscosity parameter time derivative */
+  const float alpha_dot = (const_viscosity_alpha_min - p->alpha) / tau +
+    (const_viscosity_alpha_max - p->alpha) * S;
+
+  /* Update particle's viscosity paramter */
+  p->alpha += alpha_dot * (p->t_end - p->t_begin);
+}
+
+
+/**
+ * @brief Reset acceleration fields of a particle
+ *
+ * Resets all hydro acceleration and time derivative fields in preparation
+ * for the sums taking place in the variaous force tasks
+ *
+ * @param p The particle to act upon
+ */
+__attribute__((always_inline)) INLINE static void hydro_reset_acceleration(struct part* p) {
+
+  /* Reset the acceleration. */
+  p->a[0] = 0.0f;
+  p->a[1] = 0.0f;
+  p->a[2] = 0.0f;
+  
+  /* Reset the time derivatives. */
+  p->force.u_dt = 0.0f;
+  p->force.h_dt = 0.0f;
+  p->force.v_sig = 0.0f;
+}
+
+/**
+ * @brief Predict additional particle fields forward in time when drifting
+ *
+ * @param p The particle
+ * @param xp The extended data of the particle
+ * @param dt The time-step over which to drift
+ */
+__attribute__((always_inline)) INLINE static void hydro_predict_extra(struct part* p,
+								      struct xpart* xp,
+								      float dt) {
+  float u, w;
+
+  /* Predict internal energy */
+  w = p->force.u_dt / p->u * dt;
+  if (fabsf(w) < 0.01f) /* 1st order expansion of exp(w) */
+    u = p->u *=
+      1.0f +
+      w * (1.0f + w * (0.5f + w * (1.0f / 6.0f + 1.0f / 24.0f * w))); 
+  else
+    u = p->u *= expf(w);
+  
+  /* Predict gradient term */
+  p->force.POrho2 = u * (const_hydro_gamma - 1.0f) / (p->rho * xp->omega);
+ 
+}
+
+/**
+ * @brief Finishes the force calculation. 
+ *
+ * Multiplies the forces and accelerationsby the appropiate constants 
+ *
+ * @param p The particle to act upon
+ */
+__attribute__((always_inline)) INLINE static void hydro_end_force(struct part* p) {
+
+  p->force.h_dt *= p->h * 0.333333333f;	
+  
+}

src/hydro/Gadget2/hydro.h

+/*******************************************************************************
+ * This file is part of SWIFT.
+ * Coypright (c) 2015 Matthieu Schaller (matthieu.schaller@durham.ac.uk)
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Lesser General Public License as published
+ * by the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ *
+ ******************************************************************************/
+
+/**
+ * @brief Computes the hydro time-step of a given particle
+ *
+ * @param p Pointer to the particle data
+ * @param xp Pointer to the extended particle data
+ *
+ */
+__attribute__((always_inline)) INLINE static float hydro_compute_timestep(
+    struct part* p, struct xpart* xp) {
+
+  /* CFL condition */
+  float dt_cfl = const_cfl * p->h / p->force.v_sig;
+
+  /* Limit change in h */
+  float dt_h_change = (p->force.h_dt != 0.0f)
+                          ? fabsf(const_ln_max_h_change * p->h / p->force.h_dt)
+                          : FLT_MAX;
+
+  /* Limit change in u */
+  float dt_u_change = (p->force.u_dt != 0.0f)
+                          ? fabsf(const_max_u_change * p->u / p->force.u_dt)
+                          : FLT_MAX;
+
+  return fminf(dt_cfl, fminf(dt_h_change, dt_u_change));
+}
+
+
+
+/**
+ * @brief Prepares a particle for the density calculation.
+ *
+ * Zeroes all the relevant arrays in preparation for the sums taking place in 
+ * the variaous density tasks
+ *
+ * @param p The particle to act upon
+ */
+__attribute__((always_inline)) INLINE static void hydro_init_part(struct part* p) {
+  p->density.wcount = 0.f;
+  p->density.wcount_dh = 0.f;
+  p->rho = 0.f;
+  p->rho_dh = 0.f;
+  p->density.div_v = 0.f;
+  p->density.curl_v[0] = 0.f;
+  p->density.curl_v[1] = 0.f;
+  p->density.curl_v[2] = 0.f;
+}
+
+/**
+ * @brief Finishes the density calculation. 
+ *
+ * Multiplies the density and number of neighbours by the appropiate constants 
+ * and add the self-contribution term.
+ *
+ * @param p The particle to act upon
+ */
+__attribute__((always_inline)) INLINE static void hydro_end_density(struct part* p) {
+
+  /* Some smoothing length multiples. */
+  const float h = p->h;
+  const float ih = 1.0f / h;
+  const float ih2 = ih * ih;
+  const float ih4 = ih2 * ih2;
+  
+  /* Final operation on the density. */
+  p->rho = ih * ih2 * (p->rho + p->mass * kernel_root);
+  p->rho_dh = (p->rho_dh - 3.0f * p->mass * kernel_root) * ih4;
+  p->density.wcount = (p->density.wcount + kernel_root) *
+    (4.0f / 3.0 * M_PI * kernel_gamma3);
+  p->density.wcount_dh = p->density.wcount_dh * ih * (4.0f / 3.0 * M_PI * kernel_gamma3);
+
+}
+
+
+/**
+ * @brief Prepare a particle for the force calculation.
+ *
+ * Computes viscosity term, conduction term and smoothing length gradient terms.
+ *
+ * @param p The particle to act upon
+ * @param xp The extended particle data to act upon
+ */
+__attribute__((always_inline)) INLINE static void hydro_prepare_force(struct part* p,
+								      struct xpart* xp) {
+
+  /* Some smoothing length multiples. */
+  const float h = p->h;
+  const float ih = 1.0f / h;
+  const float ih2 = ih * ih;
+  const float ih4 = ih2 * ih2;
+
+  
+  /* Pre-compute some stuff for the balsara switch. */
+  const float normDiv_v = fabs(p->density.div_v / p->rho * ih4);
+  const float normCurl_v = sqrtf(p->density.curl_v[0] * p->density.curl_v[0] +
+				 p->density.curl_v[1] * p->density.curl_v[1] +
+				 p->density.curl_v[2] * p->density.curl_v[2]) /
+    p->rho * ih4;
+
+  /* Compute this particle's sound speed. */
+  const float u = p->u;
+  const float fc = p->force.c = 
+    sqrtf(const_hydro_gamma * (const_hydro_gamma - 1.0f) * u);
+  
+  /* Compute the P/Omega/rho2. */
+  xp->omega = 1.0f + 0.3333333333f * h * p->rho_dh / p->rho;
+  p->force.POrho2 = u * (const_hydro_gamma - 1.0f) / (p->rho * xp->omega);
+  
+  /* Balsara switch */
+  p->force.balsara =
+    normDiv_v / (normDiv_v + normCurl_v + 0.0001f * fc * ih);
+}
+
+
+
+/**
+ * @brief Reset acceleration fields of a particle
+ *
+ * Resets all hydro acceleration and time derivative fields in preparation
+ * for the sums taking place in the variaous force tasks
+ *
+ * @param p The particle to act upon
+ */
+__attribute__((always_inline)) INLINE static void hydro_reset_acceleration(struct part* p) {
+
+  /* Reset the acceleration. */
+  p->a[0] = 0.0f;
+  p->a[1] = 0.0f;
+  p->a[2] = 0.0f;
+  
+  /* Reset the time derivatives. */
+  p->force.u_dt = 0.0f;
+  p->force.h_dt = 0.0f;
+  p->force.v_sig = 0.0f;
+}
+
+
+/**
+ * @brief Predict additional particle fields forward in time when drifting
+ *
+ * @param p The particle
+ * @param xp The extended data of the particle
+ * @param dt The time-step over which to drift
+ */
+__attribute__((always_inline)) INLINE static void hydro_predict_extra(struct part* p,
+								      struct xpart* xp,
+								      float dt) {
+  float u, w;
+
+  /* Predict internal energy */
+  w = p->force.u_dt / p->u * dt;
+  if (fabsf(w) < 0.01f) /* 1st order expansion of exp(w) */
+    u = p->u *=
+      1.0f +
+      w * (1.0f + w * (0.5f + w * (1.0f / 6.0f + 1.0f / 24.0f * w))); 
+  else
+    u = p->u *= expf(w);
+  
+  /* Predict gradient term */
+  p->force.POrho2 = u * (const_hydro_gamma - 1.0f) / (p->rho * xp->omega);
+ 
+}
+
+
+
+/**
+ * @brief Finishes the force calculation. 
+ *
+ * Multiplies the forces and accelerationsby the appropiate constants 
+ *
+ * @param p The particle to act upon
+ */
+__attribute__((always_inline)) INLINE static void hydro_end_force(struct part* p) {
+
+  p->force.h_dt *= p->h * 0.333333333f;	
+  
+}

src/runner.c

@@ -42,7 +42,8 @@
 #include "space.h"
 #include "task.h"
 #include "timers.h"
-#include "timestep.h"
+#include "hydro.h"
+#include "gravity.h"
 
 /* Include the right variant of the SPH interactions */
 #ifdef LEGACY_GADGET2_SPH
@@ -518,14 +519,7 @@ void runner_doinit(struct runner *r, struct cell *c) {
     if (p->t_end <= t_end) {
 
       /* Get ready for a density calculation */
-      p->density.wcount = 0.f;
-      p->density.wcount_dh = 0.f;
-      p->rho = 0.f;
-      p->rho_dh = 0.f;
-      p->density.div_v = 0.f;
-      p->density.curl_v[0] = 0.f;
-      p->density.curl_v[1] = 0.f;
-      p->density.curl_v[2] = 0.f;
+      hydro_init_part(p);
     }
   }
 }
@@ -544,11 +538,7 @@ void runner_doghost(struct runner *r, struct cell *c) {
   struct cell *finger;
   int redo, count = c->count;
   int *pid;
-  float h_corr, rho, wcount, rho_dh, wcount_dh, u, fc;
-  float normDiv_v, normCurl_v;
-#ifndef LEGACY_GADGET2_SPH
-  float alpha_dot, tau, S;
-#endif
+  float h_corr;
   float t_end = r->e->time;
 
   TIMER_TIC;
@@ -582,102 +572,50 @@ void runner_doghost(struct runner *r, struct cell *c) {
       /* Is this part within the timestep? */
       if (p->t_end <= t_end) {
 
-        /* Some smoothing length multiples. */
-        const float h = p->h;
-        const float ih = 1.0f / h;
-	const float ih2 = ih * ih;
-	const float ih4 = ih2 * ih2;
-
-        /* Final operation on the density. */
-        p->rho = rho = ih * ih2 * (p->rho + p->mass * kernel_root);
-        p->rho_dh = rho_dh = (p->rho_dh - 3.0f * p->mass * kernel_root) * ih4;
-        p->density.wcount = wcount = (p->density.wcount + kernel_root) *
-                                     (4.0f / 3.0 * M_PI * kernel_gamma3);
-        wcount_dh =
-            p->density.wcount_dh * ih * (4.0f / 3.0 * M_PI * kernel_gamma3);
+	/* Finish the density calculation */
+	hydro_end_density(p);
 
         /* If no derivative, double the smoothing length. */
-        if (wcount_dh == 0.0f) h_corr = p->h;
+        if (p->density.wcount_dh == 0.0f) h_corr = p->h;
 
         /* Otherwise, compute the smoothing length update (Newton step). */
         else {
-          h_corr = (kernel_nwneigh - wcount) / wcount_dh;
+          h_corr = (kernel_nwneigh - p->density.wcount) / p->density.wcount_dh;
 
           /* Truncate to the range [ -p->h/2 , p->h ]. */
-          h_corr = fminf(h_corr, h);
-          h_corr = fmaxf(h_corr, -h / 2.f);
+          h_corr = fminf(h_corr, p->h);
+          h_corr = fmaxf(h_corr, -p->h * 0.5f);
         }
 
         /* Did we get the right number density? */
-        if (wcount > kernel_nwneigh + const_delta_nwneigh ||
-            wcount < kernel_nwneigh - const_delta_nwneigh) {
+        if (p->density.wcount > kernel_nwneigh + const_delta_nwneigh ||
+           p->density. wcount < kernel_nwneigh - const_delta_nwneigh) {
 
           /* Ok, correct then */
           p->h += h_corr;
 
-          /* And flag for another round of fun */
+          /* Flag for another round of fun */
           pid[redo] = pid[i];
           redo += 1;
-          p->density.wcount = 0.f;
-          p->density.wcount_dh = 0.f;
-          p->rho = 0.f;
-          p->rho_dh = 0.f;
-          p->density.div_v = 0.f;
-          p->density.curl_v[0] = 0.f;
-	  p->density.curl_v[1] = 0.f;
-	  p->density.curl_v[2] = 0.f;
+
+	  /* Re-initialise everything */
+	  hydro_init_part(p);
+
+	  /* Off we go ! */
           continue;
         }
 
         /* We now have a particle whose smoothing length has converged */
 
-        /* Pre-compute some stuff for the balsara switch. */
-        normDiv_v = fabs(p->density.div_v / rho * ih4);
-        normCurl_v = sqrtf(p->density.curl_v[0] * p->density.curl_v[0] +
-                           p->density.curl_v[1] * p->density.curl_v[1] +
-                           p->density.curl_v[2] * p->density.curl_v[2]) /
-                     rho * ih4;
-
         /* As of here, particle force variables will be set. Do _NOT_
            try to read any particle density variables! */
 
-        /* Compute this particle's sound speed. */
-        u = p->u;
-        p->force.c = fc =
-            sqrtf(const_hydro_gamma * (const_hydro_gamma - 1.0f) * u);
-
-        /* Compute the P/Omega/rho2. */
-        xp->omega = 1.0f + 0.3333333333f * h * p->rho_dh / p->rho;
-        p->force.POrho2 = u * (const_hydro_gamma - 1.0f) / (rho * xp->omega);
-
-        /* Balsara switch */
-        p->force.balsara =
-            normDiv_v / (normDiv_v + normCurl_v + 0.0001f * fc * ih);
-
-#ifndef LEGACY_GADGET2_SPH
-        /* Viscosity parameter decay time */
-        tau = h / (2.f * const_viscosity_length * p->force.c);
-
-        /* Viscosity source term */
-        S = fmaxf(-normDiv_v, 0.f);
-
-        /* Compute the particle's viscosity parameter time derivative */
-        alpha_dot = (const_viscosity_alpha_min - p->alpha) / tau +
-                    (const_viscosity_alpha_max - p->alpha) * S;
-
-        /* Update particle's viscosity paramter */
-        p->alpha += alpha_dot * (p->t_end - p->t_begin);
-#endif
-
-        /* Reset the acceleration. */
-	p->a[0] = 0.0f;
-	p->a[1] = 0.0f;
-	p->a[2] = 0.0f;
+	/* Compute variables required for the force loop */
+	hydro_prepare_force(p, xp);
 	
-        /* Reset the time derivatives. */
-        p->force.u_dt = 0.0f;
-        p->force.h_dt = 0.0f;
-        p->force.v_sig = 0.0f;
+	/* Prepare the particle for the force loop over neighbours */
+	hydro_reset_acceleration(p);
+
       }
     }
 
@@ -750,10 +688,10 @@ void runner_dodrift(struct runner *r, struct cell *c, int timer) {
 
   const int nr_parts = c->count;
   const float dt = r->e->time - r->e->timeOld;
-  float u, w, rho;
   struct part *restrict p, *restrict parts = c->parts;
   struct xpart *restrict xp, *restrict xparts = c->xparts;
-
+  float w;
+  
   TIMER_TIC
 
   /* No children? */
@@ -780,38 +718,27 @@ void runner_dodrift(struct runner *r, struct cell *c, int timer) {
       p->v[1] += p->a[1] * dt;
       p->v[2] += p->a[2] * dt;
 
-      /* Predict internal energy */
-      w = p->force.u_dt / p->u * dt;
-      if (fabsf(w) < 0.01f)
-	u = p->u *=
-            1.0f +
-            w * (1.0f + w * (0.5f + w * (1.0f / 6.0f +
-                                         1.0f / 24.0f * w))); /* 1st order
-                                                                 expansion of
-                                                                 exp(w) */
-      else
-        u = p->u *= expf(w);
-
       /* Predict smoothing length */
       w = p->force.h_dt * ih * dt;
-      if (fabsf(w) < 0.01f)
-      	p->h *=
-      	  1.0f +
-      	  w * (1.0f + w * (0.5f + w * (1.0f / 6.0f + 1.0f / 24.0f * w)));
+      if (fabsf(w) < 0.01f) /* 1st order expansion of exp(w) */
+	p->h *=
+	  1.0f +
+	  w * (1.0f + w * (0.5f + w * (1.0f / 6.0f + 1.0f / 24.0f * w)));
       else
-      	p->h *= expf(w);
+	p->h *= expf(w);
       
       /* Predict density */
       w = -3.0f * p->force.h_dt * ih * dt;
       if (fabsf(w) < 0.1f)
-        rho = p->rho *=
-            1.0f +
-            w * (1.0f + w * (0.5f + w * (1.0f / 6.0f + 1.0f / 24.0f * w)));
+	p->rho *=
+	  1.0f +
+	  w * (1.0f + w * (0.5f + w * (1.0f / 6.0f + 1.0f / 24.0f * w)));
       else
-        rho = p->rho *= expf(w);
+	p->rho *= expf(w);
+      
+      /* Predict the values of the extra fields */
+      hydro_predict_extra(p, xp, dt);
 
-      /* Predict gradient term */
-      p->force.POrho2 = u * (const_hydro_gamma - 1.0f) / (rho * xp->omega);
     }
   }
 
@@ -877,8 +804,8 @@ void runner_dokick(struct runner *r, struct cell *c, int timer) {
       if ( is_fixdt || p->t_end <= t_current ) {
 
         /* First, finish the force loop */
-        p->force.h_dt *= p->h * 0.333333333f;
-
+	hydro_end_force(p);
+	  
         /* Recover the current timestep */
         current_dt = p->t_end - p->t_begin;
 
@@ -890,8 +817,8 @@ void runner_dokick(struct runner *r, struct cell *c, int timer) {
 	} else {
 	
 	  /* Compute the next timestep */
-	  new_dt_hydro = compute_timestep_hydro(p, xp);
-	  new_dt_grav = compute_timestep_grav(p, xp);
+	  new_dt_hydro = hydro_compute_timestep(p, xp);
+	  new_dt_grav = gravity_compute_timestep(p, xp);
 	  
 	  new_dt = fminf(new_dt_hydro, new_dt_grav);
 	  
@@ -931,7 +858,8 @@ void runner_dokick(struct runner *r, struct cell *c, int timer) {
 
       /* Now collect quantities for statistics */
 
-      v_full[0] = xp->v_full[0], v_full[1] = xp->v_full[1],
+      v_full[0] = xp->v_full[0];
+      v_full[1] = xp->v_full[1];
       v_full[2] = xp->v_full[2];
 
       /* Collect momentum */