ENH: radiation scheme family and microphysics scheme model attrs (col…

…umncolab#126)

* ENH: radiation scheme family and microphysics scheme model attrs

* MNT: refactoring of terminal velocity parameterization selection

* TST: update test `Model` classes to include mcphys and rad scheme attrs

emc2/core/model.py

@@ -93,8 +93,14 @@ class Model():
     lon_dim: str
         Name of the longitude dimension in the model (relevant for regional output)
     mcphys_scheme: str
-        Name of the microphysics scheme to use for models with 
-        multiple microphysics schemes.
+        Name of the microphysics scheme used by the model. Current options are:
+          1. "MG2", "MG", "Morrison" - essentially treated the same.
+          2. "NSSL"
+          3. "P3"
+    rad_scheme_family: str
+        Radiation scheme family. Many models share the same radiation scheme or bulk scattering LUT
+        ancestor (inc. same PSD parameters, ice habit description, etc.). This parameters determines
+        the single particle or bulk LUT to use given the Model class.
     stacked_time_dim: str or None
         This attribute becomes a string of the original time dimension name only if
         stacking was required to enable EMC2 to processes a domain output (time x lat x lon).
@@ -138,7 +144,8 @@ def __init__(self):
         self.height_dim = "height"
         self.lat_dim = "lat"
         self.lon_dim = "lon"
-        self.mcphys_scheme = ""
+        self.mcphys_scheme = None
+        self.rad_scheme_family = None
         self.stacked_time_dim = None
         self.process_conv = True
         self.model_name = ""
@@ -598,6 +605,8 @@ def __init__(self, file_path, time_range=None, load_processed=False):
         self.q_names_convective = {'cl': 'QCLmc', 'ci': 'QCImc', 'pl': 'QPLmc', 'pi': 'QPImc'}
         self.q_names_stratiform = {'cl': 'qcl', 'ci': 'qci', 'pl': 'qpl', 'pi': 'qpi'}
         self.hyd_types = ["cl", "ci", "pl", "pi"]
+        self.mcphys_scheme = "MG2"  # per GM2015 (J. Clim.)
+        self.rad_scheme_family = "ModelE"
 
         if load_processed:
             self.ds = xr.Dataset()
@@ -627,7 +636,8 @@ def __init__(self, file_path, time_range=None, load_processed=False):
 
 class E3SMv1(Model):
     def __init__(self, file_path, time_range=None, load_processed=False, time_dim="time", appended_str=False,
-                 all_appended_in_lat=False, single_ice_class=False, include_rain_in_rt=False):
+                 all_appended_in_lat=False, single_ice_class=False, include_rain_in_rt=False,
+                 mcphys_scheme="MG2"):
         """
         This loads an E3SMv1 simulation output with all of the necessary parameters for EMC^2 to run.
 
@@ -655,6 +665,8 @@ def __init__(self, file_path, time_range=None, load_processed=False, time_dim="t
             If True, including the rain class (`pl`) in the forward calculations.
             By default, set to False given that the rain class is excluded from the E3SM radiative scheme
             calculations.
+        mcphys_scheme: str
+            Name of the microphysics scheme used by the model. Current options are:
 
         """
         super().__init__()
@@ -705,6 +717,9 @@ def __init__(self, file_path, time_range=None, load_processed=False, time_dim="t
                 eval(f"self.{attr}")['pi'] = 'zeros_cf'
         else:
             self.hyd_types += ["pi"]
+        self.mcphys_scheme = mcphys_scheme
+        self.rad_scheme_family = "CESM"  # E3SMv1 uses the same bulk LUTs as CESM (CAM 5.0)
+
         self.process_conv = False
         if load_processed:
             self.ds = xr.Dataset()
@@ -817,7 +832,7 @@ def __init__(self, file_path, z_range=None, time_range=None,
             raise ModuleNotFoundError("wrf-python must be installed.")
 
         super().__init__()
-        if mcphys_scheme.lower() == "morrison":
+        if mcphys_scheme.lower() in ["mg2", "mg", "morrison"]:
             self.hyd_types = ["cl", "ci", "pl", "pi"]
             self.Rho_hyd = {'cl': 1000. * ureg.kg / (ureg.m**3),
                             'ci': 500. * ureg.kg / (ureg.m**3),
@@ -1024,6 +1039,7 @@ def __init__(self, file_path, z_range=None, time_range=None,
         self.lat_name = "XLAT"
         self.lon_name = "XLONG"
         self.mcphys_scheme = mcphys_scheme
+        self.rad_scheme_family = "ModelE"  # NOTE: adaptive implementation needed (as in mcphys_scheme) per WRF run
         self.process_conv = False
         wrfin.close()
         for keys in self.ds.keys():
@@ -1136,8 +1152,8 @@ def __init__(self, file_path, time_range=None, time_dim="dom_col", single_pi_cla
             self.q_names_convective.update({'pir': 'conv_dat'})
             self.q_names_stratiform.update({'pir': 'qpir'})
             self.hyd_types.append("pir")
-
-        self.process_conv = False
+        self.mcphys_scheme="MG"  # MG2008
+        self.rad_scheme_family = "ModelE"  # Similar implementation to ModelE3
 
         if load_processed:
             self.ds = xr.Dataset()
@@ -1246,6 +1262,9 @@ def __init__(self):
                                         'pl': 'cldmcpl', 'pi': 'cldmcpi'}
         self.strat_frac_names_for_rad = {'cl': 'cldsscl', 'ci': 'cldssci',
                                          'pl': 'cldsspl', 'pi': 'cldsspi'}
+
+        self.mcphys_scheme = "MG2"  # required to prevent errors from being raised.
+        self.rad_scheme_family = "E3SM"  # required to prevent errors from being raised.
         self.ds = my_ds
         self.height_dim = "height"
         self.time_dim = "time"
@@ -1385,6 +1404,8 @@ def __init__(self):
         self.strat_frac_names = {'cl': 'cldsscl', 'ci': 'cldssci', 'pl': 'cldsspl', 'pi': 'cldsspi'}
         self.conv_frac_names_for_rad = {'cl': 'cldmccl', 'ci': 'cldmcci', 'pl': 'cldmcpl', 'pi': 'cldmcpi'}
         self.strat_frac_names_for_rad = {'cl': 'cldsscl', 'ci': 'cldssci', 'pl': 'cldsspl', 'pi': 'cldsspi'}
+        self.mcphys_scheme = "MG2"  # required to prevent errors from being raised.
+        self.rad_scheme_family = "ModelE"  # required to prevent errors from being raised.
         self.ds = my_ds
         self.height_dim = "height"
         self.time_dim = "time"
@@ -1525,6 +1546,8 @@ def __init__(self):
         self.strat_frac_names = {'cl': 'cldsscl', 'ci': 'cldssci', 'pl': 'cldsspl', 'pi': 'cldsspi'}
         self.conv_frac_names_for_rad = {'cl': 'cldmccl', 'ci': 'cldmcci', 'pl': 'cldmcpl', 'pi': 'cldmcpi'}
         self.strat_frac_names_for_rad = {'cl': 'cldsscl', 'ci': 'cldssci', 'pl': 'cldsspl', 'pi': 'cldsspi'}
+        self.mcphys_scheme = "MG2"  # required to prevent errors from being raised.
+        self.rad_scheme_family = "ModelE"  # required to prevent errors from being raised.
         self.ds = my_ds
         self.height_dim = "height"
         self.time_dim = "time"
@@ -1645,6 +1668,8 @@ def __init__(self):
         self.conv_frac_names_for_rad = {'cl': 'cldmccl', 'ci': 'cldmcci', 'pl': 'cldmcpl', 'pi': 'cldmcpi'}
         self.strat_frac_names_for_rad = {'cl': 'cldsscl', 'ci': 'cldssci', 'pl': 'cldsspl', 'pi': 'cldsspi'}
         self.hyd_types = ["cl", "ci", "pl", "pi"]
+        self.mcphys_scheme = "MG2"  # required to prevent errors from being raised.
+        self.rad_scheme_family = "ModelE"  # required to prevent errors from being raised.
         self.ds = my_ds
         self.height_dim = "height"
         self.time_dim = "time"

emc2/simulator/lidar_moments.py

@@ -286,7 +286,7 @@ def calc_lidar_empirical(instrument, model, is_conv, p_values, t_values, z_value
 
 
 def calc_lidar_bulk(instrument, model, is_conv, p_values, z_values, OD_from_sfc=True,
-                    hyd_types=None, mie_for_ice=False, cesm_rad_family=None, **kwargs):
+                    hyd_types=None, mie_for_ice=False, **kwargs):
     """
     Calculates the lidar stratiform or convective backscatter, extinction, and
     optical depth in a sub-columns using bulk scattering LUTs assuming geometric
@@ -312,8 +312,6 @@ def calc_lidar_bulk(instrument, model, is_conv, p_values, z_values, OD_from_sfc=
     mie_for_ice: bool
         If True, using bulk mie caculation LUTs. Otherwise, currently using the bulk C6
         scattering LUTs for 8-column severly roughned aggregate.
-    cesm_rad_family: list or None
-        list of model classes implementing CESM-type radiation scheme (same LUT PSD assumptions, etc.).
     Additonal keyword arguments are passed into
     :py:func:`emc2.simulator.lidar_moments.accumulate_OD`.
 
@@ -322,8 +320,6 @@ def calc_lidar_bulk(instrument, model, is_conv, p_values, z_values, OD_from_sfc=
     model: :func:`emc2.core.Model`
         The model with the added simulated lidar parameters.
     """
-    if cesm_rad_family is None:
-        cesm_rad_family = ["E3SMv1", "CESM2"]
     hyd_types = model.set_hyd_types(hyd_types)
 
     optional_ice_classes = ["ci", "pi", "sn", "gr", "ha", "pir", "pid", "pif"]
@@ -337,14 +333,16 @@ def calc_lidar_bulk(instrument, model, is_conv, p_values, z_values, OD_from_sfc=
 
     n_subcolumns = model.num_subcolumns
 
-    if model.model_name in cesm_rad_family:
+    if model.rad_scheme_family == "CESM":
         bulk_ice_lut = "CESM_ice"
         bulk_mie_ice_lut = "mie_ice_CESM_PSD"
         bulk_liq_lut = "CESM_liq"
-    else:
+    elif model.rad_scheme_family == "ModelE":
         bulk_ice_lut = "E3_ice"
         bulk_mie_ice_lut = "mie_ice_E3_PSD"
         bulk_liq_lut = "E3_liq"
+    else:
+        raise ValueError(f"Unknown radiation scheme family: {model.rad_scheme_family}")
 
     Dims = model.ds["%s_q_subcolumns_cl" % cloud_str].shape
     model.ds['sub_col_beta_p_tot_%s' % cloud_str] = xr.DataArray(
@@ -390,16 +388,19 @@ def calc_lidar_bulk(instrument, model, is_conv, p_values, z_values, OD_from_sfc=
                 Qback_bulk = instrument.bulk_table[bulk_ice_lut]["Q_back"].values
                 Qext_bulk = instrument.bulk_table[bulk_ice_lut]["Q_ext"].values
         else:
-            if model.model_name in cesm_rad_family:
+            if model.rad_scheme_family in ["CESM"]:  # rad families that have bulk LUTs vs. lambda and mu
                 mu_b = np.tile(instrument.bulk_table[bulk_liq_lut]["mu"].values,
                                (instrument.bulk_table[bulk_liq_lut]["lambdas"].size)).flatten()
                 lambda_b = instrument.bulk_table[bulk_liq_lut]["lambda"].values.flatten()
-            else:
+            else:  # rad families that have bulk LUTs vs. r_eff
                 r_eff_bulk = instrument.bulk_table[bulk_liq_lut]["r_e"].values
             Qback_bulk = instrument.bulk_table[bulk_liq_lut]["Q_back"].values
             Qext_bulk = instrument.bulk_table[bulk_liq_lut]["Q_ext"].values
 
-        if np.logical_and(np.isin(hyd_type, ["cl", "pl"]), model.model_name in cesm_rad_family):
+        # The following if condition is to determine if the LUTs are vs. the PSD lambda and mu parameters
+        # The alternative default is to use LUTs vs. r_eff
+        # ===========================================================
+        if np.logical_and(np.isin(hyd_type, ["cl", "pl"]), model.rad_scheme_family in ["CESM"]):
             print("2-D interpolation of bulk liq lidar backscattering using mu-lambda values")
             rel_locs = model.ds[model.q_names_stratiform[hyd_type]].values > 0.
             back_tmp = np.ones_like(model.ds[model.q_names_stratiform[hyd_type]].values, dtype=float) * np.nan
@@ -439,7 +440,7 @@ def calc_lidar_bulk(instrument, model, is_conv, p_values, z_values, OD_from_sfc=
 
 def calc_lidar_micro(instrument, model, z_values, OD_from_sfc=True,
                      hyd_types=None, mie_for_ice=False, parallel=True, chunk=None,
-                     cesm_rad_family=None, **kwargs):
+                     **kwargs):
     """
     Calculates the lidar backscatter, extinction, and optical depth
     in a given column for the given lidar using the microphysics (MG2) logic.
@@ -466,8 +467,6 @@ def calc_lidar_micro(instrument, model, z_values, OD_from_sfc=True,
         the entries to the Dask worker queue at once. Sometimes, Dask will freeze if
         too many tasks are sent at once due to memory issues, so adjusting this number
         might be needed if that happens.
-    cesm_rad_family: list or None
-        list of model classes implementing CESM-type radiation scheme (same LUT PSD assumptions, etc.).
     Additonal keyword arguments are passed into
     :py:func:`emc2.psd.calc_mu_lambda`.
     :py:func:`emc2.simulator.lidar_moments.accumulate_OD`.
@@ -477,18 +476,18 @@ def calc_lidar_micro(instrument, model, z_values, OD_from_sfc=True,
     model: :func:`emc2.core.Model`
         The model with the added simulated lidar parameters.
     """
-    if cesm_rad_family is None:
-        cesm_rad_family = ["E3SMv1", "CESM2"]
     hyd_types = model.set_hyd_types(hyd_types)
 
     optional_ice_classes = ["ci", "pi", "sn", "gr", "ha", "pir", "pid", "pif"]
 
-    if model.model_name in cesm_rad_family:
+    if model.rad_scheme_family == "CESM":
         ice_lut = "CESM_ice"
         ice_diam_var = "p_diam"
-    else:
+    elif model.rad_scheme_family == "ModelE":
         ice_lut = "E3_ice"
         ice_diam_var = "p_diam_eq_V"
+    else:
+        raise ValueError(f"Unknown radiation scheme family: {model.rad_scheme_family}")
 
     Dims = model.ds["strat_q_subcolumns_cl"].values.shape
     for hyd_type in hyd_types:
@@ -505,7 +504,16 @@ def calc_lidar_micro(instrument, model, z_values, OD_from_sfc=True,
             np.zeros(Dims), dims=model.ds.strat_q_subcolumns_cl.dims)
         model.ds["sub_col_alpha_p_%s_strat" % hyd_type] = xr.DataArray(
             np.zeros(Dims), dims=model.ds.strat_q_subcolumns_cl.dims)
-        fits_ds = calc_mu_lambda(model, hyd_type, subcolumns=True, **kwargs).ds
+        if hyd_type in ["cl", "pl"]:  # liquid classes
+            if model.mcphys_scheme.lower() in ["mg2", "mg", "morrison", "nssl"]:
+                fits_ds = calc_mu_lambda(model, hyd_type, subcolumns=True, **kwargs).ds
+            else:
+                raise ValueError(f"no liquid PSD calulation method implemented for scheme {model.mcphys_scheme}")
+        else:  # ice classes
+            if model.mcphys_scheme.lower() in ["mg2", "mg", "morrison", "nssl"]:  # NOTE: NSSL PSD assumed like MG
+                fits_ds = calc_mu_lambda(model, hyd_type, subcolumns=True, **kwargs).ds
+            else:
+                raise ValueError(f"no ice PSD calulation method implemented for scheme {model.mcphys_scheme}")
         N_columns = len(model.ds["subcolumn"])
         total_hydrometeor = np.round(model.ds[frac_names].values * N_columns).astype(int)
         N_0 = fits_ds["N_0"].values
@@ -571,7 +579,7 @@ def calc_lidar_micro(instrument, model, z_values, OD_from_sfc=True,
 
 def calc_lidar_moments(instrument, model, is_conv,
                        OD_from_sfc=True, hyd_types=None, parallel=True, eta=1, chunk=None, mie_for_ice=False,
-                       use_rad_logic=True, use_empiric_calc=False, cesm_rad_family=None, **kwargs):
+                       use_rad_logic=True, use_empiric_calc=False, **kwargs):
     """
     Calculates the lidar backscatter, extinction, and optical depth
     in a given column for the given lidar.
@@ -623,8 +631,6 @@ def calc_lidar_moments(instrument, model, is_conv,
     use_empirical_calc: bool
         When True using empirical relations from literature for the fwd calculations
         (the cloud fraction still follows the scheme logic set by use_rad_logic).
-    cesm_rad_family: list or None
-        list of model classes implementing CESM-type radiation scheme (same LUT PSD assumptions, etc.).
     Additonal keyword arguments are passed into
     :py:func:`emc2.psd.calc_mu_lambda`.
     :py:func:`emc2.simulator.lidar_moments.accumulate_OD`.
@@ -637,8 +643,6 @@ def calc_lidar_moments(instrument, model, is_conv,
     model: :func:`emc2.core.Model`
         The model dataset with the added simulated lidar parameters.
     """
-    if cesm_rad_family is None:
-        cesm_rad_family = ["E3SMv1", "CESM2"]
     hyd_types = model.set_hyd_types(hyd_types)
 
     if is_conv:
@@ -660,10 +664,12 @@ def calc_lidar_moments(instrument, model, is_conv,
     elif mie_for_ice:
         scat_str = "Mie"
     else:
-        if model.model_name in cesm_rad_family:
+        if model.rad_scheme_family == "CESM":
             scat_str = "m-D_A-D (D. Mitchell)"
-        else:
+        elif model.rad_scheme_family == "ModelE":
             scat_str = "C6"
+        else:
+            raise ValueError(f"Unknown radiation scheme family: {model.rad_scheme_family}")
 
     if not instrument.instrument_class.lower() == "lidar":
         raise ValueError("Instrument must be a lidar!")

emc2/simulator/psd.py

@@ -48,6 +48,12 @@ def calc_mu_lambda(model, hyd_type="cl",
                    subcolumns=False, is_conv=False, **kwargs):
 
     """
+    This method calculated the Gamma PSD parameters following Morrison and Gettelman (2008).
+    Note that the dispersion cacluation from MG2008 is used in all models implementing this
+    parameterization except for ModelE and DHARMA, which use a fixed definition.
+    Note #2: `subcolumns` are hardwired as `True` in `radar_moments.py` and `lidar_moments.py`,
+    which means that we assume that the PSD mu and lambda definition applies to the SGS. This
+    defintion might be under future discussion (whether to apply this assumption or not).
     Calculates the :math:`\mu` and :math:`\lambda` of the gamma PSD given :math:`N_{0}`.
     The gamma size distribution takes the form:
 
@@ -59,7 +65,7 @@ def calc_mu_lambda(model, hyd_type="cl",
 
     Note: this method only accepts the microphysical cloud fraction in order to maintain
     consistency because the PSD calculation is necessarily related only to the MG2 scheme
-    without assumption related to the radiation logic
+    without assumption related to the radiation logic.
 
     Parameters
     ----------
@@ -77,7 +83,7 @@ def calc_mu_lambda(model, hyd_type="cl",
         The lower and upper bounds for the :math:`\mu` parameter.
     subcolumns: bool
         If True, the fit parameters will be generated using the generated subcolumns
-        rather than using the "raw" model output.
+        (in-cloud) q and N quantities) rather than using the "raw" (grid-cell mean) output.
     is_conv: bool
         If True, calculate from convective properties. IF false, do stratiform.
 
@@ -98,10 +104,10 @@ def calc_mu_lambda(model, hyd_type="cl",
     """
 
     if calc_dispersion is None:
-        if model.model_name in ["E3SM", "CESM2", "WRF"]:
-            calc_dispersion = True
-        else:
+        if model.model_name in ["ModelE", "DHARMA"]:
             calc_dispersion = False
+        else:
+            calc_dispersion = True
     if not subcolumns:
         N_name = model.N_field[hyd_type]
         if not is_conv: