Allow for non standard cosmologies, and improve cosmological time integration

.gitignore

@@ -11,6 +11,7 @@ bin/*
 *.dll
 *.exe
 *.o
+*.swp
 *.so
 *.mod
 *.a

builds/make.host.frontier

@@ -5,9 +5,9 @@ CXX               = CC
 GPUCXX           ?= hipcc
 
 CXXFLAGS_DEBUG    = -g -O0 -std=c++17
-CXXFLAGS_OPTIMIZE = -g -Ofast -std=c++17 -Wno-unused-result
+CXXFLAGS_OPTIMIZE = -g -Ofast -std=c++17 -Wno-unused-result --gcc-toolchain=${GCC_PATH}/snos
 
-GPUFLAGS_OPTIMIZE = -std=c++17 --offload-arch=gfx90a -Wall -Wno-unused-result
+GPUFLAGS_OPTIMIZE = -std=c++17 --offload-arch=gfx90a -Wall -Wno-unused-result --gcc-toolchain=${GCC_PATH}/snos
 GPUFLAGS_DEBUG    = -g -O0 -std=c++17 --offload-arch=gfx90a -Wall -Wno-unused-result
 HIPCONFIG	  = -I$(ROCM_PATH)/include $(shell hipconfig -C) # workaround for Rocm 5.2 warnings
 #HIPCONFIG	  = $(shell hipconfig -C)

builds/make.host.lux

@@ -11,8 +11,8 @@ OMP_NUM_THREADS = 10
 #-- Library
 CUDA_ROOT    = /cm/shared/apps/cuda11.2/toolkit/current
 HDF5_ROOT    = /cm/shared/apps/hdf5/1.10.6
-FFTW_ROOT    = /home/brvillas/code/fftw-3.3.8
-PFFT_ROOT    = /data/groups/comp-astro/bruno/code_mpi_local/pfft
+FFTW_ROOT    = /data/groups/comp-astro/cholla/code/fftw-3.3.10
+PFFT_ROOT    = /data/groups/comp-astro/cholla/code/pfft
 GRACKLE_ROOT = /home/brvillas/code/grackle
 
 #Paris does not do GPU_MPI transfers

builds/setup.frontier.cce.sh

@@ -9,6 +9,7 @@ module load rocm/5.7.1
 module load craype-accel-amd-gfx90a
 module load cray-hdf5 cray-fftw
 module load googletest/1.10.0
+module load gcc-mixed
 
 #-- GPU-aware MPI
 export MPICH_GPU_SUPPORT_ENABLED=1

src/analysis/io_analysis.cpp

@@ -9,6 +9,9 @@
 
 // #define OUTPUT_SKEWERS_TRANSMITTED_FLUX
 
+using std::ifstream;
+using std::string;
+
   #ifdef OUTPUT_SKEWERS
 void Grid3D::Output_Skewers_File(struct Parameters *P)
 {
@@ -73,6 +76,18 @@ void Grid3D::Write_Skewers_Header_HDF5(hid_t file_id)
   status       = H5Aclose(attribute_id);
   attribute_id = H5Acreate(file_id, "Omega_b", H5T_IEEE_F64BE, dataspace_id, H5P_DEFAULT, H5P_DEFAULT);
   status       = H5Awrite(attribute_id, H5T_NATIVE_DOUBLE, &Cosmo.Omega_b);
+  status       = H5Aclose(attribute_id);
+  attribute_id = H5Acreate(file_id, "Omega_K", H5T_IEEE_F64BE, dataspace_id, H5P_DEFAULT, H5P_DEFAULT);
+  status       = H5Awrite(attribute_id, H5T_NATIVE_DOUBLE, &Cosmo.Omega_K);
+  status       = H5Aclose(attribute_id);
+  attribute_id = H5Acreate(file_id, "Omega_R", H5T_IEEE_F64BE, dataspace_id, H5P_DEFAULT, H5P_DEFAULT);
+  status       = H5Awrite(attribute_id, H5T_NATIVE_DOUBLE, &Cosmo.Omega_R);
+  status       = H5Aclose(attribute_id);
+  attribute_id = H5Acreate(file_id, "w0", H5T_IEEE_F64BE, dataspace_id, H5P_DEFAULT, H5P_DEFAULT);
+  status       = H5Awrite(attribute_id, H5T_NATIVE_DOUBLE, &Cosmo.w0);
+  status       = H5Aclose(attribute_id);
+  attribute_id = H5Acreate(file_id, "wa", H5T_IEEE_F64BE, dataspace_id, H5P_DEFAULT, H5P_DEFAULT);
+  status       = H5Awrite(attribute_id, H5T_NATIVE_DOUBLE, &Cosmo.wa);
   status       = H5Aclose(attribute_id);
     #endif
 
@@ -525,6 +540,18 @@ void Grid3D::Write_Analysis_Header_HDF5(hid_t file_id)
   attribute_id = H5Acreate(file_id, "Omega_b", H5T_IEEE_F64BE, dataspace_id, H5P_DEFAULT, H5P_DEFAULT);
   status       = H5Awrite(attribute_id, H5T_NATIVE_DOUBLE, &Cosmo.Omega_b);
   status       = H5Aclose(attribute_id);
+  attribute_id = H5Acreate(file_id, "Omega_K", H5T_IEEE_F64BE, dataspace_id, H5P_DEFAULT, H5P_DEFAULT);
+  status       = H5Awrite(attribute_id, H5T_NATIVE_DOUBLE, &Cosmo.Omega_K);
+  status       = H5Aclose(attribute_id);
+  attribute_id = H5Acreate(file_id, "Omega_R", H5T_IEEE_F64BE, dataspace_id, H5P_DEFAULT, H5P_DEFAULT);
+  status       = H5Awrite(attribute_id, H5T_NATIVE_DOUBLE, &Cosmo.Omega_R);
+  status       = H5Aclose(attribute_id);
+  attribute_id = H5Acreate(file_id, "w0", H5T_IEEE_F64BE, dataspace_id, H5P_DEFAULT, H5P_DEFAULT);
+  status       = H5Awrite(attribute_id, H5T_NATIVE_DOUBLE, &Cosmo.w0);
+  status       = H5Aclose(attribute_id);
+  attribute_id = H5Acreate(file_id, "wa", H5T_IEEE_F64BE, dataspace_id, H5P_DEFAULT, H5P_DEFAULT);
+  status       = H5Awrite(attribute_id, H5T_NATIVE_DOUBLE, &Cosmo.wa);
+  status       = H5Aclose(attribute_id);
   #endif
 
   status = H5Sclose(dataspace_id);

src/analysis/lya_statistics.cpp

@@ -13,6 +13,7 @@
 
 // #define PRINT_ANALYSIS_LOG
 
+using std::max;
 void AnalysisModule::Transfer_Skewers_Global_Axis(int axis)
 {
   bool am_I_root;

src/cosmology/cosmology.cpp

@@ -15,8 +15,11 @@ void Cosmology::Initialize(struct Parameters *P, Grav3D &Grav, Particles3D &Part
   H0 /= 1000;  //[km/s / kpc]
   Omega_M = P->Omega_M;
   Omega_L = P->Omega_L;
-  Omega_K = 1 - (Omega_M + Omega_L);
+  Omega_R = P->Omega_R;
+  Omega_K = 1 - (Omega_M + Omega_L + Omega_R);
   Omega_b = P->Omega_b;
+  w0      = P->w0;
+  wa      = P->wa;
 
   if (strcmp(P->init, "Read_Grid") == 0) {
     // Read scale factor value from Particles
@@ -42,8 +45,34 @@ void Cosmology::Initialize(struct Parameters *P, Grav3D &Grav, Particles3D &Part
   delta_a     = max_delta_a;
 
   // Initialize Time and set the time conversion
-  t_secs          = 0;
   time_conversion = KPC;
+  t_secs          = 0;
+
+  // The following code computes the universal time
+  // at the scale factor of the ICs or restart file
+  // Pick a small scale factor step for integrating the universal time
+  Real da_t_sec = 1.0e-2 * current_a;
+  // Pick a small but non-zero starting scale factor for the integral
+  Real a_t_sec = 1.0e-6;
+  // Step for the time integral, corresponding to da_t_sec
+  Real dt_physical;
+
+  // Advance a_t_sec until it matches current_a, and integrate time
+  while (a_t_sec < current_a) {
+    // Limit the scale a_t_sec factor to current_a
+    if (a_t_sec + da_t_sec > current_a) {
+      da_t_sec = current_a - a_t_sec;
+    }
+
+    // Compute the time step
+    dt_physical = Get_dt_from_da(da_t_sec, a_t_sec);
+
+    // Advance the time in seconds and the scale factor
+    t_secs += dt_physical * time_conversion;
+    a_t_sec += da_t_sec;
+    // chprintf(" Revised a_t_sec %f da_t_sec %f t_start : %e Myr\n", a_t_sec, da_t_sec, t_secs / MYR);
+  }
+  chprintf(" Revised t_start : %f Myr\n", t_secs / MYR);
 
   // Set Normalization factors
   r_0_dm          = P->xlen / P->nx;

src/cosmology/cosmology.h

@@ -19,6 +19,9 @@ class Cosmology
   Real Omega_L;
   Real Omega_K;
   Real Omega_b;
+  Real Omega_R;
+  Real w0;
+  Real wa;
 
   Real cosmo_G;
   Real cosmo_h;
@@ -64,8 +67,14 @@ class Cosmology
 
   Real Get_Hubble_Parameter(Real a);
 
+  Real dtda_cosmo(Real da, Real a);
+  Real Get_dt_from_da_rk(Real da, Real a);
   Real Get_da_from_dt(Real dt);
-  Real Get_dt_from_da(Real da);
+  Real Get_dt_from_da(Real da, Real a);
+
+  // write expansion history log file
+  void Create_Expansion_History_File(struct Parameters *P);
+  void Write_Expansion_History_Entry(void);
 };
 
   #endif

src/cosmology/cosmology_functions.cpp

@@ -1,4 +1,5 @@
 #ifdef COSMOLOGY
+  #include <fstream>
 
   #include "../global/global.h"
   #include "../grid/grid3D.h"
@@ -10,8 +11,11 @@ void Grid3D::Initialize_Cosmology(struct Parameters *P)
   chprintf("Initializing Cosmology... \n");
   Cosmo.Initialize(P, Grav, Particles);
 
+  // Create expansion history log file
+  Cosmo.Create_Expansion_History_File(P);
+
   // Change to comoving Cosmological System
-  Change_Cosmological_Frame_Sytem(true);
+  Change_Cosmological_Frame_System(true);
 
   if (fabs(Cosmo.current_a - Cosmo.next_output) < 1e-5) {
     H.Output_Now = true;
@@ -20,29 +24,74 @@ void Grid3D::Initialize_Cosmology(struct Parameters *P)
   chprintf("Cosmology Successfully Initialized. \n\n");
 }
 
+/* Computes dt/da * da */
+Real Cosmology::dtda_cosmo(Real da, Real a)
+{
+  Real a2     = a * a;
+  Real fac_de = pow(a, -3 * (1 + w0 + wa)) * exp(-3 * wa * (1 - current_a));
+  Real a_dot  = sqrt(Omega_R / a2 + Omega_M / a + a2 * Omega_L * fac_de + Omega_K) * H0;
+  return da / a_dot;
+}
+
+/* Compute dt/da * da. dt/da is computed with a Runge-Kutta integration step */
+Real Cosmology::Get_dt_from_da_rk(Real da, Real a)
+{
+  Real a3 = 0.3;
+  Real a4 = 0.6;
+  Real a5 = 1.0;
+  Real a6 = 0.875;
+  Real c1 = 37.0 / 378.0;
+  Real c3 = 250.0 / 621.0;
+  Real c4 = 125.0 / 594.0;
+  Real c6 = 512.0 / 1771.0;
+
+  // compute RK average derivatives
+  Real ak1 = dtda_cosmo(da, a);
+  Real ak3 = dtda_cosmo(da, a + a3 * da);
+  Real ak4 = dtda_cosmo(da, a + a4 * da);
+  Real ak6 = dtda_cosmo(da, a + a6 * da);
+
+  // compute timestep
+  Real dt = (c1 * ak1 + c3 * ak3 + c4 * ak4 + c6 * ak6);
+
+  // return timestep
+  return dt;
+}
+
 Real Cosmology::Get_da_from_dt(Real dt)
 {
-  Real a2    = current_a * current_a;
-  Real a_dot = sqrt(Omega_M / current_a + a2 * Omega_L + Omega_K) * H0;
+  Real a2     = current_a * current_a;
+  Real fac_de = pow(current_a, -3 * (1 + w0 + wa)) * exp(-3 * wa * (1 - current_a));
+  Real a_dot  = sqrt(Omega_R / a2 + Omega_M / current_a + a2 * Omega_L * fac_de + Omega_K) * H0;
   return a_dot * dt;
 }
 
-Real Cosmology::Get_dt_from_da(Real da)
+Real Cosmology::Get_dt_from_da(Real da, Real a)
 {
-  Real a2    = current_a * current_a;
-  Real a_dot = sqrt(Omega_M / current_a + a2 * Omega_L + Omega_K) * H0;
-  return da / a_dot;
+  return Get_dt_from_da_rk(da, a);
+
+  /* The following commented code was the original Euler
+     integrator for computing time from the scale factor.
+     This has been left here temporarily to ease comparison
+     with the Runge-Kutta integrator, but it can be removed
+     eventually. */
+  /* Real a2     = a * a;
+  Real fac_de = pow(a, -3 * (1 + w0 + wa)) * exp(-3 * wa * (1 - a));
+  Real a_dot  = sqrt(Omega_R / a2 + Omega_M / a + a2 * Omega_L * fac_de + Omega_K) * H0;
+  return da / a_dot; */
 }
 
 Real Cosmology::Get_Hubble_Parameter(Real a)
 {
   Real a2     = a * a;
   Real a3     = a2 * a;
-  Real factor = (Omega_M / a3 + Omega_K / a2 + Omega_L);
+  Real a4     = a2 * a2;
+  Real fac_de = pow(a, -3 * (1 + w0 + wa)) * exp(-3 * wa * (1 - a));
+  Real factor = (Omega_R / a4 + Omega_M / a3 + Omega_K / a2 + Omega_L * fac_de);
   return H0 * sqrt(factor);
 }
 
-void Grid3D::Change_Cosmological_Frame_Sytem(bool forward)
+void Grid3D::Change_Cosmological_Frame_System(bool forward)
 {
   if (forward) {
     chprintf(" Converting to Cosmological Comoving System\n");
@@ -130,4 +179,54 @@ void Grid3D::Change_GAS_Frame_System(bool forward)
   }
 }
 
+/* create the file for recording the expansion history */
+void Cosmology::Create_Expansion_History_File(struct Parameters *P)
+{
+  if (not Is_Root_Proc()) {
+    return;
+  }
+
+  std::string file_name(EXPANSION_HISTORY_FILE_NAME);
+  chprintf("\nCreating Expansion History File: %s \n\n", file_name.c_str());
+
+  bool file_exists = false;
+  if (FILE *file = fopen(file_name.c_str(), "r")) {
+    file_exists = true;
+    chprintf("  File exists, appending values: %s \n\n", file_name.c_str());
+    fclose(file);
+  }
+
+  // current date/time based on current system
+  time_t now = time(0);
+  // convert now to string form
+  char *dt = ctime(&now);
+
+  std::string message = "# H0 OmegaM Omega_b OmegaL w0 wa Omega_R Omega_K\n";
+  message += "# " + std::to_string(H0 * 1e3) + " " + std::to_string(Omega_M);
+  message += " " + std::to_string(Omega_b);
+  message += " " + std::to_string(Omega_L) + " " + std::to_string(w0) + " " + std::to_string(wa);
+  message += " " + std::to_string(Omega_R) + " " + std::to_string(Omega_K);
+
+  std::ofstream out_file;
+  out_file.open(file_name.c_str(), std::ios::app);
+  out_file << "# Run date: " << dt;
+  out_file << message.c_str() << std::endl;
+  out_file.close();
+}
+
+/* Write the current entry to the expansion history file */
+void Cosmology::Write_Expansion_History_Entry(void)
+{
+  if (not Is_Root_Proc()) {
+    return;
+  }
+
+  std::string message = std::to_string(t_secs / MYR) + " " + std::to_string(current_a);
+  std::string file_name(EXPANSION_HISTORY_FILE_NAME);
+  std::ofstream out_file;
+  out_file.open(file_name.c_str(), std::ios::app);
+  out_file << message.c_str() << std::endl;
+  out_file.close();
+}
+
 #endif

src/global/global.cpp

@@ -102,10 +102,11 @@ char *Trim(char *s)
 }
 
 // NOLINTNEXTLINE(cert-err58-cpp)
+// NOLINTNEXTLINE(*)
 const std::set<std::string> optionalParams = {
-    "flag_delta",   "ddelta_dt",   "n_delta",  "Lz",       "Lx",      "phi",     "theta",
-    "delta",        "nzr",         "nxr",      "H0",       "Omega_M", "Omega_L", "Init_redshift",
-    "End_redshift", "tile_length", "n_proc_x", "n_proc_y", "n_proc_z"};
+    "flag_delta", "ddelta_dt",     "n_delta",      "Lz",          "Lx",       "phi",      "theta",   "delta",
+    "nzr",        "nxr",           "H0",           "Omega_M",     "Omega_L",  "Omega_R",  "Omega_K", "w0",
+    "wa",         "Init_redshift", "End_redshift", "tile_length", "n_proc_x", "n_proc_y", "n_proc_z"};  // NOLINT
 
 bool Old_Style_Parse_Param(const char *name, const char *value, struct Parameters *parms);
 
@@ -368,6 +369,12 @@ bool Old_Style_Parse_Param(const char *name, const char *value, struct Parameter
     parms->Omega_L = atof(value);
   } else if (strcmp(name, "Omega_b") == 0) {
     parms->Omega_b = atof(value);
+  } else if (strcmp(name, "Omega_R") == 0) {
+    parms->Omega_R = atof(value);
+  } else if (strcmp(name, "w0") == 0) {
+    parms->w0 = atof(value);
+  } else if (strcmp(name, "wa") == 0) {
+    parms->wa = atof(value);
 #endif  // COSMOLOGY
 #ifdef TILED_INITIAL_CONDITIONS
   } else if (strcmp(name, "tile_length") == 0) {

src/global/global.h

@@ -321,11 +321,16 @@ struct Parameters {
   Real Omega_M;
   Real Omega_L;
   Real Omega_b;
+  Real Omega_R;
+  Real w0;
+  Real wa;
   Real Init_redshift;
   Real End_redshift;
-  char scale_outputs_file[MAXLEN];  // File for the scale_factor output values
-                                    // for cosmological simulations
-#endif                              // COSMOLOGY
+
+  // File for the scale_factor output values for cosmological simulations
+  char scale_outputs_file[MAXLEN];
+  #define EXPANSION_HISTORY_FILE_NAME "expansion_history.txt"
+#endif  // COSMOLOGY
 #ifdef TILED_INITIAL_CONDITIONS
   Real tile_length;
 #endif  // TILED_INITIAL_CONDITIONS