From e79d3cfa13bafb0f701ab36c27596d648c74cf7b Mon Sep 17 00:00:00 2001
From: Xylar Asay-Davis <xylarstorm@gmail.com>
Date: Thu, 17 Oct 2024 18:07:36 -0500
Subject: [PATCH] Update variable mapping methods

This merge also adds or updates some attribute docstrings
---
 polaris/ocean/model/ocean_model_step.py | 56 ++++++++++++++++---------
 1 file changed, 37 insertions(+), 19 deletions(-)

diff --git a/polaris/ocean/model/ocean_model_step.py b/polaris/ocean/model/ocean_model_step.py
index 0bb6cf85f..69d53d8c4 100644
--- a/polaris/ocean/model/ocean_model_step.py
+++ b/polaris/ocean/model/ocean_model_step.py
@@ -17,7 +17,14 @@ class OceanModelStep(ModelStep):
         ``min_tasks``) are computed dynamically from the number of cells
         in the mesh
 
-    map : dict
+    mpaso_to_omega_var_map : dict
+        A map from MPAS-Ocean variable names to their Omega equivalents
+
+    omega_to_mpaso_var_map : dict
+        A map from Omega variable names to their MPAS-Ocean equivalents, the
+        inverse of ``mpaso_to_omega_var_map``
+
+    config_map : dict
         A nested dictionary that maps from MPAS-Ocean to Omega model config
         options
 
@@ -104,7 +111,8 @@ def __init__(self, component, name, subdir=None, indir=None, ntasks=None,
 
         self.dynamic_ntasks = (ntasks is None and min_tasks is None)
 
-        self.var_map: Union[None, Dict[str, str]] = None
+        self.mpaso_to_omega_var_map: Union[None, Dict[str, str]] = None
+        self.omega_to_mpaso_var_map: Union[None, Dict[str, str]] = None
         self.config_map: Union[None, List[Dict[str, Dict[str, str]]]] = None
         self.graph_target = graph_target
 
@@ -116,11 +124,11 @@ def setup(self):
         """
         config = self.config
         model = config.get('ocean', 'model')
-        self._read_var_map()
         if model == 'omega':
             self.make_yaml = True
             self.config_models = ['ocean', 'Omega']
             self.yaml = 'omega.yml'
+            self._read_var_map()
             self._read_config_map()
             self.partition_graph = False
         elif model == 'mpas-ocean':
@@ -160,10 +168,11 @@ def compute_cell_count(self):
         """
         return None
 
-    def rename_vars_mpaso_to_omega(self, ds):
+    def map_input_dataset(self, ds):
         """
-        Rename variables in a dataset from their MPAS-Ocean names to the Omega
-        equivalent
+        If the model is Omega, rename variables in a dataset from their
+        MPAS-Ocean names to the Omega equivalent (appropriate for input
+        datasets like an initial condition)
 
         Parameters
         ----------
@@ -173,31 +182,37 @@ def rename_vars_mpaso_to_omega(self, ds):
         Returns
         -------
         ds : xarray.Dataset
-            The same dataset with variables renamed to their Omega equivalents
+            The same dataset with variables renamed as appropriate for the
+            ocean model being run
         """
-        assert self.var_map is not None
-        ds = ds.rename(self.var_map)
+        config = self.config
+        model = config.get('ocean', 'model')
+        if model == 'omega':
+            assert self.mpaso_to_omega_var_map is not None
+            ds = ds.rename(self.mpaso_to_omega_var_map)
         return ds
 
-    def rename_vars_omega_to_mpaso(self, ds):
+    def map_output_dataset(self, ds):
         """
-        Rename variables in a dataset from their Omega names to the MPAS-Ocean
-        equivalent
+        If the model is Omega, rename variables in a dataset from their
+        Omega names to the MPAS-Ocean equivalent (appropriate for datasets
+        that are output from the model)
 
         Parameters
         ----------
         ds : xarray.Dataset
-            A dataset containing Omega variable names
+            A dataset containing variable names native to either ocean model
 
         Returns
         -------
         ds : xarray.Dataset
-            The same dataset with variables renamed to their MPAS-Ocean
-            equivalents
+            The same dataset with variables named as expected in MPAS-Ocean
         """
-        assert self.var_map is not None
-        inv_map = {v: k for k, v in self.var_map.items()}
-        ds = ds.rename(inv_map)
+        config = self.config
+        model = config.get('ocean', 'model')
+        if model == 'omega':
+            assert self.omega_to_mpaso_var_map is not None
+            ds = ds.rename(self.omega_to_mpaso_var_map)
         return ds
 
     def map_yaml_options(self, options, config_model):
@@ -324,7 +339,10 @@ def _read_var_map(self):
 
         yaml_data = YAML(typ='rt')
         nested_dict = yaml_data.load(text)
-        self.var_map = nested_dict['variables']
+        self.mpaso_to_omega_var_map = nested_dict['variables']
+        assert self.mpaso_to_omega_var_map is not None
+        self.omega_to_mpaso_var_map = {
+            v: k for k, v in self.mpaso_to_omega_var_map.items()}
 
     def _read_config_map(self):
         """