Skip to content

Cache gmt accessor info for non-virtualfile grid/image inputs #4008

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 2 commits into
base: main
Choose a base branch
from
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 27 additions & 0 deletions pygmt/tests/test_xarray_backend.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
import numpy.testing as npt
import pytest
import xarray as xr
from pygmt.clib import Session
from pygmt.enums import GridRegistration, GridType
from pygmt.exceptions import GMTValueError
from pygmt.helpers import GMTTempFile
Expand Down Expand Up @@ -38,6 +39,32 @@ def test_xarray_backend_load_dataarray():
dataarray.to_netcdf(tmpfile.name)


def test_xarray_backend_load_dataarray_temp_nc_grid():
"""
Check that xarray.load_dataarray works to read a temporary netCDF grid, and ensure
that GMTDataArrayAccessor information is retained after original file is deleted.

This is a regression test for
https://github.com/GenericMappingTools/pygmt/issues/4005
"""

with Session() as lib:
with GMTTempFile(suffix=".nc") as tmpfile:
args = [
"@earth_relief_01d_g",
"-T", # change from gridline to pixel registration
f"-G{tmpfile.name}",
]
lib.call_module(module="grdedit", args=args)
dataarray = xr.load_dataarray(
tmpfile.name, engine="gmt", raster_kind="grid"
)

# Ensure GMTDataArrayAccessor info is preserved after tempfile is deleted
assert dataarray.gmt.registration is GridRegistration.PIXEL
assert dataarray.gmt.gtype is GridType.CARTESIAN


def test_xarray_backend_gmt_open_nc_grid():
"""
Ensure that passing engine='gmt' to xarray.open_dataarray works to open a netCDF
Expand Down
1 change: 1 addition & 0 deletions pygmt/xarray/accessor.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@


@xr.register_dataarray_accessor("gmt")
@xr.register_dataset_accessor("gmt")
class GMTDataArrayAccessor:
"""
GMT accessor for :class:`xarray.DataArray`.
Expand Down
4 changes: 3 additions & 1 deletion pygmt/xarray/backend.py
Original file line number Diff line number Diff line change
Expand Up @@ -150,4 +150,6 @@ def open_dataset( # type: ignore[override]
raster.encoding["source"] = (
sorted(source)[0] if isinstance(source, list) else source
)
return raster.to_dataset()
raster_dataset: xr.Dataset = raster.to_dataset()
_ = raster_dataset.gmt # Load GMTDataArray accessor information
return raster_dataset
Comment on lines +153 to +155
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Originally posted by @seisman in #3932 (comment)

Actually, it's likely that the accessor information will be lost when converting via to_dataset.

So we can set the accessor info at the xr.Dataset level, but not sure if it's a good idea because then xr.Dataset would have GMT-functions enabled...

Copy link
Member

@seisman seisman Jul 23, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The test is still failing. I guess that the accessor information is lost in both dataarray->dataset and dataset->dataarray conversions.

When calling xr.open_dataarray with the "gmt" engine, it actually calls the GMTBackendEntry's open_dataset method first, so the conversion from dataset to dataarray cannot be avoided.

Copy link
Member

@seisman seisman Jul 23, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

>>> import pygmt
>>> import xarray as xr

>>> grid = xr.open_dataarray("@static_earth_relief.nc", engine="gmt", raster_kind="grid")

>>> grid.encoding["source"]
np.str_('/home/seisman/.gmt/cache/static_earth_relief.nc')

>>> grid.gmt.registration, grid.gmt.gtype
(<GridRegistration.PIXEL: 1>, <GridType.GEOGRAPHIC: 1>)

>>> # Set source encoding to a fake path, similar to a temporary file being deleted
>>> grid.encoding["source"] = "/path/to/nothing"  
>>> grid.gmt.registration, grid.gmt.gtype
(<GridRegistration.PIXEL: 1>, <GridType.GEOGRAPHIC: 1>)

>>> # Convert a dataarray to a dataset. The source encoding is not kept.
>>> da = grid.to_dataset()
>>> da.encoding
{}

>>> # Convert a dataset to a dataarray.
>>> # Xref: https://github.com/pydata/xarray/blob/fb49a3b0a83272467aa01ac189b95e9e7ebe01aa/xarray/backends/api.py#L950
>>> (ds, ) = da.data_vars.values()
>>> ds.encoding
{'source': '/path/to/nothing'}

>>> ds.gmt.registration, ds.gmt.gtype
(<GridRegistration.GRIDLINE: 0>, <GridType.CARTESIAN: 0>)

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As shown above, the accessor information is lost during the dataset->dataarray conversion, and we rely on the source encoding to decide the registration and gtype, which default to 0 for temporary files being deleted.

I don't think we have any solutions here, and need to document the limitation and recommend:

  1. Avoiding using temporary files if possible
  2. If temporary files cannot be avoided, use _ = grid.gmt before deleting the temporary file

Copy link
Member Author

@weiji14 weiji14 Jul 23, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, I found this upstream issue - pydata/xarray#3268 that accessors are not meant to cache a state. I do wish though that the @xr.register_dataarray_accessor decorator could work by us returning an xr.DataArray directly (instead of having to return an xr.Dataset that then gets converted to an xr.DataArray).

Agree that we should document that temporary files should be avoided and/or trigger caching using _ = grid.gmt, then close #4005 as wontfix (with a workaround).

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I do wish though that the @xr.register_dataarray_accessor decorator could work by us returning an xr.DataArray directly (instead of having to return an xr.Dataset that then gets converted to an xr.DataArray).

Opened an upstream issue on xarray for this -> pydata/xarray#10562

Loading