Merge pull request #250 from xylar/fix-manufactured-solution-viz

Fix manufactured solution viz step

Author: xylar

polaris/ocean/convergence/analysis.py

@@ -121,6 +121,7 @@ def setup(self):
         Add input files based on resolutions, which may have been changed by
         user config options
         """
+        super().setup()
         config = self.config
         dependencies = self.dependencies_dict
         if self.refinement == 'time':

polaris/ocean/tasks/manufactured_solution/__init__.py

@@ -119,10 +119,12 @@ def __init__(self, component, config, refinement='both'):
 
         self.add_step(Analysis(component=component,
                                subdir=f'{self.subdir}/analysis',
-                               refinement=refinement,
-                               dependencies=analysis_dependencies))
-        self.add_step(Viz(component=component, resolutions=resolutions,
-                          taskdir=self.subdir),
+                               dependencies=analysis_dependencies,
+                               refinement=refinement))
+        self.add_step(Viz(component=component,
+                          dependencies=analysis_dependencies,
+                          taskdir=self.subdir,
+                          refinement=refinement),
                       run_by_default=False)
         config.add_from_package('polaris.ocean.convergence',
                                 'convergence.cfg')

polaris/ocean/tasks/manufactured_solution/analysis.py

@@ -8,11 +8,6 @@ class Analysis(ConvergenceAnalysis):
     """
     A step for analysing the output from the manufactured solution
     test case
-
-    Attributes
-    ----------
-    resolutions : list of float
-        The resolutions of the meshes that have been run
     """
     def __init__(self, component, subdir, dependencies, refinement='both'):
         """

polaris/ocean/tasks/manufactured_solution/viz.py

@@ -4,8 +4,11 @@
 import matplotlib.pyplot as plt
 import numpy as np
 
+from polaris.ocean.convergence import (
+    get_resolution_for_task,
+    get_timestep_for_task,
+)
 from polaris.ocean.model import OceanIOStep
-from polaris.ocean.resolution import resolution_to_subdir
 from polaris.ocean.tasks.manufactured_solution.exact_solution import (
     ExactSolution,
 )
@@ -19,10 +22,31 @@ class Viz(OceanIOStep):
 
     Attributes
     ----------
-    resolutions : list of float
-        The resolutions of the meshes that have been run
+    dependencies_dict : dict of dict of polaris.Steps
+        The dependencies of this step must be given as separate keys in the
+        dict:
+
+            mesh : dict of polaris.Steps
+                Keys of the dict correspond to `refinement_factors`
+                Values of the dict are polaris.Steps, which must have the
+                attribute `path`, the path to `base_mesh.nc` of that
+                resolution
+            init : dict of polaris.Steps
+                Keys of the dict correspond to `refinement_factors`
+                Values of the dict are polaris.Steps, which must have the
+                attribute `path`, the path to `initial_state.nc` of that
+                resolution
+            forward : dict of polaris.Steps
+                Keys of the dict correspond to `refinement_factors`
+                Values of the dict are polaris.Steps, which must have the
+                attribute `path`, the path to `forward.nc` of that
+                resolution
+
+    refinement : str
+        Refinement type. One of 'space', 'time' or 'both' indicating both
+        space and time
     """
-    def __init__(self, component, resolutions, taskdir):
+    def __init__(self, component, dependencies, taskdir, refinement='both'):
         """
         Create the step
 
@@ -31,37 +55,82 @@ def __init__(self, component, resolutions, taskdir):
         component : polaris.Component
             The component the step belongs to
 
-        resolutions : list of float
-            The resolutions of the meshes that have been run
+        dependencies : dict of dict of polaris.Steps
+            The dependencies of this step must be given as separate keys in the
+            dict:
+
+                mesh : dict of polaris.Steps
+                    Keys of the dict correspond to `refinement_factors`
+                    Values of the dict are polaris.Steps, which must have the
+                    attribute `path`, the path to `base_mesh.nc` of that
+                    resolution
+                init : dict of polaris.Steps
+                    Keys of the dict correspond to `refinement_factors`
+                    Values of the dict are polaris.Steps, which must have the
+                    attribute `path`, the path to `initial_state.nc` of that
+                    resolution
+                forward : dict of polaris.Steps
+                    Keys of the dict correspond to `refinement_factors`
+                    Values of the dict are polaris.Steps, which must have the
+                    attribute `path`, the path to `forward.nc` of that
+                    resolution
 
         taskdir : str
             The subdirectory that the task belongs to
+
+        refinement : str, optional
+            Refinement type. One of 'space', 'time' or 'both' indicating both
+            space and time
         """
         super().__init__(component=component, name='viz', indir=taskdir)
-        self.resolutions = resolutions
+        self.dependencies_dict = dependencies
+        self.refinement = refinement
+        self.add_output_file('comparison.png')
 
-        for resolution in resolutions:
-            mesh_name = resolution_to_subdir(resolution)
+    def setup(self):
+        """
+        Add input files based on resolutions, which may have been changed by
+        user config options
+        """
+        super().setup()
+        config = self.config
+        dependencies = self.dependencies_dict
+
+        if self.refinement == 'time':
+            option = 'refinement_factors_time'
+        else:
+            option = 'refinement_factors_space'
+        refinement_factors = config.getlist('convergence', option,
+                                            dtype=float)
+
+        for refinement_factor in refinement_factors:
+            base_mesh = dependencies['mesh'][refinement_factor]
+            init = dependencies['init'][refinement_factor]
+            forward = dependencies['forward'][refinement_factor]
             self.add_input_file(
-                filename=f'mesh_{mesh_name}.nc',
-                target=f'../init/{mesh_name}/culled_mesh.nc')
+                filename=f'mesh_r{refinement_factor:02g}.nc',
+                work_dir_target=f'{base_mesh.path}/base_mesh.nc')
             self.add_input_file(
-                filename=f'init_{mesh_name}.nc',
-                target=f'../init/{mesh_name}/initial_state.nc')
+                filename=f'init_r{refinement_factor:02g}.nc',
+                work_dir_target=f'{init.path}/initial_state.nc')
             self.add_input_file(
-                filename=f'output_{mesh_name}.nc',
-                target=f'../forward/{mesh_name}/output.nc')
-
-        self.add_output_file('comparison.png')
+                filename=f'output_r{refinement_factor:02g}.nc',
+                work_dir_target=f'{forward.path}/output.nc')
 
     def run(self):
         """
         Run this step of the test case
         """
         plt.switch_backend('Agg')
         config = self.config
-        resolutions = self.resolutions
-        nres = len(resolutions)
+        if self.refinement == 'time':
+            option = 'refinement_factors_time'
+        else:
+            option = 'refinement_factors_space'
+        refinement_factors = config.getlist('convergence', option,
+                                            dtype=float)
+
+        nres = len(refinement_factors)
 
         section = config['manufactured_solution']
         eta0 = section.getfloat('ssh_amplitude')
@@ -70,11 +139,14 @@ def run(self):
         fig, axes = plt.subplots(nrows=nres, ncols=3, figsize=(12, 2 * nres))
         rmse = []
         error_range = None
-        for i, res in enumerate(resolutions):
-            mesh_name = resolution_to_subdir(res)
-            ds_mesh = self.open_model_dataset(f'mesh_{mesh_name}.nc')
-            ds_init = self.open_model_dataset(f'init_{mesh_name}.nc')
-            ds = self.open_model_dataset(f'output_{mesh_name}.nc')
+
+        for i, refinement_factor in enumerate(refinement_factors):
+            ds_mesh = self.open_model_dataset(
+                f'mesh_r{refinement_factor:02g}.nc')
+            ds_init = self.open_model_dataset(
+                f'init_r{refinement_factor:02g}.nc')
+            ds = self.open_model_dataset(
+                f'output_r{refinement_factor:02g}.nc')
             exact = ExactSolution(config, ds_init)
 
             t0 = datetime.datetime.strptime(ds.xtime.values[0].decode(),
@@ -110,8 +182,13 @@ def run(self):
         axes[0, 2].set_title('Error (Numerical - Analytical)')
 
         pad = 5
-        for ax, res in zip(axes[:, 0], resolutions):
-            ax.annotate(f'{res}km', xy=(0, 0.5),
+        for ax, refinement_factor in zip(axes[:, 0], refinement_factors):
+            timestep, _ = get_timestep_for_task(
+                config, refinement_factor, refinement=self.refinement)
+            resolution = get_resolution_for_task(
+                config, refinement_factor, refinement=self.refinement)
+
+            ax.annotate(f'{resolution}km\n{timestep}s', xy=(0, 0.5),
                         xytext=(-ax.yaxis.labelpad - pad, 0),
                         xycoords=ax.yaxis.label, textcoords='offset points',
                         size='large', ha='right', va='center')