From 46b290f6e826b359a6c33b10383c87f4b6712790 Mon Sep 17 00:00:00 2001
From: Cristian C Lalescu <Cristian.Lalescu@ds.mpg.de>
Date: Mon, 11 Nov 2019 17:08:02 +0100
Subject: [PATCH] fix field io mechanism for Kraichnan model

---
 TurTLE/DNS.py                     | 39 +++++++++-----
 TurTLE/_code.py                   |  3 +-
 cpp/field.cpp                     |  1 +
 cpp/full_code/kraichnan_field.cpp | 89 ++++++++++++++++++++++++++++---
 cpp/full_code/kraichnan_field.hpp | 16 ++----
 5 files changed, 114 insertions(+), 34 deletions(-)

diff --git a/TurTLE/DNS.py b/TurTLE/DNS.py
index deabe148..8bfdebe9 100644
--- a/TurTLE/DNS.py
+++ b/TurTLE/DNS.py
@@ -680,6 +680,7 @@ class DNS(_code):
             dns_type):
         pars = {}
         if dns_type == 'kraichnan_field':
+            pars['output_velocity'] = int(1)
             pars['field_random_seed'] = int(1)
             pars['spectrum_slope'] = float(-5./3)
             pars['spectrum_k_cutoff'] = float(16)
@@ -1026,19 +1027,25 @@ class DNS(_code):
         # take care of fields' initial condition
         # first, check if initial field exists
         need_field = False
-        if not os.path.exists(self.get_checkpoint_0_fname()):
-            need_field = True
-        else:
-            f = h5py.File(self.get_checkpoint_0_fname(), 'r')
-            try:
-                dset = f['vorticity/complex/0']
-                need_field = (dset.shape == (self.parameters['ny'],
-                                             self.parameters['nz'],
-                                             self.parameters['nx']//2+1,
-                                             3))
-            except:
+        if self.dns_type in [
+                'static_field',
+                'NSVEparticles',
+                'NSVEcomplex_particles',
+                'NSVEparticles_no_output',
+                'NSVEp_extra_sampling']:
+            if not os.path.exists(self.get_checkpoint_0_fname()):
                 need_field = True
-            f.close()
+            else:
+                f = h5py.File(self.get_checkpoint_0_fname(), 'r')
+                try:
+                    dset = f['vorticity/complex/0']
+                    need_field = (dset.shape == (self.parameters['ny'],
+                                                 self.parameters['nz'],
+                                                 self.parameters['nx']//2+1,
+                                                 3))
+                except:
+                    need_field = True
+                f.close()
         if need_field:
             f = h5py.File(self.get_checkpoint_0_fname(), 'a')
             if len(opt.src_simname) > 0:
@@ -1065,6 +1072,11 @@ class DNS(_code):
                        amplitude = 0.05)
                 f['vorticity/complex/{0}'.format(0)] = data
             f.close()
+        if self.dns_type == 'kraichnan_field':
+            if not os.path.exists(self.get_checkpoint_0_fname()):
+                f = h5py.File(self.get_checkpoint_0_fname(), 'a')
+                f.create_group('velocity/real')
+                f.close()
         # now take care of particles' initial condition
         if self.dns_type in [
                 'kraichnan_field',
@@ -1081,6 +1093,9 @@ class DNS(_code):
         if not os.path.exists(self.get_data_file_name()):
             self.generate_initial_condition(opt = opt)
             self.write_par()
+        if (('test' in self.dns_type) or
+            (self.dns_type in ['kraichnan_field'])):
+            self.check_current_vorticity_exists = False
         self.run(
                 nb_processes = opt.nb_processes,
                 nb_threads_per_process = opt.nb_threads_per_process,
diff --git a/TurTLE/_code.py b/TurTLE/_code.py
index 05b36518..be898972 100644
--- a/TurTLE/_code.py
+++ b/TurTLE/_code.py
@@ -176,6 +176,7 @@ class _code(_base):
                 """
         self.host_info = host_info
         self.main = ''
+        self.check_current_vorticity_exists = True
         return None
     def write_src(self):
         with open(self.name + '.cpp', 'w') as outfile:
@@ -278,7 +279,7 @@ class _code(_base):
         if os.path.exists(os.path.join(self.work_dir, 'stop_' + self.simname)):
             warnings.warn("Found stop_<simname> file, I will remove it.")
             os.remove(os.path.join(self.work_dir, 'stop_' + self.simname))
-        if not 'test' in self.name:
+        if self.check_current_vorticity_exists:
             with h5py.File(os.path.join(self.work_dir, self.simname + '.h5'), 'r') as data_file:
                 iter0 = data_file['iteration'][...]
                 checkpoint0 = data_file['checkpoint'][()]
diff --git a/cpp/field.cpp b/cpp/field.cpp
index fdd8c2b4..3f9b9938 100644
--- a/cpp/field.cpp
+++ b/cpp/field.cpp
@@ -2052,6 +2052,7 @@ int make_gaussian_random_field(
     rgen.resize(omp_get_max_threads());
     // seed random number generators such that no seed is ever repeated if we change the value of rseed.
     // basically use a multi-dimensional array indexing technique to come up with actual seed.
+    // Note: this method IS NOT MPI/OpenMP-invariant!
     for (int thread_id=0; thread_id < omp_get_max_threads(); thread_id++)
     {
         int current_seed = (
diff --git a/cpp/full_code/kraichnan_field.cpp b/cpp/full_code/kraichnan_field.cpp
index 459c9d5f..212ba9b8 100644
--- a/cpp/full_code/kraichnan_field.cpp
+++ b/cpp/full_code/kraichnan_field.cpp
@@ -65,7 +65,7 @@ int kraichnan_field<rnumber>::initialize(void)
             nx, ny, nz,
             this->comm,
 	        fftw_planner_string_to_flag[this->fftw_plan_rigor]);
-    this->velocity->real_space_representation = false;
+    this->velocity->real_space_representation = true;
 
     // the velocity layout should be the same as the vorticity one
     this->kk = new kspace<FFTW, SMOOTH>(
@@ -104,30 +104,60 @@ int kraichnan_field<rnumber>::initialize(void)
     DEBUG_MSG("Coefficient: %g\n",
             this->spectrum_coefficient);
 
+    // is this the first iteration?
+    if ((this->iteration == 0) and (this->output_velocity == 1))
+    {
+        // if yes, generate random field and save it
+        this->generate_random_velocity();
+        this->velocity->io(
+                this->get_current_fname(),
+                "velocity",
+                this->iteration,
+                false);
+    }
+    else
+    {
+        // if not, read the random field that exists in the checkpoint file
+        this->velocity->io(
+                this->get_current_fname(),
+                "velocity",
+                this->iteration,
+                true);
+    }
+
     return EXIT_SUCCESS;
 }
 
 template <typename rnumber>
-int kraichnan_field<rnumber>::step(void)
+int kraichnan_field<rnumber>::generate_random_velocity(void)
 {
-    TIMEZONE("kraichnan_file::step");
-
+    TIMEZONE("kraichnan_file::make_random_field");
+    // generate Gaussian random field
     make_gaussian_random_field(
         this->kk,
         this->velocity,
-        this->iteration,
+        this->iteration, // not an ideal choice because resulting field sequence depends on MPI/OpenMP configuration
         this->spectrum_slope,
         this->spectrum_k_cutoff,
         this->spectrum_coefficient * 3./2.); // incompressibility projection factor
-    this->velocity->ift();
-
+    // this->velocity is now in Fourier space
+    // project divfree, requires field in Fourier space
     DEBUG_MSG("L2Norm before: %g\n",
             this->velocity->L2norm(this->kk));
     this->kk->template project_divfree<rnumber>(this->velocity->get_cdata());
     DEBUG_MSG("L2Norm after: %g\n",
             this->velocity->L2norm(this->kk));
+    // transform velocity to real space representation
+    this->velocity->ift();
+    return EXIT_SUCCESS;
+}
 
+template <typename rnumber>
+int kraichnan_field<rnumber>::step(void)
+{
+    TIMEZONE("kraichnan_file::step");
     this->ps->completeLoop(sqrt(this->dt));
+    this->generate_random_velocity();
     this->iteration++;
     return EXIT_SUCCESS;
 }
@@ -136,7 +166,19 @@ template <typename rnumber>
 int kraichnan_field<rnumber>::write_checkpoint(void)
 {
     TIMEZONE("kraichnan_file::write_checkpoint");
-    this->io_checkpoint(false);
+    this->update_checkpoint();
+    // output field
+    if (this->output_velocity == 1)
+    {
+        assert(this->velocity->real_space_representation);
+        std::string fname = this->get_current_fname();
+        this->velocity->io(
+                fname,
+                "velocity",
+                this->iteration,
+                false);
+    }
+    // output particles
     this->particles_output_writer_mpi->open_file(this->get_current_fname());
     this->particles_output_writer_mpi->template save<3>(
             this->ps->getParticlesState(),
@@ -172,6 +214,32 @@ template <typename rnumber>
 int kraichnan_field<rnumber>::do_stats()
 {
     TIMEZONE("kraichnan_field::do_stats");
+    /// compute field statistics
+    if (this->iteration % this->niter_stat == 0)
+    {
+        hid_t stat_group;
+        if (this->myrank == 0)
+            stat_group = H5Gopen(
+                    this->stat_file,
+                    "statistics",
+                    H5P_DEFAULT);
+        else
+            stat_group = 0;
+        field<rnumber, FFTW, THREE> *tvelocity = new field<rnumber, FFTW, THREE>(
+                this->nx, this->ny, this->nz,
+                this->comm,
+	            fftw_planner_string_to_flag[this->fftw_plan_rigor]);
+        *tvelocity = *this->velocity;
+        tvelocity->compute_stats(
+            this->kk,
+            stat_group,
+            "velocity",
+            this->iteration / this->niter_stat,
+            this->max_velocity_estimate/sqrt(3));
+        if (this->myrank == 0)
+            H5Gclose(stat_group);
+        delete tvelocity;
+    }
     /// either one of two conditions suffices to compute statistics:
     /// 1) current iteration is a multiple of niter_part
     /// 2) we are within niter_part_fine_duration/2 of a multiple of niter_part_fine_period
@@ -235,6 +303,8 @@ int kraichnan_field<rnumber>::read_parameters(void)
     this->spectrum_k_cutoff = hdf5_tools::read_value<double>(parameter_file, "parameters/spectrum_k_cutoff");
     this->spectrum_coefficient = hdf5_tools::read_value<double>(parameter_file, "parameters/spectrum_coefficient");
     this->field_random_seed = hdf5_tools::read_value<int>(parameter_file, "parameters/field_random_seed");
+    this->output_velocity = hdf5_tools::read_value<int>(parameter_file, "parameters/output_velocity");
+    this->max_velocity_estimate = hdf5_tools::read_value<double>(parameter_file, "parameters/max_velocity_estimate");
     H5Fclose(parameter_file);
     return EXIT_SUCCESS;
 }
@@ -242,6 +312,8 @@ int kraichnan_field<rnumber>::read_parameters(void)
 template <class rnumber>
 void kraichnan_field<rnumber>::update_checkpoint()
 {
+    TIMEZONE("kraichnan_field::update_checkpoint");
+    DEBUG_MSG("entered kraichan_field::update_checkpoint\n");
     std::string fname = this->get_current_fname();
     if (this->kk->layout->myrank == 0)
     {
@@ -297,6 +369,7 @@ void kraichnan_field<rnumber>::update_checkpoint()
         }
     }
     MPI_Bcast(&this->checkpoint, 1, MPI_INT, 0, this->kk->layout->comm);
+    DEBUG_MSG("exiting kraichan_field::update_checkpoint\n");
 }
 
 template class kraichnan_field<float>;
diff --git a/cpp/full_code/kraichnan_field.hpp b/cpp/full_code/kraichnan_field.hpp
index 992724fb..d7fb215e 100644
--- a/cpp/full_code/kraichnan_field.hpp
+++ b/cpp/full_code/kraichnan_field.hpp
@@ -55,6 +55,7 @@ class kraichnan_field: public direct_numerical_simulation
         int tracers0_integration_steps;
         int tracers0_neighbours;
         int tracers0_smoothness;
+        int output_velocity;
 
         /* other stuff */
         kspace<FFTW, SMOOTH> *kk;
@@ -62,6 +63,7 @@ class kraichnan_field: public direct_numerical_simulation
         double spectrum_slope;
         double spectrum_k_cutoff;
         double spectrum_coefficient;
+        double max_velocity_estimate;
         int field_random_seed;
 
         /* other stuff */
@@ -87,20 +89,8 @@ class kraichnan_field: public direct_numerical_simulation
         int write_checkpoint(void);
         int do_stats(void);
 
-        
+        int generate_random_velocity(void);
         void update_checkpoint(void);
-        inline void io_checkpoint(bool read = true)
-        {
-            assert(this->velocity->real_space_representation);
-            if (!read)
-                this->update_checkpoint();
-            std::string fname = this->get_current_fname();
-            this->velocity->io(
-                    fname,
-                    "velocity",
-                    this->iteration,
-                    read);
-        }
 };
 
 #endif//KRAICHNAN_FIELD_HPP
-- 
GitLab