Merge pull request #1060 from PCMDI/feature/lee1043-mov-modularize

Feature/lee1043 mov modularize
PCMDI · Feb 22, 2024 · 84e5f9d · 84e5f9d
2 parents f2c1593 + 6bc7232
commit 84e5f9d
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 19 deletions.
diff --git a/pcmdi_metrics/variability_mode/lib/eof_analysis.py b/pcmdi_metrics/variability_mode/lib/eof_analysis.py
@@ -396,10 +396,10 @@ def get_anomaly_timeseries(ds: xr.Dataset, data_var: str, season: str) -> xr.Dat
         raise TypeError(
             "The first parameter of get_anomaly_timeseries must be an xarray Dataset"
         )
-    # Get anomaly field
+    # Get anomaly field by removing annual cycle
+    ds_anomaly = ds.temporal.departures(data_var, freq="month", weighted=True)
+    # Get temporal average if needed
     if season == "yearly":
-        # remove seasonal cycle
-        ds_anomaly = ds.temporal.departures(data_var, freq="month", weighted=True)
         # yearly time series
         ds_anomaly = ds_anomaly.temporal.group_average(
             data_var, freq="year", weighted=True
@@ -410,27 +410,26 @@ def get_anomaly_timeseries(ds: xr.Dataset, data_var: str, season: str) -> xr.Dat
         ds_ave = ds_anomaly.temporal.average(data_var)
         # anomaly
         ds_anomaly[data_var] = ds_anomaly[data_var] - ds_ave[data_var]
-    else:
-        # Remove annual cycle
-        ds_anomaly = ds.temporal.departures(data_var, freq="month", weighted=True)
-        if season != "monthly":
-            ds_anomaly_all_seasons = ds_anomaly.temporal.departures(
-                data_var,
-                freq="season",
-                weighted=True,
-                season_config={"dec_mode": "DJF", "drop_incomplete_djf": True},
-            )
-            ds_anomaly = select_by_season(ds_anomaly_all_seasons, season)
+    elif season.upper() in ["DJF", "MAM", "JJA", "SON"]:
+        ds_anomaly_all_seasons = ds_anomaly.temporal.departures(
+            data_var,
+            freq="season",
+            weighted=True,
+            season_config={"dec_mode": "DJF", "drop_incomplete_djf": True},
+        )
+        ds_anomaly = select_by_season(ds_anomaly_all_seasons, season)
     # return result
     return ds_anomaly
 
 
 def select_by_season(ds: xr.Dataset, season: str) -> xr.Dataset:
     time_key = get_time_key(ds)
-    lat_bnds_key = get_latitude_bounds_key(ds)
-    lon_bnds_key = get_longitude_bounds_key(ds)
     ds_subset = ds.where(ds[time_key].dt.season == season, drop=True)
     # Preserve original spatial bounds info
+    # Extract original bounds
+    lat_bnds_key = get_latitude_bounds_key(ds)
+    lon_bnds_key = get_longitude_bounds_key(ds)
+    # Assign back to the new dataset
     ds_subset[lat_bnds_key] = get_latitude_bounds(ds)
     ds_subset[lon_bnds_key] = get_longitude_bounds(ds)
     return ds_subset

diff --git a/pcmdi_metrics/variability_mode/variability_modes_driver.py b/pcmdi_metrics/variability_mode/variability_modes_driver.py
@@ -304,10 +304,10 @@
         debug=debug,
     )
 
-    # Save global grid information for regrid below
+    # Get global grid information for later use: regrid
     ref_grid_global = get_grid(obs_timeseries)
 
-    # Declare dictionary variables to keep information from observation
+    # Set dictionary variables to keep information from observation in the memory during the season and model loop
     eof_obs = {}
     pc_obs = {}
     frac_obs = {}
@@ -317,7 +317,7 @@
     stdv_pc_obs = {}
     obs_timeseries_season_dict = {}
 
-    # Dictonary for json archive
+    # Set dictonary for json archive
     if "obs" not in result_dict["REF"]:
         result_dict["REF"]["obs"] = {}
     if "defaultReference" not in result_dict["REF"]["obs"]: