Merge remote-tracking branch 'origin/master' into refactor/fv3net-onl…

…ine-modules

CONTRIBUTING.md

@@ -0,0 +1,17 @@
+# Contributions Guides and Standards
+
+
+## fv3net
+
+- Only imports to `vcm`, `vcm.cubedsphere`, and `vcm.cloud` are allowed. No
+  deeper nesting (e.g. `vcm.cloud.fsspec`) or imports to other modules are
+  allowed.
+
+
+##  vcm
+
+- The external interfaces are the modules `vcm`, `vcm.cubedsphere`, and
+  `vcm.cloud`. All routines to be used externally should be imported into one
+   of these namespaces. This rule could change pending future changes to the vcm API.
+
+

Makefile

@@ -3,7 +3,7 @@
 #################################################################################
 # GLOBALS                                                                       #
 #################################################################################
-VERSION = 0.1.0
+VERSION = v0.1.0
 ENVIRONMENT_SCRIPTS = .environment-scripts
 PROJECT_DIR := $(shell dirname $(realpath $(lastword $(MAKEFILE_LIST))))
 BUCKET = [OPTIONAL] your-bucket-for-syncing-data (do not include 's3://')

external/vcm/tests/test_visualize.py

@@ -269,7 +269,7 @@ def test_plot_cube_axes(sample_dataset, plotting_function):
     "plotting_function", [("pcolormesh"), ("contour"), ("contourf")]
 )
 def test_plot_cube_with_facets(sample_dataset, plotting_function):
-    f, axes, hs, cbar = plot_cube(
+    f, axes, hs, cbar, facet_grid = plot_cube(
         mappable_var(sample_dataset, "t2m"),
         col="time",
         plotting_function=plotting_function,
@@ -281,7 +281,7 @@ def test_plot_cube_with_facets(sample_dataset, plotting_function):
 )
 def test_plot_cube_on_axis(sample_dataset, plotting_function):
     ax = plt.axes(projection=ccrs.Robinson())
-    f, axes, hs, cbar = plot_cube(
+    f, axes, hs, cbar, facet_grid = plot_cube(
         mappable_var(sample_dataset, "t2m").isel(time=0),
         plotting_function=plotting_function,
         ax=ax,

external/vcm/vcm/__init__.py

@@ -11,6 +11,14 @@
     parse_timestep_str_from_path,
     parse_datetime_from_str,
 )
+from .calc import mass_integrate
+from .calc.thermo import (
+    net_heating,
+    net_precipitation,
+    pressure_at_midpoint_log,
+    potential_temperature,
+)
 from .coarsen import coarsen_restarts_on_pressure, coarsen_restarts_on_sigma
+from .select import mask_to_surface_type
 from .visualize import plot_cube, mappable_var, plot_cube_axes
 from .xarray_loaders import open_tiles, open_delayed

external/vcm/vcm/calc/thermo.py

@@ -2,6 +2,7 @@
 import xarray as xr
 from ..cubedsphere.constants import COORD_Z_CENTER, COORD_Z_OUTER
 
+
 # following are defined as in FV3GFS model (see FV3/fms/constants/constants.f90)
 _GRAVITY = 9.80665  # m /s2
 _RDGAS = 287.05  # J / K / kg
@@ -19,6 +20,8 @@
 _REFERENCE_SURFACE_PRESSURE = 100000  # reference pressure for potential temp [Pa]
 _REVERSE = slice(None, None, -1)
 
+_SEC_PER_DAY = 86400
+
 
 def potential_temperature(P, T):
     return T * (_REFERENCE_SURFACE_PRESSURE / P) ** _POISSON_CONST
@@ -196,8 +199,7 @@ def net_heating(
     """
 
     lv = latent_heat_vaporization(surface_temperature)
-
-    return (
+    da = (
         -dlw_sfc
         - dsw_sfc
         + ulw_sfc
@@ -208,6 +210,8 @@ def net_heating(
         + shf
         + surface_rain_rate * lv
     )
+    da.attrs = {"long_name": "net heating from model physics", "units": "W/m^2"}
+    return da
 
 
 def latent_heat_flux_to_evaporation(
@@ -226,35 +230,9 @@ def latent_heat_flux_to_evaporation(
     return lhf / latent_heat_vaporization(surface_temperature)
 
 
-def net_heating_from_dataset(ds: xr.Dataset, suffix: str = None) -> xr.DataArray:
-    """Compute the net heating from a dataset of diagnostic output
-
-    This should be equivalent to the vertical integral (i.e. <>) of Q1::
-
-        cp <Q1>
-
-    Args:
-        ds: a datasets with the names for the heat fluxes and precipitation used
-            by the ML pipeline
-        suffix: (optional) suffix of flux data vars if applicable. Will add "_" before
-            appending to variable names if not already in suffix.
-
-    Returns:
-        the total net heating, the rate of change of the dry enthalpy <c_p T>
-    """
-    if suffix and suffix[0] != "_":
-        suffix = "_" + suffix
-    elif not suffix or suffix == "":
-        suffix = ""
-    fluxes = (
-        ds["DLWRFsfc" + suffix],
-        ds["DSWRFsfc" + suffix],
-        ds["ULWRFsfc" + suffix],
-        ds["ULWRFtoa" + suffix],
-        ds["USWRFsfc" + suffix],
-        ds["USWRFtoa" + suffix],
-        ds["DSWRFtoa" + suffix],
-        ds["SHTFLsfc" + suffix],
-        ds["PRATEsfc" + suffix],
+def net_precipitation(lhf, prate):
+    da = (prate - latent_heat_flux_to_evaporation(lhf)) * _SEC_PER_DAY
+    da.attrs = (
+        {"long_name": "net precipitation from model physics", "units": "mm/day"},
     )
-    return net_heating(*fluxes)
+    return da

external/vcm/vcm/visualize/plot_cube.py

@@ -105,10 +105,13 @@ def plot_cube(
         cbar (obj):
             `plt.colorbar` object handle associated with figure, if `colorbar`
             arg is True, else None.
+        facet_grid (xarray.plot.facetgrid):
+            xarray plotting facetgrid for multi-axes case. In single-axes case,
+            retunrs None.
 
     Example:
         # plot diag winds at two times
-        fig, axes, hs, cbar = plot_cube(
+        fig, axes, hs, cbar, facet_grid = plot_cube(
             mappable_var(diag_ds, 'VGRD850').isel(time = slice(2, 4)),
             plotting_function = "contourf",
             col = "time",
@@ -165,6 +168,7 @@ def plot_cube(
         handle = _plot_func_short(array, ax=ax)
         axes = np.array(ax)
         handles = [handle]
+        facet_grid = None
 
     if coastlines:
         coastlines_kwargs = dict() if not coastlines_kwargs else coastlines_kwargs
@@ -186,7 +190,7 @@ def plot_cube(
     else:
         cbar = None
 
-    return fig, axes, handles, cbar
+    return fig, axes, handles, cbar, facet_grid
 
 
 def mappable_var(ds: xr.Dataset, var_name: str):

external/vcm/vcm/visualize/plot_diagnostics.py

@@ -82,7 +82,9 @@ def plot_diag_var_single_map(da, grid, var_name, plot_cube_kwargs=None):
     """
     da = da.rename(var_name)
     ds_mappable = mappable_var(xr.merge([da, grid]), var_name)
-    fig, axes, handles, cbar = plot_cube(ds_mappable, **(plot_cube_kwargs or {}))
+    fig, axes, handles, cbar, facet_grid = plot_cube(
+        ds_mappable, **(plot_cube_kwargs or {})
+    )
     return fig
 
 

fv3net/diagnostics/data_funcs.py

@@ -2,6 +2,7 @@
 import pandas as pd
 import xarray as xr
 
+import vcm
 from vcm.select import drop_nondim_coords, get_latlon_grid_coords
 
 # give as [lat, lon]
@@ -12,10 +13,11 @@
     "central_canada": [55.0, 258.0],
     "tropical_west_pacific": [-5.0, 165.0],
 }
+_KG_M2S_TO_MM_DAY = 86400  # kg/m2/s same as mm/s. Using 1000 km/m3 for H20 density
 
 
 def merge_comparison_datasets(
-    var, datasets, dataset_labels, grid, additional_dataset=None
+    data_vars, datasets, dataset_labels, grid, additional_dataset=None
 ):
     """ Makes a comparison dataset out of multiple datasets that all have a common
     data variable. They are concatenated with a new dim "dataset" that can be used
@@ -39,7 +41,7 @@ def merge_comparison_datasets(
     src_dim_index = pd.Index(dataset_labels, name="dataset")
     datasets = [drop_nondim_coords(ds) for ds in datasets]
     datasets_to_merge = [
-        xr.concat([ds[var].squeeze(drop=True) for ds in datasets], src_dim_index),
+        xr.concat([ds[data_vars].squeeze(drop=True) for ds in datasets], src_dim_index),
         grid,
     ]
     if additional_dataset is not None:
@@ -90,3 +92,37 @@ def _conditions(d):
 
     cond = np.vectorize(_conditions)
     return cond(phase)
+
+
+def net_heating_from_dataset(ds: xr.Dataset, suffix: str = None) -> xr.DataArray:
+    """Compute the net heating from a dataset of diagnostic output
+
+    This should be equivalent to the vertical integral (i.e. <>) of Q1::
+
+        cp <Q1>
+
+    Args:
+        ds: a datasets with the names for the heat fluxes and precipitation used
+            by the ML pipeline
+        suffix: (optional) suffix of flux data vars if applicable. Will add '_' before
+            appending to variable names if not already in suffix.
+
+    Returns:
+        the total net heating, the rate of change of the dry enthalpy <c_p T>
+    """
+    if suffix and suffix[0] != "_":
+        suffix = "_" + suffix
+    elif not suffix or suffix == "":
+        suffix = ""
+    fluxes = (
+        ds["DLWRFsfc" + suffix],
+        ds["DSWRFsfc" + suffix],
+        ds["ULWRFsfc" + suffix],
+        ds["ULWRFtoa" + suffix],
+        ds["USWRFsfc" + suffix],
+        ds["USWRFtoa" + suffix],
+        ds["DSWRFtoa" + suffix],
+        ds["SHTFLsfc" + suffix],
+        ds["PRATEsfc" + suffix],
+    )
+    return vcm.net_heating(*fluxes)

fv3net/diagnostics/sklearn_model_performance/__main__.py

@@ -36,8 +36,9 @@
         help="Output dir to write results to. Can be local or a GCS path.",
     )
     parser.add_argument(
-        "num_test_zarrs",
+        "--num_test_zarrs",
         type=int,
+        default=4,
         help="Number of zarrs to concat together for use as test set.",
     )
     parser.add_argument(
@@ -61,6 +62,7 @@
         help="Factor by which to downsample test set time steps",
     )
     args = parser.parse_args()
+    args.test_data_path = os.path.join(args.test_data_path, "test")
 
     # if output path is remote GCS location, save results to local output dir first
     proto = get_protocol(args.output_path)

fv3net/diagnostics/sklearn_model_performance/data_funcs_sklearn.py

@@ -1,14 +1,15 @@
-from scipy.interpolate import UnivariateSpline
 import os
+from scipy.interpolate import UnivariateSpline
 import xarray as xr
 
 import fv3net
+from ..data_funcs import net_heating_from_dataset
 from fv3net.pipelines.create_training_data import (
     SUFFIX_COARSE_TRAIN_DIAG,
     VAR_Q_HEATING_ML,
     VAR_Q_MOISTENING_ML,
 )
-from vcm.calc import mass_integrate, thermo
+import vcm
 from vcm.cloud.fsspec import get_fs
 from vcm.convenience import round_time
 from vcm.cubedsphere.constants import (
@@ -23,12 +24,24 @@
     kg_m2_to_mm,
     SPECIFIC_HEAT_CONST_PRESSURE,
     GRAVITY,
-    SEC_PER_DAY,
 )
 
 SAMPLE_DIM = "sample"
 STACK_DIMS = ["tile", INIT_TIME_DIM, COORD_X_CENTER, COORD_Y_CENTER]
 
+THERMO_DATA_VAR_ATTRS = {
+    "net_precipitation": {"long_name": "net column precipitation", "units": "mm/day"},
+    "net_heating": {"long_name": "net column heating", "units": "W/m^2"},
+    "net_precipitation_ml": {
+        "long_name": "residual P-E predicted by ML model",
+        "units": "mm/day",
+    },
+    "net_heating_ml": {
+        "long_name": "residual heating predicted by ML model",
+        "units": "W/m^2",
+    },
+}
+
 
 def predict_on_test_data(
     test_data_path,
@@ -79,11 +92,11 @@ def load_high_res_diag_dataset(coarsened_hires_diags_path, init_times):
             f"are not matched in high res dataset."
         )
 
-    evaporation = thermo.latent_heat_flux_to_evaporation(ds_hires["LHTFLsfc_coarse"])
-    ds_hires["P-E_total"] = SEC_PER_DAY * (ds_hires["PRATEsfc_coarse"] - evaporation)
-    ds_hires["heating_total"] = thermo.net_heating_from_dataset(
-        ds_hires, suffix="coarse"
+    ds_hires["net_precipitation"] = vcm.net_precipitation(
+        ds_hires[f"LHTFLsfc_coarse"], ds_hires[f"PRATEsfc_coarse"]
     )
+    ds_hires["net_heating"] = net_heating_from_dataset(ds_hires, suffix="coarse")
+
     return ds_hires
 
 
@@ -95,17 +108,28 @@ def add_column_heating_moistening(ds):
         ds (xarray dataset): train/test or prediction dataset
             that has dQ1, dQ2, delp, precip and LHF data variables
     """
-    ds["P-E_total"] = (
-        mass_integrate(-ds[VAR_Q_MOISTENING_ML], ds.delp)
-        - thermo.latent_heat_flux_to_evaporation(
-            ds[f"LHTFLsfc_{SUFFIX_COARSE_TRAIN_DIAG}"]
-        )
-        + ds[f"PRATEsfc_{SUFFIX_COARSE_TRAIN_DIAG}"]
-    ) * kg_m2s_to_mm_day
 
-    ds["heating_total"] = SPECIFIC_HEAT_CONST_PRESSURE * mass_integrate(
+    ds["net_precipitation_ml"] = (
+        vcm.mass_integrate(-ds[VAR_Q_MOISTENING_ML], ds.delp) * kg_m2s_to_mm_day
+    )
+    ds["net_precipitation_physics"] = vcm.net_precipitation(
+        ds[f"LHTFLsfc_{SUFFIX_COARSE_TRAIN_DIAG}"],
+        ds[f"PRATEsfc_{SUFFIX_COARSE_TRAIN_DIAG}"],
+    )
+
+    ds["net_precipitation"] = (
+        ds["net_precipitation_ml"] + ds["net_precipitation_physics"]
+    )
+
+    ds["net_heating_ml"] = SPECIFIC_HEAT_CONST_PRESSURE * vcm.mass_integrate(
         ds[VAR_Q_HEATING_ML], ds.delp
-    ) + thermo.net_heating_from_dataset(ds, suffix=SUFFIX_COARSE_TRAIN_DIAG)
+    )
+    ds["net_heating_physics"] = net_heating_from_dataset(
+        ds, suffix=SUFFIX_COARSE_TRAIN_DIAG
+    )
+    ds["net_heating"] = ds["net_heating_ml"] + ds["net_heating_physics"]
+    for data_var, data_attrs in THERMO_DATA_VAR_ATTRS.items():
+        ds[data_var].attrs = data_attrs
 
 
 def integrate_for_Q(P, sphum, lower_bound=55000, upper_bound=85000):
@@ -114,7 +138,7 @@ def integrate_for_Q(P, sphum, lower_bound=55000, upper_bound=85000):
 
 
 def lower_tropospheric_stability(ds):
-    pressure = thermo.pressure_at_midpoint_log(ds.delp)
+    pressure = vcm.pressure_at_midpoint_log(ds.delp)
     T_at_700mb = (
         regrid_to_shared_coords(
             ds["T"],
@@ -126,5 +150,5 @@ def lower_tropospheric_stability(ds):
         .squeeze()
         .drop("p700mb")
     )
-    theta_700mb = thermo.potential_temperature(70000, T_at_700mb)
+    theta_700mb = vcm.potential_temperature(70000, T_at_700mb)
     return theta_700mb - ds["tsea"]