Compute and provide layer thickness

monetio/models/_wrfchem_mm.py

@@ -89,23 +89,15 @@ def open_mfdataset(
         # need to calculate surface pressure and dp and optionally dz here.
 
         # Additional defaults for satellite analysis
-        var_list.append("PH")
-        var_list.append("PHB")
-        var_list.append("PB")
-        var_list.append("P")
-        var_list.append("T")
+        var_list.append("zstag")
 
     var_wrf_list = []
     for var in var_list:
         if var == "pres":  # Insert special versions.
             var_wrf = getvar(
                 wrflist, var, timeidx=ALL_TIMES, method="cat", squeeze=False, units="Pa"
             )
-        elif var == "height":
-            var_wrf = getvar(
-                wrflist, var, timeidx=ALL_TIMES, method="cat", squeeze=False, units="m"
-            )
-        elif var == "height_agl":
+        elif var in {"height", "height_agl", "zstag"}:
             var_wrf = getvar(
                 wrflist, var, timeidx=ALL_TIMES, method="cat", squeeze=False, units="m"
             )
@@ -115,6 +107,13 @@ def open_mfdataset(
 
     dset = xr.merge(var_wrf_list)
 
+    if not surf_only_nc:
+        # Compute layer thickness
+        dset["dz"] = (
+            dset["zstag"].diff("bottom_top_stag").swap_dims({"bottom_top_stag": "bottom_top"})
+        )
+        dset["dz"] = dset["dz"].assign_attrs({"units": "m", "long_name": "layer thickness"})
+
     # Add global attributes needed
     a_truelat1 = extract_global_attrs(wrflist[0], "TRUELAT1")
     a_truelat2 = extract_global_attrs(wrflist[0], "TRUELAT2")
@@ -210,6 +209,7 @@ def open_mfdataset(
                 "temp": "temperature_k",
                 "height": "alt_msl_m_mid",
                 "height_agl": "alt_agl_m_mid",
+                "dz": "dz_m",
                 "PSFC": "surfpres_pa",
                 "pressure": "pres_pa_mid",
             }