From 92044dcee0598602bfd040dc55ea8060a4f292b0 Mon Sep 17 00:00:00 2001 From: Henri Drake Date: Tue, 24 Feb 2026 05:51:18 -0800 Subject: [PATCH 1/5] Add minimal test of full water mass budget In the process of writing this test, I uncovered several places where there is unnecessary friction in applying the package to very simple examples. This test should be used to expand testing coverage to all options of the keyword arguments in order to find remaining bugs! --- .github/workflows/ci.yml | 50 ++++++++++++++++++ ci/environment.yml | 2 +- xwmb/budget.py | 2 + xwmb/tests/test_budget.py | 108 ++++++++++++++++++++++++++++++++++++++ 4 files changed, 161 insertions(+), 1 deletion(-) create mode 100644 .github/workflows/ci.yml create mode 100644 xwmb/tests/test_budget.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..7bb6844 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,50 @@ +name: CI + +on: + push: + branches: + - main + pull_request: + branches: + - main + +jobs: + build: + + runs-on: ubuntu-latest + defaults: + run: + shell: bash -l {0} + strategy: + fail-fast: false + matrix: + python-version: ['3.11', '3.12', '3.13', '3.14'] + + steps: + - name: Cancel previous runs + uses: styfle/cancel-workflow-action@0.7.0 + with: + access_token: ${{ github.token }} + - name: Checkout source + uses: actions/checkout@v2 + + - name: Conda setup + uses: conda-incubator/setup-miniconda@v2 + with: + channels: conda-forge + mamba-version: '*' + python-version: ${{ matrix.python-version }} + activate-environment: test_env_xwmb + auto-activate-base: false + + - name: Set up conda environment + run: | + mamba env update -f ci/environment.yml + python -m pip install -e . + - name: Conda list information + run: | + conda env list + conda list + - name: Test with pytest + run: | + pytest diff --git a/ci/environment.yml b/ci/environment.yml index cc3256a..0c54158 100644 --- a/ci/environment.yml +++ b/ci/environment.yml @@ -3,7 +3,7 @@ channels: - conda-forge - nodefaults dependencies: - - python>=3.12 + - python>=3.11 - black - netcdf4 - pylint diff --git a/xwmb/budget.py b/xwmb/budget.py index 3d77a7d..566fa5a 100644 --- a/xwmb/budget.py +++ b/xwmb/budget.py @@ -161,6 +161,8 @@ def mass_budget(self, lambda_name, greater_than=False, integrate=True, along_sec raise TypeError(f"Boolean or list of 3 numbers expected, got {type(default_bins).__name__}") elif isinstance(bins, np.ndarray): self.add_bins_gridcoords(lambda_name, bins) + elif isinstance(bins, xr.DataArray): + self.add_bins_gridcoords(lambda_name, bins.values) elif bins is None: pass else: diff --git a/xwmb/tests/test_budget.py b/xwmb/tests/test_budget.py new file mode 100644 index 0000000..a86ed4c --- /dev/null +++ b/xwmb/tests/test_budget.py @@ -0,0 +1,108 @@ +import xgcm +import xwmb +import xarray as xr +import numpy as np + +def synthetic_dataset(): + x_f = np.array([-0.5, 0.5]) + x_c = 0.5*(x_f[:-1] + x_f[1:]) + + y_f = np.array([0.5, 1.5, 2.5]) + y_c = 0.5*(y_f[:-1] + y_f[1:]) + + lam_f = np.array([0, 1, 2]) + lam_c = 0.5*(lam_f[:-1] + lam_f[1:]) + + t_f = np.array([0, 1]) + t_c = 0.5*(t_f[:-1] + t_f[1:]) + + coords = { + "x_c": x_c, + "x_f": x_f, + "y_c": y_c, + "y_f": y_f, + "lam_c": lam_c, + "lam_f": lam_f, + "t_c": t_c, + "t_f": t_f + } + ds = xr.Dataset(coords=coords) + ds = ds.assign_coords({ + "geolon": xr.broadcast(ds.x_c, ds.y_c)[0], + "geolat": xr.broadcast(ds.x_c, ds.y_c)[0], + "geolon_c": xr.broadcast(ds.x_f, ds.y_f)[0], + "geolat_c": xr.broadcast(ds.x_f, ds.y_f)[1], + }) + + # Grid cell area + ds["area"] = xr.ones_like(xr.broadcast(ds.x_c, ds.y_c)[0]) + + # Time-averaged size and contours of water mass + ds["lam"] = ds.lam_c * xr.ones_like(xr.broadcast(ds.t_c, ds.lam_c, ds.y_c, ds.x_c)[0]) + ds["thickness"] = xr.ones_like(ds.lam) + + # Bounding snapshots of size and contours of water mass + ds["lam_bounds"] = ds.lam_c * xr.ones_like(xr.broadcast(ds.t_f, ds.lam_c, ds.y_c, ds.x_c)[0]) + ds["thickness_bounds"] = xr.ones_like(ds.lam_bounds) + + # Lateral mass transport + ds["umo"] = xr.zeros_like(xr.broadcast(ds.t_c, ds.lam_c, ds.y_c, ds.x_f)[0]) + ds["vmo"] = xr.where( + ds.y_f == 1.5, + ds.lam_c - 1., + xr.zeros_like(xr.broadcast(ds.t_c, ds.lam_c, ds.y_f, ds.x_c)[0]) + ) + + # Volume-integrated tendency + ds["tend"] = xr.where( + ds.y_c == 2.0, + 1., + xr.zeros_like(xr.broadcast(ds.t_c, ds.lam_c, ds.y_c, ds.x_c)[0]) + ) + return ds + +def synthetic_grid(): + + ds = synthetic_dataset() + + # Placeholder until https://github.com/hdrake/xbudget/issues/21 + ds = ds.rename({"t_c":"time", "t_f":"time_bounds"}) + + coords = { + "X": {"center":"x_c", "outer":"x_f"}, + "Y": {"center":"y_c", "outer":"y_f"}, + "Z": {"center":"lam_c", "outer":"lam_f"}, + "T": {"center":"time", "outer":"time_bounds"} + } + grid = xgcm.Grid( + ds, + coords = coords, + boundary = {"X": "extend", "Y":"extend", "Z":"extend", "T":"extend"}, + metrics = {("X","Y"): "area"}, + autoparse_metadata=False + ) + + return grid + +def test_mass_budget(): + grid = synthetic_grid() + + xbudget_dict = { + "mass": { + "thickness": "thickness", + "rhs": {"sum": {"advection": {"sum": {"lateral": {"sum": { + "zonal_convergence": {"product": {"zonal_divergence": {"difference": {"zonal_mass_transport": "umo"}}}}, + "meridional_convergence": {"product": {"meridional_divergence": {"difference": {"meridional_mass_transport": "vmo"}}}} + }}}}}} + }, + "tracer": {"lambda": "lam", "rhs": {"sum": {"tendency": {"var": "tend"}}}} + } + + wmb = xwmb.WaterMassBudget( + grid, + xbudget_dict, + rebin=False, + rho_ref = 1. + ) + + wmb.mass_budget("tracer", bins=grid._ds.lam_f).compute() \ No newline at end of file From ceb2e6b1d9bc7465ab0ff9f35a4977dd7f1af31e Mon Sep 17 00:00:00 2001 From: Henri Drake Date: Tue, 24 Feb 2026 05:53:43 -0800 Subject: [PATCH 2/5] Add testing suite --- xwmb/version.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xwmb/version.py b/xwmb/version.py index b4f94f4..b8005f9 100644 --- a/xwmb/version.py +++ b/xwmb/version.py @@ -1,3 +1,3 @@ """xwmb: version information""" -__version__ = "0.5.6" +__version__ = "0.5.7" From 1c43cddf6fb9ef27408ca7390e17367bf5c7743a Mon Sep 17 00:00:00 2001 From: Henri Drake Date: Tue, 24 Feb 2026 05:54:11 -0800 Subject: [PATCH 3/5] Requires xwmt > 0.2.0 because uses a generic tracer --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index ffaf2a3..9bf4531 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,7 +18,7 @@ classifiers = [ ] dependencies = [ "regionate >= 0.5.0", - "xwmt >= 0.1.0", + "xwmt >= 0.2.0", ] [project.urls] From a89766ba4e7fa08904731934b8edad570a7bbe1f Mon Sep 17 00:00:00 2001 From: Henri Drake Date: Tue, 24 Feb 2026 15:48:09 -0800 Subject: [PATCH 4/5] Don't test on 3.14 yet and make CI workflow more robust --- .github/workflows/ci.yml | 4 ++-- ci/environment.yml | 1 - 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 7bb6844..1b816c6 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -18,7 +18,7 @@ jobs: strategy: fail-fast: false matrix: - python-version: ['3.11', '3.12', '3.13', '3.14'] + python-version: ['3.11', '3.12', '3.13'] steps: - name: Cancel previous runs @@ -39,7 +39,7 @@ jobs: - name: Set up conda environment run: | - mamba env update -f ci/environment.yml + mamba env update -n test_env_xwmb -f ci/environment.yml --prune python -m pip install -e . - name: Conda list information run: | diff --git a/ci/environment.yml b/ci/environment.yml index 0c54158..f9b6dd4 100644 --- a/ci/environment.yml +++ b/ci/environment.yml @@ -3,7 +3,6 @@ channels: - conda-forge - nodefaults dependencies: - - python>=3.11 - black - netcdf4 - pylint From e655e129ae3030f684e155def72125576c9f2d4c Mon Sep 17 00:00:00 2001 From: Henri Drake Date: Tue, 24 Feb 2026 15:55:12 -0800 Subject: [PATCH 5/5] Remove prune to avoid build issues --- .github/workflows/ci.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 1b816c6..dc273e8 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -39,7 +39,7 @@ jobs: - name: Set up conda environment run: | - mamba env update -n test_env_xwmb -f ci/environment.yml --prune + mamba env update -n test_env_xwmb -f ci/environment.yml python -m pip install -e . - name: Conda list information run: |