Merge PR #402 (Optimize get_diff_of_diffs and add_lumped_species_to_dataset)

This merge brings PR #402 (Optimize get_diff_of_diffs and add_
lumped_species_to_dataset, by @yantosca) into the GCPy 1.7.0
development stream.

PR #402 does the following:

1. In function "make_benchmark_conc_plots":
   - Optimizes the algorithm that computes the list of variables
     for each benchmark category
   - Passes truncated datasets to compare_single_level and compare_zonal_mean.

2. In routine "add_lumped_species_to_dataset"
   - Vectorizes the summing of species
   - Merges new species into the Dataset in a single operation.

3. In routine "get_diff_of_diffs":
   - Optimized the algorithm to compute diff-of-diffs.
   - Aligns cubed-sphere grids before computing diff-of-diffs

4. In routines "create_regridders", "compare_single_level", and
   "compare_zonal_mean":
   - Delete regridder objects once we no longer need them
   - Manually call gc.collect() to force garbage collection.

Signed-off-by: Bob Yantosca <yantosca@seas.harvard.edu>

Author: yantosca

CHANGELOG.md

@@ -19,6 +19,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Added functions `get_molwt_from_metadata` and `read_species_metadata` to `gcpy/util.py`
 - Added function `get_species_database_files` to `gcpy/benchmark/modules/benchmark_utils.py`
 - Added constant `SPECIES_DATABASE` to `gcpy/benchmark/modules/benchmark_utils.py`
+- Added manual garbage collection in `create_regridders`, `compare_single_level`, and `compare_zonal_mean` functions.
+- Added helpful tips to the `gcpy/benchmark/benchmark.slurm.sh` script
 
 ### Changed
 - Modified criteria for terminating read of log files in `benchmark_scrape_gcclassic_timers.py` to avoid being spoofed by  output that is attached by Intel VTune
@@ -44,6 +46,9 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Replaced `get_species_database_dir` with `get_species_database_files` in `gcpy/benchmark/modules/benchmark_funcs.py`
 - Updated `gcpy/benchmark/modules/benchmark_scrape_gchp_timers.py` to look for GCHP timers in `allPEs.log` if not found in the log file
 - Updated routine `make_benchmark_aerosol_tables` to include all dust species in the aerosol burdens table
+- Optimized function `get_diff_of_diffs` (in `gcpy/util.py`) for performance
+- Optimized function `add_lumped_species_to_dataset` (in `gcpy/benchmark/modules/benchmark_utils.py`) for performance
+- Optimized the algorithm to generate `varlist` in `make_benchmark_conc_plots`.  Also truncated datasets to only contain varibales in `varlist`.
 
 ### Fixed
 - Fixed grid area calculation scripts of `grid_area` in `gcpy/gcpy/cstools.py`
@@ -59,7 +64,6 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Removed `.load()` statements from xarray Datasets to improve performance
 - Removed `paths:spcdb_dir` YAML tag in benchmark configuration files
 - Removed `st_Ox` from `benchmark_categories.yml`; this species is no longer used in TransportTracers simulations
-- Removed special data handling for files generated with MAPL versions prior to 1.0.0 in function `get_diff_of_diffs` (located in `gcpy/util.py`)
 
 ## [1.6.2] - 2025-06-12
 ### Added

gcpy/benchmark/benchmark_slurm.sh

@@ -2,9 +2,9 @@
 
 #SBATCH -c 8
 #SBATCH -N 1
-#SBATCH -t 0-4:00
-#SBATCH -p seas_compute,shared
-#SBATCH --mem=100000
+#SBATCH -t 0-6:00
+#SBATCH -p sapphire,huce_cascade,seas_compute,shared
+#SBATCH --mem=180000
 #SBATCH --mail-type=END
 
 #============================================================================
@@ -13,7 +13,17 @@
 #
 # You can modify the SLURM parameters above for your setup.
 #
-# Tip: Using less cores can reduce the amount of memory required.
+# Tips:
+# -----
+# (1) Use fewer cores to reduce the memory footprint. This may prevent
+#     your job from running out of memory.  Python under Linux seems
+#     to have an issue where not all memory is released back to the OS.
+#
+# (2) We recommend that you generate only one benchmark comparison
+#     (GCC vs GCC, GCHP vs GCC, GCHP vs GCC, or diff of diffs)
+#     at a time.  Otherwise your job will probaly run out of memory.
+#
+# (3) For diff-of-diffs plots, we recommend using 6 cores.
 #============================================================================
 
 # Apply all bash initialization settings

gcpy/benchmark/modules/benchmark_funcs.py

@@ -1277,7 +1277,7 @@ def make_benchmark_conc_plots(
         devds = add_lumped_species_to_dataset(devds)
 
         if diff_of_diffs:
-            print("-->Adding lumped species to dev datasets")
+            print("-->Adding lumped species to second ref and dev datasets")
             second_refds = add_lumped_species_to_dataset(second_refds)
             second_devds = add_lumped_species_to_dataset(second_devds)
 
@@ -1340,19 +1340,34 @@ def createplots(filecat):
             if not os.path.isdir(catdir):
                 os.mkdir(catdir)
 
-        varlist = []
-        warninglist = []
-        for subcat in catdict[filecat]:
-            for spc in catdict[filecat][subcat]:
-                varname = coll_prefix + spc
-                if varname not in refds.data_vars or \
-                   varname not in devds.data_vars:
-                    warninglist.append(varname)
-                    continue
-                varlist.append(varname)
-        if warninglist:
-            msg = f"\n\nWarning: variables in {filecat} category not in dataset: {warninglist}"
-            print(msg)
+        # Get the list of variables in both Ref and Dev for each category
+        # (this is computationally efficient)
+        ref_vars = set(refds.data_vars)
+        dev_vars = set(devds.data_vars)
+        candidates = [
+            coll_prefix + spc
+            for subcat in catdict[filecat]
+            for spc in catdict[filecat][subcat]
+        ]
+        varlist = \
+            [var for var in candidates \
+             if var in ref_vars and var in dev_vars
+        ]
+        warninglist = \
+            [var for var in candidates \
+             if var not in ref_vars or var not in dev_vars
+        ]
+
+        # Create new datasets containing only the variables for a
+        # given category, as this will optimize performance.
+        refds_cat = refds[varlist]
+        devds_cat = devds[varlist]
+        second_refds_cat = None
+        if second_refds is not None:
+            second_refds_cat = second_refds[varlist]
+        second_devds_cat = None
+        if second_devds is not None:
+            second_devds_cat = second_devds[varlist]
 
         # -----------------------
         # Surface plots
@@ -1373,9 +1388,9 @@ def createplots(filecat):
 
             diff_sfc = []
             compare_single_level(
-                refds,
+                refds_cat,
                 refstr,
-                devds,
+                devds_cat,
                 devstr,
                 varlist=varlist,
                 ilev=0,
@@ -1390,8 +1405,8 @@ def createplots(filecat):
                 sigdiff_list=diff_sfc,
                 weightsdir=weightsdir,
                 convert_to_ugm3=convert_to_ugm3,
-                second_ref=second_refds,
-                second_dev=second_devds,
+                second_ref=second_refds_cat,
+                second_dev=second_devds_cat,
                 n_job=n_job,
                 spcdb_files=spcdb_files,
             )
@@ -1412,7 +1427,8 @@ def createplots(filecat):
 
             if subdst is not None:
                 pdfname = os.path.join(
-                    catdir, f"{filecat}_500hPa_{subdst}.pdf"
+                    catdir,
+                    f"{filecat}_500hPa_{subdst}.pdf"
                 )
             else:
                 pdfname = os.path.join(
@@ -1422,9 +1438,9 @@ def createplots(filecat):
 
             diff_500 = []
             compare_single_level(
-                refds,
+                refds_cat,
                 refstr,
-                devds,
+                devds_cat,
                 devstr,
                 varlist=varlist,
                 ilev=22,
@@ -1439,8 +1455,8 @@ def createplots(filecat):
                 sigdiff_list=diff_500,
                 weightsdir=weightsdir,
                 convert_to_ugm3=convert_to_ugm3,
-                second_ref=second_refds,
-                second_dev=second_devds,
+                second_ref=second_refds_cat,
+                second_dev=second_devds_cat,
                 n_job=n_job,
                 spcdb_files=spcdb_files
             )
@@ -1473,9 +1489,9 @@ def createplots(filecat):
 
             diff_zm = []
             compare_zonal_mean(
-                refds,
+                refds_cat,
                 refstr,
-                devds,
+                devds_cat,
                 devstr,
                 varlist=varlist,
                 refmet=refmetds,
@@ -1488,8 +1504,8 @@ def createplots(filecat):
                 sigdiff_list=diff_zm,
                 weightsdir=weightsdir,
                 convert_to_ugm3=convert_to_ugm3,
-                second_ref=second_refds,
-                second_dev=second_devds,
+                second_ref=second_refds_cat,
+                second_dev=second_devds_cat,
                 n_job=n_job,
                 spcdb_files=spcdb_files
             )
@@ -1518,9 +1534,9 @@ def createplots(filecat):
                 )
 
             compare_zonal_mean(
-                refds,
+                refds_cat,
                 refstr,
-                devds,
+                devds_cat,
                 devstr,
                 varlist=varlist,
                 refmet=refmetds,
@@ -1534,8 +1550,8 @@ def createplots(filecat):
                 normalize_by_area=normalize_by_area,
                 convert_to_ugm3=convert_to_ugm3,
                 weightsdir=weightsdir,
-                second_ref=second_refds,
-                second_dev=second_devds,
+                second_ref=second_refds_cat,
+                second_dev=second_devds_cat,
                 n_job=n_job,
                 spcdb_files=spcdb_files
             )

gcpy/benchmark/modules/benchmark_utils.py

@@ -366,7 +366,8 @@ def add_lumped_species_to_dataset(
     as a dictionary or a path to a yaml file. If neither is passed then
     the lumped species yaml file stored in gcpy is used. This file is
     customized for use with benchmark simuation SpeciesConc diagnostic
-    collection output.
+    collection output.  The algorithm has been optimized by AI to
+    improve performance.
 
     Args
     dset      : xr.Dataset : Data prior to adding lumped species
@@ -378,11 +379,22 @@ def add_lumped_species_to_dataset(
 
     Returns
     dset      : xr.Dataset : Original species plus added lumped species
+
+    Remarks
+    -------
+    Key Improvements:
+    1. Vectorized summation: Uses sum(to_sum) instead of incremental +=
+    2. Lazy evaluation: Operations remain lazy until actual computation
+    3. Single merge: Uses .assign() instead of merging many DataArrays
+    4. Cleaner logic: More Pythonic dictionary iteration
+
+    Performance Impact:
+    Original: O(n_lumped × n_constituents) individual array operations
+    Optimized: O(n_lumped) vectorized operations
     """
 
-    # Default is to add all benchmark lumped species.
-    # Can overwrite by passing a dictionary
-    # or a yaml file path containing one
+    # Default is to add all benchmark lumped species.  Can overwrite
+    # by passing a dictionary or a yaml file path containing one.
     assert not (
         lspc_dict is not None and lspc_yaml != ""
     ), "Cannot pass both lspc_dict and lspc_yaml. Choose one only."
@@ -394,67 +406,64 @@ def add_lumped_species_to_dataset(
     # Make sure attributes are transferred when copying dataset / dataarrays
     with xr.set_options(keep_attrs=True):
 
-        # Get a dummy DataArray to use for initialization
-        dummy_darr = None
-        for var in dset.data_vars:
-            if prefix in var or prefix.replace("VV", "") in var:
-                dummy_darr = dset[var]
-                dummy_type = dummy_darr.dtype
-                dummy_shape = dummy_darr.shape
-                break
-        if dummy_darr is None:
-            msg = "Invalid prefix: " + prefix
-            raise ValueError(msg)
-
-        # Create a list with a copy of the dummy DataArray object
-        n_lumped_spc = len(lspc_dict)
-        lumped_spc = [None] * n_lumped_spc
-        for var, spcname in enumerate(lspc_dict):
-            lumped_spc[var] = dummy_darr.copy(deep=False)
-            lumped_spc[var].name = prefix + spcname
-            lumped_spc[var].values = np.full(dummy_shape, 0.0, dtype=dummy_type)
-
-        # Loop over lumped species list
-        for var, lspc in enumerate(lumped_spc):
-
-            # Search key for lspc_dict is lspc.name minus the prefix
-            cidx = lspc.name.find("_")
-            key = lspc.name[cidx+1:]
+        # Dictionary to store new lumped species
+        new_vars = {}
+
+        # Loop over lumped species
+        for lspc_name, constituents in lspc_dict.items():
+            full_name = prefix + lspc_name
 
             # Check if overlap with existing species
-            if lspc.name in dset.data_vars and overwrite:
-                dset.drop(lspc.name)
-            else:
-                assert(lspc.name not in dset.data_vars), \
-                    f"{lspc.name} already in dataset. To overwrite pass overwrite=True."
+            if full_name in dset.data_vars:
+                if overwrite:
+                    if verbose:
+                        print(f"Overwriting existing {full_name}")
+                else:
+                    raise ValueError(
+                        f"{full_name} already in dataset. "
+                        "To overwrite pass overwrite=True."
+                    )
 
-            # Verbose prints
             if verbose:
-                print(f"Creating {lspc.name}")
+                print(f"Creating {full_name}")
 
-            # Loop over and sum constituent species values
-            num_spc = 0
-            for _, spcname in enumerate(lspc_dict[key]):
+            # Collect all constituent species that exist
+            to_sum = []
+            for spcname, scale_factor in constituents.items():
                 varname = prefix + spcname
                 if varname not in dset.data_vars:
                     if verbose:
-                        print(f"Warning: {varname} needed for {lspc_dict[key][spcname]} not in dataset")
+                        print(f"Warning: {varname} needed for {scale_factor} not in dataset")
                     continue
+
                 if verbose:
-                    print(f" -> adding {varname} with scale {lspc_dict[key][spcname]}")
-                lspc.values += dset[varname].values * lspc_dict[key][spcname]
-                num_spc += 1
+                    print(f" -> adding {varname} with scale {scale_factor}")
+
+                # Build list of scaled species (lazy operations)
+                to_sum.append(dset[varname] * scale_factor)
 
-            # Replace values with NaN if no species found in dataset
-            if num_spc == 0:
+            # Vectorized sum of all constituents at once
+            if len(to_sum) > 0:
+                new_vars[full_name] = sum(to_sum)
+            else:
+                # Create NaN array matching first species shape
                 if verbose:
                     print("No constituent species found! Setting to NaN.")
-                lspc.values = np.full(lspc.shape, np.nan)
-
-        # Insert the DataSet into the list of DataArrays
-        # so that we can only do the merge operation once
-        lumped_spc.insert(0, dset)
-        dset = xr.merge(lumped_spc)
+                template_var = next(
+                    (var for key, var in dset.data_vars.items()
+                     if prefix in key or prefix.replace("VV", "") in key),
+                    None
+                )
+                if template_var is not None:
+                    new_vars[full_name] = template_var.copy(deep=False) * np.nan
+
+        # Single merge operation
+        if overwrite:
+            dset = dset.drop_vars(
+                [key for key in new_vars.keys() if key in dset.data_vars],
+                errors='ignore'
+            )
+        dset = dset.assign(new_vars)
 
     return dset