Skip to content

Commit

Permalink
pygmt.grdcut: Support both grid and image output
Browse files Browse the repository at this point in the history
  • Loading branch information
seisman committed Apr 16, 2024
1 parent 48e6505 commit 84ce139
Show file tree
Hide file tree
Showing 4 changed files with 52 additions and 18 deletions.
24 changes: 19 additions & 5 deletions pygmt/clib/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,14 +25,15 @@
vectors_to_arrays,
)
from pygmt.clib.loading import load_libgmt
from pygmt.datatypes import _GMT_DATASET, _GMT_GRID
from pygmt.datatypes import _GMT_DATASET, _GMT_GRID, _GMT_IMAGE
from pygmt.exceptions import (
GMTCLibError,
GMTCLibNoSessionError,
GMTInvalidInput,
GMTVersionError,
)
from pygmt.helpers import (
GMTTempFile,
data_kind,
fmt_docstring,
tempfile_from_geojson,
Expand Down Expand Up @@ -1649,7 +1650,9 @@ def virtualfile_in( # noqa: PLR0912

@contextlib.contextmanager
def virtualfile_out(
self, kind: Literal["dataset", "grid"] = "dataset", fname: str | None = None
self,
kind: Literal["dataset", "grid", "image"] = "dataset",
fname: str | None = None,
):
r"""
Create a virtual file or an actual file for storing output data.
Expand Down Expand Up @@ -1706,8 +1709,11 @@ def virtualfile_out(
family, geometry = {
"dataset": ("GMT_IS_DATASET", "GMT_IS_PLP"),
"grid": ("GMT_IS_GRID", "GMT_IS_SURFACE"),
"image": ("GMT_IS_IMAGE", "GMT_IS_SURFACE"),
}[kind]
with self.open_virtualfile(family, geometry, "GMT_OUT", None) as vfile:
with self.open_virtualfile(
family, geometry, "GMT_OUT|GMT_IS_REFERENCE", None
) as vfile:
yield vfile

def inquire_virtualfile(self, vfname: str) -> int:
Expand Down Expand Up @@ -1801,9 +1807,13 @@ def read_virtualfile(
# _GMT_DATASET).
if kind is None: # Return the ctypes void pointer
return pointer
if kind in ["image", "cube"]:
if kind == "cube":
raise NotImplementedError(f"kind={kind} is not supported yet.")
dtype = {"dataset": _GMT_DATASET, "grid": _GMT_GRID}[kind]
dtype = {
"dataset": _GMT_DATASET,
"grid": _GMT_GRID,
"image": _GMT_IMAGE,
}[kind]
return ctp.cast(pointer, ctp.POINTER(dtype))

def virtualfile_to_dataset(
Expand Down Expand Up @@ -2013,6 +2023,10 @@ def virtualfile_to_raster(
self["GMT_IS_IMAGE"]: "image",
self["GMT_IS_CUBE"]: "cube",
}[family]
if kind == "image":
with GMTTempFile(suffix=".tif") as tmpfile:
self.call_module("write", f"{vfname} {tmpfile.name} -Ti")
return xr.load_dataarray(tmpfile.name)
return self.read_virtualfile(vfname, kind=kind).contents.to_dataarray()

def extract_region(self):
Expand Down
1 change: 1 addition & 0 deletions pygmt/datatypes/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,3 +4,4 @@

from pygmt.datatypes.dataset import _GMT_DATASET
from pygmt.datatypes.grid import _GMT_GRID
from pygmt.datatypes.image import _GMT_IMAGE
9 changes: 9 additions & 0 deletions pygmt/datatypes/image.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,9 @@
"""
Wrapper for the GMT_GRID data type.
"""

import ctypes as ctp


class _GMT_IMAGE(ctp.Structure): # noqa: N801
pass
36 changes: 23 additions & 13 deletions pygmt/src/grdcut.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,20 +4,19 @@

from pygmt.clib import Session
from pygmt.helpers import (
GMTTempFile,
build_arg_string,
data_kind,
fmt_docstring,
kwargs_to_strings,
use_alias,
)
from pygmt.io import load_dataarray
from pygmt.src.which import which

__doctest_skip__ = ["grdcut"]


@fmt_docstring
@use_alias(
G="outgrid",
R="region",
J="projection",
N="extend",
Expand All @@ -27,7 +26,7 @@
f="coltypes",
)
@kwargs_to_strings(R="sequence")
def grdcut(grid, **kwargs):
def grdcut(grid, outgrid: str | None = None, **kwargs):
r"""
Extract subregion from a grid.
Expand Down Expand Up @@ -99,13 +98,24 @@ def grdcut(grid, **kwargs):
>>> # 12° E to 15° E and a latitude range of 21° N to 24° N
>>> new_grid = pygmt.grdcut(grid=grid, region=[12, 15, 21, 24])
"""
with GMTTempFile(suffix=".nc") as tmpfile:
with Session() as lib:
with lib.virtualfile_in(check_kind="raster", data=grid) as vingrd:
if (outgrid := kwargs.get("G")) is None:
kwargs["G"] = outgrid = tmpfile.name # output to tmpfile
lib.call_module(
module="grdcut", args=build_arg_string(kwargs, infile=vingrd)
)
inkind = data_kind(grid)
outkind = "image" if inkind == "image" else "grid"
if inkind == "file":
realpath = which(str(grid))
if isinstance(realpath, list):
realpath = realpath[0]
if realpath.endswith(".tif"):
outkind = "image"

return load_dataarray(outgrid) if outgrid == tmpfile.name else None
with Session() as lib:
with (
lib.virtualfile_in(check_kind="raster", data=grid) as vingrd,
lib.virtualfile_out(kind=outkind, fname=outgrid) as voutgrd,
):
kwargs["G"] = voutgrd
lib.call_module(
module="grdcut", args=build_arg_string(kwargs, infile=vingrd)
)
return lib.virtualfile_to_raster(
outgrid=outgrid, kind=outkind, vfname=voutgrd
)

0 comments on commit 84ce139

Please sign in to comment.