Prognostic run diagnostics assume 3D verification data is interpolated to pressure levels

[spencerkclark]

In working on a lightweight end-to-end example of our corrective ML workflow I have been starting from a C48 resolution "fine-grid" simulation coarsened to C12 resolution. To generate 3D verification data I started by simply outputting 3D coarse diagnostics of temperature, specific humidity, horizontal winds, etc. These are naturally output on model levels. It seems our prognostic run diagnostics workflow currently cannot handle this (even though it handles 3D outputs from our coarse prognostic runs on model levels just fine), and does not provide an informative error. It instead fails in the report generation step with this obscure message:

ValueError: cannot use non-scalar arrays in a slice for xarray indexing: <xarray.DataArray 'latitude' (z: 79)>
    raise ValueError(
  File "/opt/conda/envs/fv3net/lib/python3.8/site-packages/xarray/core/indexes.py", line 73, in _sanitize_slice_element
    _sanitize_slice_element(label.start),
  File "/opt/conda/envs/fv3net/lib/python3.8/site-packages/xarray/core/indexes.py", line 92, in _query_slice
    indexer = _query_slice(index, label, coord_name, method, tolerance)
  File "/opt/conda/envs/fv3net/lib/python3.8/site-packages/xarray/core/indexes.py", line 198, in query
    idxr, new_idx = index.query(labels, method=method, tolerance=tolerance)
  File "/opt/conda/envs/fv3net/lib/python3.8/site-packages/xarray/core/indexing.py", line 117, in remap_label_indexers
    pos_indexers, new_indexes = indexing.remap_label_indexers(
  File "/opt/conda/envs/fv3net/lib/python3.8/site-packages/xarray/core/coordinates.py", line 421, in remap_label_indexers
    pos_indexers, new_indexes = remap_label_indexers(
  File "/opt/conda/envs/fv3net/lib/python3.8/site-packages/xarray/core/dataset.py", line 2474, in sel
    ascent_region_diags = zonal_mean_diags.sel(latitude=slice(lat_min, lat_max))
  File "/home/jovyan/fv3net/workflows/diagnostics/fv3net/diagnostics/prognostic_run/metrics.py", line 229, in tropical_ascent_region_mean
    return name, func(*args, **kwargs).load()
  File "/home/jovyan/fv3net/workflows/diagnostics/fv3net/diagnostics/_shared/registry.py", line 37, in load
    return [func(*args, **kwargs)
  File "/opt/conda/envs/fv3net/lib/python3.8/site-packages/joblib/parallel.py", line 262, in <listcomp>
    return [func(*args, **kwargs)
  File "/opt/conda/envs/fv3net/lib/python3.8/site-packages/joblib/parallel.py", line 262, in __call__
    self.results = batch()
  File "/opt/conda/envs/fv3net/lib/python3.8/site-packages/joblib/_parallel_backends.py", line 572, in __init__
    result = ImmediateResult(func)
  File "/opt/conda/envs/fv3net/lib/python3.8/site-packages/joblib/_parallel_backends.py", line 208, in apply_async
    job = self._backend.apply_async(batch, callback=cb)
  File "/opt/conda/envs/fv3net/lib/python3.8/site-packages/joblib/parallel.py", line 779, in _dispatch
    self._dispatch(tasks)
  File "/opt/conda/envs/fv3net/lib/python3.8/site-packages/joblib/parallel.py", line 861, in dispatch_one_batch
    while self.dispatch_one_batch(iterator):
  File "/opt/conda/envs/fv3net/lib/python3.8/site-packages/joblib/parallel.py", line 1046, in __call__
    computed_outputs = Parallel(n_jobs=n_jobs, verbose=True)(
  File "/home/jovyan/fv3net/workflows/diagnostics/fv3net/diagnostics/_shared/registry.py", line 28, in compute
    metrics[ds.run] = metrics_registry.compute(ds, n_jobs=1)
  File "/home/jovyan/fv3net/workflows/diagnostics/fv3net/diagnostics/prognostic_run/computed_diagnostics.py", line 270, in load_metrics_from_diagnostics
    return RunMetrics(load_metrics_from_diagnostics(self.folders))
  File "/home/jovyan/fv3net/workflows/diagnostics/fv3net/diagnostics/prognostic_run/computed_diagnostics.py", line 87, in load_metrics_from_diagnostics
    metrics = computed_diagnostics.load_metrics_from_diagnostics()
  File "/home/jovyan/fv3net/workflows/diagnostics/fv3net/diagnostics/prognostic_run/views/static_report.py", line 653, in make_report
    make_report(computed_diagnostics, args.output)
  File "/home/jovyan/fv3net/workflows/diagnostics/fv3net/diagnostics/prognostic_run/views/static_report.py", line 733, in main_new
    args.func(dissoc(args, "func"))
  File "/home/jovyan/fv3net/workflows/diagnostics/fv3net/diagnostics/prognostic_run/cli.py", line 47, in main
    sys.exit(load_entry_point('fv3net-diagnostics', 'console_scripts', 'prognostic_run_diags')())
  File "/opt/conda/envs/fv3net/bin/prognostic_run_diags", line 33, in <module>
Traceback (most recent call last):
array([ -5., -25.,  -5.,   5., -29., -29., -29., -29., -15., -15., -15.,
        -3.,   3.,   5., -13., -29., -29., -29.,  -9.,  -9.,  -9.,  -9.,
        -9., -13., -13., -13., -13., -13., -13., -13., -13., -13., -13.,
       -13., -13., -13., -13., -13., -13., -13., -13., -13.,  -5., -25.,
       -25., -25., -25., -25., -25., -25., -25., -25., -25., -25., -25.,
       -25., -25., -25., -25., -25., -25., -25., -25., -29., -29., -29.,
       -29., -29., -29., -29., -29., -29., -29., -29., -29., -29., -29.,
       -29., -29.])
Coordinates:
  * z        (z) float64 4.514 8.301 12.45 16.74 ... 983.8 989.5 994.4 998.3
Error: exit status 1

It would great if the prognostic run diagnostics would also automatically interpolate 3D verification data to levels of constant pressure if needed as well.

[oliverwm1]

Thanks for submitting the issue. Requiring the 3D verification data to be on pressure levels was an intentional design choice, the idea being that verification datasets are re-used many times and so the time/compute savings of having this data available on pressure levels would be significant. But I'm open to doing interpolation on the fly if necessary. And at the very least we should give a better error message.

[spencerkclark]

Yes, that's a good rationale for preferring they be pre-interpolated. One could also argue that could be enforced upon reviewing a PR to add verification entries to the catalog (which is where we put frequently used datasets) rather than disallowing it in the workflow.

The main motivation for me here is (for ease of testing purposes) to try and remove as many manual steps as possible between a fine-resolution run and our corrective ML workflow.

[spencerkclark]

If we wanted to be extra safe though, we could guard it behind a flag so that you had to explicitly opt-in to interpolating verification data on the fly (otherwise an informative error message could be raised).

[spencerkclark]

(No need to address this immediately though. It would probably be good for me to outline a little more what an end-to-end test would look like before assessing which minor inconveniences would be useful to address in existing aspects of our workflow, and which we might be OK dealing with manually in an end-to-end test).

[spencerkclark]

Apologies for missing that this was at least documented here.