From 058c874b1b33f4c5bb108fd69cd974fb4915812b Mon Sep 17 00:00:00 2001 From: tomchor Date: Sun, 20 Apr 2025 12:51:09 -0600 Subject: [PATCH 01/13] start updating names to new nomenclature --- pynanigans/grids.py | 4 ++-- pynanigans/pnplot.py | 2 +- pynanigans/utils.py | 10 +++++----- 3 files changed, 8 insertions(+), 8 deletions(-) diff --git a/pynanigans/grids.py b/pynanigans/grids.py index b845536..9061630 100644 --- a/pynanigans/grids.py +++ b/pynanigans/grids.py @@ -4,8 +4,8 @@ def get_coords(ds, topology="PPN",): Constructs the coords dict for ds to be passed to xgcm.Grid Flat dimensions (F) are treated the same as Periodic ones (P) """ - per = dict(left='xF', center='xC') - nper = dict(outer='xF', center='xC') + per = dict(left='xF', center='x_caa') + nper = dict(outer='xF', center='x_caa') per = { dim : dict(left=f"{dim}F", center=f"{dim}C") for dim in "xyz" } nper = { dim : dict(outer=f"{dim}F", center=f"{dim}C") for dim in "xyz" } coords = { dim : per[dim] if top in "FP" else nper[dim] for dim, top in zip("xyz", topology) } diff --git a/pynanigans/pnplot.py b/pynanigans/pnplot.py index 695be0e..8a90558 100644 --- a/pynanigans/pnplot.py +++ b/pynanigans/pnplot.py @@ -5,7 +5,7 @@ def pnplot(darray, surjection=surjection, **kwargs): """ Bijects darray to rename the dimensions before calling plot(). - This makes plot easier as, instead of calling, `ds.u.plot(x='xF', y='zC')`, + This makes plot easier as, instead of calling, `ds.u.plot(x='xF', y='z_aac')`, you can call `ds.pnplot(x='x', y='z')` """ return biject(darray, surjection=surjection).plot(**kwargs) diff --git a/pynanigans/utils.py b/pynanigans/utils.py index ec1ad38..9d8dcf9 100644 --- a/pynanigans/utils.py +++ b/pynanigans/utils.py @@ -1,11 +1,11 @@ import xarray as xr from .grids import get_grid -surjection = dict(xC='x', +surjection = dict(x_caa='x', xF='x', - yC='y', + y_aca='y', yF='y', - zC='z', + z_aac='z', zF='z', ) @@ -17,7 +17,7 @@ def biject(darray, *args, surjection=surjection): If `*args` is provided, only those dimensions will be renamed. If not, `x`, `y` and `z` will be automatically renamed. - This makes calling functions easier as instead of calling `darray.u.plot(x='xF', y='zC')`, + This makes calling functions easier as instead of calling `darray.u.plot(x='xF', y='z_aac')`, you can call `darray.pnplot(x='x', y='z')` """ da_dims = darray.dims @@ -60,7 +60,7 @@ def downsample(darray, round_func=round, **dim_limits): Downsamples `darray` based on dimensions given in dim_limits dim_limits should be of the form: - dim_limits = dict(yC=1000, zF=2048) + dim_limits = dict(y_aca=1000, zF=2048) """ for dim, dim_limit in dim_limits.items(): dim_length = len(darray[dim]) From 4fde9ee76f6c9d8bafe6c526413ab94e1d16b79f Mon Sep 17 00:00:00 2001 From: tomchor Date: Sun, 20 Apr 2025 12:52:22 -0600 Subject: [PATCH 02/13] change names of face locations --- pynanigans/grids.py | 4 ++-- pynanigans/pnplot.py | 2 +- pynanigans/utils.py | 10 +++++----- 3 files changed, 8 insertions(+), 8 deletions(-) diff --git a/pynanigans/grids.py b/pynanigans/grids.py index 9061630..e0472ba 100644 --- a/pynanigans/grids.py +++ b/pynanigans/grids.py @@ -4,8 +4,8 @@ def get_coords(ds, topology="PPN",): Constructs the coords dict for ds to be passed to xgcm.Grid Flat dimensions (F) are treated the same as Periodic ones (P) """ - per = dict(left='xF', center='x_caa') - nper = dict(outer='xF', center='x_caa') + per = dict(left='x_faa', center='x_caa') + nper = dict(outer='x_faa', center='x_caa') per = { dim : dict(left=f"{dim}F", center=f"{dim}C") for dim in "xyz" } nper = { dim : dict(outer=f"{dim}F", center=f"{dim}C") for dim in "xyz" } coords = { dim : per[dim] if top in "FP" else nper[dim] for dim, top in zip("xyz", topology) } diff --git a/pynanigans/pnplot.py b/pynanigans/pnplot.py index 8a90558..ff1ec90 100644 --- a/pynanigans/pnplot.py +++ b/pynanigans/pnplot.py @@ -5,7 +5,7 @@ def pnplot(darray, surjection=surjection, **kwargs): """ Bijects darray to rename the dimensions before calling plot(). - This makes plot easier as, instead of calling, `ds.u.plot(x='xF', y='z_aac')`, + This makes plot easier as, instead of calling, `ds.u.plot(x='x_faa', y='z_aac')`, you can call `ds.pnplot(x='x', y='z')` """ return biject(darray, surjection=surjection).plot(**kwargs) diff --git a/pynanigans/utils.py b/pynanigans/utils.py index 9d8dcf9..0545d64 100644 --- a/pynanigans/utils.py +++ b/pynanigans/utils.py @@ -2,11 +2,11 @@ from .grids import get_grid surjection = dict(x_caa='x', - xF='x', + x_faa='x', y_aca='y', - yF='y', + y_afa='y', z_aac='z', - zF='z', + z_aaf='z', ) @@ -17,7 +17,7 @@ def biject(darray, *args, surjection=surjection): If `*args` is provided, only those dimensions will be renamed. If not, `x`, `y` and `z` will be automatically renamed. - This makes calling functions easier as instead of calling `darray.u.plot(x='xF', y='z_aac')`, + This makes calling functions easier as instead of calling `darray.u.plot(x='x_faa', y='z_aac')`, you can call `darray.pnplot(x='x', y='z')` """ da_dims = darray.dims @@ -60,7 +60,7 @@ def downsample(darray, round_func=round, **dim_limits): Downsamples `darray` based on dimensions given in dim_limits dim_limits should be of the form: - dim_limits = dict(y_aca=1000, zF=2048) + dim_limits = dict(y_aca=1000, z_aaf=2048) """ for dim, dim_limit in dim_limits.items(): dim_length = len(darray[dim]) From 0d13fa5fb6fb3089f3c44c1b839523df780decc9 Mon Sep 17 00:00:00 2001 From: tomaschor Date: Wed, 14 May 2025 18:40:38 -0300 Subject: [PATCH 03/13] remove get_distances --- pynanigans/grids.py | 56 ++++++++++++++------------------------------- 1 file changed, 17 insertions(+), 39 deletions(-) diff --git a/pynanigans/grids.py b/pynanigans/grids.py index e0472ba..44e1c44 100644 --- a/pynanigans/grids.py +++ b/pynanigans/grids.py @@ -1,57 +1,35 @@ -def get_coords(ds, topology="PPN",): +def get_coords(topology="PPN"): """ - Constructs the coords dict for ds to be passed to xgcm.Grid + Constructs the coords dict to be passed to xgcm.Grid Flat dimensions (F) are treated the same as Periodic ones (P) """ - per = dict(left='x_faa', center='x_caa') - nper = dict(outer='x_faa', center='x_caa') - per = { dim : dict(left=f"{dim}F", center=f"{dim}C") for dim in "xyz" } - nper = { dim : dict(outer=f"{dim}F", center=f"{dim}C") for dim in "xyz" } - coords = { dim : per[dim] if top in "FP" else nper[dim] for dim, top in zip("xyz", topology) } - - return coords + x_per = dict(left='x_faa', center='x_caa') + y_per = dict(left='y_afa', center='y_aca') + z_per = dict(left='z_aaf', center='z_aac') + x_nper = dict(outer='x_faa', center='x_caa') + y_nper = dict(outer='y_afa', center='y_aca') + z_nper = dict(outer='z_aaf', center='z_aac') -def get_distances(ds, dim="x", topology="P"): - """ - Get distance metrics for Center and Face points of one specific dimension ξ. - If the topology of this dimension is periodic, len(ξC)==len(ξF), but if it - is nonperiodic, then len(ξC)+1==len(ξF). + per = dict(x = x_per, y = y_per, z = z_per) + nper = dict(x = x_nper, y = y_nper, z = z_nper) - Currently does not deal with stretched domains where ΔξC!=ΔξF in the interior. - """ - import numpy as np - import xarray as xr + coords = { dim : per[dim] if top in "FP" else nper[dim] for dim, top in zip("xyz", topology) } - Δξ_mean = ds[dim+"C"].diff(dim+"C").mean().item() - ΔξC = xr.DataArray(np.ones(len(ds[dim+"C"])), dims=[dim+'C']) - if topology=="P" or topology=="F": - ΔξF = xr.DataArray(np.ones(len(ds[dim+"F"])), dims=[dim+'F']) - elif topology=="N": - if len(ds[dim+"F"]) != 1: - interior = np.ones(len(ds[dim+"F"])-2) - ΔξF = xr.DataArray(np.hstack([0.5, interior, 0.5]), dims=[dim+'F']) - else: # Especial case of a slice in a non-periodic dimension - ΔξF = xr.DataArray([1], dims=[dim+'F']) - return Δξ_mean * xr.Dataset({f"Δ{dim}C" : ΔξC, f"Δ{dim}F" : ΔξF}) + return coords def get_metrics(ds, topology="PPN"): """ - Constructs the metric dict for ds. + Constructs the metric dict for `ds`. (Not sure if the metrics are correct at the boundary points """ - for ξ, top in zip('xyz', topology): - ξdist = get_distances(ds, dim=ξ, topology=top) - ds.coords[f"Δ{ξ}C"] = ξdist[f"Δ{ξ}C"] - ds.coords[f"Δ{ξ}F"] = ξdist[f"Δ{ξ}F"] - metrics = { - ('x',): ['ΔxC', 'ΔxF'], # X distances - ('y',): ['ΔyC', 'ΔyF'], # Y distances - ('z',): ['ΔzC', 'ΔzF'], # Z distances + ('x',): ['Δx_caa', 'Δx_faa'], # X distances + ('y',): ['Δy_aca', 'Δy_afa'], # Y distances + ('z',): ['Δz_aac', 'Δz_aaf'], # Z distances } return metrics @@ -62,7 +40,7 @@ def get_grid(ds, coords=None, metrics=None, topology="PPN", **kwargs): import xgcm as xg if coords is None: - coords = get_coords(ds, topology=topology) + coords = get_coords(topology) if metrics is None: metrics = get_metrics(ds, topology=topology) From 69873cd971188fe3551777effaa344ed448ce1d9 Mon Sep 17 00:00:00 2001 From: tomaschor Date: Wed, 14 May 2025 18:42:07 -0300 Subject: [PATCH 04/13] bugfix --- pynanigans/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pynanigans/__init__.py b/pynanigans/__init__.py index 0f825fd..e0f62ec 100644 --- a/pynanigans/__init__.py +++ b/pynanigans/__init__.py @@ -1,6 +1,6 @@ __version__ = "0.1.0" -from .grids import get_distances, get_metrics, get_coords, get_grid +from .grids import get_metrics, get_coords, get_grid from .utils import * from . import pnplot from . import utils From ca15be698b78e6c2922190e8c06599c66fbd111f Mon Sep 17 00:00:00 2001 From: tomaschor Date: Wed, 14 May 2025 18:43:43 -0300 Subject: [PATCH 05/13] single quotes -> double quotes --- pynanigans/grids.py | 18 +++++++++--------- pynanigans/pnplot.py | 4 ++-- pynanigans/utils.py | 18 +++++++++--------- 3 files changed, 20 insertions(+), 20 deletions(-) diff --git a/pynanigans/grids.py b/pynanigans/grids.py index 44e1c44..39519e5 100644 --- a/pynanigans/grids.py +++ b/pynanigans/grids.py @@ -4,13 +4,13 @@ def get_coords(topology="PPN"): Constructs the coords dict to be passed to xgcm.Grid Flat dimensions (F) are treated the same as Periodic ones (P) """ - x_per = dict(left='x_faa', center='x_caa') - y_per = dict(left='y_afa', center='y_aca') - z_per = dict(left='z_aaf', center='z_aac') + x_per = dict(left="x_faa", center="x_caa") + y_per = dict(left="y_afa", center="y_aca") + z_per = dict(left="z_aaf", center="z_aac") - x_nper = dict(outer='x_faa', center='x_caa') - y_nper = dict(outer='y_afa', center='y_aca') - z_nper = dict(outer='z_aaf', center='z_aac') + x_nper = dict(outer="x_faa", center="x_caa") + y_nper = dict(outer="y_afa", center="y_aca") + z_nper = dict(outer="z_aaf", center="z_aac") per = dict(x = x_per, y = y_per, z = z_per) nper = dict(x = x_nper, y = y_nper, z = z_nper) @@ -27,9 +27,9 @@ def get_metrics(ds, topology="PPN"): """ metrics = { - ('x',): ['Δx_caa', 'Δx_faa'], # X distances - ('y',): ['Δy_aca', 'Δy_afa'], # Y distances - ('z',): ['Δz_aac', 'Δz_aaf'], # Z distances + ("x",): ["Δx_caa", "Δx_faa"], # X distances + ("y",): ["Δy_aca", "Δy_afa"], # Y distances + ("z",): ["Δz_aac", "Δz_aaf"], # Z distances } return metrics diff --git a/pynanigans/pnplot.py b/pynanigans/pnplot.py index ff1ec90..8893598 100644 --- a/pynanigans/pnplot.py +++ b/pynanigans/pnplot.py @@ -5,8 +5,8 @@ def pnplot(darray, surjection=surjection, **kwargs): """ Bijects darray to rename the dimensions before calling plot(). - This makes plot easier as, instead of calling, `ds.u.plot(x='x_faa', y='z_aac')`, - you can call `ds.pnplot(x='x', y='z')` + This makes plot easier as, instead of calling, `ds.u.plot(x="x_faa", y="z_aac")`, + you can call `ds.pnplot(x="x", y="z")` """ return biject(darray, surjection=surjection).plot(**kwargs) xr.DataArray.pnplot = pnplot diff --git a/pynanigans/utils.py b/pynanigans/utils.py index 0545d64..c5ab23c 100644 --- a/pynanigans/utils.py +++ b/pynanigans/utils.py @@ -1,12 +1,12 @@ import xarray as xr from .grids import get_grid -surjection = dict(x_caa='x', - x_faa='x', - y_aca='y', - y_afa='y', - z_aac='z', - z_aaf='z', +surjection = dict(x_caa="x", + x_faa="x", + y_aca="y", + y_afa="y", + z_aac="z", + z_aaf="z", ) @@ -44,9 +44,9 @@ def normalize_time_by(darray, seconds=1, new_units="seconds"): object while normalizing it by number of seconds `seconds`. """ import numpy as np - if darray.time.dtype == ' Date: Wed, 14 May 2025 18:44:05 -0300 Subject: [PATCH 06/13] update version --- pynanigans/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pynanigans/__init__.py b/pynanigans/__init__.py index e0f62ec..4049212 100644 --- a/pynanigans/__init__.py +++ b/pynanigans/__init__.py @@ -1,4 +1,4 @@ -__version__ = "0.1.0" +__version__ = "0.1.1" from .grids import get_metrics, get_coords, get_grid from .utils import * From 2691e8abe7202b2503879ee01173573585af97d9 Mon Sep 17 00:00:00 2001 From: tomaschor Date: Wed, 14 May 2025 18:54:26 -0300 Subject: [PATCH 07/13] change folding point syntax --- pynanigans/__init__.py | 2 +- pynanigans/utils.py | 12 ++++++------ 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/pynanigans/__init__.py b/pynanigans/__init__.py index 4049212..e181608 100644 --- a/pynanigans/__init__.py +++ b/pynanigans/__init__.py @@ -1,4 +1,4 @@ -__version__ = "0.1.1" +__version__ = "0.2.0" from .grids import get_metrics, get_coords, get_grid from .utils import * diff --git a/pynanigans/utils.py b/pynanigans/utils.py index c5ab23c..5f3047d 100644 --- a/pynanigans/utils.py +++ b/pynanigans/utils.py @@ -132,25 +132,25 @@ def open_simulation(fname, `kwargs` are passed to `xarray.open_dataset()` and `grid_kwargs` are passed to `pynanigans.get_grid()`. """ - #++++ Open dataset and create grid before squeezing + #+++ Open dataset and create grid before squeezing if load: ds = xr.load_dataset(fname, **kwargs) else: ds = xr.open_dataset(fname, **kwargs) grid_ds = get_grid(ds, topology=topology, **grid_kwargs) - #---- + #--- - #++++ Squeeze? + #+++ Squeeze? if squeeze: ds = ds.squeeze() - #---- + #--- - #++++ Returning only unique times. Useful if simulation was restarted and there's overlap in time + #+++ Returning only unique times. Useful if simulation was restarted and there's overlap in time if unique: import numpy as np _, index = np.unique(ds["time"], return_index=True) if verbose and (len(index)!=len(ds.time)): print("Cleaning non-unique indices") ds = ds.isel(time=index) - #---- + #--- return grid_ds, ds From 6d9c3bb98ac8329b0053c61a453b585cdde33d6a Mon Sep 17 00:00:00 2001 From: tomaschor Date: Thu, 15 May 2025 10:20:11 -0300 Subject: [PATCH 08/13] update test_grids --- tests/test_grids.py | 56 +++++++++++++++++++++++++++++++++++++++------ 1 file changed, 49 insertions(+), 7 deletions(-) diff --git a/tests/test_grids.py b/tests/test_grids.py index 3103580..d8b0a29 100644 --- a/tests/test_grids.py +++ b/tests/test_grids.py @@ -5,14 +5,56 @@ def test_get_coords(): # Test periodic coordinates - coords = get_coords(None, topology="PPP") + coords = get_coords(topology="PPP") assert "x" in coords assert "y" in coords assert "z" in coords - assert coords["x"]["left"] == "xF" - assert coords["x"]["center"] == "xC" - + assert coords["x"]["left"] == "x_faa" + assert coords["x"]["center"] == "x_caa" + # Test non-periodic coordinates - coords = get_coords(None, topology="NNN") - assert coords["x"]["outer"] == "xF" - assert coords["x"]["center"] == "xC" + coords = get_coords(topology="NNN") + assert coords["x"]["outer"] == "x_faa" + assert coords["x"]["center"] == "x_caa" + +def test_get_metrics(): + # Create a test dataset + data = np.random.rand(10, 10, 10) + dims = ['x_caa', 'y_aca', 'z_aac'] + coords = { + 'x_caa': np.linspace(0, 1, 10), + 'y_aca': np.linspace(0, 1, 10), + 'z_aac': np.linspace(0, 1, 10) + } + ds = xr.Dataset( + data_vars={'u': (dims, data)}, + coords=coords + ) + + # Test metrics + metrics = get_metrics(ds) + assert ('x',) in metrics + assert ('y',) in metrics + assert ('z',) in metrics + assert "Δx_caa" in metrics[('x',)] + assert "Δx_faa" in metrics[('x',)] + +def test_get_grid(): + # Create a test dataset + data = np.random.rand(10, 10, 10) + dims = ['x_caa', 'y_aca', 'z_aac'] + coords = { + 'x_caa': np.linspace(0, 1, 10), + 'y_aca': np.linspace(0, 1, 10), + 'z_aac': np.linspace(0, 1, 10) + } + ds = xr.Dataset( + data_vars={'u': (dims, data)}, + coords=coords + ) + + # Test grid creation + grid = get_grid(ds, topology="PPP") + assert grid is not None + assert hasattr(grid, 'coords') + assert hasattr(grid, 'metrics') From ec60c9b4db86993ad9bc815b9725aaadc4f435c2 Mon Sep 17 00:00:00 2001 From: tomaschor Date: Thu, 15 May 2025 10:31:04 -0300 Subject: [PATCH 09/13] change notation in tests --- tests/test_pnplot.py | 40 ++++++++++++++++++++-------------------- 1 file changed, 20 insertions(+), 20 deletions(-) diff --git a/tests/test_pnplot.py b/tests/test_pnplot.py index 9235ca2..d02a9e9 100644 --- a/tests/test_pnplot.py +++ b/tests/test_pnplot.py @@ -6,10 +6,10 @@ def test_pnplot(): # Create a test dataset data = np.random.rand(10, 10) - dims = ['xC', 'yC'] + dims = ['x_caa', 'y_aca'] coords = { - 'xC': np.linspace(0, 1, 10), - 'yC': np.linspace(0, 1, 10) + 'x_caa': np.linspace(0, 1, 10), + 'y_aca': np.linspace(0, 1, 10) } ds = xr.Dataset( data_vars={'u': (dims, data)}, @@ -17,16 +17,16 @@ def test_pnplot(): ) # Test plotting - plot = pnplot(ds.u, x='x', y='y') + plot = pnplot(ds.u, x='x_caa', y='y_aca') assert plot is not None def test_imshow(): # Create a test dataset data = np.random.rand(10, 10) - dims = ['xC', 'yC'] + dims = ['x_caa', 'y_aca'] coords = { - 'xC': np.linspace(0, 1, 10), - 'yC': np.linspace(0, 1, 10) + 'x_caa': np.linspace(0, 1, 10), + 'y_aca': np.linspace(0, 1, 10) } ds = xr.Dataset( data_vars={'u': (dims, data)}, @@ -34,16 +34,16 @@ def test_imshow(): ) # Test imshow - plot = _imshow(ds.u, x='x', y='y') + plot = _imshow(ds.u, x='x_caa', y='y_aca') assert plot is not None def test_pcolormesh(): # Create a test dataset data = np.random.rand(10, 10) - dims = ['xC', 'yC'] + dims = ['x_caa', 'y_aca'] coords = { - 'xC': np.linspace(0, 1, 10), - 'yC': np.linspace(0, 1, 10) + 'x_caa': np.linspace(0, 1, 10), + 'y_aca': np.linspace(0, 1, 10) } ds = xr.Dataset( data_vars={'u': (dims, data)}, @@ -51,16 +51,16 @@ def test_pcolormesh(): ) # Test pcolormesh - plot = _pcolormesh(ds.u, x='x', y='y') + plot = _pcolormesh(ds.u, x='x_caa', y='y_aca') assert plot is not None def test_contour(): # Create a test dataset data = np.random.rand(10, 10) - dims = ['xC', 'yC'] + dims = ['x_caa', 'y_aca'] coords = { - 'xC': np.linspace(0, 1, 10), - 'yC': np.linspace(0, 1, 10) + 'x_caa': np.linspace(0, 1, 10), + 'y_aca': np.linspace(0, 1, 10) } ds = xr.Dataset( data_vars={'u': (dims, data)}, @@ -68,16 +68,16 @@ def test_contour(): ) # Test contour - plot = _contour(ds.u, x='x', y='y') + plot = _contour(ds.u, x='x_caa', y='y_aca') assert plot is not None def test_contourf(): # Create a test dataset data = np.random.rand(10, 10) - dims = ['xC', 'yC'] + dims = ['x_caa', 'y_aca'] coords = { - 'xC': np.linspace(0, 1, 10), - 'yC': np.linspace(0, 1, 10) + 'x_caa': np.linspace(0, 1, 10), + 'y_aca': np.linspace(0, 1, 10) } ds = xr.Dataset( data_vars={'u': (dims, data)}, @@ -85,5 +85,5 @@ def test_contourf(): ) # Test contourf - plot = _contourf(ds.u, x='x', y='y') + plot = _contourf(ds.u, x='x_caa', y='y_aca') assert plot is not None \ No newline at end of file From 881f075a2c6dc9834e18b2790122f4a2e6ed9923 Mon Sep 17 00:00:00 2001 From: tomaschor Date: Thu, 15 May 2025 10:33:41 -0300 Subject: [PATCH 10/13] update tests to catch some more bugs --- tests/test_pnplot.py | 32 ++++++++++++++++++++++++++------ tests/test_utils.py | 34 +++++++++++++++++----------------- 2 files changed, 43 insertions(+), 23 deletions(-) diff --git a/tests/test_pnplot.py b/tests/test_pnplot.py index d02a9e9..146515c 100644 --- a/tests/test_pnplot.py +++ b/tests/test_pnplot.py @@ -1,7 +1,7 @@ import pytest import xarray as xr import numpy as np -from pynanigans.pnplot import pnplot, _imshow, _pcolormesh, _contour, _contourf +from pynanigans import pnplot def test_pnplot(): # Create a test dataset @@ -20,6 +20,10 @@ def test_pnplot(): plot = pnplot(ds.u, x='x_caa', y='y_aca') assert plot is not None + # Test error case - invalid dimension + with pytest.raises(ValueError): + pnplot(ds.u, x='invalid_dim', y='y_aca') + def test_imshow(): # Create a test dataset data = np.random.rand(10, 10) @@ -34,9 +38,13 @@ def test_imshow(): ) # Test imshow - plot = _imshow(ds.u, x='x_caa', y='y_aca') + plot = ds.u.pnimshow(x='x_caa', y='y_aca') assert plot is not None + # Test error case - invalid dimension + with pytest.raises(ValueError): + ds.u.pnimshow(x='invalid_dim', y='y_aca') + def test_pcolormesh(): # Create a test dataset data = np.random.rand(10, 10) @@ -51,9 +59,13 @@ def test_pcolormesh(): ) # Test pcolormesh - plot = _pcolormesh(ds.u, x='x_caa', y='y_aca') + plot = ds.u.pnpcolormesh(x='x_caa', y='y_aca') assert plot is not None + # Test error case - invalid dimension + with pytest.raises(ValueError): + ds.u.pnpcolormesh(x='invalid_dim', y='y_aca') + def test_contour(): # Create a test dataset data = np.random.rand(10, 10) @@ -68,9 +80,13 @@ def test_contour(): ) # Test contour - plot = _contour(ds.u, x='x_caa', y='y_aca') + plot = ds.u.pncontour(x='x_caa', y='y_aca') assert plot is not None + # Test error case - invalid dimension + with pytest.raises(ValueError): + ds.u.pncontour(x='invalid_dim', y='y_aca') + def test_contourf(): # Create a test dataset data = np.random.rand(10, 10) @@ -85,5 +101,9 @@ def test_contourf(): ) # Test contourf - plot = _contourf(ds.u, x='x_caa', y='y_aca') - assert plot is not None \ No newline at end of file + plot = ds.u.pncontourf(x='x_caa', y='y_aca') + assert plot is not None + + # Test error case - invalid dimension + with pytest.raises(ValueError): + ds.u.pncontourf(x='invalid_dim', y='y_aca') \ No newline at end of file diff --git a/tests/test_utils.py b/tests/test_utils.py index 4b9a22a..2bff3ea 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -6,11 +6,11 @@ def test_biject(): # Create a test dataset data = np.random.rand(10, 10, 10) - dims = ['xC', 'yC', 'zC'] + dims = ['x_caa', 'y_aca', 'z_aac'] coords = { - 'xC': np.linspace(0, 1, 10), - 'yC': np.linspace(0, 1, 10), - 'zC': np.linspace(0, 1, 10) + 'x_caa': np.linspace(0, 1, 10), + 'y_aca': np.linspace(0, 1, 10), + 'z_aac': np.linspace(0, 1, 10) } ds = xr.Dataset( data_vars={'u': (dims, data)}, @@ -22,9 +22,9 @@ def test_biject(): assert 'x' in result.dims assert 'y' in result.dims assert 'z' in result.dims - assert 'xC' not in result.dims - assert 'yC' not in result.dims - assert 'zC' not in result.dims + assert 'x_caa' not in result.dims + assert 'y_aca' not in result.dims + assert 'z_aac' not in result.dims def test_normalize_time_by(): # Create a test dataset with time @@ -43,10 +43,10 @@ def test_normalize_time_by(): def test_downsample(): # Create a test dataset data = np.random.rand(100, 100) - dims = ['xC', 'yC'] + dims = ['x_caa', 'y_aca'] coords = { - 'xC': np.linspace(0, 1, 100), - 'yC': np.linspace(0, 1, 100) + 'x_caa': np.linspace(0, 1, 100), + 'y_aca': np.linspace(0, 1, 100) } ds = xr.Dataset( data_vars={'u': (dims, data)}, @@ -54,20 +54,20 @@ def test_downsample(): ) # Test downsampling - result = downsample(ds.u, xC=50, yC=50) - assert len(result.xC) == 50 - assert len(result.yC) == 50 + result = downsample(ds.u, x_caa=50, y_aca=50) + assert len(result.x_caa) == 50 + assert len(result.y_aca) == 50 def test_pnchunk(): # Create a test dataset with time data = np.random.rand(100, 10, 10, 10) time = np.array([np.timedelta64(i, 'ns') for i in range(100)]) - dims = ['time', 'xC', 'yC', 'zC'] + dims = ['time', 'x_caa', 'y_aca', 'z_aac'] coords = { 'time': time, - 'xC': np.linspace(0, 1, 10), - 'yC': np.linspace(0, 1, 10), - 'zC': np.linspace(0, 1, 10) + 'x_caa': np.linspace(0, 1, 10), + 'y_aca': np.linspace(0, 1, 10), + 'z_aac': np.linspace(0, 1, 10) } ds = xr.Dataset( data_vars={'u': (dims, data)}, From 5598102c2cb4485ed60c1f51f79c684946c7538a Mon Sep 17 00:00:00 2001 From: tomaschor Date: Thu, 15 May 2025 10:40:01 -0300 Subject: [PATCH 11/13] remove test_get_grid --- tests/test_grids.py | 20 -------------------- 1 file changed, 20 deletions(-) diff --git a/tests/test_grids.py b/tests/test_grids.py index d8b0a29..ec4e1cd 100644 --- a/tests/test_grids.py +++ b/tests/test_grids.py @@ -38,23 +38,3 @@ def test_get_metrics(): assert ('z',) in metrics assert "Δx_caa" in metrics[('x',)] assert "Δx_faa" in metrics[('x',)] - -def test_get_grid(): - # Create a test dataset - data = np.random.rand(10, 10, 10) - dims = ['x_caa', 'y_aca', 'z_aac'] - coords = { - 'x_caa': np.linspace(0, 1, 10), - 'y_aca': np.linspace(0, 1, 10), - 'z_aac': np.linspace(0, 1, 10) - } - ds = xr.Dataset( - data_vars={'u': (dims, data)}, - coords=coords - ) - - # Test grid creation - grid = get_grid(ds, topology="PPP") - assert grid is not None - assert hasattr(grid, 'coords') - assert hasattr(grid, 'metrics') From 019f75a48446f3355070bc800542bb1e1ce009f0 Mon Sep 17 00:00:00 2001 From: tomaschor Date: Thu, 15 May 2025 10:42:16 -0300 Subject: [PATCH 12/13] fix some pnplot tests --- tests/test_pnplot.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/tests/test_pnplot.py b/tests/test_pnplot.py index 146515c..294f4c5 100644 --- a/tests/test_pnplot.py +++ b/tests/test_pnplot.py @@ -17,12 +17,12 @@ def test_pnplot(): ) # Test plotting - plot = pnplot(ds.u, x='x_caa', y='y_aca') + plot = ds.u.pnplot(x='x', y='y') assert plot is not None # Test error case - invalid dimension with pytest.raises(ValueError): - pnplot(ds.u, x='invalid_dim', y='y_aca') + ds.u.pnplot(x='invalid_dim', y='y') def test_imshow(): # Create a test dataset @@ -38,7 +38,7 @@ def test_imshow(): ) # Test imshow - plot = ds.u.pnimshow(x='x_caa', y='y_aca') + plot = ds.u.pnimshow(x='x', y='y') assert plot is not None # Test error case - invalid dimension @@ -59,12 +59,12 @@ def test_pcolormesh(): ) # Test pcolormesh - plot = ds.u.pnpcolormesh(x='x_caa', y='y_aca') + plot = ds.u.pnpcolormesh(x='x', y='y') assert plot is not None # Test error case - invalid dimension with pytest.raises(ValueError): - ds.u.pnpcolormesh(x='invalid_dim', y='y_aca') + ds.u.pnpcolormesh(x='invalid_dim', y='y') def test_contour(): # Create a test dataset @@ -80,12 +80,12 @@ def test_contour(): ) # Test contour - plot = ds.u.pncontour(x='x_caa', y='y_aca') + plot = ds.u.pncontour(x='x', y='y') assert plot is not None # Test error case - invalid dimension with pytest.raises(ValueError): - ds.u.pncontour(x='invalid_dim', y='y_aca') + ds.u.pncontour(x='invalid_dim', y='y') def test_contourf(): # Create a test dataset @@ -101,9 +101,9 @@ def test_contourf(): ) # Test contourf - plot = ds.u.pncontourf(x='x_caa', y='y_aca') + plot = ds.u.pncontourf(x='x', y='y') assert plot is not None # Test error case - invalid dimension with pytest.raises(ValueError): - ds.u.pncontourf(x='invalid_dim', y='y_aca') \ No newline at end of file + ds.u.pncontourf(x='invalid_dim', y='y') From cc3030d3d8c6072d95cc0afd81492e1ee85291c8 Mon Sep 17 00:00:00 2001 From: tomaschor Date: Thu, 15 May 2025 10:47:47 -0300 Subject: [PATCH 13/13] remove white space --- pynanigans/grids.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pynanigans/grids.py b/pynanigans/grids.py index 8ed9a02..4ab61c9 100644 --- a/pynanigans/grids.py +++ b/pynanigans/grids.py @@ -1,6 +1,6 @@ def get_coords(topology="PPN"): - """ + """ Constructs the coords dict to be passed to xgcm.Grid Flat dimensions (F) are treated the same as Periodic ones (P) """