Create `xtorch_harmonics` package #2267

spencerkclark · 2023-06-30T21:25:01Z

This PR adds a package containing an xarray wrapper of a spherical harmonic transform roundtrip that can be completed with NVIDIA's torch_harmonics package. The roundtrip transform code itself is derived from this notebook of @oliverwm1. The xarray wrapper is based on that developed in this notebook of mine.

Some updates from that development notebook include:

Code to ensure that DataArray and Dataset level attributes are preserved in round trips.
A split of the grid parameter into forward_grid and inverse_grid components to allow for more flexibility.
Optional grid validation code to ensure that the coordinates of the grid match the expectation of the forward_grid specified. This can be disabled by setting unsafe=True.
Updating of the latitude coordinate in the case that the forward_grid does not equal the inverse_grid.

Added public API:

xtorch_harmonics.roundtrip: a function to apply a roundtrip spherical harmonic transform to all horizontally varying variables in a Dataset or to a single DataArray.

Requirement changes:

torch-harmonics@git+https://github.com/NVIDIA/torch-harmonics.git@8826246cacf6c37b600cdd63fde210815ba238fd
Install git and git-lfs in the fv3net Docker image to enable installation of torch_harmonics from GitHub.
Pin pip to version 20.2.4 (consistent with what we do in environment.yml in fv3net) in the lint build in CircleCI to avoid errors like this.
Ran make lock_deps/lock_pip following these instructions
Add PR review with license info for any additions to constraints.txt
(example)
Tests added

Coverage reports (updated automatically):

spencerkclark · 2023-06-30T21:27:10Z

constraints.txt

@@ -288,6 +288,7 @@ tf-estimator-nightly==2.8.0.dev2021122109
 threadpoolctl==3.1.0
 toml==0.10.1
 toolz==0.10.0
+git+https://github.com/NVIDIA/torch-harmonics.git@8826246cacf6c37b600cdd63fde210815ba238fd


torch_harmonics uses a BSD-3-Clause license (see here).

oliverwm1

In general looks great! The coordinate validation is tricky... In general I think it's fine to leave in what you have but it is a bit unclear to me whether a shifted/flipped coordinate should be valid.

docker/fv3net/Dockerfile

external/xtorch_harmonics/xtorch_harmonics/xtorch_harmonics.py

spencerkclark

Thanks for the initial review @oliverwm1. I tried to address all of your comments; see my responses below.

docker/fv3net/Dockerfile

external/xtorch_harmonics/xtorch_harmonics/xtorch_harmonics.py

oliverwm1

LGTM, thanks for the changes.

Not sure if I've done everything necessary to hook into CI. We'll see.

…itude

@frodre

#2267 added `torch_harmonics` via a GitHub URL to the `pip-requirements.txt` file, which subsequently led to it being used in the `constraints.txt` file. This broke the dataflow image build as newer versions of `pip` no longer allow this: #2304. This PR implements the fix suggested in a discussion between @frodre and @oliverwm1 in Slack a while back: > Could we just avoid having it in the constraints entirely? It would still get installed in the image right? Indeed it seems like even if older versions of `pip` allow GitHub URLs to appear in constraints files, they do not have the desired effect. In other words if one does: ``` $ pip install -c constraints.txt torch_harmonics ``` in an environment without it already installed, `pip` will download the latest version off of PyPI instead of fetching it from the GitHub URL in the constraints file (perhaps this is not surprising since it is unnamed and `pip` is not set up to infer the name from the URL). Thus I suppose it is somewhat misleading to include GitHub URLs in the constraints to begin with, and sufficient / better that we specify `torch_harmonics` from GitHub in the `setup.py` in `xtorch_harmonics`: https://github.com/ai2cm/fv3net/blob/314edd50a1d9f1fe31d14a418dbe981d16c63f1d/external/xtorch_harmonics/setup.py#L16 All that being said, I have implemented this in a way such to retain `git+https://github.com/NVIDIA/torch-harmonics.git@8826246cacf6c37b600cdd63fde210815ba238fd` in `pip-requirements.txt` since it appears `pip-compile` can still resolve dependencies from packages on GitHub and use that to inform the versions of other packages pinned in the `constraints.txt` file, which is still useful. Requirement changes: - [x] Ran `make lock_deps/lock_pip` following these [instructions](https://vulcanclimatemodeling.com/docs/fv3net/dependency_management.html#dependency-management) Resolves #2304 Coverage reports (updated automatically): --------- Co-authored-by: Oliver Watt-Meyer <oliverwatt@gmail.com>

spencerkclark commented Jun 30, 2023

View reviewed changes

spencerkclark changed the title ~~Create xtorch_harmonics package~~ Create xtorch_harmonics package Jun 30, 2023

spencerkclark force-pushed the xtorch_harmonics branch 4 times, most recently from 5bdc807 to 5cbc355 Compare July 6, 2023 00:39

oliverwm1 reviewed Jul 20, 2023

View reviewed changes

spencerkclark commented Jul 24, 2023

View reviewed changes

spencerkclark marked this pull request as ready for review July 24, 2023 15:35

oliverwm1 approved these changes Jul 24, 2023

View reviewed changes

spencerkclark added 15 commits July 25, 2023 12:55

Add initial code and tests

a5272b7

Not sure if I've done everything necessary to hook into CI. We'll see.

Add tox.ini file similar to vcm

5dadef7

Add scipy to testing requirements

1fa0739

Use same version of pip in lint container as fv3net requires

f60788d

Install git and git-lfs in build_default CI workflow

4075281

Install git and git-lfs in fv3net image instead

f2b726c

Pin torch_harmonics in setup.py

96def4c

Simplify tests some

37d0cd7

Explicitly note rtol needed for equiangular grid

1f40ab4

Rename main file

c15a7db

Address all but dimension name default comments

1c84fda

Remove default values for lat_dim and lon_dim

07d3997

lint

62ac78e

Ensure replaced latitude has orientation consistent with original lat…

edf8874

…itude

Add xtorch_harmonics to mypy pre-commit hook; fix type errors

0e43658

spencerkclark force-pushed the xtorch_harmonics branch from 58befda to 0e43658 Compare July 25, 2023 12:55

spencerkclark merged commit 345f07b into master Jul 25, 2023

spencerkclark deleted the xtorch_harmonics branch July 25, 2023 17:22

spencerkclark mentioned this pull request Jul 28, 2023

Constraints from GitHub break dataflow image build #2304

Closed

spencerkclark mentioned this pull request Nov 3, 2023

Remove constraints from GitHub in constraints.txt #2364

Merged

1 task

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Create `xtorch_harmonics` package #2267

Create `xtorch_harmonics` package #2267

spencerkclark commented Jun 30, 2023 •

edited

Loading

spencerkclark Jun 30, 2023

oliverwm1 left a comment

spencerkclark left a comment

oliverwm1 left a comment

Create xtorch_harmonics package #2267

Create xtorch_harmonics package #2267

Conversation

spencerkclark commented Jun 30, 2023 • edited Loading

spencerkclark Jun 30, 2023

Choose a reason for hiding this comment

oliverwm1 left a comment

Choose a reason for hiding this comment

spencerkclark left a comment

Choose a reason for hiding this comment

oliverwm1 left a comment

Choose a reason for hiding this comment

Create `xtorch_harmonics` package #2267

Create `xtorch_harmonics` package #2267

spencerkclark commented Jun 30, 2023 •

edited

Loading