Adding previously removed restart file support.

pressel · Oct 13, 2015 · 765d277 · 765d277
1 parent eaab8ec
commit 765d277
Show file tree

Hide file tree

Showing 17 changed files with 411 additions and 74 deletions.
diff --git a/Grid.pxd b/Grid.pxd
@@ -1,4 +1,5 @@
 cimport ParallelMPI
+cimport Restart
 
 cdef extern from "grid.h":
     struct DimStruct:
@@ -49,6 +50,7 @@ cdef class Grid:
 
     cpdef extract_local_ghosted(self, double [:] global_array, int dim)
 
+    cpdef restart(self, Restart.Restart Re)
 
 
 
diff --git a/Grid.pyx b/Grid.pyx
@@ -6,6 +6,7 @@
 
 cimport mpi4py.mpi_c as mpi
 cimport ParallelMPI
+cimport Restart
 cimport numpy as np
 import numpy as np
 import time
@@ -178,4 +179,44 @@ cdef class Grid:
         cdef int start = self.dims.indx_lo_g[dim]
         cdef int end = self.dims.indx_lo_g[dim] + self.dims.nlg[dim]
         #Force a copy with the return statement
-        return np.array(global_array[start:end],dtype=np.double)
+        return np.array(global_array[start:end],dtype=np.double)
+
+    cpdef restart(self, Restart.Restart Re):
+        Re.restart_data['Gr'] = {}
+        Re.restart_data['Gr']['dims'] = self.dims.dims
+        Re.restart_data['Gr']['n'] = np.array([self.dims.n[0],
+                                               self.dims.n[1],
+                                               self.dims.n[2]])
+        Re.restart_data['Gr']['ng'] =  np.array([self.dims.ng[0],
+                                                 self.dims.ng[1],
+                                                 self.dims.ng[2]])
+        Re.restart_data['Gr']['nl'] = np.array([self.dims.nl[0],
+                                                self.dims.nl[1],
+                                                self.dims.nl[2]])
+        Re.restart_data['Gr']['nlg'] = np.array([self.dims.nlg[0],
+                                                  self.dims.nlg[1],
+                                                  self.dims.nlg[2]])
+        Re.restart_data['Gr']['indx_lo_g'] = np.array([self.dims.indx_lo_g[0],
+                                                  self.dims.indx_lo_g[1],
+                                                  self.dims.indx_lo_g[2]])
+        Re.restart_data['Gr']['indx_lo'] = np.array([self.dims.indx_lo[0],
+                                                  self.dims.indx_lo[1],
+                                                  self.dims.indx_lo[2]])
+        Re.restart_data['Gr']['npd'] = self.dims.npd
+        Re.restart_data['Gr']['npl'] = self.dims.npl
+        Re.restart_data['Gr']['npg'] = self.dims.npg
+        Re.restart_data['Gr']['gw'] = self.dims.gw
+        Re.restart_data['Gr']['nbuffer']  = np.array([self.dims.nbuffer[0],
+                                                  self.dims.nbuffer[1],
+                                                  self.dims.nbuffer[2]])
+        Re.restart_data['Gr']['nbuffer']  = np.array([self.dims.ghosted_stride[0],
+                                                  self.dims.ghosted_stride[1],
+                                                  self.dims.ghosted_stride[2]])
+        Re.restart_data['Gr']['dx'] = np.array([self.dims.dx[0],
+                                                self.dims.dx[1],
+                                                self.dims.dx[2]])
+        Re.restart_data['Gr']['dxi'] = np.array([self.dims.dxi[0],
+                                                self.dims.dxi[1],
+                                                self.dims.dxi[2]])
+
+        return
diff --git a/NetCDFIO.pyx b/NetCDFIO.pyx
@@ -27,18 +27,9 @@ cdef class NetCDFIO_Stats:
         self.frequency = namelist['stats_io']['frequency']
 
         # Setup the statistics output path
-        outpath = str(os.path.join(namelist['output'][
-                      'output_root'] + 'Output.' + namelist['meta']['simname'] + '.' + self.uuid[-5:]))
-        self.stats_path = str(
-            os.path.join(
-                outpath,
-                namelist['stats_io']['stats_dir']))
-        self.path_plus_file = str(
-            self.stats_path +
-            '/' +
-            'Stats.' +
-            namelist['meta']['simname'] +
-            '.nc')
+        outpath = str(os.path.join(namelist['output']['output_root'] + 'Output.' + namelist['meta']['simname'] + '.' + self.uuid[-5:]))
+        self.stats_path = str( os.path.join(outpath, namelist['stats_io']['stats_dir']))
+        self.path_plus_file = str( self.stats_path + '/' + 'Stats.' + namelist['meta']['simname'] + '.nc')
         if Pa.rank == 0:
             try:
                 os.mkdir(outpath)
@@ -50,14 +41,8 @@ cdef class NetCDFIO_Stats:
                 pass
 
             shutil.copyfile(
-                os.path.join(
-                    './',
-                    namelist['meta']['simname'] +
-                    '.in'),
-                os.path.join(
-                    outpath,
-                    namelist['meta']['simname'] +
-                    '.in'))
+                os.path.join( './', namelist['meta']['simname'] + '.in'),
+                os.path.join( outpath, namelist['meta']['simname'] + '.in'))
             self.setup_stats_file(Gr, Pa)
         return
 
@@ -191,12 +176,8 @@ cdef class NetCDFIO_Fields:
         self.diagnostic_fields = namelist['fields_io']['diagnostic_fields']
 
         # Setup the statistics output path
-        outpath = str(os.path.join(namelist['output'][
-                      'output_root'] + 'Output.' + namelist['meta']['simname'] + '.' + self.uuid[-5:]))
-        self.fields_path = str(
-            os.path.join(
-                outpath,
-                namelist['fields_io']['fields_dir']))
+        outpath = str(os.path.join(namelist['output']['output_root'] + 'Output.' + namelist['meta']['simname'] + '.' + self.uuid[-5:]))
+        self.fields_path = str(os.path.join(outpath, namelist['fields_io']['fields_dir']))
         if Pa.rank == 0:
             try:
                 os.mkdir(outpath)
@@ -207,15 +188,8 @@ cdef class NetCDFIO_Fields:
             except:
                 pass
 
-            shutil.copyfile(
-                os.path.join(
-                    './',
-                    namelist['meta']['simname'] +
-                    '.in'),
-                os.path.join(
-                    outpath,
-                    namelist['meta']['simname'] +
-                    '.in'))
+            shutil.copyfile( os.path.join('./', namelist['meta']['simname'] + '.in'),
+                             os.path.join( outpath, namelist['meta']['simname'] + '.in'))
         return
 
     cpdef update(self, Grid.Grid Gr, PrognosticVariables.PrognosticVariables PV, DiagnosticVariables.DiagnosticVariables DV, TimeStepping.TimeStepping TS, ParallelMPI.ParallelMPI Pa):

diff --git a/ParallelMPI.pxd b/ParallelMPI.pxd
@@ -22,7 +22,7 @@ cdef class ParallelMPI:
         int sub_z_rank
 
         void barrier(self)
-        void kill(self)
+
         void create_sub_communicators(self)
 
         double domain_scalar_sum(self, double local_value)
@@ -39,6 +39,7 @@ cdef class ParallelMPI:
         double [:] HorizontalMeanofSquaresConditional(self,Grid.Grid Gr, double* values1, double* values2, double* mask)
 
     cpdef root_print(self, txt_output)
+    cpdef void kill(self)
 
 cdef class Pencil:
 

diff --git a/ParallelMPI.pyx b/ParallelMPI.pyx
@@ -75,7 +75,7 @@ cdef class ParallelMPI:
             print(txt_output)
         return
 
-    cdef void kill(self):
+    cpdef void kill(self):
         '''
         Call MPI_Abort.
         :return:
@@ -1065,3 +1065,4 @@ cdef class Pencil:
                         ijk_no_gw = ishift_nogw + jshift_nogw+ (k-dims.gw)*kstride_nogw
                         data[ijk] = recv_buffer[ijk_no_gw]
         return
+
diff --git a/PrognosticVariables.pxd b/PrognosticVariables.pxd
@@ -2,6 +2,7 @@ from NetCDFIO cimport NetCDFIO_Stats
 cimport Grid
 cimport ParallelMPI
 cimport ReferenceState
+cimport Restart
 
 cdef extern from "prognostic_variables.h":
     struct VelocityDofs:
@@ -47,3 +48,5 @@ cdef class PrognosticVariables:
     cpdef val_nan(self,PA,message)
     cpdef val_bounds(self,var_name,Grid.Grid Gr)
     cpdef stats_io(self, Grid.Grid Gr, ReferenceState.ReferenceState RS, NetCDFIO_Stats NS, ParallelMPI.ParallelMPI Pa)
+    cpdef restart(self, Grid.Grid Gr, Restart.Restart Re)
+    cpdef init_from_restart(self, Grid.Grid Gr, Restart.Restart Re)
diff --git a/PrognosticVariables.pyx b/PrognosticVariables.pyx
@@ -13,7 +13,7 @@ from NetCDFIO cimport NetCDFIO_Stats
 cimport Grid
 cimport ParallelMPI
 cimport ReferenceState
-
+cimport Restart
 
 cdef class PrognosticVariables:
     def __init__(self, Grid.Grid Gr):
@@ -250,3 +250,93 @@ cdef class PrognosticVariables:
     cpdef val_bounds(self,var_name,Grid.Grid Gr):
         var_array = self.get_variable_array(var_name, Gr)
         return np.amin(var_array), np.amax(var_array)
+
+
+    cpdef restart(self, Grid.Grid Gr, Restart.Restart Re):
+
+        Re.restart_data['PV'] = {}
+        Re.restart_data['PV']['name_index'] = self.name_index
+        Re.restart_data['PV']['units'] = self.units
+        Re.restart_data['PV']['index_name'] = self.index_name
+        Re.restart_data['PV']['nv'] = self.nv
+        Re.restart_data['PV']['nv_scalars'] = self.nv_scalars
+        Re.restart_data['PV']['nv_velocities'] = self.nv_velocities
+        Re.restart_data['PV']['bc_type'] = np.array(self.bc_type)
+        Re.restart_data['PV']['var_type'] = np.array(self.var_type)
+        Re.restart_data['PV']['velocity_directions'] = np.array(self.velocity_directions)
+        Re.restart_data['PV']['velocity_names_directional'] = self.velocity_names_directional
+
+
+        cdef:
+            double [:] values = np.empty((self.nv * Gr.dims.npl),dtype=np.double,order='c')
+            Py_ssize_t imin = Gr.dims.gw
+            Py_ssize_t jmin = Gr.dims.gw
+            Py_ssize_t kmin = Gr.dims.gw
+            Py_ssize_t imax = Gr.dims.nlg[0] - Gr.dims.gw
+            Py_ssize_t jmax = Gr.dims.nlg[1] - Gr.dims.gw
+            Py_ssize_t kmax = Gr.dims.nlg[2] - Gr.dims.gw
+            Py_ssize_t i, j, k, count, ijk, n, v_shift
+            Py_ssize_t ishift, jshift
+            Py_ssize_t istride = Gr.dims.nlg[1] * Gr.dims.nlg[2]
+            Py_ssize_t jstride = Gr.dims.nlg[2]
+
+        with nogil:
+            count = 0
+            for n in xrange(self.nv):
+                v_shift = Gr.dims.nlg[0] * Gr.dims.nlg[1] * Gr.dims.nlg[2] * n
+                for i in xrange(imin, imax):
+                    ishift = istride * i
+                    for j in xrange(jmin, jmax):
+                        jshift = jstride * j
+                        for k in xrange(kmin, kmax):
+                            ijk = v_shift + ishift + jshift + k
+                            values[count] = self.values[ijk]
+                            count += 1
+
+        Re.restart_data['PV']['values'] = np.array(values)
+
+        return
+
+
+    cpdef init_from_restart(self, Grid.Grid Gr, Restart.Restart Re):
+
+        self.name_index = Re.restart_data['PV']['name_index']
+        self.units  = Re.restart_data['PV']['units']
+        self.index_name =  Re.restart_data['PV']['index_name']
+        self.nv = Re.restart_data['PV']['nv']
+        self.nv_scalars = Re.restart_data['PV']['nv_scalars']
+        self.nv_velocities = Re.restart_data['PV']['nv_velocities']
+        self.bc_type = Re.restart_data['PV']['bc_type']
+        self.var_type = Re.restart_data['PV']['var_type']
+        self.velocity_directions = Re.restart_data['PV']['velocity_directions']
+        self.velocity_names_directional = Re.restart_data['PV']['velocity_names_directional']
+
+
+        cdef:
+            double [:] values = Re.restart_data['PV']['values']
+            Py_ssize_t imin = Gr.dims.gw
+            Py_ssize_t jmin = Gr.dims.gw
+            Py_ssize_t kmin = Gr.dims.gw
+            Py_ssize_t imax = Gr.dims.nlg[0] - Gr.dims.gw
+            Py_ssize_t jmax = Gr.dims.nlg[1] - Gr.dims.gw
+            Py_ssize_t kmax = Gr.dims.nlg[2] - Gr.dims.gw
+            Py_ssize_t i, j, k, count, ijk, n
+            Py_ssize_t ishift, jshift, v_shift
+            Py_ssize_t istride = Gr.dims.nlg[1] * Gr.dims.nlg[2]
+            Py_ssize_t jstride = Gr.dims.nlg[2]
+
+
+        with nogil:
+            count = 0
+            for n in xrange(self.nv):
+                v_shift = Gr.dims.nlg[0] * Gr.dims.nlg[1] * Gr.dims.nlg[2] * n
+                for i in xrange(imin, imax):
+                    ishift = istride * i
+                    for j in xrange(jmin, jmax):
+                        jshift = jstride * j
+                        for k in xrange(kmin, kmax):
+                            ijk = v_shift + ishift + jshift + k
+                            self.values[ijk] =  values[count]
+                            count += 1
+
+        return
diff --git a/ReferenceState.pxd b/ReferenceState.pxd
@@ -1,4 +1,5 @@
 cimport Grid
+cimport Restart
 cdef class ReferenceState:
     cdef:
         public double [:] p0
@@ -7,6 +8,7 @@ cdef class ReferenceState:
         public double [:] alpha0_half
         public double [:] rho0
         public double [:] rho0_half
+
         double sg
 
     cdef public:
@@ -17,6 +19,6 @@ cdef class ReferenceState:
         double u0 #u velocity removed in Galilean transformation
         double v0 #v velocity removed in Galilean transformation
 
-
-
+    cpdef restart(self, Grid.Grid Gr, Restart.Restart Re)
+    cpdef init_from_restart(self, Grid.Grid Gr, Restart.Restart Re)
 
diff --git a/ReferenceState.pyx b/ReferenceState.pyx
@@ -5,6 +5,7 @@
 # cython: cdivision=True
 
 cimport Grid
+cimport Restart
 cimport numpy as np
 import numpy as np
 from NetCDFIO cimport NetCDFIO_Stats
@@ -134,3 +135,34 @@ cdef class ReferenceState:
         NS.write_reference_profile('qi0', qi_half[Gr.dims.gw:-Gr.dims.gw], Pa)
 
         return
+
+    cpdef restart(self, Grid.Grid Gr, Restart.Restart Re):
+        Re.restart_data['Ref'] = {}
+
+        Re.restart_data['Ref']['p0'] = np.array(self.p0)[Gr.dims.gw:-Gr.dims.gw]
+        Re.restart_data['Ref']['p0_half'] = np.array(self.p0_half)[Gr.dims.gw:-Gr.dims.gw]
+        Re.restart_data['Ref']['alpha0'] = np.array(self.alpha0)[Gr.dims.gw:-Gr.dims.gw]
+        Re.restart_data['Ref']['alpha0_half'] = np.array(self.alpha0_half)[Gr.dims.gw:-Gr.dims.gw]
+
+        Re.restart_data['Ref']['Tg'] = self.Tg
+        Re.restart_data['Ref']['Pg'] = self.Pg
+        Re.restart_data['Ref']['sg'] = self.sg
+        Re.restart_data['Ref']['qtg'] = self.qtg
+        Re.restart_data['Ref']['u0'] = self.u0
+        Re.restart_data['Ref']['v0'] = self.v0
+
+        return
+
+
+    cpdef init_from_restart(self, Grid.Grid Gr, Restart.Restart Re):
+
+
+        #self.Tg = Re.restart_data['Ref']['Tg']
+        #self.Pg = Re.restart_data['Ref']['Pg']
+        #self.sg = Re.restart_data['Ref']['sg']
+        #self.qtg = Re.restart_data['Ref']['qtg']
+        #self.u0 = Re.restart_data['Ref']['u0']
+        #self.v0 = Re.restart_data['Ref']['v0']
+
+
+        return
diff --git a/Restart.pxd b/Restart.pxd
@@ -0,0 +1,16 @@
+cimport ParallelMPI
+
+cdef class Restart:
+    cdef:
+        public dict restart_data
+        str restart_path
+        str input_path
+        public bint is_restart_run
+        str uuid
+        public double last_restart_time
+        public double frequency
+
+    cpdef initialize(self)
+    cpdef write(self, ParallelMPI.ParallelMPI Pa)
+    cpdef read(self, ParallelMPI.ParallelMPI Pa)
+    cpdef free_memory(self)