Skip to content

Commit

Permalink
cleaned up version and mpl_plottin gof sgrids (not finished yet)
Browse files Browse the repository at this point in the history
  • Loading branch information
ChrisBarker-NOAA committed Oct 9, 2024
1 parent 00a8395 commit f123650
Show file tree
Hide file tree
Showing 5 changed files with 70 additions and 57 deletions.
4 changes: 2 additions & 2 deletions pixi.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

105 changes: 56 additions & 49 deletions tests/test_grids/test_sgrid.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,12 +2,17 @@

import numpy as np
import xarray as xr
import pytest

import xarray_subset_grid.accessor # noqa: F401
from tests.test_utils import get_test_file_dir
from xarray_subset_grid.grids.sgrid import _get_location_info_from_topology

# open dataset as zarr object using fsspec reference file system and xarray
try:
import fsspec
except ImportError:
fsspec = None


test_dir = get_test_file_dir()
Expand All @@ -33,56 +38,58 @@ def test_grid_topology_location_parse():
'coords': ['lon_rho', 'lat_rho'],
'padding': {'xi_rho': 'both', 'eta_rho': 'both'}}

@pytest.mark.online
def test_polygon_subset():
'''
This is a basic integration test for the subsetting of a ROMS sgrid dataset using a polygon.
'''
if fsspec is None:
raise ImportError("Must have fsspec installed to run --online tests")
fs = fsspec.filesystem(
"reference",
fo="s3://nextgen-dmac-cloud-ingest/nos/wcofs/nos.wcofs.2ds.best.nc.zarr",
remote_protocol="s3",
remote_options={"anon": True},
target_protocol="s3",
target_options={"anon": True},
)
m = fs.get_mapper("")

ds = xr.open_dataset(
m, engine="zarr", backend_kwargs=dict(consolidated=False), chunks={}
)

polygon = np.array(
[
[-122.38488806417945, 34.98888604471138],
[-122.02425311530737, 33.300351211467074],
[-120.60402628930146, 32.723214427630836],
[-116.63789131284673, 32.54346959375448],
[-116.39346090873218, 33.8541384965596],
[-118.83845767505964, 35.257586401855164],
[-121.34541503969862, 35.50073821008141],
[-122.38488806417945, 34.98888604471138],
]
)
ds_temp = ds.xsg.subset_vars(['temp_sur'])
ds_subset = ds_temp.xsg.subset_polygon(polygon)

# def test_polygon_subset():
# '''
# This is a basic integration test for the subsetting of a ROMS sgrid dataset using a polygon.
# '''
# fs = fsspec.filesystem(
# "reference",
# fo="s3://nextgen-dmac-cloud-ingest/nos/wcofs/nos.wcofs.2ds.best.nc.zarr",
# remote_protocol="s3",
# remote_options={"anon": True},
# target_protocol="s3",
# target_options={"anon": True},
# )
# m = fs.get_mapper("")

# ds = xr.open_dataset(
# m, engine="zarr", backend_kwargs=dict(consolidated=False), chunks={}
# )

# polygon = np.array(
# [
# [-122.38488806417945, 34.98888604471138],
# [-122.02425311530737, 33.300351211467074],
# [-120.60402628930146, 32.723214427630836],
# [-116.63789131284673, 32.54346959375448],
# [-116.39346090873218, 33.8541384965596],
# [-118.83845767505964, 35.257586401855164],
# [-121.34541503969862, 35.50073821008141],
# [-122.38488806417945, 34.98888604471138],
# ]
# )
# ds_temp = ds.xsg.subset_vars(['temp_sur'])
# ds_subset = ds_temp.xsg.subset_polygon(polygon)

# #Check that the subset dataset has the correct dimensions given the original padding
# assert ds_subset.sizes['eta_rho'] == ds_subset.sizes['eta_psi'] + 1
# assert ds_subset.sizes['eta_u'] == ds_subset.sizes['eta_psi'] + 1
# assert ds_subset.sizes['eta_v'] == ds_subset.sizes['eta_psi']
# assert ds_subset.sizes['xi_rho'] == ds_subset.sizes['xi_psi'] + 1
# assert ds_subset.sizes['xi_u'] == ds_subset.sizes['xi_psi']
# assert ds_subset.sizes['xi_v'] == ds_subset.sizes['xi_psi'] + 1

# #Check that the subset rho/psi/u/v positional relationsip makes sense aka psi point is
# #'between' it's neighbor rho points
# #Note that this needs to be better generalized; it's not trivial to write a test that
# #works in all potential cases.
# assert (ds_subset['lon_rho'][0,0] < ds_subset['lon_psi'][0,0]
# and ds_subset['lon_rho'][0,1] > ds_subset['lon_psi'][0,0])

# #ds_subset.temp_sur.isel(ocean_time=0).plot(x="lon_rho", y="lat_rho")
#Check that the subset dataset has the correct dimensions given the original padding
assert ds_subset.sizes['eta_rho'] == ds_subset.sizes['eta_psi'] + 1
assert ds_subset.sizes['eta_u'] == ds_subset.sizes['eta_psi'] + 1
assert ds_subset.sizes['eta_v'] == ds_subset.sizes['eta_psi']
assert ds_subset.sizes['xi_rho'] == ds_subset.sizes['xi_psi'] + 1
assert ds_subset.sizes['xi_u'] == ds_subset.sizes['xi_psi']
assert ds_subset.sizes['xi_v'] == ds_subset.sizes['xi_psi'] + 1

#Check that the subset rho/psi/u/v positional relationsip makes sense aka psi point is
#'between' it's neighbor rho points
#Note that this needs to be better generalized; it's not trivial to write a test that
#works in all potential cases.
assert (ds_subset['lon_rho'][0,0] < ds_subset['lon_psi'][0,0]
and ds_subset['lon_rho'][0,1] > ds_subset['lon_psi'][0,0])

#ds_subset.temp_sur.isel(ocean_time=0).plot(x="lon_rho", y="lat_rho")

def test_polygon_subset_2():
ds = xr.open_dataset(sample_sgrid_file, decode_times=False)
Expand Down
12 changes: 6 additions & 6 deletions tests/test_visualization/test_mpl_plotting.py
Original file line number Diff line number Diff line change
Expand Up @@ -68,15 +68,15 @@ def test_plot_ugrid_start_index_1():
# SGRID tests
#############

def test_plot_sgrid_only_grid():
import cftime
ds = xr.open_dataset(EXAMPLE_DATA / "wcofs_small_subset.nc")
# def test_plot_sgrid_only_grid():
# import cftime
# ds = xr.open_dataset(EXAMPLE_DATA / "wcofs_small_subset.nc", decode_times=False)

fig, axis = plt.subplots()
# fig, axis = plt.subplots()

plot_sgrid(axis, ds)
# plot_sgrid(axis, ds)

fig.savefig(OUTPUT_DIR / "sgrid_just_plot")
# fig.savefig(OUTPUT_DIR / "sgrid_just_plot")



2 changes: 2 additions & 0 deletions xarray_subset_grid/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,4 +5,6 @@
from . import accessor # noqa
from .selector import Selector

from._version import __version__

__all__ = ['Selector']
4 changes: 4 additions & 0 deletions xarray_subset_grid/visualization/mpl_plotting.py
Original file line number Diff line number Diff line change
Expand Up @@ -111,9 +111,13 @@ def plot_sgrid(axes, ds, nodes=False, rho_points=False, edge_points=False):
(where U and V are in ROMS)
"""

raise NotImplementedError("have to port ugrid code to Sgrid")

mesh_defs = ds[ds.cf.cf_roles["grid_topology"][0]].attrs
lon_var, lat_var = mesh_defs["node_coordinates"].split()
nodes_lon, nodes_lat = (ds[n] for n in mesh_defs["node_coordinates"].split())


faces = ds[mesh_defs["face_node_connectivity"]]

if faces.shape[0] == 3:
Expand Down

0 comments on commit f123650

Please sign in to comment.