write all outputs to single netcdf file

- add time dimension to netcdf output files
- add time variable to netcdf output files
- propagate simulation time to export method

thetis/exporter.py

@@ -24,6 +24,8 @@ def get_visu_space(fs):
     is_vector = len(fs.ufl_element().value_shape()) == 1
     mesh = fs.mesh()
     cell = mesh.ufl_cell()
+    if isinstance(cell, TensorProductCell):
+        cell = cell.sub_cells()[0]
     family_selector = {
         (True, triangle): 'Lagrange',
         (False, triangle): 'Discontinuous Lagrange',
@@ -63,7 +65,7 @@ def set_next_export_ix(self, next_export_ix):
         """Sets the index of next export"""
         self.next_export_ix = next_export_ix
 
-    def export(self, function):
+    def export(self, function, time):
         """Exports given function to disk"""
         raise NotImplementedError('This method must be implemented in the derived class')
 
@@ -104,8 +106,13 @@ def set_next_export_ix(self, next_export_ix):
         # FIXME hack to change correct output file count
         self.outfile.counter = itertools.count(start=self.next_export_ix)
 
-    def export(self, function):
-        """Exports given function to disk"""
+    def export(self, function, time):
+        """
+        Export function to disk.
+
+        :arg function: :class:`Function` to export
+        :arg float time: simulation time in seconds
+        """
         assert self.fs_visu.max_work_functions == 1
         tmp_proj_func = self.fs_visu.get_work_function()
         # NOTE tmp function must be invariant as the projector is built only once
@@ -160,12 +167,13 @@ def gen_filename(self, iexport):
         filename = '{0:s}_{1:05d}'.format(self.filename, iexport)
         return os.path.join(self.outputdir, filename)
 
-    def export_as_index(self, iexport, function):
+    def export_as_index(self, iexport, function, time):
         """
         Export function to disk using the specified export index number
 
         :arg int iexport: export index >= 0
         :arg function: :class:`Function` to export
+        :arg float time: simulation time in seconds
         """
         assert function.function_space() == self.function_space,\
             'Function space does not match'
@@ -176,15 +184,16 @@ def export_as_index(self, iexport, function):
             f.store(function)
         self.next_export_ix = iexport + 1
 
-    def export(self, function):
+    def export(self, function, time):
         """
         Export function to disk.
 
         Increments export index by 1.
 
         :arg function: :class:`Function` to export
+        :arg float time: simulation time in seconds
         """
-        self.export_as_index(self.next_export_ix, function)
+        self.export_as_index(self.next_export_ix, function, time)
 
     def load(self, iexport, function):
         """
@@ -232,19 +241,20 @@ def __init__(self, fs_visu, func_name, outputdir, filename_prefix,
 
     def gen_filename(self, iexport):
         """
-        Generate file name 'prefix_nnnnn.nc' for i-th export
+        Generate file name for i-th export
 
         :arg int iexport: export index >= 0
         """
-        filename = '{0:s}_{1:05d}.nc'.format(self.filename, iexport)
+        filename = '{0:s}.nc'.format(self.filename, iexport)
         return os.path.join(self.outputdir, filename)
 
-    def export_as_index(self, iexport, function):
+    def export_as_index(self, iexport, function, time):
         """
         Export function to disk using the specified export index number
 
         :arg int iexport: export index >= 0
         :arg function: :class:`Function` to export
+        :arg float time: simulation time in seconds
         """
         filename = self.gen_filename(iexport)
         if self.verbose:
@@ -264,19 +274,21 @@ def export_as_index(self, iexport, function):
                 self.cast_operators[function] = op
             op.interpolate()
 
-        write_nc_field(tmp_proj_func, filename, self.func_name)
+        write_nc_field(tmp_proj_func, time, filename, self.func_name,
+                       time_index=iexport)
         self.fs_visu.restore_work_function(tmp_proj_func)
         self.next_export_ix = iexport + 1
 
-    def export(self, function):
+    def export(self, function, time):
         """
         Export function to disk.
 
         Increments export index by 1.
 
         :arg function: :class:`Function` to export
+        :arg float time: simulation time in seconds
         """
-        self.export_as_index(self.next_export_ix, function)
+        self.export_as_index(self.next_export_ix, function, time)
 
 
 class ExportManager(object):
@@ -380,11 +392,13 @@ def set_next_export_ix(self, next_export_ix):
         for k in self.exporters:
             self.exporters[k].set_next_export_ix(next_export_ix)
 
-    def export(self):
+    def export(self, time):
         """
         Export all designated functions to disk
 
         Increments export index by 1.
+
+        :arg float time: simulation time in seconds
         """
         if self.verbose and COMM_WORLD.rank == 0:
             sys.stdout.write('Exporting: ')
@@ -396,7 +410,7 @@ def export(self):
                     sys.stdout.flush()
                 if key in self.preproc_callbacks:
                     self.preproc_callbacks[key]()
-                self.exporters[key].export(field)
+                self.exporters[key].export(field, time)
         if self.verbose and COMM_WORLD.rank == 0:
             sys.stdout.write('\n')
             sys.stdout.flush()

thetis/netcdf_io.py

@@ -1,15 +1,18 @@
 import netCDF4
 import firedrake as fd
 from .utility import get_functionspace
+import os
 
 
-def write_nc_field(function, filename, fieldname):
+def write_nc_field(function, time, filename, fieldname, time_index=0):
     """
     Write Function to disk in netCDF format
 
     :arg function: Firedrake function to export
+    :arg float time: simulation time in seconds
     :arg string filename: output filename
     :arg string fieldname: canonical field name
+    :kwarg int time_index: write to specified time index in the output file
     """
     fs = function.function_space()
     mesh = fs.mesh()
@@ -38,8 +41,7 @@ def write_nc_field(function, filename, fieldname):
     is_dg = family in ['Discontinuous Lagrange', 'DQ']
     fs_vertex = fs_p1dg if (is_dg and degree == 1) else fs_p1
     is_vector = fs.shape != ()
-    if is_vector:
-        vector_dim = fs.shape[0]
+    vector_dim = fs.shape[0] if is_vector else None
     face_nodes = elem.cell().num_vertices()
 
     gdim = mesh.geometric_dimension()
@@ -75,37 +77,68 @@ def get_lg_index(fs):
     x, y = fd.SpatialCoordinate(mesh)
     f_coords.interpolate(fd.as_vector((x, y)))
 
-    with netCDF4.Dataset(filename, 'w', parallel=True) as ncfile:
-        ncfile.createDimension('face', global_nb_cells)
-        ncfile.createDimension('face_nb_nodes', face_nodes)
-        ncfile.createDimension('vertex', global_nb_vertices)
-        ncfile.createDimension('time', None)
+    def _append_data(ncfile, time_index, function, time, fieldname,
+                     is_vector, vector_dim, l2g_ix):
 
-        mesh_prefix = 'Mesh'
-
-        var = ncfile.createVariable(
-            f'{mesh_prefix}_face_nodes', 'u8', ('face', 'face_nb_nodes'))
-        var.start_index = 0
-        var[local2global_cell_ix, :] = global_cell_nodes
-
-        for i, label in zip(range(gdim), ('x', 'y', 'z')):
-            var = ncfile.createVariable(f'{mesh_prefix}_node_{label}', 'f', ('vertex', ))
-            var[local2global_vertex_ix] = f_coords.dat.data_ro[:, i]
-
-        if data_at_cells:
-            l2g_ix = local2global_cell_ix
-            shape = ('time', 'face', )
-        else:
-            l2g_ix = local2global_vertex_ix
-            shape = ('time', 'vertex', )
-
-        time_index = 0
+        var = ncfile['time']
+        var.set_collective(True)
+        var[time_index] = time
         if is_vector:
             for i, label in zip(range(vector_dim), ('x', 'y', 'z')):
-                var = ncfile.createVariable(f'{fieldname}_{label}', 'f', shape)
-                var.set_collective(True)  # needed for unlimited time dim
+                var = ncfile[f'{fieldname}_{label}']
+                var.set_collective(True)
                 var[time_index, l2g_ix] = function.dat.data_ro[:, i]
         else:
-            var = ncfile.createVariable(fieldname, 'f', shape)
-            var.set_collective(True)  # needed for unlimited time dim
+            var = ncfile[fieldname]
+            var.set_collective(True)
             var[time_index, l2g_ix] = function.dat.data_ro
+
+    if data_at_cells:
+        l2g_ix = local2global_cell_ix
+        shape = ('time', 'face', )
+    else:
+        l2g_ix = local2global_vertex_ix
+        shape = ('time', 'vertex', )
+
+    if time_index == 0:
+        # create new file
+        with netCDF4.Dataset(filename, 'w', parallel=True) as ncfile:
+            ncfile.createDimension('face', global_nb_cells)
+            ncfile.createDimension('face_nb_nodes', face_nodes)
+            ncfile.createDimension('vertex', global_nb_vertices)
+            ncfile.createDimension('time', None)
+
+            mesh_prefix = 'Mesh'
+
+            var = ncfile.createVariable(
+                f'{mesh_prefix}_face_nodes', 'u8', ('face', 'face_nb_nodes'))
+            var.start_index = 0
+            var[local2global_cell_ix, :] = global_cell_nodes
+
+            for i, label in zip(range(gdim), ('x', 'y', 'z')):
+                var = ncfile.createVariable(f'{mesh_prefix}_node_{label}', 'f', ('vertex', ))
+                var[local2global_vertex_ix] = f_coords.dat.data_ro[:, i]
+
+            var = ncfile.createVariable('time', 'f', ('time'))
+            var.standard_name = 'time'
+            var.units = 's'
+
+            if is_vector:
+                for i, label in zip(range(vector_dim), ('x', 'y', 'z')):
+                    var = ncfile.createVariable(f'{fieldname}_{label}', 'f', shape)
+                    var.set_collective(True)  # needed for unlimited time dim
+            else:
+                var = ncfile.createVariable(fieldname, 'f', shape)
+                var.set_collective(True)  # needed for unlimited time dim
+
+            _append_data(
+                ncfile, time_index, function, time, fieldname, is_vector,
+                vector_dim, l2g_ix
+            )
+    else:
+        # append to existing file
+        with netCDF4.Dataset(filename, 'r+', parallel=True) as ncfile:
+            _append_data(
+                ncfile, time_index, function, time, fieldname, is_vector,
+                vector_dim, l2g_ix
+            )

thetis/solver.py

@@ -983,7 +983,7 @@ def export(self):
         # TODO find a cleaner way of doing this ...
         self.fields.uv_3d += self.fields.uv_dav_3d
         for e in self.exporters.values():
-            e.export()
+            e.export(self.simulation_time)
         # restore uv_3d
         self.fields.uv_3d -= self.fields.uv_dav_3d
 

thetis/solver2d.py

@@ -722,7 +722,7 @@ def export(self):
         """
         self.callbacks.evaluate(mode='export')
         for e in self.exporters.values():
-            e.export()
+            e.export(self.simulation_time)
 
     def load_state(self, i_export, outputdir=None, t=None, iteration=None):
         """