From b12ad6f8c5dd58e1c56c8d3199933ddca6674be4 Mon Sep 17 00:00:00 2001 From: Kristen Thyng Date: Wed, 27 Aug 2025 10:58:14 -0700 Subject: [PATCH 1/5] ignoring xgcm deprecation warnings and updating test versions also pinning xgcm=0.8.1 for now --- ci/environment-py3.11.yml | 5 +- ...nment-py3.9.yml => environment-py3.12.yml} | 8 +- ...ment-py3.10.yml => environment-py3.13.yml} | 8 +- docs/environment.yml | 6 +- docs/whats_new.md | 3 + environment.yml | 4 +- requirements.txt | 2 +- setup.cfg | 6 +- xroms/interp.py | 28 +- xroms/tests/test_accessor.py | 1066 ++++++++--------- xroms/utilities.py | 148 ++- xroms/xroms.py | 96 +- 12 files changed, 730 insertions(+), 650 deletions(-) rename ci/{environment-py3.9.yml => environment-py3.12.yml} (83%) rename ci/{environment-py3.10.yml => environment-py3.13.yml} (83%) diff --git a/ci/environment-py3.11.yml b/ci/environment-py3.11.yml index 4a87960..e9c737d 100644 --- a/ci/environment-py3.11.yml +++ b/ci/environment-py3.11.yml @@ -1,4 +1,4 @@ -name: test_env_intake-axds +name: test-env-xroms channels: - conda-forge dependencies: @@ -15,11 +15,12 @@ dependencies: - numpy - pooch - xarray + - xgcm=0.8.1 # - xgcm >= 0.8.1 ############## - pytest - pip: - - xgcm >= 0.8.1 + # - xgcm >= 0.8.1 - codecov - pytest-cov - coverage[toml] diff --git a/ci/environment-py3.9.yml b/ci/environment-py3.12.yml similarity index 83% rename from ci/environment-py3.9.yml rename to ci/environment-py3.12.yml index d5678ae..ff206f7 100644 --- a/ci/environment-py3.9.yml +++ b/ci/environment-py3.12.yml @@ -1,8 +1,8 @@ -name: test_env_intake-axds +name: test-env-xroms channels: - conda-forge dependencies: - - python=3.9 + - python=3.12 ############## These will have to be adjusted to your specific project - cartopy - cf_xarray @@ -15,11 +15,11 @@ dependencies: - numpy - pooch - xarray - # - xgcm >= 0.8.1 + - xgcm=0.8.1 ############## - pytest - pip: - - xgcm >= 0.8.1 + # - xgcm >= 0.8.1 - codecov - pytest-cov - coverage[toml] diff --git a/ci/environment-py3.10.yml b/ci/environment-py3.13.yml similarity index 83% rename from ci/environment-py3.10.yml rename to ci/environment-py3.13.yml index 6b59e24..764e45a 100644 --- a/ci/environment-py3.10.yml +++ b/ci/environment-py3.13.yml @@ -1,8 +1,8 @@ -name: test_env_intake-axds +name: test-env-xroms channels: - conda-forge dependencies: - - python=3.10 + - python=3.13 ############## These will have to be adjusted to your specific project - cartopy - cf_xarray @@ -15,11 +15,11 @@ dependencies: - numpy - pooch - xarray - # - xgcm >= 0.8.1 + - xgcm=0.8.1 ############## - pytest - pip: - - xgcm >= 0.8.1 + # - xgcm >= 0.8.1 - codecov - pytest-cov - coverage[toml] diff --git a/docs/environment.yml b/docs/environment.yml index 93191b4..4ca38b4 100644 --- a/docs/environment.yml +++ b/docs/environment.yml @@ -3,7 +3,7 @@ channels: - conda-forge - nodefaults dependencies: - - python=3.10 + - python=3.12 # If your docs code examples depend on other packages add them here - cartopy - cf_xarray @@ -13,14 +13,14 @@ dependencies: - esmpy=>8.5.0 # https://github.com/esmf-org/esmf/issues/140 - matplotlib-base - netcdf4 - - numpy <1.24 # https://github.com/numba/numba/issues/8615#issuecomment-1360792615 + - numpy #<1.24 # https://github.com/numba/numba/issues/8615#issuecomment-1360792615 - numba >= 0.49 # required by xgcm - pip - pooch - requests - xarray - xcmocean - - xgcm >= 0.8.1 + - xgcm=0.8.1 - xesmf # don't install if on windows - xroms # These are needed for the docs themselves diff --git a/docs/whats_new.md b/docs/whats_new.md index a73617d..5084b0b 100644 --- a/docs/whats_new.md +++ b/docs/whats_new.md @@ -1,5 +1,8 @@ # What's New +## v0.6.2 (August 27, 2025) +* catching and ignoring a bunch of warnings from `xgcm`, but still staying with `xgcm` `v0.8.1` until I can update this code to match. + ## v0.6.1 (October 28, 2024) * Correction in a few built-in calculations of u/v grid to rho-grid interpolations of u and v velocities (currently `speed` and `_uv2eastnorth`). In these cases, we need to fill nans with zeros so that the masked locations in the velocity fields are not fully brought forward into the rho mask but are instead interpolated over. By making them 0, they are calculated into the mask\_rho positions by combining them with neighboring cells. If this wasn't done, the fact that they are masked would supersede the neighboring cells and they would be masked in mask\_rho. This needs to be done anytime the velocities are moved from their native grids to the rho or other grids to preserve their locations around masked cells. diff --git a/environment.yml b/environment.yml index d9c2dcf..1dc6eaa 100644 --- a/environment.yml +++ b/environment.yml @@ -15,14 +15,14 @@ dependencies: - esmpy>=8.5.0 # https://github.com/esmf-org/esmf/issues/140 - matplotlib-base - netcdf4 - - numpy <1.24 # https://github.com/numba/numba/issues/8615#issuecomment-1360792615 + - numpy #<1.24 # https://github.com/numba/numba/issues/8615#issuecomment-1360792615 - numba >= 0.49 # required by xgcm - pip - pooch - requests - xarray - xcmocean - - xgcm >= 0.8.1 + - xgcm=0.8.1 - xesmf # don't install if on windows # - pip: # install from github to get recent PRs I contributed # - git@github.com:xarray-contrib/cf-xarray.git diff --git a/requirements.txt b/requirements.txt index 0d0415a..370011c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -8,4 +8,4 @@ numba >= 0.49 numpy pooch xarray -xgcm >= 0.8.1 +xgcm == 0.8.1 diff --git a/setup.cfg b/setup.cfg index a76e534..df20499 100644 --- a/setup.cfg +++ b/setup.cfg @@ -58,9 +58,9 @@ classifiers = Operating System :: OS Independent Programming Language :: Python Programming Language :: Python :: 3 - Programming Language :: Python :: 3.9 - Programming Language :: Python :: 3.10 Programming Language :: Python :: 3.11 + Programming Language :: Python :: 3.12 + Programming Language :: Python :: 3.13 # Dont change this one License :: OSI Approved :: MIT License @@ -80,7 +80,7 @@ install_requires = numpy pooch xarray - xgcm >= 0.8.1 + xgcm == 0.8.1 setup_requires= setuptools_scm python_requires = >=3.6 diff --git a/xroms/interp.py b/xroms/interp.py index d25a6b8..e50ae72 100644 --- a/xroms/interp.py +++ b/xroms/interp.py @@ -261,7 +261,9 @@ def isoslice(var, iso_values, xgrid, iso_array=None, axis="Z"): key = "z" # perform interpolation - transformed = xgrid.transform(var, axis, iso_values, target_data=iso_array) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + transformed = xgrid.transform(var, axis, iso_values, target_data=iso_array) if key not in transformed.coords: transformed = transformed.assign_coords({key: iso_array}) @@ -287,9 +289,11 @@ def isoslice(var, iso_values, xgrid, iso_array=None, axis="Z"): iso_array = iso_array.cf.isel(Z=0).drop_vars( iso_array.cf["Z"].name, errors="ignore" ) - transformedlon = xgrid.transform( - var[lonkey], axis, iso_values, target_data=iso_array - ) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + transformedlon = xgrid.transform( + var[lonkey], axis, iso_values, target_data=iso_array + ) transformed = transformed.assign_coords({lonkey: transformedlon}) transformed[lonkey].attrs["standard_name"] = "longitude" @@ -308,9 +312,11 @@ def isoslice(var, iso_values, xgrid, iso_array=None, axis="Z"): iso_array = iso_array.cf.isel(Z=0).drop_vars( iso_array.cf["Z"].name, errors="ignore" ) - transformedlat = xgrid.transform( - var[latkey], axis, iso_values, target_data=iso_array - ) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + transformedlat = xgrid.transform( + var[latkey], axis, iso_values, target_data=iso_array + ) transformed = transformed.assign_coords({latkey: transformedlat}) transformed[latkey].attrs["standard_name"] = "latitude" @@ -319,9 +325,11 @@ def isoslice(var, iso_values, xgrid, iso_array=None, axis="Z"): zkey = var.cf["vertical"].name if zkey not in transformed.coords: - transformedZ = xgrid.transform( - var[zkey], axis, iso_values, target_data=iso_array - ) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + transformedZ = xgrid.transform( + var[zkey], axis, iso_values, target_data=iso_array + ) transformed = transformed.assign_coords({zkey: transformedZ}) transformed[zkey].attrs["positive"] = "up" diff --git a/xroms/tests/test_accessor.py b/xroms/tests/test_accessor.py index 3debff6..7fe5d54 100644 --- a/xroms/tests/test_accessor.py +++ b/xroms/tests/test_accessor.py @@ -74,548 +74,548 @@ def test_grid(): assert isinstance(ds.xroms.xgrid, xgcm.grid.Grid) -def test_speed(): - - acc = ds.xroms.speed - - assert np.allclose(acc, xroms.speed(ds.u, ds.v, grid)) - - # also check attributes - assert acc.name == acc.attrs["name"] - # cf-xarray: make sure all Axes and Coordinates available in output - hcoord = "rho" - scoord = "s_rho" - dims = dim_dict[hcoord][scoord] - axes = axesTZYX - coords = coord_dict[hcoord][scoord] - coordnames = coordnamesTZYX - for ax, dim in zip(axes, dims): - assert acc.cf[ax].name == dim - for coordname, coord in zip(coordnames, coords): - assert acc.cf[coordname].name == coord - - -def test_KE(): - - s = xroms.speed(ds.u, ds.v, grid) - acc = ds.xroms.KE - - assert np.allclose(acc, xroms.KE(ds.rho0, s)) - - # also check attributes - assert acc.name == acc.attrs["name"] - # cf-xarray: make sure all Axes and Coordinates available in output - hcoord = "rho" - scoord = "s_rho" - dims = dim_dict[hcoord][scoord] - axes = axesTZYX - coords = coord_dict[hcoord][scoord] - coordnames = coordnamesTZYX - for ax, dim in zip(axes, dims): - assert acc.cf[ax].name == dim - for coordname, coord in zip(coordnames, coords): - assert acc.cf[coordname].name == coord - - -def test_uv_geostrophic(): - - acc = ds.xroms.ug - assert np.allclose(acc, xroms.uv_geostrophic(ds.zeta, ds.f, grid, which="xi")) - # also check attributes - assert ( - acc.name == acc.attrs["name"] - ) # cf-xarray: make sure all Axes and Coordinates available in output - hcoord = "u" - scoord = None - dims = dim_dict[hcoord][scoord] - axes = axesTYX - coords = coord_dict[hcoord][scoord] - coordnames = coordnamesTYX - for ax, dim in zip(axes, dims): - assert acc.cf[ax].name == dim - for coordname, coord in zip(coordnames, coords): - assert acc.cf[coordname].name == coord - - acc = ds.xroms.vg - assert np.allclose(acc, xroms.uv_geostrophic(ds.zeta, ds.f, grid, which="eta")) - assert acc.name == acc.attrs["name"] - hcoord = "v" - scoord = None - dims = dim_dict[hcoord][scoord] - axes = axesTYX - coords = coord_dict[hcoord][scoord] - coordnames = coordnamesTYX - for ax, dim in zip(axes, dims): - assert acc.cf[ax].name == dim - for coordname, coord in zip(coordnames, coords): - assert acc.cf[coordname].name == coord - - -def test_EKE(): - - acc = ds.xroms.EKE - xug, xvg = xroms.uv_geostrophic(ds.zeta, ds.f, grid, which="both") - assert np.allclose(acc, xroms.EKE(xug, xvg, grid)) - assert acc.name == acc.attrs["name"] - hcoord = "rho" - scoord = None - dims = dim_dict[hcoord][scoord] - axes = axesTYX - coords = coord_dict[hcoord][scoord] - coordnames = coordnamesTYX - for ax, dim in zip(axes, dims): - assert acc.cf[ax].name == dim - for coordname, coord in zip(coordnames, coords): - assert acc.cf[coordname].name == coord - - -def test_dudz(): - acc = ds.xroms.dudz - assert np.allclose(acc, xroms.dudz(ds.u, grid)) - assert acc.name == acc.attrs["name"] - hcoord = "u" - scoord = "s_w" - dims = dim_dict[hcoord][scoord] - axes = axesTZYX - coords = coord_dict[hcoord][scoord] - coordnames = coordnamesTZYX - for ax, dim in zip(axes, dims): - assert acc.cf[ax].name == dim - for coordname, coord in zip(coordnames, coords): - assert acc.cf[coordname].name == coord - - -def test_dvdz(): - acc = ds.xroms.dvdz - assert np.allclose(acc, xroms.dvdz(ds.v, grid)) - assert acc.name == acc.attrs["name"] - hcoord = "v" - scoord = "s_w" - dims = dim_dict[hcoord][scoord] - axes = axesTZYX - coords = coord_dict[hcoord][scoord] - coordnames = coordnamesTZYX - for ax, dim in zip(axes, dims): - assert acc.cf[ax].name == dim - for coordname, coord in zip(coordnames, coords): - assert acc.cf[coordname].name == coord - - -def test_vertical_shear(): - xdudz = ds.xroms.dudz - xdvdz = ds.xroms.dvdz - acc = ds.xroms.vertical_shear - assert np.allclose(acc, xroms.vertical_shear(xdudz, xdvdz, grid)) - assert acc.name == acc.attrs["name"] - hcoord = "rho" - scoord = "s_w" - dims = dim_dict[hcoord][scoord] - axes = axesTZYX - coords = coord_dict[hcoord][scoord] - coordnames = coordnamesTZYX - for ax, dim in zip(axes, dims): - assert acc.cf[ax].name == dim - for coordname, coord in zip(coordnames, coords): - assert acc.cf[coordname].name == coord - - -def test_relative_vorticity(): - acc = ds.xroms.vort - assert np.allclose(acc, 0) - assert acc.name == acc.attrs["name"] - hcoord = "psi" - scoord = "s_w" - dims = dim_dict[hcoord][scoord] - axes = axesTZYX - coords = coord_dict[hcoord][scoord] - coordnames = coordnamesTZYX - for ax, dim in zip(axes, dims): - assert acc.cf[ax].name == dim - for coordname, coord in zip(coordnames, coords): - assert acc.cf[coordname].name == coord - - -def test_convergence(): - acc = ds.xroms.convergence - assert np.allclose(acc, xroms.convergence(ds["u"], ds["v"], grid)) - - -def test_convergence_norm(): - acc = ds.xroms.convergence_norm - assert np.allclose( - acc, xroms.convergence(ds["u"], ds["v"], grid).cf.isel(Z=-1) / ds["f"] - ) - - -def test_ertel(): - acc = ds.xroms.ertel - xsig0 = xroms.potential_density(ds.temp, ds.salt) - xbuoy = xroms.buoyancy(xsig0) - assert np.allclose(acc, xroms.ertel(xbuoy, ds.u, ds.v, ds.f, grid)) - assert acc.name == acc.attrs["name"] - hcoord = "rho" - scoord = "s_rho" - dims = dim_dict[hcoord][scoord] - axes = axesTZYX - coords = coord_dict[hcoord][scoord] - coordnames = coordnamesTZYX - for ax, dim in zip(axes, dims): - assert acc.cf[ax].name == dim - for coordname, coord in zip(coordnames, coords): - assert acc.cf[coordname].name == coord - - -def test_w(): - # VRX - pass - - -# acc = ds.xroms.w -# assert np.allclose(acc, xroms.w(ds.u, ds.v, grid)) -# acc.name == acc.attrs['name'] -# acc.attrs['grid'] == ds.xroms.grid -# items = ['T','X','Y','Z','longitude','latitude','vertical','time'] -# assert set(items).issubset(acc.cf.get_valid_keys()) - - -def test_omega(): - # VRX - pass - - -# acc = ds.xroms.omega -# assert np.allclose(acc, xroms.omega(ds.u, ds.v, grid)) -# acc.name == acc.attrs['name'] -# acc.attrs['grid'] == ds.xroms.grid -# items = ['T','X','Y','Z','longitude','latitude','vertical','time'] -# assert set(items).issubset(acc.cf.get_valid_keys()) - - -def test_rho(): - acc = ds.xroms.rho - assert np.allclose(acc, xroms.density(ds.temp, ds.salt, ds.z_rho)) - assert acc.name == acc.attrs["name"] - hcoord = "rho" - scoord = "s_rho" - dims = dim_dict[hcoord][scoord] - axes = axesTZYX - coords = coord_dict[hcoord][scoord] - coordnames = coordnamesTZYX - for ax, dim in zip(axes, dims): - assert acc.cf[ax].name == dim - for coordname, coord in zip(coordnames, coords): - assert acc.cf[coordname].name == coord - - -def test_sig0(): - acc = ds.xroms.sig0 - assert np.allclose(acc, xroms.potential_density(ds.temp, ds.salt, 0)) - assert acc.name == acc.attrs["name"] - hcoord = "rho" - scoord = "s_rho" - dims = dim_dict[hcoord][scoord] - axes = axesTZYX - coords = coord_dict[hcoord][scoord] - coordnames = coordnamesTZYX - for ax, dim in zip(axes, dims): - assert acc.cf[ax].name == dim - for coordname, coord in zip(coordnames, coords): - assert acc.cf[coordname].name == coord - - -def test_buoyancy(): - acc = ds.xroms.buoyancy - xsig0 = xroms.potential_density(ds.temp, ds.salt) - assert np.allclose(acc, xroms.buoyancy(xsig0)) - assert acc.name == acc.attrs["name"] - hcoord = "rho" - scoord = "s_rho" - dims = dim_dict[hcoord][scoord] - axes = axesTZYX - coords = coord_dict[hcoord][scoord] - coordnames = coordnamesTZYX - for ax, dim in zip(axes, dims): - assert acc.cf[ax].name == dim - for coordname, coord in zip(coordnames, coords): - assert acc.cf[coordname].name == coord - - -def test_N2(): - acc = ds.xroms.N2 - xrho = xroms.density(ds.temp, ds.salt, ds.z_rho) - assert np.allclose(acc, xroms.N2(xrho, grid), equal_nan=True) - assert acc.name == acc.attrs["name"] - hcoord = "rho" - scoord = "s_w" - dims = dim_dict[hcoord][scoord] - axes = axesTZYX - coords = coord_dict[hcoord][scoord] - coordnames = coordnamesTZYX - for ax, dim in zip(axes, dims): - assert acc.cf[ax].name == dim - for coordname, coord in zip(coordnames, coords): - assert acc.cf[coordname].name == coord - - -def test_M2(): - acc = ds.xroms.M2 - xrho = xroms.density(ds.temp, ds.salt, ds.z_rho) - assert np.allclose(acc, xroms.M2(xrho, grid), equal_nan=True) - assert acc.name == acc.attrs["name"] - hcoord = "rho" - scoord = "s_w" - dims = dim_dict[hcoord][scoord] - axes = axesTZYX - coords = coord_dict[hcoord][scoord] - coordnames = coordnamesTZYX - for ax, dim in zip(axes, dims): - assert acc.cf[ax].name == dim - for coordname, coord in zip(coordnames, coords): - assert acc.cf[coordname].name == coord - - -def test_mld(): - acc = ds.xroms.mld(thresh=0.03) - sig0 = xroms.potential_density(ds.temp, ds.salt, 0) - assert np.allclose(acc, xroms.mld(sig0, grid, ds.h, ds.mask_rho), equal_nan=True) - assert acc.name == acc.attrs["name"] - hcoord = "rho" - scoord = None - dims = dim_dict[hcoord][scoord] - axes = axesTYX - coords = coord_dict[hcoord][scoord] - coordnames = coordnamesTYX - for ax, dim in zip(axes, dims): - assert acc.cf[ax].name == dim - for coordname, coord in zip(coordnames, coords): - assert acc.cf[coordname].name == coord - - -def test_ddxi(): - testvars = ["salt", "u", "v"] - for testvar in testvars: - with pytest.raises(KeyError): - acc = ds[testvar].xroms.ddxi(grid) - - if testvar == "salt": - hcoord = "u" - scoord = "s_w" - elif testvar == "u": - hcoord = "rho" - scoord = "s_w" - elif testvar == "v": - hcoord = "psi" - scoord = "s_w" - dims = dim_dict[hcoord][scoord] - axes = axesTZYX - coords = coord_dict[hcoord][scoord] - coordnames = coordnamesTZYX - - acc = ds.xroms.ddxi(testvar) - assert np.allclose(acc, xroms.ddxi(ds[testvar], grid)) - assert acc.name == acc.attrs["name"] - for ax, dim in zip(axes, dims): - assert acc.cf[ax].name == dim - for coordname, coord in zip(coordnames, coords): - assert acc.cf[coordname].name == coord - - -def test_ddeta(): - testvars = ["salt", "u", "v"] - for testvar in testvars: - with pytest.raises(KeyError): - acc = ds[testvar].xroms.ddeta(grid) - - if testvar == "salt": - hcoord = "v" - scoord = "s_w" - elif testvar == "u": - hcoord = "psi" - scoord = "s_w" - elif testvar == "v": - hcoord = "rho" - scoord = "s_w" - dims = dim_dict[hcoord][scoord] - axes = axesTZYX - coords = coord_dict[hcoord][scoord] - coordnames = coordnamesTZYX - - acc = ds.xroms.ddeta(testvar) - assert np.allclose(acc, xroms.ddeta(ds[testvar], grid)) - assert acc.name == acc.attrs["name"] - for ax, dim in zip(axes, dims): - assert acc.cf[ax].name == dim - for coordname, coord in zip(coordnames, coords): - assert acc.cf[coordname].name == coord - - -def test_ddz(): - testvars = ["salt", "u", "v"] - for testvar in testvars: - with pytest.raises(KeyError): - acc = ds[testvar].xroms.ddz(grid) - dims = list(ds[testvar].dims) - axes = axesTZYX - coords = [ds[testvar].cf[coordname].name for coordname in coordnamesTZYX] - coordnames = coordnamesTZYX - # correct dim and coord in derivative direction - # import pdb; pdb.set_trace() - if grid._get_dims_from_axis(ds[testvar], "Z")[0] == "s_rho": - # if grid.axes["Z"]._get_axis_coord(ds[testvar])[1] == "s_rho": - dims[1] = "s_w" - coords[1] = coords[1].replace("rho", "w") - else: - dims[1] = "s_rho" - coords[1] = coords[1].replace("w", "rho") - - acc = ds.xroms.ddz(testvar) - assert np.allclose(acc, xroms.ddz(ds[testvar], grid)) - assert acc.name == acc.attrs["name"] - for ax, dim in zip(axes, dims): - assert acc.cf[ax].name == dim - for coordname, coord in zip(coordnames, coords): - assert acc.cf[coordname].name == coord - - -def test_to_grid(): - - testvars = ["salt", "u", "v"] - for testvar in testvars: - for scoord in ["s_w", "s_rho"]: - for hcoord in ["rho", "u", "v", "psi"]: - acc = ds.xroms.to_grid(testvar, hcoord=hcoord, scoord=scoord) - # acc = ds[testvar].xroms.to_grid(grid, hcoord=hcoord, scoord=scoord) - assert np.allclose( - acc, xroms.to_grid(ds[testvar], grid, hcoord=hcoord, scoord=scoord) - ) - assert acc.name == acc.attrs["name"] - dims = dim_dict[hcoord][scoord] - axes = axesTZYX - coords = coord_dict[hcoord][scoord] - coordnames = coordnamesTZYX - for ax, dim in zip(axes, dims): - assert acc.cf[ax].name == dim - for coordname, coord in zip(coordnames, coords): - assert acc.cf[coordname].name == coord - - acc = ds.xroms.to_grid(testvar, hcoord=hcoord, scoord=scoord) - assert np.allclose( - acc, xroms.to_grid(ds[testvar], grid, hcoord=hcoord, scoord=scoord) - ) - assert acc.name == acc.attrs["name"] - for ax, dim in zip(axes, dims): - assert acc.cf[ax].name == dim - for coordname, coord in zip(coordnames, coords): - assert acc.cf[coordname].name == coord - - -# can't figure out what is wrong here, will have to come back -# def test_sel2d(): -# lon0, lat0 = -94.8, 28.0 +# def test_speed(): + +# acc = ds.xroms.speed + +# assert np.allclose(acc, xroms.speed(ds.u, ds.v, grid)) + +# # also check attributes +# assert acc.name == acc.attrs["name"] +# # cf-xarray: make sure all Axes and Coordinates available in output +# hcoord = "rho" +# scoord = "s_rho" +# dims = dim_dict[hcoord][scoord] +# axes = axesTZYX +# coords = coord_dict[hcoord][scoord] +# coordnames = coordnamesTZYX +# for ax, dim in zip(axes, dims): +# assert acc.cf[ax].name == dim +# for coordname, coord in zip(coordnames, coords): +# assert acc.cf[coordname].name == coord + + +# def test_KE(): + +# s = xroms.speed(ds.u, ds.v, grid) +# acc = ds.xroms.KE + +# assert np.allclose(acc, xroms.KE(ds.rho0, s)) + +# # also check attributes +# assert acc.name == acc.attrs["name"] +# # cf-xarray: make sure all Axes and Coordinates available in output +# hcoord = "rho" +# scoord = "s_rho" +# dims = dim_dict[hcoord][scoord] +# axes = axesTZYX +# coords = coord_dict[hcoord][scoord] +# coordnames = coordnamesTZYX +# for ax, dim in zip(axes, dims): +# assert acc.cf[ax].name == dim +# for coordname, coord in zip(coordnames, coords): +# assert acc.cf[coordname].name == coord + + +# def test_uv_geostrophic(): + +# acc = ds.xroms.ug +# assert np.allclose(acc, xroms.uv_geostrophic(ds.zeta, ds.f, grid, which="xi")) +# # also check attributes +# assert ( +# acc.name == acc.attrs["name"] +# ) # cf-xarray: make sure all Axes and Coordinates available in output +# hcoord = "u" +# scoord = None +# dims = dim_dict[hcoord][scoord] +# axes = axesTYX +# coords = coord_dict[hcoord][scoord] +# coordnames = coordnamesTYX +# for ax, dim in zip(axes, dims): +# assert acc.cf[ax].name == dim +# for coordname, coord in zip(coordnames, coords): +# assert acc.cf[coordname].name == coord + +# acc = ds.xroms.vg +# assert np.allclose(acc, xroms.uv_geostrophic(ds.zeta, ds.f, grid, which="eta")) +# assert acc.name == acc.attrs["name"] +# hcoord = "v" +# scoord = None +# dims = dim_dict[hcoord][scoord] +# axes = axesTYX +# coords = coord_dict[hcoord][scoord] +# coordnames = coordnamesTYX +# for ax, dim in zip(axes, dims): +# assert acc.cf[ax].name == dim +# for coordname, coord in zip(coordnames, coords): +# assert acc.cf[coordname].name == coord + + +# def test_EKE(): + +# acc = ds.xroms.EKE +# xug, xvg = xroms.uv_geostrophic(ds.zeta, ds.f, grid, which="both") +# assert np.allclose(acc, xroms.EKE(xug, xvg, grid)) +# assert acc.name == acc.attrs["name"] +# hcoord = "rho" +# scoord = None +# dims = dim_dict[hcoord][scoord] +# axes = axesTYX +# coords = coord_dict[hcoord][scoord] +# coordnames = coordnamesTYX +# for ax, dim in zip(axes, dims): +# assert acc.cf[ax].name == dim +# for coordname, coord in zip(coordnames, coords): +# assert acc.cf[coordname].name == coord + + +# def test_dudz(): +# acc = ds.xroms.dudz +# assert np.allclose(acc, xroms.dudz(ds.u, grid)) +# assert acc.name == acc.attrs["name"] +# hcoord = "u" +# scoord = "s_w" +# dims = dim_dict[hcoord][scoord] +# axes = axesTZYX +# coords = coord_dict[hcoord][scoord] +# coordnames = coordnamesTZYX +# for ax, dim in zip(axes, dims): +# assert acc.cf[ax].name == dim +# for coordname, coord in zip(coordnames, coords): +# assert acc.cf[coordname].name == coord + + +# def test_dvdz(): +# acc = ds.xroms.dvdz +# assert np.allclose(acc, xroms.dvdz(ds.v, grid)) +# assert acc.name == acc.attrs["name"] +# hcoord = "v" +# scoord = "s_w" +# dims = dim_dict[hcoord][scoord] +# axes = axesTZYX +# coords = coord_dict[hcoord][scoord] +# coordnames = coordnamesTZYX +# for ax, dim in zip(axes, dims): +# assert acc.cf[ax].name == dim +# for coordname, coord in zip(coordnames, coords): +# assert acc.cf[coordname].name == coord + + +# def test_vertical_shear(): +# xdudz = ds.xroms.dudz +# xdvdz = ds.xroms.dvdz +# acc = ds.xroms.vertical_shear +# assert np.allclose(acc, xroms.vertical_shear(xdudz, xdvdz, grid)) +# assert acc.name == acc.attrs["name"] +# hcoord = "rho" +# scoord = "s_w" +# dims = dim_dict[hcoord][scoord] +# axes = axesTZYX +# coords = coord_dict[hcoord][scoord] +# coordnames = coordnamesTZYX +# for ax, dim in zip(axes, dims): +# assert acc.cf[ax].name == dim +# for coordname, coord in zip(coordnames, coords): +# assert acc.cf[coordname].name == coord + + +# def test_relative_vorticity(): +# acc = ds.xroms.vort +# assert np.allclose(acc, 0) +# assert acc.name == acc.attrs["name"] +# hcoord = "psi" +# scoord = "s_w" +# dims = dim_dict[hcoord][scoord] +# axes = axesTZYX +# coords = coord_dict[hcoord][scoord] +# coordnames = coordnamesTZYX +# for ax, dim in zip(axes, dims): +# assert acc.cf[ax].name == dim +# for coordname, coord in zip(coordnames, coords): +# assert acc.cf[coordname].name == coord + + +# def test_convergence(): +# acc = ds.xroms.convergence +# assert np.allclose(acc, xroms.convergence(ds["u"], ds["v"], grid)) + + +# def test_convergence_norm(): +# acc = ds.xroms.convergence_norm +# assert np.allclose( +# acc, xroms.convergence(ds["u"], ds["v"], grid).cf.isel(Z=-1) / ds["f"] +# ) + + +# def test_ertel(): +# acc = ds.xroms.ertel +# xsig0 = xroms.potential_density(ds.temp, ds.salt) +# xbuoy = xroms.buoyancy(xsig0) +# assert np.allclose(acc, xroms.ertel(xbuoy, ds.u, ds.v, ds.f, grid)) +# assert acc.name == acc.attrs["name"] +# hcoord = "rho" +# scoord = "s_rho" +# dims = dim_dict[hcoord][scoord] +# axes = axesTZYX +# coords = coord_dict[hcoord][scoord] +# coordnames = coordnamesTZYX +# for ax, dim in zip(axes, dims): +# assert acc.cf[ax].name == dim +# for coordname, coord in zip(coordnames, coords): +# assert acc.cf[coordname].name == coord + + +# def test_w(): +# # VRX +# pass + + +# # acc = ds.xroms.w +# # assert np.allclose(acc, xroms.w(ds.u, ds.v, grid)) +# # acc.name == acc.attrs['name'] +# # acc.attrs['grid'] == ds.xroms.grid +# # items = ['T','X','Y','Z','longitude','latitude','vertical','time'] +# # assert set(items).issubset(acc.cf.get_valid_keys()) + + +# def test_omega(): +# # VRX +# pass + + +# # acc = ds.xroms.omega +# # assert np.allclose(acc, xroms.omega(ds.u, ds.v, grid)) +# # acc.name == acc.attrs['name'] +# # acc.attrs['grid'] == ds.xroms.grid +# # items = ['T','X','Y','Z','longitude','latitude','vertical','time'] +# # assert set(items).issubset(acc.cf.get_valid_keys()) + + +# def test_rho(): +# acc = ds.xroms.rho +# assert np.allclose(acc, xroms.density(ds.temp, ds.salt, ds.z_rho)) +# assert acc.name == acc.attrs["name"] +# hcoord = "rho" +# scoord = "s_rho" +# dims = dim_dict[hcoord][scoord] +# axes = axesTZYX +# coords = coord_dict[hcoord][scoord] +# coordnames = coordnamesTZYX +# for ax, dim in zip(axes, dims): +# assert acc.cf[ax].name == dim +# for coordname, coord in zip(coordnames, coords): +# assert acc.cf[coordname].name == coord + + +# def test_sig0(): +# acc = ds.xroms.sig0 +# assert np.allclose(acc, xroms.potential_density(ds.temp, ds.salt, 0)) +# assert acc.name == acc.attrs["name"] +# hcoord = "rho" +# scoord = "s_rho" +# dims = dim_dict[hcoord][scoord] +# axes = axesTZYX +# coords = coord_dict[hcoord][scoord] +# coordnames = coordnamesTZYX +# for ax, dim in zip(axes, dims): +# assert acc.cf[ax].name == dim +# for coordname, coord in zip(coordnames, coords): +# assert acc.cf[coordname].name == coord + + +# def test_buoyancy(): +# acc = ds.xroms.buoyancy +# xsig0 = xroms.potential_density(ds.temp, ds.salt) +# assert np.allclose(acc, xroms.buoyancy(xsig0)) +# assert acc.name == acc.attrs["name"] +# hcoord = "rho" +# scoord = "s_rho" +# dims = dim_dict[hcoord][scoord] +# axes = axesTZYX +# coords = coord_dict[hcoord][scoord] +# coordnames = coordnamesTZYX +# for ax, dim in zip(axes, dims): +# assert acc.cf[ax].name == dim +# for coordname, coord in zip(coordnames, coords): +# assert acc.cf[coordname].name == coord + + +# def test_N2(): +# acc = ds.xroms.N2 +# xrho = xroms.density(ds.temp, ds.salt, ds.z_rho) +# assert np.allclose(acc, xroms.N2(xrho, grid), equal_nan=True) +# assert acc.name == acc.attrs["name"] +# hcoord = "rho" +# scoord = "s_w" +# dims = dim_dict[hcoord][scoord] +# axes = axesTZYX +# coords = coord_dict[hcoord][scoord] +# coordnames = coordnamesTZYX +# for ax, dim in zip(axes, dims): +# assert acc.cf[ax].name == dim +# for coordname, coord in zip(coordnames, coords): +# assert acc.cf[coordname].name == coord + + +# def test_M2(): +# acc = ds.xroms.M2 +# xrho = xroms.density(ds.temp, ds.salt, ds.z_rho) +# assert np.allclose(acc, xroms.M2(xrho, grid), equal_nan=True) +# assert acc.name == acc.attrs["name"] +# hcoord = "rho" +# scoord = "s_w" +# dims = dim_dict[hcoord][scoord] +# axes = axesTZYX +# coords = coord_dict[hcoord][scoord] +# coordnames = coordnamesTZYX +# for ax, dim in zip(axes, dims): +# assert acc.cf[ax].name == dim +# for coordname, coord in zip(coordnames, coords): +# assert acc.cf[coordname].name == coord + + +# def test_mld(): +# acc = ds.xroms.mld(thresh=0.03) +# sig0 = xroms.potential_density(ds.temp, ds.salt, 0) +# assert np.allclose(acc, xroms.mld(sig0, grid, ds.h, ds.mask_rho), equal_nan=True) +# assert acc.name == acc.attrs["name"] +# hcoord = "rho" +# scoord = None +# dims = dim_dict[hcoord][scoord] +# axes = axesTYX +# coords = coord_dict[hcoord][scoord] +# coordnames = coordnamesTYX +# for ax, dim in zip(axes, dims): +# assert acc.cf[ax].name == dim +# for coordname, coord in zip(coordnames, coords): +# assert acc.cf[coordname].name == coord + + +# def test_ddxi(): # testvars = ["salt", "u", "v"] # for testvar in testvars: -# acc = ds[testvar].xroms.sel2d(lon0, lat0) -# out = xroms.sel2d( -# ds[testvar], -# ds[testvar].cf["longitude"], -# ds[testvar].cf["latitude"], -# lon0, -# lat0, -# ) -# assert np.allclose(acc, out) -# assert acc.name == testvar -# dims = ds[testvar].dims +# with pytest.raises(KeyError): +# acc = ds[testvar].xroms.ddxi(grid) + +# if testvar == "salt": +# hcoord = "u" +# scoord = "s_w" +# elif testvar == "u": +# hcoord = "rho" +# scoord = "s_w" +# elif testvar == "v": +# hcoord = "psi" +# scoord = "s_w" +# dims = dim_dict[hcoord][scoord] +# axes = axesTZYX +# coords = coord_dict[hcoord][scoord] +# coordnames = coordnamesTZYX + +# acc = ds.xroms.ddxi(testvar) +# assert np.allclose(acc, xroms.ddxi(ds[testvar], grid)) +# assert acc.name == acc.attrs["name"] +# for ax, dim in zip(axes, dims): +# assert acc.cf[ax].name == dim +# for coordname, coord in zip(coordnames, coords): +# assert acc.cf[coordname].name == coord + + +# def test_ddeta(): +# testvars = ["salt", "u", "v"] +# for testvar in testvars: +# with pytest.raises(KeyError): +# acc = ds[testvar].xroms.ddeta(grid) + +# if testvar == "salt": +# hcoord = "v" +# scoord = "s_w" +# elif testvar == "u": +# hcoord = "psi" +# scoord = "s_w" +# elif testvar == "v": +# hcoord = "rho" +# scoord = "s_w" +# dims = dim_dict[hcoord][scoord] +# axes = axesTZYX +# coords = coord_dict[hcoord][scoord] +# coordnames = coordnamesTZYX + +# acc = ds.xroms.ddeta(testvar) +# assert np.allclose(acc, xroms.ddeta(ds[testvar], grid)) +# assert acc.name == acc.attrs["name"] +# for ax, dim in zip(axes, dims): +# assert acc.cf[ax].name == dim +# for coordname, coord in zip(coordnames, coords): +# assert acc.cf[coordname].name == coord + + +# def test_ddz(): +# testvars = ["salt", "u", "v"] +# for testvar in testvars: +# with pytest.raises(KeyError): +# acc = ds[testvar].xroms.ddz(grid) +# dims = list(ds[testvar].dims) # axes = axesTZYX # coords = [ds[testvar].cf[coordname].name for coordname in coordnamesTZYX] # coordnames = coordnamesTZYX +# # correct dim and coord in derivative direction # # import pdb; pdb.set_trace() +# if grid._get_dims_from_axis(ds[testvar], "Z")[0] == "s_rho": +# # if grid.axes["Z"]._get_axis_coord(ds[testvar])[1] == "s_rho": +# dims[1] = "s_w" +# coords[1] = coords[1].replace("rho", "w") +# else: +# dims[1] = "s_rho" +# coords[1] = coords[1].replace("w", "rho") + +# acc = ds.xroms.ddz(testvar) +# assert np.allclose(acc, xroms.ddz(ds[testvar], grid)) +# assert acc.name == acc.attrs["name"] # for ax, dim in zip(axes, dims): # assert acc.cf[ax].name == dim # for coordname, coord in zip(coordnames, coords): # assert acc.cf[coordname].name == coord -def test_argsel2d(): - lon0, lat0 = -94.8, 28.0 - testvars = ["salt", "u", "v"] - for testvar in testvars: - inds = ds[testvar].xroms.argsel2d(lon0, lat0) - outinds = xroms.argsel2d( - ds[testvar].cf["longitude"], ds[testvar].cf["latitude"], lon0, lat0 - ) - assert np.allclose(inds, outinds) - - -def test_gridmean(): - testvars = ["salt", "u", "v"] - for testvar in testvars: - for axis in ["Z", "Y", "X"]: - var1 = ds[testvar].xroms.gridmean(grid, axis) - var2 = xroms.gridmean(ds[testvar], grid, axis) - assert np.allclose(var1, var2) - - -def test_gridsum(): - testvars = ["salt", "u", "v"] - for testvar in testvars: - for axis in ["Z", "Y", "X"]: - var1 = ds[testvar].xroms.gridsum(grid, axis) - var2 = xroms.gridsum(ds[testvar], grid, axis) - assert np.allclose(var1, var2) - - -def test_interpll(): - XESMF_AVAILABLE = xroms.XESMF_AVAILABLE - xroms.XESMF_AVAILABLE = False - - with pytest.raises(ModuleNotFoundError): - ie, ix = 2, 3 - indexer = {"eta_rho": [ie], "xi_rho": [ix]} - testvars = ["salt", "u", "v"] - for testvar in testvars: - var1 = xroms.interpll( - ds[testvar], ds.lon_rho.isel(indexer), ds.lat_rho.isel(indexer) - ) - var2 = ds[testvar].xroms.interpll( - ds.lon_rho.isel(indexer), ds.lat_rho.isel(indexer) - ) - assert np.allclose(var1, var2) - - # put back the way it was for testing - xroms.XESMF_AVAILABLE = XESMF_AVAILABLE - - -def test_zslice(): - testvars = ["salt", "u", "v"] - for testvar in testvars: - varin = ds[testvar] - depths = np.asarray(ds[testvar].cf["vertical"][0, :, 0, 0].values) - varout = xroms.isoslice(varin, depths, grid, axis="Z") - varcomp = ds[testvar].xroms.zslice(grid, depths) - # varcomp = ds[testvar].xroms.isoslice(grid, depths, axis="Z") - assert np.allclose( - varout.cf.isel(T=0, Y=0, X=0), varcomp.cf.isel(T=0, Y=0, X=0) - ) - - varcompds = ds.xroms.zslice(testvar, depths) - # varcomp = ds[testvar].xroms.isoslice(grid, depths, axis="Z") - assert np.allclose( - varout.cf.isel(T=0, Y=0, X=0), varcompds.cf.isel(T=0, Y=0, X=0) - ) - - -def test_find_horizontal_velocities(): - uname, vname = ds.xroms.find_horizontal_velocities() - assert uname == "u" - assert vname == "v" - - # have to delete variables in the xroms accessor for this - # test to work - ds.xroms.ds["u_eastward"] = ds.xroms.ds["u"].copy() - del ds.xroms.ds["u"] - ds.xroms.ds["v_northward"] = ds.xroms.ds["v"].copy() - del ds.xroms.ds["v"] - - uname, vname = ds.xroms.find_horizontal_velocities() - assert uname == "u_eastward" - assert vname == "v_northward" +# def test_to_grid(): + +# testvars = ["salt", "u", "v"] +# for testvar in testvars: +# for scoord in ["s_w", "s_rho"]: +# for hcoord in ["rho", "u", "v", "psi"]: +# acc = ds.xroms.to_grid(testvar, hcoord=hcoord, scoord=scoord) +# # acc = ds[testvar].xroms.to_grid(grid, hcoord=hcoord, scoord=scoord) +# assert np.allclose( +# acc, xroms.to_grid(ds[testvar], grid, hcoord=hcoord, scoord=scoord) +# ) +# assert acc.name == acc.attrs["name"] +# dims = dim_dict[hcoord][scoord] +# axes = axesTZYX +# coords = coord_dict[hcoord][scoord] +# coordnames = coordnamesTZYX +# for ax, dim in zip(axes, dims): +# assert acc.cf[ax].name == dim +# for coordname, coord in zip(coordnames, coords): +# assert acc.cf[coordname].name == coord + +# acc = ds.xroms.to_grid(testvar, hcoord=hcoord, scoord=scoord) +# assert np.allclose( +# acc, xroms.to_grid(ds[testvar], grid, hcoord=hcoord, scoord=scoord) +# ) +# assert acc.name == acc.attrs["name"] +# for ax, dim in zip(axes, dims): +# assert acc.cf[ax].name == dim +# for coordname, coord in zip(coordnames, coords): +# assert acc.cf[coordname].name == coord + + +# # can't figure out what is wrong here, will have to come back +# # def test_sel2d(): +# # lon0, lat0 = -94.8, 28.0 +# # testvars = ["salt", "u", "v"] +# # for testvar in testvars: +# # acc = ds[testvar].xroms.sel2d(lon0, lat0) +# # out = xroms.sel2d( +# # ds[testvar], +# # ds[testvar].cf["longitude"], +# # ds[testvar].cf["latitude"], +# # lon0, +# # lat0, +# # ) +# # assert np.allclose(acc, out) +# # assert acc.name == testvar +# # dims = ds[testvar].dims +# # axes = axesTZYX +# # coords = [ds[testvar].cf[coordname].name for coordname in coordnamesTZYX] +# # coordnames = coordnamesTZYX +# # # import pdb; pdb.set_trace() +# # for ax, dim in zip(axes, dims): +# # assert acc.cf[ax].name == dim +# # for coordname, coord in zip(coordnames, coords): +# # assert acc.cf[coordname].name == coord + + +# def test_argsel2d(): +# lon0, lat0 = -94.8, 28.0 +# testvars = ["salt", "u", "v"] +# for testvar in testvars: +# inds = ds[testvar].xroms.argsel2d(lon0, lat0) +# outinds = xroms.argsel2d( +# ds[testvar].cf["longitude"], ds[testvar].cf["latitude"], lon0, lat0 +# ) +# assert np.allclose(inds, outinds) + + +# def test_gridmean(): +# testvars = ["salt", "u", "v"] +# for testvar in testvars: +# for axis in ["Z", "Y", "X"]: +# var1 = ds[testvar].xroms.gridmean(grid, axis) +# var2 = xroms.gridmean(ds[testvar], grid, axis) +# assert np.allclose(var1, var2) + + +# def test_gridsum(): +# testvars = ["salt", "u", "v"] +# for testvar in testvars: +# for axis in ["Z", "Y", "X"]: +# var1 = ds[testvar].xroms.gridsum(grid, axis) +# var2 = xroms.gridsum(ds[testvar], grid, axis) +# assert np.allclose(var1, var2) + + +# def test_interpll(): +# XESMF_AVAILABLE = xroms.XESMF_AVAILABLE +# xroms.XESMF_AVAILABLE = False + +# with pytest.raises(ModuleNotFoundError): +# ie, ix = 2, 3 +# indexer = {"eta_rho": [ie], "xi_rho": [ix]} +# testvars = ["salt", "u", "v"] +# for testvar in testvars: +# var1 = xroms.interpll( +# ds[testvar], ds.lon_rho.isel(indexer), ds.lat_rho.isel(indexer) +# ) +# var2 = ds[testvar].xroms.interpll( +# ds.lon_rho.isel(indexer), ds.lat_rho.isel(indexer) +# ) +# assert np.allclose(var1, var2) + +# # put back the way it was for testing +# xroms.XESMF_AVAILABLE = XESMF_AVAILABLE + + +# def test_zslice(): +# testvars = ["salt", "u", "v"] +# for testvar in testvars: +# varin = ds[testvar] +# depths = np.asarray(ds[testvar].cf["vertical"][0, :, 0, 0].values) +# varout = xroms.isoslice(varin, depths, grid, axis="Z") +# varcomp = ds[testvar].xroms.zslice(grid, depths) +# # varcomp = ds[testvar].xroms.isoslice(grid, depths, axis="Z") +# assert np.allclose( +# varout.cf.isel(T=0, Y=0, X=0), varcomp.cf.isel(T=0, Y=0, X=0) +# ) + +# varcompds = ds.xroms.zslice(testvar, depths) +# # varcomp = ds[testvar].xroms.isoslice(grid, depths, axis="Z") +# assert np.allclose( +# varout.cf.isel(T=0, Y=0, X=0), varcompds.cf.isel(T=0, Y=0, X=0) +# ) + + +# def test_find_horizontal_velocities(): +# uname, vname = ds.xroms.find_horizontal_velocities() +# assert uname == "u" +# assert vname == "v" + +# # have to delete variables in the xroms accessor for this +# # test to work +# ds.xroms.ds["u_eastward"] = ds.xroms.ds["u"].copy() +# del ds.xroms.ds["u"] +# ds.xroms.ds["v_northward"] = ds.xroms.ds["v"].copy() +# del ds.xroms.ds["v"] + +# uname, vname = ds.xroms.find_horizontal_velocities() +# assert uname == "u_eastward" +# assert vname == "v_northward" diff --git a/xroms/utilities.py b/xroms/utilities.py index ab217e8..deacfa6 100644 --- a/xroms/utilities.py +++ b/xroms/utilities.py @@ -53,9 +53,11 @@ def grid_interp(xgrid, da, dim, which_xgcm_function="interp", **kwargs): chunk = list(da.chunks[i_chunk_dim]) # to interpolate, first remove chunking to 1 chunk - new_coord = getattr(xgrid, which_xgcm_function)( - da.chunk({dim_name: -1}), dim, **kwargs - ) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + new_coord = getattr(xgrid, which_xgcm_function)( + da.chunk({dim_name: -1}), dim, **kwargs + ) # new_coord = xgrid.interp(da.chunk({dim_name: -1}), dim, **kwargs) if ( @@ -87,7 +89,9 @@ def grid_interp(xgrid, da, dim, which_xgcm_function="interp", **kwargs): raise ValueError("chunks probably are not being dealt with properly") else: - new_coord = getattr(xgrid, which_xgcm_function)(da, dim, **kwargs) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + new_coord = getattr(xgrid, which_xgcm_function)(da, dim, **kwargs) # new_coord = xgrid.interp(da, dim, **kwargs) return new_coord @@ -214,35 +218,39 @@ def hgrad( if which in ["both", "xi"]: if is3D: - dqdx = xgrid.interp( - xgrid.derivative(q, "X", boundary=hboundary, fill_value=hfill_value), - "Z", - boundary=sboundary, - fill_value=sfill_value, - ) - dqdz = xgrid.interp( - xgrid.derivative(q, "Z", boundary=sboundary, fill_value=sfill_value), - "X", - boundary=hboundary, - fill_value=hfill_value, - ) - dzdx = xgrid.interp( - xgrid.derivative(z, "X", boundary=hboundary, fill_value=hfill_value), - "Z", - boundary=sboundary, - fill_value=sfill_value, - ) - dzdz = xgrid.interp( - xgrid.derivative(z, "Z", boundary=sboundary, fill_value=sfill_value), - "X", - boundary=hboundary, - fill_value=hfill_value, - ) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + dqdx = xgrid.interp( + xgrid.derivative(q, "X", boundary=hboundary, fill_value=hfill_value), + "Z", + boundary=sboundary, + fill_value=sfill_value, + ) + dqdz = xgrid.interp( + xgrid.derivative(q, "Z", boundary=sboundary, fill_value=sfill_value), + "X", + boundary=hboundary, + fill_value=hfill_value, + ) + dzdx = xgrid.interp( + xgrid.derivative(z, "X", boundary=hboundary, fill_value=hfill_value), + "Z", + boundary=sboundary, + fill_value=sfill_value, + ) + dzdz = xgrid.interp( + xgrid.derivative(z, "Z", boundary=sboundary, fill_value=sfill_value), + "X", + boundary=hboundary, + fill_value=hfill_value, + ) dqdxi = dqdx * dzdz - dqdz * dzdx else: # 2D variables - dqdxi = xgrid.derivative(q, "X", boundary=hboundary, fill_value=hfill_value) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + dqdxi = xgrid.derivative(q, "X", boundary=hboundary, fill_value=hfill_value) if attrs is None and isinstance(q, xr.DataArray): attrs = q.attrs.copy() @@ -266,37 +274,41 @@ def hgrad( if which in ["both", "eta"]: if is3D: - dqdy = xgrid.interp( - xgrid.derivative(q, "Y", boundary=hboundary, fill_value=hfill_value), - "Z", - boundary=sboundary, - fill_value=sfill_value, - ) - dqdz = xgrid.interp( - xgrid.derivative(q, "Z", boundary=sboundary, fill_value=sfill_value), - "Y", - boundary=hboundary, - fill_value=hfill_value, - ) - dzdy = xgrid.interp( - xgrid.derivative(z, "Y", boundary=hboundary, fill_value=hfill_value), - "Z", - boundary=sboundary, - fill_value=sfill_value, - ) - dzdz = xgrid.interp( - xgrid.derivative(z, "Z", boundary=sboundary, fill_value=sfill_value), - "Y", - boundary=hboundary, - fill_value=hfill_value, - ) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + dqdy = xgrid.interp( + xgrid.derivative(q, "Y", boundary=hboundary, fill_value=hfill_value), + "Z", + boundary=sboundary, + fill_value=sfill_value, + ) + dqdz = xgrid.interp( + xgrid.derivative(q, "Z", boundary=sboundary, fill_value=sfill_value), + "Y", + boundary=hboundary, + fill_value=hfill_value, + ) + dzdy = xgrid.interp( + xgrid.derivative(z, "Y", boundary=hboundary, fill_value=hfill_value), + "Z", + boundary=sboundary, + fill_value=sfill_value, + ) + dzdz = xgrid.interp( + xgrid.derivative(z, "Z", boundary=sboundary, fill_value=sfill_value), + "Y", + boundary=hboundary, + fill_value=hfill_value, + ) dqdeta = dqdy * dzdz - dqdz * dzdy else: # 2D variables - dqdeta = xgrid.derivative( - q, "Y", boundary=hboundary, fill_value=hfill_value - ) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + dqdeta = xgrid.derivative( + q, "Y", boundary=hboundary, fill_value=hfill_value + ) if attrs is None and isinstance(q, xr.DataArray): attrs = q.attrs.copy() @@ -639,7 +651,9 @@ def ddz( "long_name", "var" ) - var = xgrid.derivative(var, "Z", boundary=sboundary, fill_value=sfill_value) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + var = xgrid.derivative(var, "Z", boundary=sboundary, fill_value=sfill_value) var = to_grid( var, xgrid, @@ -1042,9 +1056,11 @@ def to_s_rho(var, xgrid, sboundary="extend", sfill_value=None): # only change if not already on s_rho if "s_rho" not in var.dims: - var = xgrid.interp( - var, "Z", to="center", boundary=sboundary, fill_value=sfill_value - ) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + var = xgrid.interp( + var, "Z", to="center", boundary=sboundary, fill_value=sfill_value + ) return var @@ -1094,9 +1110,11 @@ def to_s_w(var, xgrid, sboundary="extend", sfill_value=None): # only change if not already on s_w if "s_w" not in var.dims: - var = xgrid.interp( - var, "Z", to="outer", boundary=sboundary, fill_value=sfill_value - ) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + var = xgrid.interp( + var, "Z", to="outer", boundary=sboundary, fill_value=sfill_value + ) return var @@ -1257,7 +1275,9 @@ def gridmean(var, xgrid, dim): attrs.setdefault("long_name", "var") + ", grid mean over dim " + dimstr ) - var = xgrid.average(var, dim) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + var = xgrid.average(var, dim) return var diff --git a/xroms/xroms.py b/xroms/xroms.py index 486ec4c..9e07d4d 100644 --- a/xroms/xroms.py +++ b/xroms/xroms.py @@ -186,7 +186,9 @@ def roms_dataset( } ) - xgrid = xgcm.Grid(ds, coords=coords, periodic=[]) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + xgrid = xgcm.Grid(ds, coords=coords, periodic=[]) if "Vtransform" in ds.variables.keys(): Vtransform = ds.Vtransform @@ -339,21 +341,27 @@ def roms_dataset( if include_Z0: ds.coords["z_rho0"] = order(z_rho0) - ds.coords["z_rho_u0"] = xgrid.interp(ds.z_rho0, "X") + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + ds.coords["z_rho_u0"] = xgrid.interp(ds.z_rho0, "X") ds.coords["z_rho_u0"].attrs = { "long_name": "depth of U-points on vertical RHO grid", "field": "z_rho_u0, scalar", "units": "m", } - ds.coords["z_rho_v0"] = xgrid.interp(ds.z_rho0, "Y") + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + ds.coords["z_rho_v0"] = xgrid.interp(ds.z_rho0, "Y") ds.coords["z_rho_v0"].attrs = { "long_name": "depth of V-points on vertical RHO grid", "field": "z_rho_v0, scalar", "units": "m", } - ds.coords["z_rho_psi0"] = xgrid.interp(ds.z_rho_u0, "Y") + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + ds.coords["z_rho_psi0"] = xgrid.interp(ds.z_rho_u0, "Y") ds.coords["z_rho_psi0"].attrs = { "long_name": "depth of PSI-points on vertical RHO grid", "field": "z_rho_psi0, scalar", @@ -361,21 +369,27 @@ def roms_dataset( } ds.coords["z_w0"] = order(z_w0) - ds.coords["z_w_u0"] = xgrid.interp(ds.z_w0, "X") + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + ds.coords["z_w_u0"] = xgrid.interp(ds.z_w0, "X") ds.coords["z_w_u0"].attrs = { "long_name": "depth of U-points on vertical W grid", "field": "z_w_u0, scalar", "units": "m", } - ds.coords["z_w_v0"] = xgrid.interp(ds.z_w0, "Y") + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + ds.coords["z_w_v0"] = xgrid.interp(ds.z_w0, "Y") ds.coords["z_w_v0"].attrs = { "long_name": "depth of V-points on vertical W grid", "field": "z_w_v0, scalar", "units": "m", } - ds.coords["z_w_psi0"] = xgrid.interp(ds.z_w_u0, "Y") + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + ds.coords["z_w_psi0"] = xgrid.interp(ds.z_w_u0, "Y") ds.coords["z_w_psi0"].attrs = { "long_name": "depth of PSI-points on vertical W grid", "field": "z_w_psi0, scalar", @@ -405,42 +419,54 @@ def roms_dataset( # just keep these local instead of saving to Dataset — doesn't look like they # are used in any functions outside of this function. - pm_v = xgrid.interp(ds.pm, "Y") + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + pm_v = xgrid.interp(ds.pm, "Y") # ds["pm_v"].attrs = { # "long_name": "curvilinear coordinate metric in XI on V grid", # "units": "meter-1", # "field": "pm_v, scalar", # } - pn_u = xgrid.interp(ds.pn, "X") + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + pn_u = xgrid.interp(ds.pn, "X") # ds["pn_u"].attrs = { # "long_name": "curvilinear coordinate metric in ETA on U grid", # "units": "meter-1", # "field": "pn_u, scalar", # } - pm_u = xgrid.interp(ds.pm, "X") + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + pm_u = xgrid.interp(ds.pm, "X") # ds["pm_u"].attrs = { # "long_name": "curvilinear coordinate metric in XI on U grid", # "units": "meter-1", # "field": "pm_u, scalar", # } - pn_v = xgrid.interp(ds.pn, "Y") + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + pn_v = xgrid.interp(ds.pn, "Y") # ds["pn_v"].attrs = { # "long_name": "curvilinear coordinate metric in ETA on V grid", # "units": "meter-1", # "field": "pn_v, scalar", # } - pm_psi = xgrid.interp(xgrid.interp(ds.pm, "Y"), "X") # at psi points (eta_v, xi_u) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + pm_psi = xgrid.interp(xgrid.interp(ds.pm, "Y"), "X") # at psi points (eta_v, xi_u) # ds["pm_psi"].attrs = { # "long_name": "curvilinear coordinate metric in XI on PSI grid", # "units": "meter-1", # "field": "pm_psi, scalar", # } - pn_psi = xgrid.interp(xgrid.interp(ds.pn, "X"), "Y") # at psi points (eta_v, xi_u) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + pn_psi = xgrid.interp(xgrid.interp(ds.pn, "X"), "Y") # at psi points (eta_v, xi_u) # ds["pn_psi"].attrs = { # "long_name": "curvilinear coordinate metric in ETA on PSI grid", # "units": "meter-1", @@ -504,7 +530,9 @@ def roms_dataset( } if ds["3d"] and include_3D_metrics: - ds["dz"] = xgrid.diff(ds.z_w.chunk({ds.z_w.cf["Z"].name: -1}), "Z") + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + ds["dz"] = xgrid.diff(ds.z_w.chunk({ds.z_w.cf["Z"].name: -1}), "Z") ds["dz"].attrs = { "long_name": "vertical layer thickness on vertical RHO grid", "time": "ocean_time", @@ -512,7 +540,9 @@ def roms_dataset( "units": "m", } - ds["dz_w"] = xgrid.diff(ds.z_rho, "Z", boundary="fill") + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + ds["dz_w"] = xgrid.diff(ds.z_rho, "Z", boundary="fill") ds["dz_w"].attrs = { "long_name": "vertical layer thickness on vertical W grid", "time": "ocean_time", @@ -577,56 +607,72 @@ def roms_dataset( if include_Z0: # also include z coordinates with mean sea level (constant over time) - ds["dz0"] = xgrid.diff(ds.z_w0, "Z") + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + ds["dz0"] = xgrid.diff(ds.z_w0, "Z") ds["dz0"].attrs = { "long_name": "vertical layer thickness on vertical RHO grid", "field": "dz0, scalar", "units": "m", } - ds["dz_w0"] = xgrid.diff(ds.z_rho0, "Z", boundary="fill") + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + ds["dz_w0"] = xgrid.diff(ds.z_rho0, "Z", boundary="fill") ds["dz_w0"].attrs = { "long_name": "vertical layer thickness on vertical W grid", "field": "dz_w0, scalar", "units": "m", } - ds["dz_u0"] = xgrid.interp(ds.dz0, "X") + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + ds["dz_u0"] = xgrid.interp(ds.dz0, "X") ds["dz_u0"].attrs = { "long_name": "vertical layer thickness on vertical RHO grid on U grid", "field": "dz_u0, scalar", "units": "m", } - ds["dz_w_u0"] = xgrid.interp(ds.dz_w0, "X") + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + ds["dz_w_u0"] = xgrid.interp(ds.dz_w0, "X") ds["dz_w_u0"].attrs = { "long_name": "vertical layer thickness on vertical W grid on U grid", "field": "dz_w_u0, scalar", "units": "m", } - ds["dz_v0"] = xgrid.interp(ds.dz0, "Y") + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + ds["dz_v0"] = xgrid.interp(ds.dz0, "Y") ds["dz_v0"].attrs = { "long_name": "vertical layer thickness on vertical RHO grid on V grid", "field": "dz_v0, scalar", "units": "m", } - ds["dz_w_v0"] = xgrid.interp(ds.dz_w0, "Y") + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + ds["dz_w_v0"] = xgrid.interp(ds.dz_w0, "Y") ds["dz_w_v0"].attrs = { "long_name": "vertical layer thickness on vertical W grid on V grid", "field": "dz_w_v0, scalar", "units": "m", } - ds["dz_psi0"] = xgrid.interp(ds.dz_v0, "X") + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + ds["dz_psi0"] = xgrid.interp(ds.dz_v0, "X") ds["dz_psi0"].attrs = { "long_name": "vertical layer thickness on vertical RHO grid on PSI grid", "field": "dz_psi0, scalar", "units": "m", } - ds["dz_w_psi0"] = xgrid.interp(ds.dz_w_v0, "X") + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + ds["dz_w_psi0"] = xgrid.interp(ds.dz_w_v0, "X") ds["dz_w_psi0"].attrs = { "long_name": "vertical layer thickness on vertical W grid on PSI grid", "field": "dz_w_psi0, scalar", @@ -792,7 +838,9 @@ def roms_dataset( ("X", "Y"): ["dA"], # Areas } - xgrid = xgcm.Grid(ds, coords=coords, metrics=metrics, periodic=[]) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + xgrid = xgcm.Grid(ds, coords=coords, metrics=metrics, periodic=[]) # # ds.attrs['grid'] = grid # causes recursion error # # also put grid into every variable with at least 2D From 353b03e0ddd8ce694691134aa06ea5b19f2fb3fd Mon Sep 17 00:00:00 2001 From: Kristen Thyng Date: Wed, 27 Aug 2025 11:07:12 -0700 Subject: [PATCH 2/5] small changes, undid a large comment --- docs/whats_new.md | 2 + setup.cfg | 2 +- xroms/tests/test_accessor.py | 1066 +++++++++++++++++----------------- 3 files changed, 536 insertions(+), 534 deletions(-) diff --git a/docs/whats_new.md b/docs/whats_new.md index 5084b0b..5706637 100644 --- a/docs/whats_new.md +++ b/docs/whats_new.md @@ -2,6 +2,8 @@ ## v0.6.2 (August 27, 2025) * catching and ignoring a bunch of warnings from `xgcm`, but still staying with `xgcm` `v0.8.1` until I can update this code to match. +* updating CI test versions to 3.11, 3.12, 3.13 + ## v0.6.1 (October 28, 2024) * Correction in a few built-in calculations of u/v grid to rho-grid interpolations of u and v velocities (currently `speed` and `_uv2eastnorth`). In these cases, we need to fill nans with zeros so that the masked locations in the velocity fields are not fully brought forward into the rho mask but are instead interpolated over. By making them 0, they are calculated into the mask\_rho positions by combining them with neighboring cells. If this wasn't done, the fact that they are masked would supersede the neighboring cells and they would be masked in mask\_rho. This needs to be done anytime the velocities are moved from their native grids to the rho or other grids to preserve their locations around masked cells. diff --git a/setup.cfg b/setup.cfg index df20499..ae4d007 100644 --- a/setup.cfg +++ b/setup.cfg @@ -80,7 +80,7 @@ install_requires = numpy pooch xarray - xgcm == 0.8.1 + xgcm==0.8.1 setup_requires= setuptools_scm python_requires = >=3.6 diff --git a/xroms/tests/test_accessor.py b/xroms/tests/test_accessor.py index 7fe5d54..3debff6 100644 --- a/xroms/tests/test_accessor.py +++ b/xroms/tests/test_accessor.py @@ -74,548 +74,548 @@ def test_grid(): assert isinstance(ds.xroms.xgrid, xgcm.grid.Grid) -# def test_speed(): - -# acc = ds.xroms.speed - -# assert np.allclose(acc, xroms.speed(ds.u, ds.v, grid)) - -# # also check attributes -# assert acc.name == acc.attrs["name"] -# # cf-xarray: make sure all Axes and Coordinates available in output -# hcoord = "rho" -# scoord = "s_rho" -# dims = dim_dict[hcoord][scoord] -# axes = axesTZYX -# coords = coord_dict[hcoord][scoord] -# coordnames = coordnamesTZYX -# for ax, dim in zip(axes, dims): -# assert acc.cf[ax].name == dim -# for coordname, coord in zip(coordnames, coords): -# assert acc.cf[coordname].name == coord - - -# def test_KE(): - -# s = xroms.speed(ds.u, ds.v, grid) -# acc = ds.xroms.KE - -# assert np.allclose(acc, xroms.KE(ds.rho0, s)) - -# # also check attributes -# assert acc.name == acc.attrs["name"] -# # cf-xarray: make sure all Axes and Coordinates available in output -# hcoord = "rho" -# scoord = "s_rho" -# dims = dim_dict[hcoord][scoord] -# axes = axesTZYX -# coords = coord_dict[hcoord][scoord] -# coordnames = coordnamesTZYX -# for ax, dim in zip(axes, dims): -# assert acc.cf[ax].name == dim -# for coordname, coord in zip(coordnames, coords): -# assert acc.cf[coordname].name == coord - - -# def test_uv_geostrophic(): - -# acc = ds.xroms.ug -# assert np.allclose(acc, xroms.uv_geostrophic(ds.zeta, ds.f, grid, which="xi")) -# # also check attributes -# assert ( -# acc.name == acc.attrs["name"] -# ) # cf-xarray: make sure all Axes and Coordinates available in output -# hcoord = "u" -# scoord = None -# dims = dim_dict[hcoord][scoord] -# axes = axesTYX -# coords = coord_dict[hcoord][scoord] -# coordnames = coordnamesTYX -# for ax, dim in zip(axes, dims): -# assert acc.cf[ax].name == dim -# for coordname, coord in zip(coordnames, coords): -# assert acc.cf[coordname].name == coord - -# acc = ds.xroms.vg -# assert np.allclose(acc, xroms.uv_geostrophic(ds.zeta, ds.f, grid, which="eta")) -# assert acc.name == acc.attrs["name"] -# hcoord = "v" -# scoord = None -# dims = dim_dict[hcoord][scoord] -# axes = axesTYX -# coords = coord_dict[hcoord][scoord] -# coordnames = coordnamesTYX -# for ax, dim in zip(axes, dims): -# assert acc.cf[ax].name == dim -# for coordname, coord in zip(coordnames, coords): -# assert acc.cf[coordname].name == coord - - -# def test_EKE(): - -# acc = ds.xroms.EKE -# xug, xvg = xroms.uv_geostrophic(ds.zeta, ds.f, grid, which="both") -# assert np.allclose(acc, xroms.EKE(xug, xvg, grid)) -# assert acc.name == acc.attrs["name"] -# hcoord = "rho" -# scoord = None -# dims = dim_dict[hcoord][scoord] -# axes = axesTYX -# coords = coord_dict[hcoord][scoord] -# coordnames = coordnamesTYX -# for ax, dim in zip(axes, dims): -# assert acc.cf[ax].name == dim -# for coordname, coord in zip(coordnames, coords): -# assert acc.cf[coordname].name == coord - - -# def test_dudz(): -# acc = ds.xroms.dudz -# assert np.allclose(acc, xroms.dudz(ds.u, grid)) -# assert acc.name == acc.attrs["name"] -# hcoord = "u" -# scoord = "s_w" -# dims = dim_dict[hcoord][scoord] -# axes = axesTZYX -# coords = coord_dict[hcoord][scoord] -# coordnames = coordnamesTZYX -# for ax, dim in zip(axes, dims): -# assert acc.cf[ax].name == dim -# for coordname, coord in zip(coordnames, coords): -# assert acc.cf[coordname].name == coord - - -# def test_dvdz(): -# acc = ds.xroms.dvdz -# assert np.allclose(acc, xroms.dvdz(ds.v, grid)) -# assert acc.name == acc.attrs["name"] -# hcoord = "v" -# scoord = "s_w" -# dims = dim_dict[hcoord][scoord] -# axes = axesTZYX -# coords = coord_dict[hcoord][scoord] -# coordnames = coordnamesTZYX -# for ax, dim in zip(axes, dims): -# assert acc.cf[ax].name == dim -# for coordname, coord in zip(coordnames, coords): -# assert acc.cf[coordname].name == coord - - -# def test_vertical_shear(): -# xdudz = ds.xroms.dudz -# xdvdz = ds.xroms.dvdz -# acc = ds.xroms.vertical_shear -# assert np.allclose(acc, xroms.vertical_shear(xdudz, xdvdz, grid)) -# assert acc.name == acc.attrs["name"] -# hcoord = "rho" -# scoord = "s_w" -# dims = dim_dict[hcoord][scoord] -# axes = axesTZYX -# coords = coord_dict[hcoord][scoord] -# coordnames = coordnamesTZYX -# for ax, dim in zip(axes, dims): -# assert acc.cf[ax].name == dim -# for coordname, coord in zip(coordnames, coords): -# assert acc.cf[coordname].name == coord - - -# def test_relative_vorticity(): -# acc = ds.xroms.vort -# assert np.allclose(acc, 0) -# assert acc.name == acc.attrs["name"] -# hcoord = "psi" -# scoord = "s_w" -# dims = dim_dict[hcoord][scoord] -# axes = axesTZYX -# coords = coord_dict[hcoord][scoord] -# coordnames = coordnamesTZYX -# for ax, dim in zip(axes, dims): -# assert acc.cf[ax].name == dim -# for coordname, coord in zip(coordnames, coords): -# assert acc.cf[coordname].name == coord - - -# def test_convergence(): -# acc = ds.xroms.convergence -# assert np.allclose(acc, xroms.convergence(ds["u"], ds["v"], grid)) - - -# def test_convergence_norm(): -# acc = ds.xroms.convergence_norm -# assert np.allclose( -# acc, xroms.convergence(ds["u"], ds["v"], grid).cf.isel(Z=-1) / ds["f"] -# ) - - -# def test_ertel(): -# acc = ds.xroms.ertel -# xsig0 = xroms.potential_density(ds.temp, ds.salt) -# xbuoy = xroms.buoyancy(xsig0) -# assert np.allclose(acc, xroms.ertel(xbuoy, ds.u, ds.v, ds.f, grid)) -# assert acc.name == acc.attrs["name"] -# hcoord = "rho" -# scoord = "s_rho" -# dims = dim_dict[hcoord][scoord] -# axes = axesTZYX -# coords = coord_dict[hcoord][scoord] -# coordnames = coordnamesTZYX -# for ax, dim in zip(axes, dims): -# assert acc.cf[ax].name == dim -# for coordname, coord in zip(coordnames, coords): -# assert acc.cf[coordname].name == coord - - -# def test_w(): -# # VRX -# pass - - -# # acc = ds.xroms.w -# # assert np.allclose(acc, xroms.w(ds.u, ds.v, grid)) -# # acc.name == acc.attrs['name'] -# # acc.attrs['grid'] == ds.xroms.grid -# # items = ['T','X','Y','Z','longitude','latitude','vertical','time'] -# # assert set(items).issubset(acc.cf.get_valid_keys()) - - -# def test_omega(): -# # VRX -# pass - - -# # acc = ds.xroms.omega -# # assert np.allclose(acc, xroms.omega(ds.u, ds.v, grid)) -# # acc.name == acc.attrs['name'] -# # acc.attrs['grid'] == ds.xroms.grid -# # items = ['T','X','Y','Z','longitude','latitude','vertical','time'] -# # assert set(items).issubset(acc.cf.get_valid_keys()) - - -# def test_rho(): -# acc = ds.xroms.rho -# assert np.allclose(acc, xroms.density(ds.temp, ds.salt, ds.z_rho)) -# assert acc.name == acc.attrs["name"] -# hcoord = "rho" -# scoord = "s_rho" -# dims = dim_dict[hcoord][scoord] -# axes = axesTZYX -# coords = coord_dict[hcoord][scoord] -# coordnames = coordnamesTZYX -# for ax, dim in zip(axes, dims): -# assert acc.cf[ax].name == dim -# for coordname, coord in zip(coordnames, coords): -# assert acc.cf[coordname].name == coord - - -# def test_sig0(): -# acc = ds.xroms.sig0 -# assert np.allclose(acc, xroms.potential_density(ds.temp, ds.salt, 0)) -# assert acc.name == acc.attrs["name"] -# hcoord = "rho" -# scoord = "s_rho" -# dims = dim_dict[hcoord][scoord] -# axes = axesTZYX -# coords = coord_dict[hcoord][scoord] -# coordnames = coordnamesTZYX -# for ax, dim in zip(axes, dims): -# assert acc.cf[ax].name == dim -# for coordname, coord in zip(coordnames, coords): -# assert acc.cf[coordname].name == coord - - -# def test_buoyancy(): -# acc = ds.xroms.buoyancy -# xsig0 = xroms.potential_density(ds.temp, ds.salt) -# assert np.allclose(acc, xroms.buoyancy(xsig0)) -# assert acc.name == acc.attrs["name"] -# hcoord = "rho" -# scoord = "s_rho" -# dims = dim_dict[hcoord][scoord] -# axes = axesTZYX -# coords = coord_dict[hcoord][scoord] -# coordnames = coordnamesTZYX -# for ax, dim in zip(axes, dims): -# assert acc.cf[ax].name == dim -# for coordname, coord in zip(coordnames, coords): -# assert acc.cf[coordname].name == coord - - -# def test_N2(): -# acc = ds.xroms.N2 -# xrho = xroms.density(ds.temp, ds.salt, ds.z_rho) -# assert np.allclose(acc, xroms.N2(xrho, grid), equal_nan=True) -# assert acc.name == acc.attrs["name"] -# hcoord = "rho" -# scoord = "s_w" -# dims = dim_dict[hcoord][scoord] -# axes = axesTZYX -# coords = coord_dict[hcoord][scoord] -# coordnames = coordnamesTZYX -# for ax, dim in zip(axes, dims): -# assert acc.cf[ax].name == dim -# for coordname, coord in zip(coordnames, coords): -# assert acc.cf[coordname].name == coord - - -# def test_M2(): -# acc = ds.xroms.M2 -# xrho = xroms.density(ds.temp, ds.salt, ds.z_rho) -# assert np.allclose(acc, xroms.M2(xrho, grid), equal_nan=True) -# assert acc.name == acc.attrs["name"] -# hcoord = "rho" -# scoord = "s_w" -# dims = dim_dict[hcoord][scoord] -# axes = axesTZYX -# coords = coord_dict[hcoord][scoord] -# coordnames = coordnamesTZYX -# for ax, dim in zip(axes, dims): -# assert acc.cf[ax].name == dim -# for coordname, coord in zip(coordnames, coords): -# assert acc.cf[coordname].name == coord - - -# def test_mld(): -# acc = ds.xroms.mld(thresh=0.03) -# sig0 = xroms.potential_density(ds.temp, ds.salt, 0) -# assert np.allclose(acc, xroms.mld(sig0, grid, ds.h, ds.mask_rho), equal_nan=True) -# assert acc.name == acc.attrs["name"] -# hcoord = "rho" -# scoord = None -# dims = dim_dict[hcoord][scoord] -# axes = axesTYX -# coords = coord_dict[hcoord][scoord] -# coordnames = coordnamesTYX -# for ax, dim in zip(axes, dims): -# assert acc.cf[ax].name == dim -# for coordname, coord in zip(coordnames, coords): -# assert acc.cf[coordname].name == coord - - -# def test_ddxi(): -# testvars = ["salt", "u", "v"] -# for testvar in testvars: -# with pytest.raises(KeyError): -# acc = ds[testvar].xroms.ddxi(grid) - -# if testvar == "salt": -# hcoord = "u" -# scoord = "s_w" -# elif testvar == "u": -# hcoord = "rho" -# scoord = "s_w" -# elif testvar == "v": -# hcoord = "psi" -# scoord = "s_w" -# dims = dim_dict[hcoord][scoord] -# axes = axesTZYX -# coords = coord_dict[hcoord][scoord] -# coordnames = coordnamesTZYX - -# acc = ds.xroms.ddxi(testvar) -# assert np.allclose(acc, xroms.ddxi(ds[testvar], grid)) -# assert acc.name == acc.attrs["name"] -# for ax, dim in zip(axes, dims): -# assert acc.cf[ax].name == dim -# for coordname, coord in zip(coordnames, coords): -# assert acc.cf[coordname].name == coord - - -# def test_ddeta(): -# testvars = ["salt", "u", "v"] -# for testvar in testvars: -# with pytest.raises(KeyError): -# acc = ds[testvar].xroms.ddeta(grid) - -# if testvar == "salt": -# hcoord = "v" -# scoord = "s_w" -# elif testvar == "u": -# hcoord = "psi" -# scoord = "s_w" -# elif testvar == "v": -# hcoord = "rho" -# scoord = "s_w" -# dims = dim_dict[hcoord][scoord] -# axes = axesTZYX -# coords = coord_dict[hcoord][scoord] -# coordnames = coordnamesTZYX - -# acc = ds.xroms.ddeta(testvar) -# assert np.allclose(acc, xroms.ddeta(ds[testvar], grid)) -# assert acc.name == acc.attrs["name"] -# for ax, dim in zip(axes, dims): -# assert acc.cf[ax].name == dim -# for coordname, coord in zip(coordnames, coords): -# assert acc.cf[coordname].name == coord - - -# def test_ddz(): +def test_speed(): + + acc = ds.xroms.speed + + assert np.allclose(acc, xroms.speed(ds.u, ds.v, grid)) + + # also check attributes + assert acc.name == acc.attrs["name"] + # cf-xarray: make sure all Axes and Coordinates available in output + hcoord = "rho" + scoord = "s_rho" + dims = dim_dict[hcoord][scoord] + axes = axesTZYX + coords = coord_dict[hcoord][scoord] + coordnames = coordnamesTZYX + for ax, dim in zip(axes, dims): + assert acc.cf[ax].name == dim + for coordname, coord in zip(coordnames, coords): + assert acc.cf[coordname].name == coord + + +def test_KE(): + + s = xroms.speed(ds.u, ds.v, grid) + acc = ds.xroms.KE + + assert np.allclose(acc, xroms.KE(ds.rho0, s)) + + # also check attributes + assert acc.name == acc.attrs["name"] + # cf-xarray: make sure all Axes and Coordinates available in output + hcoord = "rho" + scoord = "s_rho" + dims = dim_dict[hcoord][scoord] + axes = axesTZYX + coords = coord_dict[hcoord][scoord] + coordnames = coordnamesTZYX + for ax, dim in zip(axes, dims): + assert acc.cf[ax].name == dim + for coordname, coord in zip(coordnames, coords): + assert acc.cf[coordname].name == coord + + +def test_uv_geostrophic(): + + acc = ds.xroms.ug + assert np.allclose(acc, xroms.uv_geostrophic(ds.zeta, ds.f, grid, which="xi")) + # also check attributes + assert ( + acc.name == acc.attrs["name"] + ) # cf-xarray: make sure all Axes and Coordinates available in output + hcoord = "u" + scoord = None + dims = dim_dict[hcoord][scoord] + axes = axesTYX + coords = coord_dict[hcoord][scoord] + coordnames = coordnamesTYX + for ax, dim in zip(axes, dims): + assert acc.cf[ax].name == dim + for coordname, coord in zip(coordnames, coords): + assert acc.cf[coordname].name == coord + + acc = ds.xroms.vg + assert np.allclose(acc, xroms.uv_geostrophic(ds.zeta, ds.f, grid, which="eta")) + assert acc.name == acc.attrs["name"] + hcoord = "v" + scoord = None + dims = dim_dict[hcoord][scoord] + axes = axesTYX + coords = coord_dict[hcoord][scoord] + coordnames = coordnamesTYX + for ax, dim in zip(axes, dims): + assert acc.cf[ax].name == dim + for coordname, coord in zip(coordnames, coords): + assert acc.cf[coordname].name == coord + + +def test_EKE(): + + acc = ds.xroms.EKE + xug, xvg = xroms.uv_geostrophic(ds.zeta, ds.f, grid, which="both") + assert np.allclose(acc, xroms.EKE(xug, xvg, grid)) + assert acc.name == acc.attrs["name"] + hcoord = "rho" + scoord = None + dims = dim_dict[hcoord][scoord] + axes = axesTYX + coords = coord_dict[hcoord][scoord] + coordnames = coordnamesTYX + for ax, dim in zip(axes, dims): + assert acc.cf[ax].name == dim + for coordname, coord in zip(coordnames, coords): + assert acc.cf[coordname].name == coord + + +def test_dudz(): + acc = ds.xroms.dudz + assert np.allclose(acc, xroms.dudz(ds.u, grid)) + assert acc.name == acc.attrs["name"] + hcoord = "u" + scoord = "s_w" + dims = dim_dict[hcoord][scoord] + axes = axesTZYX + coords = coord_dict[hcoord][scoord] + coordnames = coordnamesTZYX + for ax, dim in zip(axes, dims): + assert acc.cf[ax].name == dim + for coordname, coord in zip(coordnames, coords): + assert acc.cf[coordname].name == coord + + +def test_dvdz(): + acc = ds.xroms.dvdz + assert np.allclose(acc, xroms.dvdz(ds.v, grid)) + assert acc.name == acc.attrs["name"] + hcoord = "v" + scoord = "s_w" + dims = dim_dict[hcoord][scoord] + axes = axesTZYX + coords = coord_dict[hcoord][scoord] + coordnames = coordnamesTZYX + for ax, dim in zip(axes, dims): + assert acc.cf[ax].name == dim + for coordname, coord in zip(coordnames, coords): + assert acc.cf[coordname].name == coord + + +def test_vertical_shear(): + xdudz = ds.xroms.dudz + xdvdz = ds.xroms.dvdz + acc = ds.xroms.vertical_shear + assert np.allclose(acc, xroms.vertical_shear(xdudz, xdvdz, grid)) + assert acc.name == acc.attrs["name"] + hcoord = "rho" + scoord = "s_w" + dims = dim_dict[hcoord][scoord] + axes = axesTZYX + coords = coord_dict[hcoord][scoord] + coordnames = coordnamesTZYX + for ax, dim in zip(axes, dims): + assert acc.cf[ax].name == dim + for coordname, coord in zip(coordnames, coords): + assert acc.cf[coordname].name == coord + + +def test_relative_vorticity(): + acc = ds.xroms.vort + assert np.allclose(acc, 0) + assert acc.name == acc.attrs["name"] + hcoord = "psi" + scoord = "s_w" + dims = dim_dict[hcoord][scoord] + axes = axesTZYX + coords = coord_dict[hcoord][scoord] + coordnames = coordnamesTZYX + for ax, dim in zip(axes, dims): + assert acc.cf[ax].name == dim + for coordname, coord in zip(coordnames, coords): + assert acc.cf[coordname].name == coord + + +def test_convergence(): + acc = ds.xroms.convergence + assert np.allclose(acc, xroms.convergence(ds["u"], ds["v"], grid)) + + +def test_convergence_norm(): + acc = ds.xroms.convergence_norm + assert np.allclose( + acc, xroms.convergence(ds["u"], ds["v"], grid).cf.isel(Z=-1) / ds["f"] + ) + + +def test_ertel(): + acc = ds.xroms.ertel + xsig0 = xroms.potential_density(ds.temp, ds.salt) + xbuoy = xroms.buoyancy(xsig0) + assert np.allclose(acc, xroms.ertel(xbuoy, ds.u, ds.v, ds.f, grid)) + assert acc.name == acc.attrs["name"] + hcoord = "rho" + scoord = "s_rho" + dims = dim_dict[hcoord][scoord] + axes = axesTZYX + coords = coord_dict[hcoord][scoord] + coordnames = coordnamesTZYX + for ax, dim in zip(axes, dims): + assert acc.cf[ax].name == dim + for coordname, coord in zip(coordnames, coords): + assert acc.cf[coordname].name == coord + + +def test_w(): + # VRX + pass + + +# acc = ds.xroms.w +# assert np.allclose(acc, xroms.w(ds.u, ds.v, grid)) +# acc.name == acc.attrs['name'] +# acc.attrs['grid'] == ds.xroms.grid +# items = ['T','X','Y','Z','longitude','latitude','vertical','time'] +# assert set(items).issubset(acc.cf.get_valid_keys()) + + +def test_omega(): + # VRX + pass + + +# acc = ds.xroms.omega +# assert np.allclose(acc, xroms.omega(ds.u, ds.v, grid)) +# acc.name == acc.attrs['name'] +# acc.attrs['grid'] == ds.xroms.grid +# items = ['T','X','Y','Z','longitude','latitude','vertical','time'] +# assert set(items).issubset(acc.cf.get_valid_keys()) + + +def test_rho(): + acc = ds.xroms.rho + assert np.allclose(acc, xroms.density(ds.temp, ds.salt, ds.z_rho)) + assert acc.name == acc.attrs["name"] + hcoord = "rho" + scoord = "s_rho" + dims = dim_dict[hcoord][scoord] + axes = axesTZYX + coords = coord_dict[hcoord][scoord] + coordnames = coordnamesTZYX + for ax, dim in zip(axes, dims): + assert acc.cf[ax].name == dim + for coordname, coord in zip(coordnames, coords): + assert acc.cf[coordname].name == coord + + +def test_sig0(): + acc = ds.xroms.sig0 + assert np.allclose(acc, xroms.potential_density(ds.temp, ds.salt, 0)) + assert acc.name == acc.attrs["name"] + hcoord = "rho" + scoord = "s_rho" + dims = dim_dict[hcoord][scoord] + axes = axesTZYX + coords = coord_dict[hcoord][scoord] + coordnames = coordnamesTZYX + for ax, dim in zip(axes, dims): + assert acc.cf[ax].name == dim + for coordname, coord in zip(coordnames, coords): + assert acc.cf[coordname].name == coord + + +def test_buoyancy(): + acc = ds.xroms.buoyancy + xsig0 = xroms.potential_density(ds.temp, ds.salt) + assert np.allclose(acc, xroms.buoyancy(xsig0)) + assert acc.name == acc.attrs["name"] + hcoord = "rho" + scoord = "s_rho" + dims = dim_dict[hcoord][scoord] + axes = axesTZYX + coords = coord_dict[hcoord][scoord] + coordnames = coordnamesTZYX + for ax, dim in zip(axes, dims): + assert acc.cf[ax].name == dim + for coordname, coord in zip(coordnames, coords): + assert acc.cf[coordname].name == coord + + +def test_N2(): + acc = ds.xroms.N2 + xrho = xroms.density(ds.temp, ds.salt, ds.z_rho) + assert np.allclose(acc, xroms.N2(xrho, grid), equal_nan=True) + assert acc.name == acc.attrs["name"] + hcoord = "rho" + scoord = "s_w" + dims = dim_dict[hcoord][scoord] + axes = axesTZYX + coords = coord_dict[hcoord][scoord] + coordnames = coordnamesTZYX + for ax, dim in zip(axes, dims): + assert acc.cf[ax].name == dim + for coordname, coord in zip(coordnames, coords): + assert acc.cf[coordname].name == coord + + +def test_M2(): + acc = ds.xroms.M2 + xrho = xroms.density(ds.temp, ds.salt, ds.z_rho) + assert np.allclose(acc, xroms.M2(xrho, grid), equal_nan=True) + assert acc.name == acc.attrs["name"] + hcoord = "rho" + scoord = "s_w" + dims = dim_dict[hcoord][scoord] + axes = axesTZYX + coords = coord_dict[hcoord][scoord] + coordnames = coordnamesTZYX + for ax, dim in zip(axes, dims): + assert acc.cf[ax].name == dim + for coordname, coord in zip(coordnames, coords): + assert acc.cf[coordname].name == coord + + +def test_mld(): + acc = ds.xroms.mld(thresh=0.03) + sig0 = xroms.potential_density(ds.temp, ds.salt, 0) + assert np.allclose(acc, xroms.mld(sig0, grid, ds.h, ds.mask_rho), equal_nan=True) + assert acc.name == acc.attrs["name"] + hcoord = "rho" + scoord = None + dims = dim_dict[hcoord][scoord] + axes = axesTYX + coords = coord_dict[hcoord][scoord] + coordnames = coordnamesTYX + for ax, dim in zip(axes, dims): + assert acc.cf[ax].name == dim + for coordname, coord in zip(coordnames, coords): + assert acc.cf[coordname].name == coord + + +def test_ddxi(): + testvars = ["salt", "u", "v"] + for testvar in testvars: + with pytest.raises(KeyError): + acc = ds[testvar].xroms.ddxi(grid) + + if testvar == "salt": + hcoord = "u" + scoord = "s_w" + elif testvar == "u": + hcoord = "rho" + scoord = "s_w" + elif testvar == "v": + hcoord = "psi" + scoord = "s_w" + dims = dim_dict[hcoord][scoord] + axes = axesTZYX + coords = coord_dict[hcoord][scoord] + coordnames = coordnamesTZYX + + acc = ds.xroms.ddxi(testvar) + assert np.allclose(acc, xroms.ddxi(ds[testvar], grid)) + assert acc.name == acc.attrs["name"] + for ax, dim in zip(axes, dims): + assert acc.cf[ax].name == dim + for coordname, coord in zip(coordnames, coords): + assert acc.cf[coordname].name == coord + + +def test_ddeta(): + testvars = ["salt", "u", "v"] + for testvar in testvars: + with pytest.raises(KeyError): + acc = ds[testvar].xroms.ddeta(grid) + + if testvar == "salt": + hcoord = "v" + scoord = "s_w" + elif testvar == "u": + hcoord = "psi" + scoord = "s_w" + elif testvar == "v": + hcoord = "rho" + scoord = "s_w" + dims = dim_dict[hcoord][scoord] + axes = axesTZYX + coords = coord_dict[hcoord][scoord] + coordnames = coordnamesTZYX + + acc = ds.xroms.ddeta(testvar) + assert np.allclose(acc, xroms.ddeta(ds[testvar], grid)) + assert acc.name == acc.attrs["name"] + for ax, dim in zip(axes, dims): + assert acc.cf[ax].name == dim + for coordname, coord in zip(coordnames, coords): + assert acc.cf[coordname].name == coord + + +def test_ddz(): + testvars = ["salt", "u", "v"] + for testvar in testvars: + with pytest.raises(KeyError): + acc = ds[testvar].xroms.ddz(grid) + dims = list(ds[testvar].dims) + axes = axesTZYX + coords = [ds[testvar].cf[coordname].name for coordname in coordnamesTZYX] + coordnames = coordnamesTZYX + # correct dim and coord in derivative direction + # import pdb; pdb.set_trace() + if grid._get_dims_from_axis(ds[testvar], "Z")[0] == "s_rho": + # if grid.axes["Z"]._get_axis_coord(ds[testvar])[1] == "s_rho": + dims[1] = "s_w" + coords[1] = coords[1].replace("rho", "w") + else: + dims[1] = "s_rho" + coords[1] = coords[1].replace("w", "rho") + + acc = ds.xroms.ddz(testvar) + assert np.allclose(acc, xroms.ddz(ds[testvar], grid)) + assert acc.name == acc.attrs["name"] + for ax, dim in zip(axes, dims): + assert acc.cf[ax].name == dim + for coordname, coord in zip(coordnames, coords): + assert acc.cf[coordname].name == coord + + +def test_to_grid(): + + testvars = ["salt", "u", "v"] + for testvar in testvars: + for scoord in ["s_w", "s_rho"]: + for hcoord in ["rho", "u", "v", "psi"]: + acc = ds.xroms.to_grid(testvar, hcoord=hcoord, scoord=scoord) + # acc = ds[testvar].xroms.to_grid(grid, hcoord=hcoord, scoord=scoord) + assert np.allclose( + acc, xroms.to_grid(ds[testvar], grid, hcoord=hcoord, scoord=scoord) + ) + assert acc.name == acc.attrs["name"] + dims = dim_dict[hcoord][scoord] + axes = axesTZYX + coords = coord_dict[hcoord][scoord] + coordnames = coordnamesTZYX + for ax, dim in zip(axes, dims): + assert acc.cf[ax].name == dim + for coordname, coord in zip(coordnames, coords): + assert acc.cf[coordname].name == coord + + acc = ds.xroms.to_grid(testvar, hcoord=hcoord, scoord=scoord) + assert np.allclose( + acc, xroms.to_grid(ds[testvar], grid, hcoord=hcoord, scoord=scoord) + ) + assert acc.name == acc.attrs["name"] + for ax, dim in zip(axes, dims): + assert acc.cf[ax].name == dim + for coordname, coord in zip(coordnames, coords): + assert acc.cf[coordname].name == coord + + +# can't figure out what is wrong here, will have to come back +# def test_sel2d(): +# lon0, lat0 = -94.8, 28.0 # testvars = ["salt", "u", "v"] # for testvar in testvars: -# with pytest.raises(KeyError): -# acc = ds[testvar].xroms.ddz(grid) -# dims = list(ds[testvar].dims) +# acc = ds[testvar].xroms.sel2d(lon0, lat0) +# out = xroms.sel2d( +# ds[testvar], +# ds[testvar].cf["longitude"], +# ds[testvar].cf["latitude"], +# lon0, +# lat0, +# ) +# assert np.allclose(acc, out) +# assert acc.name == testvar +# dims = ds[testvar].dims # axes = axesTZYX # coords = [ds[testvar].cf[coordname].name for coordname in coordnamesTZYX] # coordnames = coordnamesTZYX -# # correct dim and coord in derivative direction # # import pdb; pdb.set_trace() -# if grid._get_dims_from_axis(ds[testvar], "Z")[0] == "s_rho": -# # if grid.axes["Z"]._get_axis_coord(ds[testvar])[1] == "s_rho": -# dims[1] = "s_w" -# coords[1] = coords[1].replace("rho", "w") -# else: -# dims[1] = "s_rho" -# coords[1] = coords[1].replace("w", "rho") - -# acc = ds.xroms.ddz(testvar) -# assert np.allclose(acc, xroms.ddz(ds[testvar], grid)) -# assert acc.name == acc.attrs["name"] # for ax, dim in zip(axes, dims): # assert acc.cf[ax].name == dim # for coordname, coord in zip(coordnames, coords): # assert acc.cf[coordname].name == coord -# def test_to_grid(): - -# testvars = ["salt", "u", "v"] -# for testvar in testvars: -# for scoord in ["s_w", "s_rho"]: -# for hcoord in ["rho", "u", "v", "psi"]: -# acc = ds.xroms.to_grid(testvar, hcoord=hcoord, scoord=scoord) -# # acc = ds[testvar].xroms.to_grid(grid, hcoord=hcoord, scoord=scoord) -# assert np.allclose( -# acc, xroms.to_grid(ds[testvar], grid, hcoord=hcoord, scoord=scoord) -# ) -# assert acc.name == acc.attrs["name"] -# dims = dim_dict[hcoord][scoord] -# axes = axesTZYX -# coords = coord_dict[hcoord][scoord] -# coordnames = coordnamesTZYX -# for ax, dim in zip(axes, dims): -# assert acc.cf[ax].name == dim -# for coordname, coord in zip(coordnames, coords): -# assert acc.cf[coordname].name == coord - -# acc = ds.xroms.to_grid(testvar, hcoord=hcoord, scoord=scoord) -# assert np.allclose( -# acc, xroms.to_grid(ds[testvar], grid, hcoord=hcoord, scoord=scoord) -# ) -# assert acc.name == acc.attrs["name"] -# for ax, dim in zip(axes, dims): -# assert acc.cf[ax].name == dim -# for coordname, coord in zip(coordnames, coords): -# assert acc.cf[coordname].name == coord - - -# # can't figure out what is wrong here, will have to come back -# # def test_sel2d(): -# # lon0, lat0 = -94.8, 28.0 -# # testvars = ["salt", "u", "v"] -# # for testvar in testvars: -# # acc = ds[testvar].xroms.sel2d(lon0, lat0) -# # out = xroms.sel2d( -# # ds[testvar], -# # ds[testvar].cf["longitude"], -# # ds[testvar].cf["latitude"], -# # lon0, -# # lat0, -# # ) -# # assert np.allclose(acc, out) -# # assert acc.name == testvar -# # dims = ds[testvar].dims -# # axes = axesTZYX -# # coords = [ds[testvar].cf[coordname].name for coordname in coordnamesTZYX] -# # coordnames = coordnamesTZYX -# # # import pdb; pdb.set_trace() -# # for ax, dim in zip(axes, dims): -# # assert acc.cf[ax].name == dim -# # for coordname, coord in zip(coordnames, coords): -# # assert acc.cf[coordname].name == coord - - -# def test_argsel2d(): -# lon0, lat0 = -94.8, 28.0 -# testvars = ["salt", "u", "v"] -# for testvar in testvars: -# inds = ds[testvar].xroms.argsel2d(lon0, lat0) -# outinds = xroms.argsel2d( -# ds[testvar].cf["longitude"], ds[testvar].cf["latitude"], lon0, lat0 -# ) -# assert np.allclose(inds, outinds) - - -# def test_gridmean(): -# testvars = ["salt", "u", "v"] -# for testvar in testvars: -# for axis in ["Z", "Y", "X"]: -# var1 = ds[testvar].xroms.gridmean(grid, axis) -# var2 = xroms.gridmean(ds[testvar], grid, axis) -# assert np.allclose(var1, var2) - - -# def test_gridsum(): -# testvars = ["salt", "u", "v"] -# for testvar in testvars: -# for axis in ["Z", "Y", "X"]: -# var1 = ds[testvar].xroms.gridsum(grid, axis) -# var2 = xroms.gridsum(ds[testvar], grid, axis) -# assert np.allclose(var1, var2) - - -# def test_interpll(): -# XESMF_AVAILABLE = xroms.XESMF_AVAILABLE -# xroms.XESMF_AVAILABLE = False - -# with pytest.raises(ModuleNotFoundError): -# ie, ix = 2, 3 -# indexer = {"eta_rho": [ie], "xi_rho": [ix]} -# testvars = ["salt", "u", "v"] -# for testvar in testvars: -# var1 = xroms.interpll( -# ds[testvar], ds.lon_rho.isel(indexer), ds.lat_rho.isel(indexer) -# ) -# var2 = ds[testvar].xroms.interpll( -# ds.lon_rho.isel(indexer), ds.lat_rho.isel(indexer) -# ) -# assert np.allclose(var1, var2) - -# # put back the way it was for testing -# xroms.XESMF_AVAILABLE = XESMF_AVAILABLE - - -# def test_zslice(): -# testvars = ["salt", "u", "v"] -# for testvar in testvars: -# varin = ds[testvar] -# depths = np.asarray(ds[testvar].cf["vertical"][0, :, 0, 0].values) -# varout = xroms.isoslice(varin, depths, grid, axis="Z") -# varcomp = ds[testvar].xroms.zslice(grid, depths) -# # varcomp = ds[testvar].xroms.isoslice(grid, depths, axis="Z") -# assert np.allclose( -# varout.cf.isel(T=0, Y=0, X=0), varcomp.cf.isel(T=0, Y=0, X=0) -# ) - -# varcompds = ds.xroms.zslice(testvar, depths) -# # varcomp = ds[testvar].xroms.isoslice(grid, depths, axis="Z") -# assert np.allclose( -# varout.cf.isel(T=0, Y=0, X=0), varcompds.cf.isel(T=0, Y=0, X=0) -# ) - - -# def test_find_horizontal_velocities(): -# uname, vname = ds.xroms.find_horizontal_velocities() -# assert uname == "u" -# assert vname == "v" - -# # have to delete variables in the xroms accessor for this -# # test to work -# ds.xroms.ds["u_eastward"] = ds.xroms.ds["u"].copy() -# del ds.xroms.ds["u"] -# ds.xroms.ds["v_northward"] = ds.xroms.ds["v"].copy() -# del ds.xroms.ds["v"] - -# uname, vname = ds.xroms.find_horizontal_velocities() -# assert uname == "u_eastward" -# assert vname == "v_northward" +def test_argsel2d(): + lon0, lat0 = -94.8, 28.0 + testvars = ["salt", "u", "v"] + for testvar in testvars: + inds = ds[testvar].xroms.argsel2d(lon0, lat0) + outinds = xroms.argsel2d( + ds[testvar].cf["longitude"], ds[testvar].cf["latitude"], lon0, lat0 + ) + assert np.allclose(inds, outinds) + + +def test_gridmean(): + testvars = ["salt", "u", "v"] + for testvar in testvars: + for axis in ["Z", "Y", "X"]: + var1 = ds[testvar].xroms.gridmean(grid, axis) + var2 = xroms.gridmean(ds[testvar], grid, axis) + assert np.allclose(var1, var2) + + +def test_gridsum(): + testvars = ["salt", "u", "v"] + for testvar in testvars: + for axis in ["Z", "Y", "X"]: + var1 = ds[testvar].xroms.gridsum(grid, axis) + var2 = xroms.gridsum(ds[testvar], grid, axis) + assert np.allclose(var1, var2) + + +def test_interpll(): + XESMF_AVAILABLE = xroms.XESMF_AVAILABLE + xroms.XESMF_AVAILABLE = False + + with pytest.raises(ModuleNotFoundError): + ie, ix = 2, 3 + indexer = {"eta_rho": [ie], "xi_rho": [ix]} + testvars = ["salt", "u", "v"] + for testvar in testvars: + var1 = xroms.interpll( + ds[testvar], ds.lon_rho.isel(indexer), ds.lat_rho.isel(indexer) + ) + var2 = ds[testvar].xroms.interpll( + ds.lon_rho.isel(indexer), ds.lat_rho.isel(indexer) + ) + assert np.allclose(var1, var2) + + # put back the way it was for testing + xroms.XESMF_AVAILABLE = XESMF_AVAILABLE + + +def test_zslice(): + testvars = ["salt", "u", "v"] + for testvar in testvars: + varin = ds[testvar] + depths = np.asarray(ds[testvar].cf["vertical"][0, :, 0, 0].values) + varout = xroms.isoslice(varin, depths, grid, axis="Z") + varcomp = ds[testvar].xroms.zslice(grid, depths) + # varcomp = ds[testvar].xroms.isoslice(grid, depths, axis="Z") + assert np.allclose( + varout.cf.isel(T=0, Y=0, X=0), varcomp.cf.isel(T=0, Y=0, X=0) + ) + + varcompds = ds.xroms.zslice(testvar, depths) + # varcomp = ds[testvar].xroms.isoslice(grid, depths, axis="Z") + assert np.allclose( + varout.cf.isel(T=0, Y=0, X=0), varcompds.cf.isel(T=0, Y=0, X=0) + ) + + +def test_find_horizontal_velocities(): + uname, vname = ds.xroms.find_horizontal_velocities() + assert uname == "u" + assert vname == "v" + + # have to delete variables in the xroms accessor for this + # test to work + ds.xroms.ds["u_eastward"] = ds.xroms.ds["u"].copy() + del ds.xroms.ds["u"] + ds.xroms.ds["v_northward"] = ds.xroms.ds["v"].copy() + del ds.xroms.ds["v"] + + uname, vname = ds.xroms.find_horizontal_velocities() + assert uname == "u_eastward" + assert vname == "v_northward" From 9e35f771cbca54db0de914682c3e7fbcc6e1a502 Mon Sep 17 00:00:00 2001 From: Kristen Thyng Date: Wed, 27 Aug 2025 11:10:40 -0700 Subject: [PATCH 3/5] updated test workflow for new ci files --- .github/workflows/test.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 989281e..7058c16 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -17,7 +17,7 @@ jobs: fail-fast: false matrix: os: ["macos-latest", "ubuntu-latest", "windows-latest"] - python-version: ["3.9", "3.10", "3.11"] + python-version: ["3.11", "3.12", "3.13"] steps: - uses: actions/checkout@v4 - name: Cache conda From 2a2f4d7e53b62546dff930c6f85c297303be366f Mon Sep 17 00:00:00 2001 From: Kristen Thyng Date: Wed, 27 Aug 2025 11:14:28 -0700 Subject: [PATCH 4/5] ran pre-commit now --- xroms/utilities.py | 36 +++++++++++++++++++++++++++--------- xroms/xroms.py | 8 ++++++-- 2 files changed, 33 insertions(+), 11 deletions(-) diff --git a/xroms/utilities.py b/xroms/utilities.py index deacfa6..a8382c2 100644 --- a/xroms/utilities.py +++ b/xroms/utilities.py @@ -221,25 +221,33 @@ def hgrad( with warnings.catch_warnings(): warnings.simplefilter("ignore") dqdx = xgrid.interp( - xgrid.derivative(q, "X", boundary=hboundary, fill_value=hfill_value), + xgrid.derivative( + q, "X", boundary=hboundary, fill_value=hfill_value + ), "Z", boundary=sboundary, fill_value=sfill_value, ) dqdz = xgrid.interp( - xgrid.derivative(q, "Z", boundary=sboundary, fill_value=sfill_value), + xgrid.derivative( + q, "Z", boundary=sboundary, fill_value=sfill_value + ), "X", boundary=hboundary, fill_value=hfill_value, ) dzdx = xgrid.interp( - xgrid.derivative(z, "X", boundary=hboundary, fill_value=hfill_value), + xgrid.derivative( + z, "X", boundary=hboundary, fill_value=hfill_value + ), "Z", boundary=sboundary, fill_value=sfill_value, ) dzdz = xgrid.interp( - xgrid.derivative(z, "Z", boundary=sboundary, fill_value=sfill_value), + xgrid.derivative( + z, "Z", boundary=sboundary, fill_value=sfill_value + ), "X", boundary=hboundary, fill_value=hfill_value, @@ -250,7 +258,9 @@ def hgrad( else: # 2D variables with warnings.catch_warnings(): warnings.simplefilter("ignore") - dqdxi = xgrid.derivative(q, "X", boundary=hboundary, fill_value=hfill_value) + dqdxi = xgrid.derivative( + q, "X", boundary=hboundary, fill_value=hfill_value + ) if attrs is None and isinstance(q, xr.DataArray): attrs = q.attrs.copy() @@ -277,25 +287,33 @@ def hgrad( with warnings.catch_warnings(): warnings.simplefilter("ignore") dqdy = xgrid.interp( - xgrid.derivative(q, "Y", boundary=hboundary, fill_value=hfill_value), + xgrid.derivative( + q, "Y", boundary=hboundary, fill_value=hfill_value + ), "Z", boundary=sboundary, fill_value=sfill_value, ) dqdz = xgrid.interp( - xgrid.derivative(q, "Z", boundary=sboundary, fill_value=sfill_value), + xgrid.derivative( + q, "Z", boundary=sboundary, fill_value=sfill_value + ), "Y", boundary=hboundary, fill_value=hfill_value, ) dzdy = xgrid.interp( - xgrid.derivative(z, "Y", boundary=hboundary, fill_value=hfill_value), + xgrid.derivative( + z, "Y", boundary=hboundary, fill_value=hfill_value + ), "Z", boundary=sboundary, fill_value=sfill_value, ) dzdz = xgrid.interp( - xgrid.derivative(z, "Z", boundary=sboundary, fill_value=sfill_value), + xgrid.derivative( + z, "Z", boundary=sboundary, fill_value=sfill_value + ), "Y", boundary=hboundary, fill_value=hfill_value, diff --git a/xroms/xroms.py b/xroms/xroms.py index 9e07d4d..2499a3e 100644 --- a/xroms/xroms.py +++ b/xroms/xroms.py @@ -457,7 +457,9 @@ def roms_dataset( with warnings.catch_warnings(): warnings.simplefilter("ignore") - pm_psi = xgrid.interp(xgrid.interp(ds.pm, "Y"), "X") # at psi points (eta_v, xi_u) + pm_psi = xgrid.interp( + xgrid.interp(ds.pm, "Y"), "X" + ) # at psi points (eta_v, xi_u) # ds["pm_psi"].attrs = { # "long_name": "curvilinear coordinate metric in XI on PSI grid", # "units": "meter-1", @@ -466,7 +468,9 @@ def roms_dataset( with warnings.catch_warnings(): warnings.simplefilter("ignore") - pn_psi = xgrid.interp(xgrid.interp(ds.pn, "X"), "Y") # at psi points (eta_v, xi_u) + pn_psi = xgrid.interp( + xgrid.interp(ds.pn, "X"), "Y" + ) # at psi points (eta_v, xi_u) # ds["pn_psi"].attrs = { # "long_name": "curvilinear coordinate metric in ETA on PSI grid", # "units": "meter-1", From 4ee24ba3e402e9e2de8bcedfdc4dbae0b96ac5ef Mon Sep 17 00:00:00 2001 From: Kristen Thyng Date: Wed, 27 Aug 2025 11:25:21 -0700 Subject: [PATCH 5/5] added scipy into docs env for readthedocs --- docs/environment.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/environment.yml b/docs/environment.yml index 4ca38b4..251c205 100644 --- a/docs/environment.yml +++ b/docs/environment.yml @@ -18,6 +18,7 @@ dependencies: - pip - pooch - requests + - scipy - xarray - xcmocean - xgcm=0.8.1