CDAT Migration Phase 1 - ReplaceCDPParameter (#638)

Co-authored-by: chengzhuzhang <[email protected]>

Author: tomvothecoder

conda-env/ci.yml

@@ -29,5 +29,4 @@ dependencies:
   - scipy
   - pytest
   - pytest-cov
-
 prefix: /opt/miniconda3/envs/e3sm_diags_ci

e3sm_diags/driver/aerosol_aeronet_driver.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import os
 from typing import TYPE_CHECKING, Optional
 
@@ -7,21 +9,22 @@
 
 import e3sm_diags
 from e3sm_diags.driver import utils
+from e3sm_diags.logger import custom_logger
 from e3sm_diags.plot.cartopy import aerosol_aeronet_plot
 
 if TYPE_CHECKING:
-    from e3sm_diags.parameter.core_parameter import CoreParameter
     from cdms2.tvariable import TransientVariable
 
-from e3sm_diags.logger import custom_logger
+    from e3sm_diags.parameter.core_parameter import CoreParameter
+
 
 logger = custom_logger(__name__)
 
 # This aerosol diagnostics scripts based on AERONET sites data was originally developed by Feng Yan and adapted and integrated in e3sm_diags by Jill Zhang.
 # Years include 2006–2015 average climatology for observation according to Feng et al. 2022:doi:10.1002/essoar.10510950.1, and Golaz et al. 2022 E3SMv2 paper.
 
 
-def run_diag(parameter: "CoreParameter") -> "CoreParameter":
+def run_diag(parameter: CoreParameter) -> CoreParameter:
     """Runs the aerosol aeronet diagnostic.
 
     :param parameter: Parameters for the run
@@ -70,7 +73,7 @@ def run_diag(parameter: "CoreParameter") -> "CoreParameter":
 
 
 def interpolate_model_output_to_obs_sites(
-    var: Optional["TransientVariable"], var_id: str
+    var: Optional[TransientVariable], var_id: str
 ):
     """Interpolate model outputs (on regular lat lon grids) to observational sites
 

e3sm_diags/driver/aerosol_budget_driver.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import csv
 import os
 from typing import TYPE_CHECKING  # , Optional
@@ -6,7 +8,6 @@
 
 if TYPE_CHECKING:
     from e3sm_diags.parameter.core_parameter import CoreParameter
-#    from cdms2.tvariable import TransientVariable
 
 import cdutil
 import numpy
@@ -89,7 +90,7 @@ def generate_metrics_dic(data, aerosol, season):
 MISSING_VALUE = 999.999
 
 
-def run_diag(parameter: "CoreParameter") -> "CoreParameter":
+def run_diag(parameter: CoreParameter) -> CoreParameter:
     """Runs the aerosol aeronet diagnostic.
 
     :param parameter: Parameters for the run

e3sm_diags/driver/annual_cycle_zonal_mean_driver.py

@@ -1,4 +1,4 @@
-from __future__ import print_function
+from __future__ import annotations
 
 from typing import TYPE_CHECKING, Any, Dict, List
 
@@ -19,10 +19,12 @@
     from cdms2.axis import TransientAxis
     from cdms2.tvariable import TransientVariable
 
-    from e3sm_diags.parameter.core_parameter import CoreParameter
+    from e3sm_diags.parameter.annual_cycle_zonal_mean_parameter import (
+        ACzonalmeanParameter,
+    )
 
 
-def run_diag(parameter: "CoreParameter"):
+def run_diag(parameter: ACzonalmeanParameter) -> ACzonalmeanParameter:
     """Runs the annual cycle zonal mean diagnostic.
 
     :param parameter: Parameters for the run
@@ -128,7 +130,7 @@ def run_diag(parameter: "CoreParameter"):
     return parameter
 
 
-def _create_annual_cycle(dataset: Dataset, variable: str) -> "TransientVariable":
+def _create_annual_cycle(dataset: Dataset, variable: str) -> TransientVariable:
     """Creates the annual climatology cycle for a dataset variable.
 
     :param dataset: Dataset
@@ -144,7 +146,7 @@ def _create_annual_cycle(dataset: Dataset, variable: str) -> "TransientVariable"
     for index, month in enumerate(month_list):
         var = dataset.get_climo_variable(variable, month)
         if month == "01":
-            var_ann_cycle: "TransientVariable" = MV2.zeros([12] + list(var.shape))
+            var_ann_cycle: TransientVariable = MV2.zeros([12] + list(var.shape))
             var_ann_cycle.id = var.id
             var_ann_cycle.long_name = var.long_name
             var_ann_cycle.units = var.units

e3sm_diags/driver/area_mean_time_series_driver.py

@@ -1,6 +1,9 @@
+from __future__ import annotations
+
 import collections
 import json
 import os
+from typing import TYPE_CHECKING
 
 import cdms2
 import cdutil
@@ -11,6 +14,12 @@
 from e3sm_diags.metrics import mean
 from e3sm_diags.plot.cartopy import area_mean_time_series_plot
 
+if TYPE_CHECKING:
+    from e3sm_diags.parameter.area_mean_time_series_parameter import (
+        AreaMeanTimeSeriesParameter,
+    )
+
+
 logger = custom_logger(__name__)
 
 RefsTestMetrics = collections.namedtuple("RefsTestMetrics", ["refs", "test", "metrics"])
@@ -24,7 +33,7 @@ def create_metrics(ref_domain):
     return {"mean": mean(ref_domain)}
 
 
-def run_diag(parameter):
+def run_diag(parameter: AreaMeanTimeSeriesParameter) -> AreaMeanTimeSeriesParameter:
     variables = parameter.variables
     regions = parameter.regions
     ref_names = parameter.ref_names

e3sm_diags/driver/arm_diags_driver.py

@@ -1,6 +1,9 @@
+from __future__ import annotations
+
 import collections
 import json
 import os
+from typing import TYPE_CHECKING
 
 import cdms2
 import numpy as np
@@ -11,6 +14,9 @@
 from e3sm_diags.logger import custom_logger
 from e3sm_diags.plot.cartopy import arm_diags_plot
 
+if TYPE_CHECKING:
+    from e3sm_diags.parameter.arm_diags_parameter import ARMDiagsParameter
+
 logger = custom_logger(__name__)
 
 RefsTestMetrics = collections.namedtuple(
@@ -50,7 +56,7 @@ def create_metrics(test, ref):
     }
 
 
-def run_diag_diurnal_cycle(parameter):
+def run_diag_diurnal_cycle(parameter: ARMDiagsParameter) -> ARMDiagsParameter:
     variables = parameter.variables
     regions = parameter.regions
     ref_name = parameter.ref_name
@@ -149,7 +155,7 @@ def run_diag_diurnal_cycle(parameter):
     return parameter
 
 
-def run_diag_diurnal_cycle_zt(parameter):
+def run_diag_diurnal_cycle_zt(parameter: ARMDiagsParameter) -> ARMDiagsParameter:
     variables = parameter.variables
     regions = parameter.regions
     ref_name = parameter.ref_name
@@ -259,7 +265,7 @@ def run_diag_diurnal_cycle_zt(parameter):
     return parameter
 
 
-def run_diag_annual_cycle(parameter):
+def run_diag_annual_cycle(parameter: ARMDiagsParameter) -> ARMDiagsParameter:
     variables = parameter.variables
     regions = parameter.regions
     ref_name = parameter.ref_name
@@ -360,7 +366,7 @@ def run_diag_annual_cycle(parameter):
     return parameter
 
 
-def run_diag_convection_onset(parameter):
+def run_diag_convection_onset(parameter: ARMDiagsParameter) -> ARMDiagsParameter:
     regions = parameter.regions
     ref_name = parameter.ref_name
     ref_path = parameter.reference_data_path
@@ -407,11 +413,11 @@ def run_diag_convection_onset(parameter):
     return parameter
 
 
-def run_diag_pdf_daily(parameter):
+def run_diag_pdf_daily(parameter: ARMDiagsParameter):
     logger.info("'run_diag_pdf_daily' is not yet implemented.")
 
 
-def run_diag(parameter):
+def run_diag(parameter: ARMDiagsParameter) -> ARMDiagsParameter:
 
     if parameter.diags_set == "annual_cycle":
         return run_diag_annual_cycle(parameter)

e3sm_diags/driver/cosp_histogram_driver.py

@@ -1,6 +1,7 @@
-from __future__ import print_function
+from __future__ import annotations
 
 import os
+from typing import TYPE_CHECKING
 
 import cdms2
 
@@ -10,6 +11,9 @@
 from e3sm_diags.metrics import corr, max_cdms, mean, min_cdms, rmse
 from e3sm_diags.plot import plot
 
+if TYPE_CHECKING:
+    from e3sm_diags.parameter.core_parameter import CoreParameter
+
 logger = custom_logger(__name__)
 
 
@@ -40,7 +44,7 @@ def create_metrics(ref, test, ref_regrid, test_regrid, diff):
     return metrics_dict
 
 
-def run_diag(parameter):
+def run_diag(parameter: CoreParameter) -> CoreParameter:
     variables = parameter.variables
     seasons = parameter.seasons
     ref_name = getattr(parameter, "ref_name", "")

e3sm_diags/driver/diurnal_cycle_driver.py

@@ -1,6 +1,7 @@
-from __future__ import print_function
+from __future__ import annotations
 
 import os
+from typing import TYPE_CHECKING
 
 import cdms2
 
@@ -11,8 +12,11 @@
 
 logger = custom_logger(__name__)
 
+if TYPE_CHECKING:
+    from e3sm_diags.parameter.diurnal_cycle_parameter import DiurnalCycleParameter
 
-def run_diag(parameter):
+
+def run_diag(parameter: DiurnalCycleParameter) -> DiurnalCycleParameter:
     variables = parameter.variables
     seasons = parameter.seasons
     ref_name = getattr(parameter, "ref_name", "")

e3sm_diags/driver/enso_diags_driver.py

@@ -1,8 +1,9 @@
-from __future__ import print_function
+from __future__ import annotations
 
 import json
 import math
 import os
+from typing import TYPE_CHECKING
 
 import cdms2
 import cdutil
@@ -16,6 +17,10 @@
 from e3sm_diags.metrics import corr, max_cdms, mean, min_cdms, rmse, std
 from e3sm_diags.plot.cartopy.enso_diags_plot import plot_map, plot_scatter
 
+if TYPE_CHECKING:
+    from e3sm_diags.parameter.enso_diags_parameter import EnsoDiagsParameter
+
+
 logger = custom_logger(__name__)
 
 
@@ -185,7 +190,7 @@ def create_metrics(ref, test, ref_regrid, test_regrid, diff):
     return metrics_dict
 
 
-def run_diag_map(parameter):
+def run_diag_map(parameter: EnsoDiagsParameter) -> EnsoDiagsParameter:
     variables = parameter.variables
     seasons = parameter.seasons
     regions = parameter.regions
@@ -378,7 +383,7 @@ def run_diag_map(parameter):
     return parameter
 
 
-def run_diag_scatter(parameter):
+def run_diag_scatter(parameter: EnsoDiagsParameter) -> EnsoDiagsParameter:
     variables = parameter.variables
     run_type = parameter.run_type
     # We will always use the same regions, so we don't do the following:
@@ -433,7 +438,7 @@ def run_diag_scatter(parameter):
     return parameter
 
 
-def run_diag(parameter):
+def run_diag(parameter: EnsoDiagsParameter) -> EnsoDiagsParameter:
     if parameter.plot_type == "map":
         return run_diag_map(parameter)
     elif parameter.plot_type == "scatter":

e3sm_diags/driver/lat_lon_driver.py

@@ -1,7 +1,8 @@
-from __future__ import print_function
+from __future__ import annotations
 
 import json
 import os
+from typing import TYPE_CHECKING
 
 import cdms2
 
@@ -13,6 +14,9 @@
 
 logger = custom_logger(__name__)
 
+if TYPE_CHECKING:
+    from e3sm_diags.parameter.core_parameter import CoreParameter
+
 
 def create_and_save_data_and_metrics(parameter, mv1_domain, mv2_domain):
     if not parameter.model_only:
@@ -107,7 +111,7 @@ def create_metrics(ref, test, ref_regrid, test_regrid, diff):
     return metrics_dict
 
 
-def run_diag(parameter):  # noqa: C901
+def run_diag(parameter: CoreParameter) -> CoreParameter:  # noqa: C901
     variables = parameter.variables
     seasons = parameter.seasons
     ref_name = getattr(parameter, "ref_name", "")