Skip to content

Commit

Permalink
Use enums GridReg and GridType in test_accessor.py
Browse files Browse the repository at this point in the history
  • Loading branch information
seisman committed Dec 18, 2024
1 parent f2e3439 commit 170a15d
Showing 1 changed file with 23 additions and 22 deletions.
45 changes: 23 additions & 22 deletions pygmt/tests/test_accessor.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
from pygmt import which
from pygmt.clib import __gmt_version__
from pygmt.datasets import load_earth_relief
from pygmt.enums import GridReg, GridType
from pygmt.exceptions import GMTInvalidInput


Expand All @@ -21,8 +22,8 @@ def test_accessor_gridline_cartesian():
"""
fname = which(fname="@test.dat.nc", download="a")
grid = xr.open_dataarray(fname)
assert grid.gmt.registration == 0 # gridline registration
assert grid.gmt.gtype == 0 # cartesian coordinate type
assert grid.gmt.registration == GridReg.GRIDLINE
assert grid.gmt.gtype == GridType.CARTESIAN


def test_accessor_pixel_geographic():
Expand All @@ -32,18 +33,18 @@ def test_accessor_pixel_geographic():
"""
fname = which(fname="@earth_relief_01d_p", download="a")
grid = xr.open_dataarray(fname, engine="netcdf4")
assert grid.gmt.registration == 1 # pixel registration
assert grid.gmt.gtype == 1 # geographic coordinate type
assert grid.gmt.registration == GridReg.PIXEL
assert grid.gmt.gtype == GridType.GEOGRAPHIC


def test_accessor_set_pixel_registration():
"""
Check that we can set a grid to be Pixel registered with a registration value of 1.
"""
grid = xr.DataArray(data=[[0.1, 0.2], [0.3, 0.4]])
assert grid.gmt.registration == 0 # default to gridline registration
grid.gmt.registration = 1 # set to pixel registration
assert grid.gmt.registration == 1 # ensure changed to pixel registration
assert grid.gmt.registration == GridReg.GRIDLINE
grid.gmt.registration = GridReg.PIXEL
assert grid.gmt.registration == GridReg.PIXEL


@pytest.mark.benchmark
Expand All @@ -53,11 +54,11 @@ def test_accessor_set_geographic_cartesian_roundtrip():
using a gtype of 1, set it to Geographic 0, and then back to Cartesian again 1.
"""
grid = xr.DataArray(data=[[0.1, 0.2], [0.3, 0.4]])
assert grid.gmt.gtype == 0 # default to cartesian coordinate type
grid.gmt.gtype = 1 # set to geographic type
assert grid.gmt.gtype == 1 # ensure changed to geographic coordinate type
grid.gmt.gtype = 0 # set back to cartesian type
assert grid.gmt.gtype == 0 # ensure changed to cartesian coordinate type
assert grid.gmt.gtype == GridType.CARTESIAN
grid.gmt.gtype = GridType.GEOGRAPHIC
assert grid.gmt.gtype == GridType.GEOGRAPHIC
grid.gmt.gtype = GridType.CARTESIAN
assert grid.gmt.gtype == GridType.CARTESIAN


def test_accessor_set_non_boolean():
Expand Down Expand Up @@ -93,8 +94,8 @@ def test_accessor_sliced_datacube():
with xr.open_dataset(fname) as dataset:
grid = dataset.sel(level=500, month=1, drop=True).z

assert grid.gmt.registration == 0 # gridline registration
assert grid.gmt.gtype == 1 # geographic coordinate type
assert grid.gmt.registration == GridReg.GRIDLINE
assert grid.gmt.gtype == GridType.GEOGRAPHIC
finally:
Path(fname).unlink()

Expand All @@ -109,19 +110,19 @@ def test_accessor_grid_source_file_not_exist():
resolution="05m", region=[0, 5, -5, 5], registration="pixel"
)
# Registration and gtype are correct
assert grid.gmt.registration == 1
assert grid.gmt.gtype == 1
assert grid.gmt.registration == GridReg.PIXEL
assert grid.gmt.gtype == GridType.GEOGRAPHIC
# The source grid file is undefined.
assert grid.encoding.get("source") is None

# For a sliced grid, fallback to default registration and gtype,
# because the source grid file doesn't exist.
sliced_grid = grid[1:3, 1:3]
assert sliced_grid.gmt.registration == 0
assert sliced_grid.gmt.gtype == 0
assert sliced_grid.gmt.registration == GridReg.GRIDLINE
assert sliced_grid.gmt.gtype == GridType.CARTESIAN

# Still possible to manually set registration and gtype
sliced_grid.gmt.registration = 1
sliced_grid.gmt.gtype = 1
assert sliced_grid.gmt.registration == 1
assert sliced_grid.gmt.gtype == 1
sliced_grid.gmt.registration = GridReg.PIXEL
sliced_grid.gmt.gtype = GridType.GEOGRAPHIC
assert sliced_grid.gmt.registration == GridReg.PIXEL
assert sliced_grid.gmt.gtype == GridType.GEOGRAPHIC

0 comments on commit 170a15d

Please sign in to comment.