Improve the check for cooling tables in the EAGLE model to prevent the...

Improve the check for cooling tables in the EAGLE model to prevent the constant re-loading of tables every step at z=0.

examples/SmallCosmoVolume_cooling/small_cosmo_volume.yml

@@ -64,7 +64,7 @@ LambdaCooling:
 
 # EAGLE cooling function
 EAGLECooling:
-  dirname:                 ./coolingtables/
+  dir_name:                 ./coolingtables/
   H_reion_z:               11.5
   He_reion_z_centre:       3.5
   He_reion_z_sigma:        0.5

src/cooling/EAGLE/cooling.c

@@ -79,32 +79,38 @@ __attribute__((always_inline)) INLINE void get_redshift_index(
     float z, int *z_index, float *dz,
     struct cooling_function_data *restrict cooling) {
 
-  /* before the earliest redshift or before hydrogen reionization, flag for
+  /* Before the earliest redshift or before hydrogen reionization, flag for
    * collisional cooling */
   if (z > cooling->H_reion_z) {
     *z_index = eagle_cooling_N_redshifts;
     *dz = 0.0;
   }
-  /* from reionization use the cooling tables */
+
+  /* From reionization use the cooling tables */
   else if (z > cooling->Redshifts[eagle_cooling_N_redshifts - 1] &&
            z <= cooling->H_reion_z) {
     *z_index = eagle_cooling_N_redshifts + 1;
     *dz = 0.0;
   }
-  /* at the end, just use the last value */
+
+  /* At the end, just use the last value */
   else if (z <= cooling->Redshifts[0]) {
     *z_index = 0;
     *dz = 0.0;
-  } else {
+  }
+
+  /* Normal case: search... */
+  else {
 
     /* start at the previous index and search */
-    for (int iz = cooling->previous_z_index; iz >= 0; iz--) {
-      if (z > cooling->Redshifts[iz]) {
+    for (int i = cooling->previous_z_index; i >= 0; i--) {
+      if (z > cooling->Redshifts[i]) {
+
+        *z_index = i;
+        cooling->previous_z_index = i;
 
-        *z_index = iz;
-        cooling->previous_z_index = iz;
-        *dz = (z - cooling->Redshifts[iz]) /
-              (cooling->Redshifts[iz + 1] - cooling->Redshifts[iz]);
+        *dz = (z - cooling->Redshifts[i]) /
+              (cooling->Redshifts[i + 1] - cooling->Redshifts[i]);
         break;
       }
     }
@@ -127,20 +133,40 @@ void cooling_update(const struct cosmology *cosmo,
   /* Current redshift */
   const float redshift = cosmo->z;
 
-  /* Get index along the redshift index of the tables */
+  /* What is the current table index along the redshift axis? */
   int z_index = -1;
   float dz = 0.f;
-  if (redshift > cooling->H_reion_z) {
-    z_index = -2;
-  } else if (redshift > cooling->Redshifts[eagle_cooling_N_redshifts - 1]) {
-    z_index = -1;
+  get_redshift_index(redshift, &z_index, &dz, cooling);
+  cooling->dz = dz;
+
+  /* Do we already have the correct tables loaded? */
+  if (cooling->z_index == z_index) return;
+
+  /* Which table should we load ? */
+  if (z_index >= eagle_cooling_N_redshifts) {
+
+    if (z_index == eagle_cooling_N_redshifts + 1) {
+
+      /* Bewtween re-ionization and first table */
+      get_redshift_invariant_table(cooling, /* photodis=*/0);
+
+    } else {
+
+      /* Above re-ionization */
+      get_redshift_invariant_table(cooling, /* photodis=*/1);
+    }
+
   } else {
-    get_redshift_index(redshift, &z_index, &dz, cooling);
+
+    /* Normal case: two tables bracketing the current z */
+    const int low_z_index = z_index;
+    const int high_z_index = z_index + 1;
+
+    get_cooling_table(cooling, low_z_index, high_z_index);
   }
-  cooling->z_index = z_index;
-  cooling->dz = dz;
 
-  eagle_check_cooling_tables(cooling, restart_flag);
+  /* Store the currently loaded index */
+  cooling->z_index = z_index;
 }
 
 /**
@@ -755,13 +781,15 @@ void cooling_init_backend(struct swift_params *parameter_file,
   cooling->S_over_Si_ratio_in_solar = parser_get_opt_param_float(
       parameter_file, "EAGLECooling:S_over_Si_in_solar", 1.f);
 
-  /* convert to cgs */
+  /* Convert to cgs (units used internally by the cooling routines) */
   cooling->He_reion_heat_cgs *=
       phys_const->const_electron_volt *
       units_cgs_conversion_factor(us, UNIT_CONV_ENERGY);
 
-  /* read in cooling table header */
+  /* Read in the list of redshifts */
   get_cooling_redshifts(cooling);
+
+  /* Read in cooling table header */
   char fname[eagle_table_path_name_length + 12];
   sprintf(fname, "%sz_0.000.hdf5", cooling->cooling_table_path);
   read_cooling_header(fname, cooling);
@@ -769,7 +797,7 @@ void cooling_init_backend(struct swift_params *parameter_file,
   /* Allocate space for cooling tables */
   allocate_cooling_tables(cooling);
 
-  /* compute conversion factors */
+  /* Compute conversion factors */
   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;
@@ -811,8 +839,10 @@ void cooling_init_backend(struct swift_params *parameter_file,
                               cooling->T_CMB_0 * cooling->T_CMB_0 *
                               cooling->T_CMB_0;
 
-  /* set low_z_index to -10 to indicate we haven't read any tables yet */
-  cooling->low_z_index = -10;
+  /* Set the redshift indices to invalid values */
+  cooling->z_index = -10;
+  cooling->previous_z_index = eagle_cooling_N_redshifts + 2;
+
   /* set previous_z_index and to last value of redshift table*/
   cooling->previous_z_index = eagle_cooling_N_redshifts - 2;
 
@@ -839,10 +869,13 @@ void cooling_restore_tables(struct cooling_function_data *cooling,
   sprintf(fname, "%sz_0.000.hdf5", cooling->cooling_table_path);
   read_cooling_header(fname, cooling);
 
-  /* Read relevant cooling tables.
-   * Third variable in cooling_update flag to mark restart*/
+  /* Allocate memory for the tables */
   allocate_cooling_tables(cooling);
-  cooling_update(cosmo, cooling, /*restart=*/1);
+
+  /* Force a re-read of the cooling tables */
+  cooling->z_index = -10;
+  cooling->previous_z_index = eagle_cooling_N_redshifts + 2;
+  cooling_update(cosmo, cooling, /*restart_flag=*/1);
 }
 
 /**

src/cooling/EAGLE/cooling_struct.h

@@ -115,17 +115,11 @@ struct cooling_function_data {
   /*! Index of the current redshift along the redshift index of the tables */
   int z_index;
 
-  /*! Index of the previous tables along the redshift index of the tables */
-  int previous_z_index;
-
-  /*! Distance between the current redshift and the table[z_index] */
+  /*! Distance between the current redshift and table[z_index] */
   float dz;
 
-  /*! Index of the table below current redshift */
-  int low_z_index;
-
-  /*! Index of the table above current redshift */
-  int high_z_index;
+  /*! Index of the previous tables along the redshift index of the tables */
+  int previous_z_index;
 
   /*! Are we doing Newton-Raphson iterations? */
   int newton_flag;

src/cooling/EAGLE/cooling_tables.c

@@ -359,9 +359,10 @@ void allocate_cooling_tables(struct cooling_function_data *restrict cooling) {
  * used to obtain temperature of particle)
  *
  * @param cooling #cooling_function_data structure
+ * @param photodis Are we loading the photo-dissociation table?
  */
-static void get_redshift_invariant_table(
-    struct cooling_function_data *restrict cooling) {
+void get_redshift_invariant_table(
+    struct cooling_function_data *restrict cooling, const int photodis) {
 #ifdef HAVE_HDF5
 
   /* Temporary tables */
@@ -390,10 +391,12 @@ static void get_redshift_invariant_table(
 
   /* Decide which high redshift table to read. Indices set in cooling_update */
   char filename[eagle_table_path_name_length + 21];
-  if (cooling->low_z_index == -1) {
-    sprintf(filename, "%sz_8.989nocompton.hdf5", cooling->cooling_table_path);
-  } else if (cooling->low_z_index == -2) {
+  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'");
   }
 
   hid_t file_id = H5Fopen(filename, H5F_ACC_RDONLY, H5P_DEFAULT);
@@ -554,8 +557,11 @@ static void get_redshift_invariant_table(
  * 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.
  */
-static void get_cooling_table(struct cooling_function_data *restrict cooling) {
+void get_cooling_table(struct cooling_function_data *restrict cooling,
+                       const int low_z_index, const int high_z_index) {
 
 #ifdef HAVE_HDF5
 
@@ -585,11 +591,10 @@ static void get_cooling_table(struct cooling_function_data *restrict cooling) {
 
   /* 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 = cooling->low_z_index; z_index <= cooling->high_z_index;
-       z_index++) {
+  for (int z_index = low_z_index; z_index <= high_z_index; z_index++) {
 
     /* Index along redhsift dimension for the subset of tables we read */
-    const int local_z_index = z_index - cooling->low_z_index;
+    const int local_z_index = z_index - low_z_index;
 
 #ifdef SWIFT_DEBUG_CHECKS
     if (local_z_index >= eagle_cooling_N_loaded_redshifts)
@@ -597,9 +602,12 @@ static void get_cooling_table(struct cooling_function_data *restrict cooling) {
 #endif
 
     /* Open table for this redshift index */
-    char fname[eagle_table_path_name_length + 32];
+    char fname[eagle_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);
 
@@ -740,48 +748,10 @@ static void get_cooling_table(struct cooling_function_data *restrict cooling) {
   free(he_electron_abundance);
 
 #ifdef SWIFT_DEBUG_CHECKS
-  message("done reading in general cooling table");
+  message("Done reading in general cooling table");
 #endif
 
 #else
   error("Need HDF5 to read cooling tables");
 #endif
 }
-
-/**
- * @brief Constructs the data structure containting the relevant cooling tables
- * for the redshift index (set in cooling_update)
- *
- * @param cooling #cooling_function_data structure
- */
-static void eagle_readtable(struct cooling_function_data *restrict cooling) {
-
-  if (cooling->z_index < 0) {
-    /* z_index is set to < 0 in cooling_update if need
-     * to read any of the high redshift tables */
-    get_redshift_invariant_table(cooling);
-  } else {
-    get_cooling_table(cooling);
-  }
-}
-
-/**
- * @brief Checks the tables that are currently loaded in memory and read
- * new ones if necessary.
- *
- * @param cooling The #cooling_function_data we play with.
- * @param restart_flag Flag indicating if we are restarting a run
- */
-void eagle_check_cooling_tables(struct cooling_function_data *restrict cooling,
-                                const int restart_flag) {
-
-  /* Do we already have the right table in memory? */
-  if (cooling->low_z_index == cooling->z_index && restart_flag != 0) return;
-
-  /* Record the table indices */
-  cooling->low_z_index = cooling->z_index;
-  cooling->high_z_index = cooling->z_index + 1;
-
-  /* Load the damn thing */
-  eagle_readtable(cooling);
-}

src/cooling/EAGLE/cooling_tables.h

@@ -57,6 +57,9 @@ void read_cooling_header(const char *fname,
 
 void allocate_cooling_tables(struct cooling_function_data *restrict cooling);
 
-void eagle_check_cooling_tables(struct cooling_function_data *restrict cooling,
-                                const int restart_flag);
+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);
+
 #endif