Conserved quantities are now computed on the fly and written to the 'energy.txt' file

src/cell.h

@@ -136,8 +136,8 @@ struct cell {
   /* Momentum of particles in cell. */
   float mom[3], ang[3];
 
-  /* Potential and kinetic energy of particles in this cell. */
-  double epot, ekin;
+  /* Mass, potential, internal  and kinetic energy of particles in this cell. */
+  double mass, e_pot, e_int, e_kin;
 
   /* Number of particles updated in this cell. */
   int updated;

src/engine.c

@@ -1617,7 +1617,7 @@ void engine_collect_kick(struct cell *c) {
 
   int updated = 0;
   float t_end_min = FLT_MAX, t_end_max = 0.0f;
-  double ekin = 0.0, epot = 0.0;
+  double e_kin = 0.0, e_int = 0.0, e_pot = 0.0;
   float mom[3] = {0.0f, 0.0f, 0.0f}, ang[3] = {0.0f, 0.0f, 0.0f};
   struct cell *cp;
 
@@ -1641,8 +1641,9 @@ void engine_collect_kick(struct cell *c) {
         t_end_min = fminf(t_end_min, cp->t_end_min);
         t_end_max = fmaxf(t_end_max, cp->t_end_max);
         updated += cp->updated;
-        ekin += cp->ekin;
-        epot += cp->epot;
+        e_kin += cp->e_kin;
+        e_int += cp->e_int;
+        e_pot += cp->e_pot;
         mom[0] += cp->mom[0];
         mom[1] += cp->mom[1];
         mom[2] += cp->mom[2];
@@ -1656,8 +1657,9 @@ void engine_collect_kick(struct cell *c) {
   c->t_end_min = t_end_min;
   c->t_end_max = t_end_max;
   c->updated = updated;
-  c->ekin = ekin;
-  c->epot = epot;
+  c->e_kin = e_kin;
+  c->e_int = e_int;
+  c->e_pot = e_pot;
   c->mom[0] = mom[0];
   c->mom[1] = mom[1];
   c->mom[2] = mom[2];
@@ -1752,7 +1754,7 @@ void engine_step(struct engine *e) {
   int k;
   int updates = 0;
   float t_end_min = FLT_MAX, t_end_max = 0.f;
-  double epot = 0.0, ekin = 0.0;
+  double e_pot = 0.0, e_int = 0.0, e_kin = 0.0;
   float mom[3] = {0.0, 0.0, 0.0};
   float ang[3] = {0.0, 0.0, 0.0};
   struct cell *c;
@@ -1764,11 +1766,16 @@ void engine_step(struct engine *e) {
   for (k = 0; k < s->nr_cells; k++)
     if (s->cells[k].nodeID == e->nodeID) {
       c = &s->cells[k];
+
+      /* Recurse */
       engine_collect_kick(c);
+
+      /* And aggregate */
       t_end_min = fminf(t_end_min, c->t_end_min);
       t_end_max = fmaxf(t_end_max, c->t_end_max);
-      ekin += c->ekin;
-      epot += c->epot;
+      e_kin += c->e_kin;
+      e_int += c->e_int;
+      e_pot += c->e_pot;
       updates += c->updated;
       mom[0] += c->mom[0];
       mom[1] += c->mom[1];
@@ -1780,7 +1787,7 @@ void engine_step(struct engine *e) {
 
 /* Aggregate the data from the different nodes. */
 #ifdef WITH_MPI
-  double in[3], out[3];
+  double in[4], out[4];
   out[0] = t_end_min;
   if (MPI_Allreduce(out, in, 1, MPI_DOUBLE, MPI_MIN, MPI_COMM_WORLD) !=
       MPI_SUCCESS)
@@ -1792,20 +1799,16 @@ void engine_step(struct engine *e) {
     error("Failed to aggregate t_end_max.");
   t_end_max = in[0];
   out[0] = updates;
-  out[1] = ekin;
-  out[2] = epot;
+  out[1] = e_kin;
+  out[2] = e_int;
+  out[3] = e_pot;
   if (MPI_Allreduce(out, in, 3, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD) !=
       MPI_SUCCESS)
     error("Failed to aggregate energies.");
   updates = in[0];
-  ekin = in[1];
-  epot = in[2];
-/* int nr_parts;
-if ( MPI_Allreduce( &s->nr_parts , &nr_parts , 1 , MPI_INT , MPI_SUM ,
-MPI_COMM_WORLD ) != MPI_SUCCESS )
-    error( "Failed to aggregate particle count." );
-if ( e->nodeID == 0 )
-    message( "nr_parts=%i." , nr_parts ); */
+  e_kin = in[1];
+  e_int = in[2];
+  e_pot = in[3];
 #endif
 
   /* Move forward in time */
@@ -1837,11 +1840,18 @@ if ( e->nodeID == 0 )
 
   TIMER_TOC2(timer_step);
 
-  /* Print some information to the screen */
   if (e->nodeID == 0) {
+
+    /* Print some information to the screen */
     printf("%d %f %f %d %.3f\n", e->step, e->time, e->timeStep, updates,
            ((double)timers[timer_count - 1]) / CPU_TPS * 1000);
     fflush(stdout);
+
+    /* Write some energy statistics */
+    fprintf(e->file_stats, "%d %f %f %f %f %f %f %f %f %f %f\n", e->step, e_kin,
+            e_int, e_pot, e_kin + e_int + e_pot, mom[0], mom[1], mom[2], ang[0],
+            ang[1], ang[2]);
+    fflush(e->file_stats);
   }
 }
 
@@ -2131,6 +2141,7 @@ void engine_init(struct engine *e, struct space *s, float dt, int nr_threads,
   e->timeStep = 0.;
   e->dt_min = dt_min;
   e->dt_max = dt_max;
+  e->file_stats = NULL;
   engine_rank = nodeID;
 
   /* Make the space link back to the engine. */
@@ -2150,6 +2161,14 @@ void engine_init(struct engine *e, struct space *s, float dt, int nr_threads,
 #endif
   }
 
+  /* Open some files */
+  if (e->nodeID == 0) {
+    e->file_stats = fopen("energy.txt", "w");
+    fprintf(e->file_stats,
+            "# Step E_kin E_int E_pot E_tot "
+            "p_x p_y p_z ang_x ang_y ang_z\n");
+  }
+
   /* Print policy */
   engine_print_policy(e);
 

src/engine.h

@@ -29,6 +29,7 @@
 
 /* Some standard headers. */
 #include <pthread.h>
+#include <stdio.h>
 
 /* Includes. */
 #include "lock.h"
@@ -109,8 +110,8 @@ struct engine {
   /* Time step */
   float timeStep;
 
-  /* The system energies from the previous step. */
-  double ekin, epot;
+  /* File for statistics */
+  FILE *file_stats;
 
   /* The current step number. */
   int step, nullstep;

src/hydro/Default/hydro.h

@@ -224,3 +224,14 @@ __attribute__((always_inline))
  */
 __attribute__((always_inline))
     INLINE static void hydro_convert_quantities(struct part* p) {}
+
+/**
+ * @brief Returns the internal energy of a particle
+ *
+ * @param p The particle of interest
+ */
+__attribute__((always_inline))
+    INLINE static float hydro_get_internal_energy(struct part* p) {
+
+  return p->u;
+}

src/hydro/Gadget2/hydro.h

@@ -238,3 +238,15 @@ __attribute__((always_inline))
   p->entropy = (const_hydro_gamma - 1.f) * p->entropy *
                powf(p->rho, -(const_hydro_gamma - 1.f));
 }
+
+/**
+ * @brief Returns the internal energy of a particle
+ *
+ * @param p The particle of interest
+ */
+__attribute__((always_inline))
+    INLINE static float hydro_get_internal_energy(struct part* p) {
+
+  return p->entropy * powf(p->rho, const_hydro_gamma - 1.f) *
+         (1.f / (const_hydro_gamma - 1.f));
+}

src/hydro/Minimal/hydro.h

@@ -218,3 +218,18 @@ __attribute__((always_inline)) INLINE static void hydro_kick_extra(
  */
 __attribute__((always_inline))
     INLINE static void hydro_convert_quantities(struct part* p) {}
+
+/**
+ * @brief Returns the internal energy of a particle
+ *
+ * For implementations where the main thermodynamic variable
+ * is not internal energy, this function computes the internal
+ * energy from the thermodynamic variable.
+ *
+ * @param p The particle of interest
+ */
+__attribute__((always_inline))
+    INLINE static float hydro_get_internal_energy(struct part* p) {
+
+  return p->u;
+}

src/runner.c

@@ -807,9 +807,10 @@ void runner_dokick(struct runner *r, struct cell *c, int timer) {
 
   int updated = 0;
   float t_end_min = FLT_MAX, t_end_max = 0.f;
-  double ekin = 0.0, epot = 0.0;
-  float mom[3] = {0.0f, 0.0f, 0.0f}, ang[3] = {0.0f, 0.0f, 0.0f};
-  float m, x[3], v_full[3];
+  double e_kin = 0.0, e_int = 0.0, e_pot = 0.0, mass = 0.0;
+  float mom[3] = {0.0f, 0.0f, 0.0f};
+  float ang[3] = {0.0f, 0.0f, 0.0f};
+  float x[3], v_full[3];
   struct part *restrict p, *restrict parts = c->parts;
   struct xpart *restrict xp, *restrict xparts = c->xparts;
 
@@ -825,7 +826,7 @@ void runner_dokick(struct runner *r, struct cell *c, int timer) {
       p = &parts[k];
       xp = &xparts[k];
 
-      m = p->mass;
+      const float m = p->mass;
       x[0] = p->x[0];
       x[1] = p->x[1];
       x[2] = p->x[2];
@@ -908,6 +909,9 @@ void runner_dokick(struct runner *r, struct cell *c, int timer) {
       v_full[1] = xp->v_full[1];
       v_full[2] = xp->v_full[2];
 
+      /* Collect mass */
+      mass += m;
+
       /* Collect momentum */
       mom[0] += m * v_full[0];
       mom[1] += m * v_full[1];
@@ -919,9 +923,10 @@ void runner_dokick(struct runner *r, struct cell *c, int timer) {
       ang[2] += m * (x[0] * v_full[1] - x[1] * v_full[0]);
 
       /* Collect total energy. */
-      ekin += 0.5 * m * (v_full[0] * v_full[0] + v_full[1] * v_full[1] +
-                         v_full[2] * v_full[2]);
-      epot += 0.f;  // MATTHIEU
+      e_kin += 0.5 * m * (v_full[0] * v_full[0] + v_full[1] * v_full[1] +
+                          v_full[2] * v_full[2]);
+      e_pot += 0.f; /* No gravitational potential thus far */
+      e_int += hydro_get_internal_energy(p);
 
       /* Minimal time for next end of time-step */
       t_end_min = fminf(p->t_end, t_end_min);
@@ -940,11 +945,16 @@ void runner_dokick(struct runner *r, struct cell *c, int timer) {
     for (int k = 0; k < 8; k++)
       if (c->progeny[k] != NULL) {
         struct cell *cp = c->progeny[k];
+
+        /* Recurse */
         runner_dokick(r, cp, 0);
 
+        /* And aggregate */
         updated += cp->updated;
-        ekin += cp->ekin;
-        epot += cp->epot;
+        e_kin += cp->e_kin;
+        e_int += cp->e_int;
+        e_pot += cp->e_pot;
+        mass += cp->mass;
         mom[0] += cp->mom[0];
         mom[1] += cp->mom[1];
         mom[2] += cp->mom[2];
@@ -958,8 +968,10 @@ void runner_dokick(struct runner *r, struct cell *c, int timer) {
 
   /* Store the values. */
   c->updated = updated;
-  c->ekin = ekin;
-  c->epot = epot;
+  c->e_kin = e_kin;
+  c->e_int = e_int;
+  c->e_pot = e_pot;
+  c->mass = mass;
   c->mom[0] = mom[0];
   c->mom[1] = mom[1];
   c->mom[2] = mom[2];