Enable vertical linear transform for spatially-varying target vertical coordinate

[NoraLoose]

This PR adds functionality to enable regridding onto a spatially varying vertical coordinate.

• Closes [DOC] Depth to depth xgcm.Grid.transform onto terrain-following depth levels #614 and closes [BUG] broken cells in transform page of docs  #573
• Documentation added
• Tests added
• Passes pre-commit run --all-files
• User visible changes (including notable bug fixes) are documented in whats-new.rst

I added a section to the existing doc/transform.ipynb notebook to document the new capability. While working on this notebook, I also fixed pre-existing issues with this notebook including:

• Insert autoparse_metadata=False into Grid() to avoid error messages
• Correct dictionary keys for pangeo cloud data library
• Remove non-existing xc coordinate from Grid object to avoid existing ValueError
• Insert a .squeeze() for plotting AMOC to avoid error

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[NoraLoose]

The changes made in this PR are not the reason for the failing checks above. The same checks fail in #641. I opened issues #643 and #644.

[NoraLoose]

Just realized that there is overlap between this PR and PR #627

[jbusecke]

Hey @NoraLoose, thank you so much for working on this and sorry for the delay. I think I got the CI issues #644 and #648 under control now. Once #649 is merged can you try to rebase this one?

Just realized that there is overlap between this PR and PR #627

Feel free to absorb and close this if needed.

[jbusecke]

Also I have restructured the linting and setup in #650, please let me know if it is ok to merge this ahead of time. I think this should not cause too big of a conflict?

[NoraLoose]

Also I have restructured the linting and setup in #650, please let me know if it is ok to merge this ahead of time. I think this should not cause too big of a conflict?

Yep, feel free to merge #650!

[NoraLoose]

Transforming onto a spatially varying vertical coordinate (e.g., terrain-following) works with the linear interpolation method, see newly added tests and the example in the documentation.

I tested this also for the conservative interpolation method by adding the following test case to the cases dictionary in xgcm/test/test_transform.py:

    "conservative_depth_depth_multidim_target": {
        "source_coord": ("z", [5, 25, 60]),
        "source_bounds_coord": ("zc", [0, 10, 50, 75]),
        "source_data": ("data", [1, 4, 0]),
        "source_additional_data_coord": (
            "zc",
            [0, 10, 50, 75],
        ),
        "source_additional_data": ("zc", [0, 10, 50, 75]),
        # 2D target
        "target_dims": ("eta_rho", "s_w"),  # eta_rho: horizontal dimension; s_w: vertical dimension
        "target_coord": ("interface_depth_rho", np.array([[0, 1, 10, 50, 80], [0, 5, 20, 30, 100]])),
        "target_data": ("interface_depth_rho", np.array([[0, 1, 10, 50, 80], [0, 5, 20, 30, 100]])),
        "expected_dims": ("eta_rho", "s_w"),
        "expected_coord": ("interface_depth_rho", np.array([[0.5, 5.5, 30, 65], [2.5, 12.5, 25, 65]])),
        "expected_data": (
            "data",
            np.array([[0.1, 0.9, 4.0, 0.0], [0.5, 1.5, 1.0, 2.0]])
        ),
        "grid_kwargs": {
            "coords": {"Z": {"center": "z", "outer": "zc"}},
            "autoparse_metadata": False,
        },
        "transform_kwargs": {"method": "conservative", "target_data": "zc", "target_dim": "s_w"},
    }

However, it fails because of this line:

xgcm/xgcm/transform.py

Line 160 in 2de04bf

assert target_theta_bins.ndim == 1

[NoraLoose]

@jbusecke I'm not sure when I have time to look into generalizing the conservative method to multiple dimensions. We could merge this PR, and open a feature request issue? If we do this, should we include more warnings (where?) to clarify that the conservative method is not yet supported for multi-dimensional vertical target coordinates?

[jbusecke]

Thank you so much @NoraLoose! This is awesome. I am pretty busy today, but will try to take some time next week to review and maybe also enable conservative interpolation.

[NoraLoose]

@jbusecke I would like to use the new functionality of this PR in a downstream package (roms-tools). Do you think we can merge this soon, and leave the conservative method stuff for a future PR?

[NoraLoose]

@jbusecke @TomNicholas I was hoping to get this PR merged because a downstream package (roms-tools) is waiting on this new functionality. Any ideas for who could review this PR if you guys are too busy?

[TomNicholas]

@NoraLoose I think at this point it's either @jbusecke or you just take over and use your own best judgement. This part of the code (the whole transform module) was @jbusecke 's work anyway. I will add you as a maintainer now regardless.

[NoraLoose]

@hdrake This PR is ready for review. It enables regridding onto a spatially varying vertical coordinate, such as a terrain-following coordinate.

I've added tests and included an example in the documentation, along with fixing some pre-existing issues in the transform.ipynb notebook.

The only remaining task is to enable conservative vertical regridding for the spatially varying coordinate, in addition to the linear method. I suggest opening an issue to track that feature, but merging this PR as is.

Let me know if you'd like to review or have any thoughts!

[hdrake]

@NoraLoose, I think fine to leave the conservative part for a later PR. I've added an explicit ValueError for the conservative case in the meantime.

I reviewed your updated notebook, which looks great! I also ran the tests (and skimmed the code, but found it a bit hard to follow at first glance) and can confirm everything passed locally.

Please feel free to merge when you're ready!

[raphaeldussin]

LGTM +1
Very useful addition!

[NoraLoose]

Thanks for your input @hdrake and @raphaeldussin!

skimmed the code, but found it a bit hard to follow at first glance

Yep, I remember it took me a while to understand the logic of the test_transform.py module as well.

I made three more modifications:

• Changed the ValueError that @hdrake added to a NonImplementedError
• Added a test that checks that the NonImplementedError is raised.
• Added some test data into the test_transform.py module that can be used to check the multi-dimensionsional conservative transform behaves correctly once that method is implemented.

[hdrake]

@NoraLoose, after adding conda install pooch jupyterlab to the doc/environment.yml environment, I was able to run your GLORYS -> ROMS example in the transform.ipynb notebook (but not the prior example because it requires access to the Google Cloud).

What environment are you using to run these example notebooks? I noticed, for example, that the notebook using intake which is not specified in either of ci/environment.yml or doc/environment.yml. Do the github actions that build the docs not actually run the notebooks? Do we need a seperate doc/environment_notebooks.yml to specify the environment that was used to create the example notebooks? We can continue discussing in another issue if there is no easy answer.

Related to your question here: #664 (comment)

[NoraLoose]

@hdrake The GitHub Actions do not execute the notebooks; they are simply rendered as provided by the user.

I recall needing to add more packages to my environment to run the notebook from start to finish, but that environment isn’t clean. We should open a new issue to create a dedicated doc/environment_notebooks.yml with the necessary dependencies.

[NoraLoose]

I'm merging this, and will open two issues:

• Feature: Implement vertical conservative tranform for spatially-varying vertical coordinate
• Create dedicated doc/environment_notebooks.yml