Baseline simulations and allPhysics configuration (#176)

Adds a bunch of edits to the workflow orchestrator to redesign its syntax and argument handling to be more generalizable and enhance the orchestrator maintainability. This includes:

general handling of arguments for both inputs and other types of arguments, with ability to specify from/location, optional arguments, and the ability to pass multiple argument values available for all arguments. This resolves #174 .
changed the argument passing syntax to reflect the merging of inputs/extra arguments; arguments are now passed in *positional_args + output_location + **{optional_arg: arg_value} order.
updated all of the current workflow step submission scripts to work with the new syntax
removed all of the intermediate orchestrator_submit.sh scripts as they are no longer needed by the orchestrator (all orchestrator calls are now directly to python), while maintaining any single-step shell script functionality in individual steps.
added the ability to launch Dataflow jobs from the orchestrator without an intermediate script using the --runner argument
These changes are included in this PR because they (among many other things) allow for the prognostic run to be called with the ML model path argument as an optional input from the model training step, to support running the ML and baseline prognostic runs together under the orchestrator.

Also introduces some changes to allow running baseline prognostic simulations with ML included.

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/pipelines/coarsen_restarts/__main__.py

@@ -14,14 +14,6 @@
         type=str,
         help="Full GCS path to input data for downloading timesteps.",
     )
-    parser.add_argument(
-        "gcs_dst_dir",
-        type=str,
-        help=(
-            "Full GCS path to output coarsened timestep data. Defaults to input path"
-            "with target resolution appended as a directory"
-        ),
-    )
     parser.add_argument(
         "gcs_grid_spec_path",
         type=str,
@@ -36,6 +28,14 @@
     parser.add_argument(
         "target_resolution", type=int, help="Target coarsening resolution to output.",
     )
+    parser.add_argument(
+        "gcs_dst_dir",
+        type=str,
+        help=(
+            "Full GCS path to output coarsened timestep data. Defaults to input path"
+            "with target resolution appended as a directory"
+        ),
+    )
     parser.add_argument(
         "--no-target-subdir",
         action="store_true",

fv3net/regression/sklearn/test.py

@@ -1,8 +1,10 @@
 from vcm.cloud import fsspec
 import joblib
 import xarray as xr
+import os
 
 from ..dataset_handler import stack_and_drop_nan_samples
+from .train import MODEL_FILENAME
 from vcm.convenience import round_time
 from vcm.cubedsphere.constants import INIT_TIME_DIM
 from fv3net.pipelines.create_training_data import (
@@ -82,7 +84,7 @@ def load_model(model_path):
     protocol = fsspec.get_protocol(model_path)
     if protocol == "gs":
         fs = fsspec.get_fs(model_path)
-        fs.get(model_path, "temp_model.pkl")
+        fs.get(os.path.join(model_path, MODEL_FILENAME), "temp_model.pkl")
         return joblib.load("temp_model.pkl")
     else:
         return joblib.load(model_path)

fv3net/regression/sklearn/train.py

@@ -160,6 +160,7 @@ def save_output(output_url, model, config):
         "remove local copy after upload.",
     )
     args = parser.parse_args()
+    args.train_data_path = os.path.join(args.train_data_path, "train")
     train_config = load_model_training_config(
         args.train_config_file, args.train_data_path
     )

workflows/coarsen_c384_diagnostics/README.md

@@ -1,7 +1,7 @@
 #### Usage
 Run 
 `workflows/coarsen_c384_diagnostics/coarsen_c384_diagnostics.sh`
-`  {INPUT_LOCATION} {OUTPUT_LOCATION} {CONFIG_LOCATION}`
+`  {INPUT_LOCATION} {CONFIG_LOCATION} {OUTPUT_LOCATION}`
 
 #### Description
 This script coarsens a subset of diagnostic variables from the high resolution runs

workflows/coarsen_c384_diagnostics/coarsen_c384_diagnostics.py

@@ -3,7 +3,6 @@
 import shutil
 import argparse
 import yaml
-import fsspec
 import xarray as xr
 
 from vcm import coarsen
@@ -18,7 +17,7 @@
     VAR_LON_OUTER,
     VAR_LAT_OUTER,
 )
-from vcm.fv3_restarts import _split_url
+from vcm.cloud.fsspec import get_fs
 from fv3net import COARSENED_DIAGS_ZARR_NAME
 
 logging.basicConfig(level=logging.INFO)
@@ -79,8 +78,7 @@ def _get_config(config_path):
 
 
 def _get_remote_diags(diags_path):
-    proto, path = _split_url(diags_path)
-    fs = fsspec.filesystem(proto)
+    fs = get_fs(diags_path)
     mapper = fs.get_mapper(diags_path)
     return xr.open_zarr(mapper)
 
@@ -90,13 +88,13 @@ def _get_remote_diags(diags_path):
     parser.add_argument(
         "input_path", type=str, help="GCS location of C384 diagnostics data zarrs."
     )
+    parser.add_argument(
+        "config_path", type=str, help="Location of diagnostics coarsening config yaml."
+    )
     parser.add_argument(
         "output_path",
         type=str,
         help="GCS location where= coarsened diagnostics zarrs will be written.",
     )
-    parser.add_argument(
-        "config_path", type=str, help="Location of diagnostics coarsening config yaml."
-    )
     args = parser.parse_args()
     coarsen_c384_diagnostics(args)

workflows/coarsen_restarts/README.md

@@ -8,26 +8,20 @@ using pressure-level coarsening defined in `vcm.coarsen`.
 ```python
 fv3net.pipelines.coarsen_restarts
 
-usage: __main__.py [-h] --gcs-src-dir GCS_SRC_DIR [--gcs-dst-dir GCS_DST_DIR]
-                   --gcs-grid-spec-path GCS_GRID_SPEC_PATH --source-resolution
-                   SOURCE_RESOLUTION --target_resolution TARGET_RESOLUTION
+usage: __main__.py [-h] GCS_SRC_DIR GCS_GRID_SPEC_PATH SOURCE_RESOLUTION 
+    TARGET_RESOLUTION GCS_DST_DIR
 
-optional arguments:
+positional arguments:
   -h, --help            show this help message and exit
-  --gcs-src-dir GCS_SRC_DIR
-                        Full GCS path to input data for downloading timesteps.
-  --gcs-dst-dir GCS_DST_DIR
-                        Full GCS path to output coarsened timestep data.
-                        Defaults to input pathwith target resolution appended
-                        as a directory
-  --gcs-grid-spec-path GCS_GRID_SPEC_PATH
-                        Full path with file wildcard 'grid_spec.tile*.nc' to
+  GCS_SRC_DIR           Full GCS path to input data for downloading timesteps
+  GCS_GRID_SPEC_PATH    Full path with file wildcard 'grid_spec.tile*.nc' to
                         select grid spec files with same resolution as the
                         source data
-  --source-resolution SOURCE_RESOLUTION
-                        Source data cubed-sphere grid resolution.
-  --target_resolution TARGET_RESOLUTION
-                        Target coarsening resolution to output
+  SOURCE_RESOLUTION     Source data cubed-sphere grid resolution.
+  TARGET_RESOLUTION     Target coarsening resolution to output
+  GCS_DST_DIR           Full GCS path to output coarsened timestep data.
+                        Defaults to input pathwith target resolution appended
+                        as a directory
 ```
 
 See `workflows/coarsen_restarts/submit_job.sh` to see an example of calling this

workflows/coarsen_restarts/submit_job.sh

@@ -1,20 +1,20 @@
 #!/bin/sh
 
 GCS_SRC="gs://vcm-ml-data/2019-12-02-40-day-X-SHiELD-simulation-C384-restart-files"
-GCS_DST="gs://vcm-ml-data/2020-01-16-X-SHiELD-2019-12-02-pressure-coarsened-rundirs/restarts"
 GCS_GRIDSPEC="gs://vcm-ml-data/2020-01-06-C384-grid-spec-with-area-dx-dy"
 SRC_RESOLUTION=384
 TARGET_RESOLUTION=48
+GCS_DST="gs://vcm-ml-data/2020-01-16-X-SHiELD-2019-12-02-pressure-coarsened-rundirs/restarts"
 
 user=$(whoami)
 user=${user,,}
 
 python -m fv3net.pipelines.coarsen_restarts\
     $GCS_SRC \
-    $GCS_DST \
     $GCS_GRIDSPEC \
     $SRC_RESOLUTION \
     $TARGET_RESOLUTION \
+    $GCS_DST \
     --runner DataflowRunner \
     --job_name coarsen-restarts-$user \
     --project vcm-ml \

workflows/coarsen_restarts/submit_job_directrunner.sh

@@ -1,15 +1,15 @@
 #!/bin/sh
 
 GCS_SRC="gs://vcm-ml-data/2019-12-02-40-day-X-SHiELD-simulation-C384-restart-files"
-GCS_DST="gs://vcm-ml-data/2020-01-16-X-SHiELD-2019-12-02-pressure-coarsened-rundirs/restarts"
 GCS_GRIDSPEC="gs://vcm-ml-data/2020-01-06-C384-grid-spec-with-area-dx-dy"
 SRC_RESOLUTION=384
 TARGET_RESOLUTION=48
+GCS_DST="gs://vcm-ml-data/2020-01-16-X-SHiELD-2019-12-02-pressure-coarsened-rundirs/restarts"
 
 python -m fv3net.pipelines.coarsen_restarts\
     $GCS_SRC \
-    $GCS_DST \
     $GCS_GRIDSPEC \
     $SRC_RESOLUTION \
     $TARGET_RESOLUTION \
+    $GCS_DST \
     --runner DirectRunner 

workflows/create_training_data/submit_job.sh

@@ -1,8 +1,8 @@
 python -m fv3net.pipelines.create_training_data \
-gs://vcm-ml-data/2020-02-28-X-SHiELD-2019-12-02-deep-and-mp-off \
-gs://vcm-ml-data/orchestration-testing/shield-coarsened-diags-2019-12-04 \
-gs://vcm-ml-data/experiments-2020-03/deep-conv-mp-off/train-test-data \
---job_name test-job-create-training-data-annak \
+gs://vcm-ml-data/2020-01-16-X-SHiELD-2019-12-02-pressure-coarsened-rundirs/one_step_output/C48 \
+gs://vcm-ml-data/2019-12-05-40-day-X-SHiELD-simulation-C384-diagnostics/C48_gfsphysics_15min_coarse.zarr \
+gs://vcm-ml-data/test-annak/2020-02-05_train_data_pipeline/ \
+--job_name test-job-create-training-data-brianh \
 --project vcm-ml \
 --region us-central1 \
 --runner DataflowRunner \