Add save_vtu_each_timestep option for VTU output

opencmp/config_functions/expanded_config_parser.py

@@ -63,6 +63,7 @@
               'restart_from': ''},
     'VISUALIZATION': {'save_to_file': False,
                       'save_type': '.sol',
+                      'save_vtu_each_timestep': False,
                       'save_frequency': ['1', 'numit'],
                       'subdivision': -1,
                       'split_components': False},

opencmp/helpers/saving.py

@@ -15,8 +15,9 @@
 # <https://www.gnu.org/licenses/>.                                                                                     #
 ########################################################################################################################
 
+import re
 from typing import Union
-from ngsolve import GridFunction, CoefficientFunction
+from ngsolve import GridFunction, CoefficientFunction, VTKOutput
 from ..models import Model
 from pathlib import Path
 
@@ -52,6 +53,15 @@ def __init__(self, model: Model, quiet: bool = False) -> None:
         self.base_filename_sol = self.save_dir_sol + model.name() + '_'
         self.base_filename_phi_sol = self.save_dir_phi_sol + 'phi' + '_'
         self.base_subdivision = model.config.get_item(['VISUALIZATION', 'subdivision'], int, quiet)
+        self.save_vtu_each_timestep = (
+            base_type == '.vtu'
+            and model.config.get_item(
+                ['VISUALIZATION', 'save_vtu_each_timestep'], bool, quiet)
+        )
+        self.model_name = model.name()
+        self.mesh = model.mesh
+        self.save_names = model.save_names
+        self.pvd_entries = {False: {}, True: {}}
 
         # Create the save dir if it doesn't exist
         Path(self.save_dir).mkdir(parents=True, exist_ok=True)
@@ -91,3 +101,43 @@ def save(self, gfu: Union[GridFunction, CoefficientFunction], timestep: float, D
 
         # Save to file
         gfu.Save(filename)
+
+        if self.save_vtu_each_timestep:
+            self._save_vtu(gfu, timestep, DIM)
+
+    def _save_vtu(self, gfu: Union[GridFunction, CoefficientFunction], timestep: float, DIM: bool) -> None:
+        """Write the just-saved checkpoint to VTU and refresh its PVD collection."""
+        if DIM:
+            vtu_dir = self.save_dir_phi_vtu
+            basename = 'phi_' + str(timestep)
+            names = ['phi']
+            pvd_filename = self.save_dir_phi + 'phi_transient.pvd'
+            relative_filename = self.model_name + '_vtu/' + basename + '.vtu'
+        else:
+            vtu_dir = self.save_dir_vtu
+            basename = self.model_name + '_' + str(timestep)
+            names = self.save_names
+            pvd_filename = self.save_dir + self.model_name + '_transient.pvd'
+            relative_filename = self.model_name + '_vtu/' + basename + '.vtu'
+
+        # On the first write of a run, keep any entries already in the .pvd so a
+        # resumed run doesn't truncate the collection to just the new timesteps.
+        if not self.pvd_entries[DIM] and Path(pvd_filename).is_file():
+            self.pvd_entries[DIM] = {float(t): f for t, f in
+                                     re.findall(r'timestep="([^"]+)".*?file="([^"]+)"',
+                                                Path(pvd_filename).read_text())}
+
+        coefs = list(gfu.components) if isinstance(gfu, GridFunction) and len(gfu.components) > 0 else [gfu]
+        VTKOutput(ma=self.mesh, coefs=coefs, names=names,
+                  filename=vtu_dir + basename, subdivision=self.base_subdivision).Do()
+
+        self.pvd_entries[DIM][float(timestep)] = relative_filename
+        with open(pvd_filename, 'w') as pvd:
+            pvd.write('<?xml version="1.0"?>\n'
+                      '<VTKFile type="Collection" version="0.1"\n'
+                      'byte_order="LittleEndian" compressor="vtkZLibDataCompressor">\n'
+                      '<Collection>\n')
+            for time, path in sorted(self.pvd_entries[DIM].items()):
+                pvd.write('<DataSet timestep="%e" group="" part="0" file="%s"/>\n'
+                          % (time, path))
+            pvd.write('</Collection>\n</VTKFile>')

opencmp/post_processing/__init__.py

@@ -44,8 +44,10 @@ def run_post_processing(config_parser: ConfigParser, solver: Solver, sol: GridFu
 
     save_output = config_parser.get_item(['VISUALIZATION', 'save_to_file'], bool, quiet=True)
     save_type = config_parser.get_item(['VISUALIZATION', 'save_type'], str, quiet=True)
+    save_vtu_each_timestep = config_parser.get_item(
+        ['VISUALIZATION', 'save_vtu_each_timestep'], bool, quiet=True)
     # Run the post-processor to convert the .sol to .vtu
-    if save_output and save_type == '.vtu':
+    if save_output and save_type == '.vtu' and not save_vtu_each_timestep:
         print('Converting saved output to VTU.')
         sol_to_vtu(config_parser, solver)
 
@@ -55,4 +57,3 @@ def run_post_processing(config_parser: ConfigParser, solver: Solver, sol: GridFu
         sol_to_components(config_parser,
                           config_parser.get_item(['OTHER', 'run_dir'], str) + '/output/',
                           solver.model)
-