Aerosol zonal plots

lib/adf_diag.py

@@ -416,7 +416,7 @@ def call_ncrcat(cmd):
                     emsg = f"Provided baseline 'cam_hist_loc' directory '{starting_location}' "
                     emsg += "not found.  Script is ending here."
                 else:
-                    emsg = "Provided 'cam_hist_loc' directory '{starting_location}' not found."
+                    emsg = f"Provided 'cam_hist_loc' directory '{starting_location}' not found."
                     emsg += " Script is ending here."
                 # End if
 
@@ -534,6 +534,18 @@ def call_ncrcat(cmd):
             vars_to_derive = []
             # create copy of var list that can be modified for derivable variables
             diag_var_list = self.diag_var_list
+
+            # Aerosol Calcs
+            #--------------
+            #Always make sure PMID is made if aerosols are desired in config file
+            if "PMID" not in diag_var_list:
+                if any(item in res["aerosol_zonal_list"] for item in diag_var_list):
+                    diag_var_list += ["PMID"]
+            if "T" not in diag_var_list:
+                if any(item in res["aerosol_zonal_list"] for item in diag_var_list):
+                    diag_var_list += ["T"]
+            #End aerosol calcs
+
             for var in diag_var_list:
                 if var not in hist_file_var_list:
                     vres = res.get(var, {})
@@ -634,7 +646,7 @@ def call_ncrcat(cmd):
 
             if vars_to_derive:
                 self.derive_variables(
-                    vars_to_derive=vars_to_derive, ts_dir=ts_dir[case_idx]
+                    res=res, vars_to_derive=vars_to_derive, ts_dir=ts_dir[case_idx]
                 )
             # End with
 
@@ -1008,7 +1020,7 @@ def setup_run_cvdp(self):
 
     #########
 
-    def derive_variables(self, vars_to_derive=None, ts_dir=None, overwrite=None):
+    def derive_variables(self, res=None, vars_to_derive=None, ts_dir=None, overwrite=None):
         """
         Derive variables acccording to steps given here.  Since derivations will depend on the
         variable, each variable to derive will need its own set of steps below.
@@ -1017,66 +1029,132 @@ def derive_variables(self, vars_to_derive=None, ts_dir=None, overwrite=None):
 
         If the file for the derived variable exists, the kwarg `overwrite` determines
         whether to overwrite the file (true) or exit with a warning message.
+
         """
 
         for var in vars_to_derive:
-            if var == "PRECT":
-                # PRECT can be found by simply adding PRECL and PRECC
-                # grab file names for the PRECL and PRECC files from the case ts directory
-                if glob.glob(os.path.join(ts_dir, "*PRECC*")) and glob.glob(
-                    os.path.join(ts_dir, "*PRECL*")
-                ):
-                    constit_files = sorted(glob.glob(os.path.join(ts_dir, "*PREC*")))
-                else:
-                    ermsg = "PRECC and PRECL were not both present; PRECT cannot be calculated."
-                    ermsg += " Please remove PRECT from diag_var_list or find the relevant CAM files."
-                    raise FileNotFoundError(ermsg)
-                # create new file name for PRECT
-                prect_file = constit_files[0].replace("PRECC", "PRECT")
-                if Path(prect_file).is_file():
-                    if overwrite:
-                        Path(prect_file).unlink()
-                    else:
-                        print(
-                            f"[{__name__}] Warning: PRECT file was found and overwrite is False. Will use existing file."
-                        )
-                        continue
-                # append PRECC to the file containing PRECL
-                os.system(f"ncks -A -v PRECC {constit_files[0]} {constit_files[1]}")
-                # create new file with the sum of PRECC and PRECL
-                os.system(
-                    f"ncap2 -s 'PRECT=(PRECC+PRECL)' {constit_files[1]} {prect_file}"
-                )
-            if var == "RESTOM":
-                # RESTOM = FSNT-FLNT
-                # Have to be more precise than with PRECT because FSNTOA, FSTNC, etc are valid variables
-                if glob.glob(os.path.join(ts_dir, "*.FSNT.*")) and glob.glob(
-                    os.path.join(ts_dir, "*.FLNT.*")
-                ):
-                    input_files = [
-                        sorted(glob.glob(os.path.join(ts_dir, f"*.{v}.*")))
-                        for v in ["FLNT", "FSNT"]
-                    ]
-                    constit_files = []
-                    for elem in input_files:
-                        constit_files += elem
-                else:
-                    ermsg = "FSNT and FLNT were not both present; RESTOM cannot be calculated."
-                    ermsg += " Please remove RESTOM from diag_var_list or find the relevant CAM files."
-                    raise FileNotFoundError(ermsg)
-                # create new file name for RESTOM
-                derived_file = constit_files[0].replace("FLNT", "RESTOM")
+            print(f"\t - deriving time series for {var}")
+
+            #Check whether there are parts to derive from and if there is an associated equation
+            vres = res.get(var, {})
+            if "derivable_from" in vres:
+                constit_list = vres['derivable_from']
+            else:
+                print("WARNING: No constituents listed in defaults config file, moving on")
+                continue
+
+            #Grab all required time series files for derived var
+            constit_files = []
+            for constit in constit_list:
+                if glob.glob(os.path.join(ts_dir, f"*.{constit}.*.nc")):
+                    constit_files.append(glob.glob(os.path.join(ts_dir, f"*.{constit}.*"))[0])
+
+            #Check if all the constituent files were found
+            if len(constit_files) != len(constit_list):
+                ermsg = f"Not all constituent files present; {var} cannot be calculated."
+                ermsg += f" Please remove {var} from diag_var_list or find the relevant CAM files."
+                print(ermsg)
+            else:
+                #Open a new dataset with all the constituent files/variables
+                ds = xr.open_mfdataset(constit_files)
+
+                # create new file name for derived variable
+                derived_file = constit_files[0].replace(constit_list[0], var)
+
+                #Check if clobber is true for file
                 if Path(derived_file).is_file():
                     if overwrite:
                         Path(derived_file).unlink()
                     else:
                         print(
-                            f"[{__name__}] Warning: RESTOM file was found and overwrite is False. Will use existing file."
+                            f"[{__name__}] Warning: '{var}' file was found and overwrite is False. Will use existing file."
                         )
                         continue
-                # append FSNT to the file containing FLNT
-                os.system(f"ncks -A -v FLNT {constit_files[0]} {constit_files[1]}")
-                # create new file with the difference of FLNT and FSNT
-                os.system(
-                    f"ncap2 -s 'RESTOM=(FSNT-FLNT)' {constit_files[1]} {derived_file}"
-                )
+
+                #NOTE: this will need to be changed when derived equations are more complex! - JR
+                if var == "RESTOM":
+                    der_val = ds["FSNT"]-ds["FLNT"]
+                else:
+                    #Loop through all constituents and sum
+                    der_val = 0
+                    for v in constit_list:
+                        der_val += ds[v]
+
+                #Set derived variable name and add to dataset
+                der_val.name = var
+                ds[var] = der_val
+
+                #Aerosol Calculations - used for zonal plots
+                #These will be multiplied by rho (density of dry air)
+                ds_pmid_done = False
+                ds_t_done = False
+                if var in res["aerosol_zonal_list"]:
+
+                    #Only calculate once for all aerosol vars
+                    if not ds_pmid_done:
+                        ds_pmid = _load_dataset(glob.glob(os.path.join(ts_dir, "*.PMID.*"))[0])
+                        ds_pmid_done = True
+                        if not ds_pmid:
+                            errmsg = f"Missing necessary files for dry air density (rho) calculation.\n"
+                            errmsg += "Please make sure 'PMID' is in the CAM run for aerosol calculations"
+                            print(errmsg)
+                            continue
+                    if not ds_t_done:
+                        ds_t = _load_dataset(glob.glob(os.path.join(ts_dir, "*.T.*"))[0])
+                        ds_t_done = True
+                        if not ds_t:
+                            errmsg = f"Missing necessary files for dry air density (rho) calculation.\n"
+                            errmsg += "Please make sure 'T' is in the CAM run for aerosol calculations"
+                            print(errmsg)
+                            continue
+
+                    #Multiply aerosol by dry air density (rho): (P/Rd*T)
+                    ds[var] = ds[var]*(ds_pmid["PMID"]/(res["Rgas"]*ds_t["T"]))
+
+                    #Sulfate conversion factor
+                    if var == "SO4":
+                        ds[var] = ds[var]*(96./115.)
+
+                #Drop all constituents from final saved dataset
+                #These are not necessary because they have their own time series files
+                ds_final = ds.drop_vars(constit_list)
+                ds_final.to_netcdf(derived_file, unlimited_dims='time', mode='w')
+
+########
+
+#Helper Function(s)
+def _load_dataset(fils):
+    """
+    This method exists to get an xarray Dataset from input file information that can be passed into the plotting methods.
+
+    Parameters
+    ----------
+    fils : list
+        strings or paths to input file(s)
+
+    Returns
+    -------
+    xr.Dataset
+
+    Notes
+    -----
+    When just one entry is provided, use `open_dataset`, otherwise `open_mfdatset`
+    """
+    import warnings  # use to warn user about missing files.
+
+    #Format warning messages:
+    def my_formatwarning(msg, *args, **kwargs):
+        """Issue `msg` as warning."""
+        return str(msg) + '\n'
+    warnings.formatwarning = my_formatwarning
+
+    if len(fils) == 0:
+        warnings.warn("Input file list is empty.")
+        return None
+    elif len(fils) > 1:
+        return xr.open_mfdataset(fils, combine='by_coords')
+    else:
+        return xr.open_dataset(fils[0])
+    #End if
+#End def
+########

lib/adf_variable_defaults.yaml

@@ -54,6 +54,9 @@
 # derivable_from -> If not present in the available output files, the variable can be derived from
 #                   other variables that are present (e.g. PRECT can be derived from PRECC and PRECL),
 #                   which are specified in this list
+#                   NOTE: this is not very flexible at the moment! It can only handle variables that
+#                         are sums of the constituents. Futher flexibility is being explored.
+#
 #
 # Final Note:  Please do not modify this file unless you plan to push your changes back to the ADF repo.
 #              If you would like to modify this file for your personal ADF runs then it is recommended
@@ -65,10 +68,16 @@
 #+++++++++++++
 # Available ADF Default Plot Types
 #+++++++++++++
-
 default_ptypes: ["Tables","LatLon","LatLon_Vector","Zonal","Meridional",
                   "NHPolar","SHPolar","Special"]
 
+#+++++++++++++
+# Constants
+#+++++++++++++
+
+#Dry Air Gas Constant:
+Rgas: 287.04 #[J/K/Kg]=8.314/0.028965
+
 #+++++++++++++
 # Category: Microphysics
 #+++++++++++++
@@ -125,6 +134,9 @@ FICE:
 # Category: Aerosols
 #+++++++++++
 
+#List of zonal areosols
+aerosol_zonal_list: ["BC","POM","SO4","SOA","NH4HSO4","DUST","SeaSalt"]
+
 AODDUST:
   category: "Aerosols"
   colormap: "Oranges"
@@ -185,6 +197,87 @@ SO2:
 SOAG:
   category: "Aerosols"
 
+BC:
+  colormap: "RdBu_r"
+  diff_colormap: "BrBG"
+  scale_factor: 1000000000
+  add_offset: 0
+  new_unit: '$\mu$g/m3'
+  mpl:
+    colorbar:
+      label : '$\mu$g/m3'
+  category: "Aerosols"
+  derivable_from: ["bc_a1", "bc_a4"]
+
+POM:
+  colormap: "RdBu_r"
+  diff_colormap: "BrBG"
+  scale_factor: 1000000000
+  add_offset: 0
+  new_unit: '$\mu$g/m3'
+  mpl:
+    colorbar:
+      label : '$\mu$g/m3'
+  category: "Aerosols"
+  derivable_from: ["pom_a1", "pom_a4"]
+
+SO4:
+  colormap: "RdBu_r"
+  diff_colormap: "BrBG"
+  scale_factor: 1000000000
+  add_offset: 0
+  new_unit: '$\mu$g/m3'
+  mpl:
+    colorbar:
+      label : '$\mu$g/m3'
+  category: "Aerosols"
+  derivable_from: ["so4_a1", "so4_a2", "so4_a3", "so4_a5"]
+
+SOA:
+  colormap: "RdBu_r"
+  diff_colormap: "BrBG"
+  scale_factor: 1000000000
+  add_offset: 0
+  new_unit: '$\mu$g/m3'
+  mpl:
+    colorbar:
+      label : '$\mu$g/m3'
+  category: "Aerosols"
+  derivable_from: ["soa_a1", "soa_a2"]
+
+DUST:
+  colormap: "RdBu_r"
+  contour_levels: [0,0.1,0.25,0.4,0.6,0.8,1.4,2,3,4,8,12,30,48,114,180]
+  non_linear: True
+  diff_colormap: "BrBG"
+  scale_factor: 1000000000
+  add_offset: 0
+  new_unit: '$\mu$g/m3'
+  mpl:
+    colorbar:
+      label : '$\mu$g/m3'
+  category: "Aerosols"
+  derivable_from: ["dst_a1", "dst_a2", "dst_a3"]
+
+SeaSalt:
+  colormap: "RdBu_r"
+  contour_levels: [0,0.05,0.075,0.2,0.3,0.4,0.7,1,1.5,2,4,6,15,24,57,90]
+  non_linear: True
+  diff_colormap: "BrBG"
+  scale_factor: 1000000000
+  add_offset: 0
+  new_unit: '$\mu$g/m3'
+  mpl:
+    colorbar:
+      label : '$\mu$g/m3'
+      ticks: [0.05,0.2,0.4,1,2,6,24,90]
+    diff_colorbar:
+      label : '$\mu$g/m3'
+      ticks: [-10,8,6,4,2,0,-2,-4,-6,-8,-10]
+  category: "Aerosols"
+  derivable_from: ["ncl_a1", "ncl_a2", "ncl_a3"]
+
+
 
 #+++++++++++++++++
 # Category: Budget
@@ -501,7 +594,6 @@ QFLX:
 # Category: Surface variables
 #+++++++++++++++++
 
-
 PBLH:
   category: "Surface variables"
   obs_file: "PBLH_ERA5_monthly_climo_197901-202112.nc"
@@ -1097,4 +1189,4 @@ utendwtem:
   obs_var_name: "utendwtem"
 
 #-----------
-#End of File
+#End of File