Feature/use onestep zarr train data

external/vcm/tests/test_calc.py

@@ -10,7 +10,7 @@
     dz_and_top_to_phis,
     _add_coords_to_interface_variable,
 )
-from vcm.calc.calc import local_time
+from vcm.calc.calc import local_time, apparent_source
 from vcm.cubedsphere.constants import COORD_Z_CENTER, COORD_Z_OUTER
 
 
@@ -85,3 +85,22 @@ def test_solar_time():
     lon = xr.DataArray([0, 180, 270, 360, 0, 270], dims=["x"], coords={"x": range(6)})
     ds_solar_test = xr.Dataset({"initialization_time": t, "lon": lon})
     assert np.allclose(local_time(ds_solar_test), [0, 12, 18, 0, 6, 0])
+
+
+def test_apparent_source():
+    coords = {
+        "initial_time": [
+            cftime.DatetimeJulian(2016, 8, 1, 0, 15, 0),
+            cftime.DatetimeJulian(2016, 8, 1, 0, 30, 0),
+        ],
+        "forecast_time": [0.0, 60.0],
+        "x": [1],
+        "y": [1],
+    }
+    T = xr.DataArray(
+        [[[[1, 2], [3, 5]]]],
+        dims=["x", "y", "initial_time", "forecast_time"],
+        coords=coords,
+    )
+    Q1 = apparent_source(T, t_dim="initial_time", s_dim="forecast_time")
+    assert Q1 == (2.0 / (15 * 60)) - (1.0 / 60)

external/vcm/vcm/calc/calc.py

@@ -46,8 +46,11 @@ def apparent_source(
     q = q.drop([t_dim, s_dim])
     dq = q.diff(t_dim)
     dq_c48 = q.diff(s_dim)
-    dt = timedelta_to_seconds(t.diff(t_dim))
-    ds = timedelta_to_seconds(s.diff(s_dim))
+
+    dt = t.diff(t_dim)
+    dt = timedelta_to_seconds(dt) if dt.dtype == "<m8[ns]" else dt
+    ds = s.diff(s_dim)
+    ds = timedelta_to_seconds(ds) if ds.dtype == "<m8[ns]" else ds
 
     tend = dq / dt
     tend_c48 = dq_c48 / ds

fv3net/pipelines/coarsen_restarts/__main__.py

@@ -23,10 +23,10 @@
         ),
     )
     parser.add_argument(
-        "source_resolution", type=int, help="Source data cubed-sphere grid resolution.",
+        "source_resolution", type=int, help="Source data cubed-sphere grid resolution."
     )
     parser.add_argument(
-        "target_resolution", type=int, help="Target coarsening resolution to output.",
+        "target_resolution", type=int, help="Target coarsening resolution to output."
     )
     parser.add_argument(
         "gcs_dst_dir",

fv3net/pipelines/common.py

@@ -134,7 +134,7 @@ def list_timesteps(path: str) -> List[str]:
 
 
 def subsample_timesteps_at_interval(
-    timesteps: List[str], sampling_interval: int,
+    timesteps: List[str], sampling_interval: int
 ) -> List[str]:
     """
     Subsample a list of timesteps at the specified interval (in minutes). Raises

fv3net/pipelines/create_training_data/__init__.py

@@ -10,10 +10,7 @@
 SUFFIX_HIRES_DIAG = "prog"
 SUFFIX_COARSE_TRAIN_DIAG = "train"
 
-DIAG_VARS = [
-    "LHTFLsfc",
-    "SHTFLsfc",
-    "PRATEsfc",
+RADIATION_VARS = [
     "DSWRFtoa",
     "DSWRFsfc",
     "USWRFtoa",
@@ -22,20 +19,30 @@
     "ULWRFtoa",
     "ULWRFsfc",
 ]
-RENAMED_PROG_DIAG_VARS = {f"{var}_coarse": f"{var}_prog" for var in DIAG_VARS}
-RENAMED_TRAIN_DIAG_VARS = {var: f"{var}_train" for var in DIAG_VARS}
 
+RENAMED_HIGH_RES_VARS = {
+    **{f"{var}_coarse": f"{var}_prog" for var in RADIATION_VARS},
+    **{
+        "LHTFLsfc_coarse": "latent_heat_flux_prog",
+        "SHTFLsfc_coarse": "sensible_heat_flux_prog",
+    },
+}
 
-RESTART_VARS = [
-    "sphum",
-    "T",
-    "delp",
-    "u",
-    "v",
-    "slmsk",
-    "phis",
-    "tsea",
-    "slope",
-    "DZ",
-    "W",
+ONE_STEP_VARS = RADIATION_VARS + [
+    "total_precipitation",
+    "surface_temperature",
+    "land_sea_mask",
+    "latent_heat_flux",
+    "sensible_heat_flux",
+    "mean_cos_zenith_angle",
+    "surface_geopotential",
+    "vertical_thickness_of_atmospheric_layer",
+    "vertical_wind",
+    "pressure_thickness_of_atmospheric_layer",
+    "specific_humidity",
+    "air_temperature",
+    "x_wind",
+    "y_wind",
 ]
+
+RENAMED_ONE_STEP_VARS = {var: f"{var}_train" for var in RADIATION_VARS}

fv3net/pipelines/create_training_data/__main__.py

@@ -38,7 +38,11 @@
         "Output zarr files will be saved in either 'train' or 'test' subdir of "
         "gcs-output-data-dir",
     )
+    parser.add_argument(
+        "--var-names-yaml",
+        type=str,
+        default=None,
+        help="optional yaml for providing data variable names",
+    )
     args, pipeline_args = parser.parse_known_args()
-    print(args)
-    """Main function"""
     run(args=args, pipeline_args=pipeline_args)

fv3net/pipelines/create_training_data/helpers.py

@@ -3,6 +3,7 @@
 import os
 import pandas as pd
 import xarray as xr
+import yaml
 
 from vcm.fv3_restarts import open_diagnostic
 from vcm.cloud.fsspec import get_fs