Updates to CESM-FV Reader

monetio/models/_cesm_fv_mm.py

@@ -46,8 +46,6 @@ def open_mfdataset(
             + "and has not been properly tested. Use at own risk."
         )
 
-    from pyresample.utils import wrap_longitudes
-
     # check that the files are netcdf format
     names, netcdf = _ensure_mfdataset_filenames(fname)
 
@@ -70,12 +68,19 @@ def open_mfdataset(
     if not surf_only:
         if "PMID" not in dset_load.keys():
             dset_load["PMID"] = _calc_pressure(dset_load)
-        var_list = var_list + ["PMID"]
+        var_list = var_list + ["pres_pa_mid"]
         if "Z3" not in dset_load.keys():
             warnings.warn("Geopotential height Z3 is not in model keys. Assuming hydrostatic runs")
             dset_load["Z3"] = _calc_hydrostatic_height(dset_load)
+        var_list = var_list + ["alt_msl_m_mid"]
+
+        # calc layer thickness if hyai and hybi exist
+        if "hyai" and "hybi" in dset_load.keys():
+            dset_load["pres_pa_int"] = _calc_pressure_i(dset_load)
+            dset_load["dz_m"] = _calc_layer_thickness_i(dset_load)
+            var_list.append("dz_m")
 
-        dset_load.rename(
+        dset_load = dset_load.rename(
             {
                 "T": "temperature_k",
                 "Z3": "alt_msl_m_mid",
@@ -89,6 +94,7 @@ def open_mfdataset(
             "description": "geopotential height above ground level",
             "units": "m",
         }
+
         var_list = var_list + [
             "temperature_k",
             "alt_msl_m_mid",
@@ -101,9 +107,12 @@ def open_mfdataset(
     # rename altitude dimension to z for monet use
     # also rename lon to x and lat to y
     dset = dset.rename_dims({"lon": "x", "lat": "y", "lev": "z"})
+    if np.any(dset["lon"] > 180):
+        dset["lon"] = (dset["lon"] + 180) % 360 - 180
+        dset = dset.sortby("lon")
 
     # convert to -180 to 180 longitude
-    lon = wrap_longitudes(dset["lon"])
+    lon = dset["lon"]
     lat = dset["lat"]
     lons, lats = meshgrid(lon, lat)
     dset["longitude"] = (("y", "x"), lons)
@@ -217,6 +226,54 @@ def _calc_pressure(dset):
     return P
 
 
+def _calc_pressure_i(dset):
+    """Calculates interface layer pressure using P0, PS, hyai, hybi
+
+    Parameters
+    ----------
+    dset : xr.Dataset
+
+    Returns
+    -------
+    xr.DataArray
+    """
+    presvars = ["PS", "hyai", "hybi"]
+    if not all(pvar in list(dset.keys()) for pvar in presvars):
+        raise KeyError(
+            "The model does not have the variables to calculate "
+            "the pressure required for satellite comparison. "
+            "If the vertical coordinate is not needed, set surface_only=True"
+        )
+    time = dset["PS"].time.values
+    vert = dset["hyai"].ilev.values
+    lat = dset["PS"].lat.values
+    lon = dset["PS"].lon.values
+    n_vert = len(vert)
+    n_time = len(time)
+    n_lat = len(lat)
+    n_lon = len(lon)
+
+    pressure_i = np.zeros((n_time, n_vert, n_lat, n_lon))
+
+    if "P0" not in dset.keys():
+        warnings.warn("P0 not in netcdf keys, assuming 100_000 Pa")
+        p0 = 100_000
+    else:
+        p0 = dset["P0"].values
+
+    for nlev in range(n_vert):
+        pressure_i[:, nlev, :, :] = (
+            dset["hyai"][0, nlev].values * p0 + dset["hybi"][0, nlev].values * dset["PS"].values
+        )
+    P_int = xr.DataArray(
+        data=pressure_i,
+        dims=["time", "ilev", "lat", "lon"],
+        coords={"time": time, "ilev": vert, "lat": lat, "lon": lon},
+        attrs={"description": "Interface layer pressure", "units": "Pa"},
+    )
+    return P_int
+
+
 def _calc_hydrostatic_height(dset):
     """Calculates midlayer height using PMID, P, PS and PHIS, T,
 
@@ -248,7 +305,6 @@ def _calc_hydrostatic_height(dset):
         raise Exception(
             "Expected default CESM behaviour:" + "pressure levels should be in decreasing order"
         )
-
     height = np.zeros((n_time, n_vert, n_lat, n_lon))
     height[:, n_vert, :, :] = dset["PHIS"].values / GRAVITY
     for nlev in range(n_vert - 1, -1, -1):
@@ -265,3 +321,88 @@ def _calc_hydrostatic_height(dset):
         attrs={"description": "Mid layer (hydrostatic) height", "units": "m"},
     )
     return z
+
+
+def _calc_hydrostatic_height_i(dset):
+    """Calculates interface layer height using pres_pa_int, PHIS, and T.
+
+    Parameters
+    ----------
+    dset : xr.Dataset
+
+    Returns
+    -------
+    xr.DataArray
+    """
+    R = 8.314  # Pa * m3 / mol K
+    M_AIR = 0.029  # kg / mol
+    GRAVITY = 9.80665  # m / s2
+
+    time = dset.time.values
+    ilev = dset.ilev.values
+    lat = dset.lat.values
+    lon = dset.lon.values
+
+    # check if the vertical levels go from highest to lowest altitude (low to high pressure)
+    # which is the default in CESM. That means, that the hybrid
+    # pressure levels should be increasing.
+    _height_decreasing = np.all(ilev[:-1] < ilev[1:])
+    if not _height_decreasing:
+        raise Exception(
+            "Expected default CESM behaviour:"
+            + "pressure levels should be in increasing order, height in decreasing order"
+        )
+    # surface geopotential height (PHIS / g)
+    height = np.zeros((len(time), len(ilev), len(lat), len(lon)))
+    height[:, -1, :, :] = dset["PHIS"].values / GRAVITY  # surface height
+
+    for nlev in range(len(ilev) - 2, -1, -1):
+        temp = dset["T"].isel(lev=nlev).values  # midlayer temp approx
+        pressure_top = dset["pres_pa_int"].isel(ilev=nlev + 1)
+        pressure = dset["pres_pa_int"].isel(ilev=nlev)
+
+        height[:, nlev, :, :] = height[:, nlev + 1, :, :] - (R * temp / (GRAVITY * M_AIR)) * np.log(
+            pressure / pressure_top
+        )
+
+    z = xr.DataArray(
+        data=height,
+        dims=["time", "ilev", "lat", "lon"],
+        coords={"time": time, "ilev": ilev, "lat": lat, "lon": lon},
+        attrs={"description": "Interface Layer (hydrostatic) Height", "units": "m"},
+    )
+    return z
+
+
+def _calc_layer_thickness_i(dset):
+    """
+    Calculates layer thickness (dz_m) from interface heights.
+    Note: This calculates based on pressure being in increasing order,
+    and altitude in decreasing order. The code flips all the variables
+    along the 'z' dimensions at the end. 'pres_pa_int' does not
+    because it has a dimension of 'ilev' instead of 'z'. Add a correction
+    for this last part.
+
+    Parameters
+    ----------
+    dset : xr.Dataset
+
+    Returns
+    ----------
+    xr.DataArray
+        Layer Thickness (m)
+    """
+    z_int = _calc_hydrostatic_height_i(dset)
+
+    # # compute layer thickness
+    dz_m = np.zeros((len(dset.time), len(dset.lev), len(dset.lat), len(dset.lon)))
+    for nlev in range(len(dset.lev)):
+        dz_m[:, nlev, :, :] = z_int[:, nlev, :, :] - z_int[:, nlev + 1, :, :]
+
+    dz_m = xr.DataArray(
+        data=dz_m,
+        dims=["time", "lev", "lat", "lon"],
+        coords={"time": dset.time, "lev": dset.lev, "lat": dset.lat, "lon": dset.lon},
+        attrs={"description": "Layer Thickness (based on interface pressure)", "units": "m"},
+    )
+    return dz_m