Merge pull request #982 from PCMDI/896_lee_meanclim_workflow

Improve mean climate workflow
PCMDI · Sep 25, 2023 · 93003a7 · 93003a7
2 parents f9dd510 + e6e7154
commit 93003a7
Show file tree

Hide file tree

Showing 7 changed files with 79 additions and 50 deletions.
diff --git a/pcmdi_metrics/mean_climate/lib/create_mean_climate_parser.py b/pcmdi_metrics/mean_climate/lib/create_mean_climate_parser.py
@@ -266,4 +266,13 @@ def create_mean_climate_parser():
         required=False,
     )
 
+    parser.add_argument(
+        "--parallel",
+        type=bool,
+        dest="parallel",
+        default=False,
+        help="Option for running code in parallel mode: True / False (default)",
+        required=False,
+    )
+
     return parser
diff --git a/pcmdi_metrics/mean_climate/lib/mean_climate_metrics_to_json.py b/pcmdi_metrics/mean_climate/lib/mean_climate_metrics_to_json.py
@@ -21,10 +21,11 @@ def mean_climate_metrics_to_json(
         for m in models_in_dict:
             if m == model:
                 for ref in list(json_dict["RESULTS"][m].keys()):
-                    runs_in_model_dict = list(json_dict["RESULTS"][m][ref].keys())
-                    for r in runs_in_model_dict:
-                        if (r != run) and (run is not None):
-                            del json_dict["RESULTS"][m][ref][r]
+                    if ref != "units":
+                        runs_in_model_dict = list(json_dict["RESULTS"][m][ref].keys())
+                        for r in runs_in_model_dict:
+                            if (r != run) and (run is not None):
+                                del json_dict["RESULTS"][m][ref][r]
             else:
                 del json_dict["RESULTS"][m]
     # Write selected dict to JSON

diff --git a/pcmdi_metrics/mean_climate/mean_climate_driver.py b/pcmdi_metrics/mean_climate/mean_climate_driver.py
@@ -50,6 +50,7 @@
 custom_obs = parameter.custom_observations
 debug = parameter.debug
 cmec = parameter.cmec
+parallel = parameter.parallel
 
 if metrics_output_path is not None:
     metrics_output_path = parameter.metrics_output_path.replace('%(case_id)', case_id)
@@ -146,6 +147,9 @@
 # -------------
 # variable loop
 # -------------
+if isinstance(vars, str):
+    vars = [vars]
+
 for var in vars:
 
     if '_' in var or '-' in var:
@@ -283,39 +287,44 @@
 
                             if debug:
                                 print('ds_test_tmp:', ds_test_tmp)
-                                ds_test_dict[region].to_netcdf('_'.join([var, 'model', model, run, region + '.nc']))
+                                ds_test_dict[region].to_netcdf('_'.join([var, 'model', model, run, region, case_id + '.nc']))
                                 if model == test_data_set[0] and run == realizations[0]:
                                     ds_ref_dict[region].to_netcdf('_'.join([var, 'ref', region + '.nc']))
 
                             # compute metrics
                             print('compute metrics start')
                             result_dict["RESULTS"][model][ref][run][region] = compute_metrics(varname, ds_test_dict[region], ds_ref_dict[region], debug=debug)
 
-                        # write individual JSON
-                        # --- single simulation, obs (need to accumulate later) / single variable
-                        json_filename_tmp = "_".join([model, var, target_grid, regrid_tool, "metrics", ref])
-                        mean_climate_metrics_to_json(
-                            os.path.join(metrics_output_path, var),
-                            json_filename_tmp,
-                            result_dict,
-                            model=model,
-                            run=run,
-                            cmec_flag=cmec,
-                            debug=debug
-                        )
+                            # write individual JSON
+                            # --- single simulation, obs (need to accumulate later) / single variable
+                            json_filename_tmp = "_".join([var, model, run, target_grid, regrid_tool, "metrics", ref, case_id])
+                            mean_climate_metrics_to_json(
+                                os.path.join(metrics_output_path, var),
+                                json_filename_tmp,
+                                result_dict,
+                                model=model,
+                                run=run,
+                                cmec_flag=cmec,
+                                debug=debug
+                            )
 
                     except Exception as e:
                         if debug:
                             raise
                         print('error occured for ', model, run)
                         print(e)
 
-    # write collective JSON --- all models / all obs / single variable
-    json_filename = "_".join([var, target_grid, regrid_tool, "metrics"])
-    mean_climate_metrics_to_json(
-        metrics_output_path,
-        json_filename,
-        result_dict,
-        cmec_flag=cmec,
-    )
-    print('pmp mean clim driver completed')
+    # ========================================================================
+    # Dictionary to JSON: collective JSON at the end of model_realization loop
+    # ------------------------------------------------------------------------
+    if not parallel:
+        # write collective JSON --- all models / all obs / single variable
+        json_filename = "_".join([var, target_grid, regrid_tool, "metrics", case_id])
+        mean_climate_metrics_to_json(
+            metrics_output_path,
+            json_filename,
+            result_dict,
+            cmec_flag=cmec,
+        )
+
+print('pmp mean clim driver completed')
diff --git a/pcmdi_metrics/mjo/lib/mjo_metric_calc.py b/pcmdi_metrics/mjo/lib/mjo_metric_calc.py
@@ -167,21 +167,35 @@ def mjo_metric_ewr_calculation(
     if plot:
         os.makedirs(outdir(output_type="graphics"), exist_ok=True)
         fout = os.path.join(outdir(output_type="graphics"), output_filename)
-        title = (
-            mip.upper()
-            + ": "
-            + model
-            + " ("
-            + run
-            + ") \n"
-            + var.capitalize()
-            + ", "
-            + season
-            + " "
-            + str(startYear)
-            + "-"
-            + str(endYear)
-        )
+        if model == 'obs':
+            title = (
+                " OBS ("
+                + run
+                + ") \n"
+                + var.capitalize()
+                + ", "
+                + season
+                + " "
+                + str(startYear)
+                + "-"
+                + str(endYear)
+            )
+        else:
+            title = (
+                mip.upper()
+                + ": "
+                + model
+                + " ("
+                + run
+                + ") \n"
+                + var.capitalize()
+                + ", "
+                + season
+                + " "
+                + str(startYear)
+                + "-"
+                + str(endYear)
+            )
         if cmmGrid:
             title += ", common grid (2.5x2.5deg)"
         plot_power(OEE, title, fout, ewr)

diff --git a/pcmdi_metrics/mjo/param/myParam_mjo.py b/pcmdi_metrics/mjo/param/myParam_mjo.py
@@ -30,7 +30,7 @@ def find_latest(path):
 # Observation
 # -------------------------------------------------
 reference_data_name = "GPCP-1-3"
-reference_data_path = "/p/user_pub/PCMDIobs/PCMDIobs2/atmos/day/pr/GPCP-1-3/gn/v20200707/pr_day_GPCP-1-3_BE_gn_v20200707_19961002-20170101.nc"  # noqa
+reference_data_path = "/p/user_pub/PCMDIobs/obs4MIPs_legacy/PCMDIobs2/atmos/day/pr/GPCP-1-3/gn/v20200924/pr_day_GPCP-1-3_BE_gn_v20200924_19961002-20170101.nc"  # noqa
 
 varOBS = "pr"
 ObsUnitsAdjust = (True, "multiply", 86400.0, "mm d-1")  # kg m-2 s-1 to mm day-1

diff --git a/pcmdi_metrics/mjo/scripts/parallel_driver.py b/pcmdi_metrics/mjo/scripts/parallel_driver.py
@@ -97,10 +97,10 @@
 # =================================================
 # Generates list of command
 # -------------------------------------------------
-param_file = "../doc/myParam_mjo.py"
+param_file = "../param/myParam_mjo.py"
 
 if debug:
-    param_file = "../doc/myParam_test.py"
+    param_file = "../param/myParam_test.py"
     print("number of models (debug mode):", len(models))
 
 cmds_list = list()
@@ -148,7 +148,6 @@
     for r, run in enumerate(runs_list):
         # command line for queue
         cmd = [
-            "python",
             "mjo_metrics_driver.py",
             "-p",
             param_file,

diff --git a/pcmdi_metrics/mjo/scripts/run.sh b/pcmdi_metrics/mjo/scripts/run.sh
@@ -1,9 +1,6 @@
 #!/bin/sh
 set -a
 
-# grim: pmp_nightly_20190628
-# gates: cdat82_20191107_py37, pmp_nightly_20190912
-
 #parallel=no
 parallel=yes
 
@@ -17,15 +14,15 @@ mkdir -p log
 if [ $parallel == no ]; then
     echo 'parallel no'
     for mip in $mips; do
-        python -u mjo_metrics_driver.py -p ../doc/myParam_mjo.py --mip ${mip} >& log/log.${mip}.txt &
+        mjo_metrics_driver.py -p ../param/myParam_mjo.py --mip ${mip} >& log/log.${mip}.txt &
         disown
     done
 elif [ $parallel == yes ]; then
     echo 'parallel yes'
     modnames="all"
     realization="all"
     for mip in $mips; do
-        python -u ./parallel_driver.py -p ../doc/myParam_mjo.py --mip ${mip} --num_workers $num_workers --modnames $modnames --realization $realization  >& log/log.parallel.${mip}.txt &
+        python -u ./parallel_driver.py -p ../param/myParam_mjo.py --mip ${mip} --num_workers $num_workers --modnames $modnames --realization $realization  >& log/log.parallel.${mip}.txt &
         disown
     done
 fi