Skip to content

Commit 046aa18

Browse files
seismanseisman
committed
Merge the WIP implementation of GMT_IMAGE
1 parent db81895 commit 046aa18

File tree

4 files changed

+199
-9
lines changed

4 files changed

+199
-9
lines changed

‎pygmt/clib/session.py

+14-8
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@
2626
vectors_to_arrays,
2727
)
2828
from pygmt.clib.loading import load_libgmt
29-
from pygmt.datatypes import _GMT_DATASET, _GMT_GRID
29+
from pygmt.datatypes import _GMT_DATASET, _GMT_GRID, _GMT_IMAGE
3030
from pygmt.exceptions import (
3131
GMTCLibError,
3232
GMTCLibNoSessionError,
@@ -1769,7 +1769,9 @@ def virtualfile_from_data(
17691769

17701770
@contextlib.contextmanager
17711771
def virtualfile_out(
1772-
self, kind: Literal["dataset", "grid"] = "dataset", fname: str | None = None
1772+
self,
1773+
kind: Literal["dataset", "grid", "image"] = "dataset",
1774+
fname: str | None = None,
17731775
):
17741776
r"""
17751777
Create a virtual file or an actual file for storing output data.
@@ -1782,8 +1784,8 @@ def virtualfile_out(
17821784
Parameters
17831785
----------
17841786
kind
1785-
The data kind of the virtual file to create. Valid values are ``"dataset"``
1786-
and ``"grid"``. Ignored if ``fname`` is specified.
1787+
The data kind of the virtual file to create. Valid values are ``"dataset"``,
1788+
``"grid"``, and ``"image"``. Ignored if ``fname`` is specified.
17871789
fname
17881790
The name of the actual file to write the output data. No virtual file will
17891791
be created.
@@ -1826,8 +1828,11 @@ def virtualfile_out(
18261828
family, geometry = {
18271829
"dataset": ("GMT_IS_DATASET", "GMT_IS_PLP"),
18281830
"grid": ("GMT_IS_GRID", "GMT_IS_SURFACE"),
1831+
"image": ("GMT_IS_IMAGE", "GMT_IS_SURFACE"),
18291832
}[kind]
1830-
with self.open_virtualfile(family, geometry, "GMT_OUT", None) as vfile:
1833+
with self.open_virtualfile(
1834+
family, geometry, "GMT_OUT|GMT_IS_REFERENCE", None
1835+
) as vfile:
18311836
yield vfile
18321837

18331838
def inquire_virtualfile(self, vfname: str) -> int:
@@ -1873,7 +1878,8 @@ def read_virtualfile(
18731878
Name of the virtual file to read.
18741879
kind
18751880
Cast the data into a GMT data container. Valid values are ``"dataset"``,
1876-
``"grid"`` and ``None``. If ``None``, will return a ctypes void pointer.
1881+
``"grid"``, ``"image"`` and ``None``. If ``None``, will return a ctypes void
1882+
pointer.
18771883
18781884
Examples
18791885
--------
@@ -1921,9 +1927,9 @@ def read_virtualfile(
19211927
# _GMT_DATASET).
19221928
if kind is None: # Return the ctypes void pointer
19231929
return pointer
1924-
if kind in {"image", "cube"}:
1930+
if kind == "cube":
19251931
raise NotImplementedError(f"kind={kind} is not supported yet.")
1926-
dtype = {"dataset": _GMT_DATASET, "grid": _GMT_GRID}[kind]
1932+
dtype = {"dataset": _GMT_DATASET, "grid": _GMT_GRID, "image": _GMT_IMAGE}[kind]
19271933
return ctp.cast(pointer, ctp.POINTER(dtype))
19281934

19291935
def virtualfile_to_dataset(

‎pygmt/datatypes/__init__.py

+1
Original file line numberDiff line numberDiff line change
@@ -4,3 +4,4 @@
44

55
from pygmt.datatypes.dataset import _GMT_DATASET
66
from pygmt.datatypes.grid import _GMT_GRID
7+
from pygmt.datatypes.image import _GMT_IMAGE

‎pygmt/datatypes/header.py

+2-1
Original file line numberDiff line numberDiff line change
@@ -203,7 +203,8 @@ def data_attrs(self) -> dict[str, Any]:
203203
Attributes for the data variable from the grid header.
204204
"""
205205
attrs: dict[str, Any] = {}
206-
attrs["Conventions"] = "CF-1.7"
206+
if self.type == 18: # Grid file format: ns = GMT netCDF format
207+
attrs["Conventions"] = "CF-1.7"
207208
attrs["title"] = self.title.decode()
208209
attrs["history"] = self.command.decode()
209210
attrs["description"] = self.remark.decode()

‎pygmt/datatypes/image.py

+182
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,182 @@
1+
"""
2+
Wrapper for the GMT_IMAGE data type.
3+
"""
4+
5+
import ctypes as ctp
6+
from typing import ClassVar
7+
8+
import numpy as np
9+
import xarray as xr
10+
from pygmt.datatypes.header import _GMT_GRID_HEADER
11+
12+
13+
class _GMT_IMAGE(ctp.Structure): # noqa: N801
14+
"""
15+
GMT image data structure.
16+
17+
Examples
18+
--------
19+
>>> from pygmt.clib import Session
20+
>>> import numpy as np
21+
>>> import xarray as xr
22+
23+
>>> with Session() as lib:
24+
... with lib.virtualfile_out(kind="image") as voutimg:
25+
... lib.call_module("read", f"@earth_day_01d {voutimg} -Ti")
26+
... # Read the image from the virtual file
27+
... image = lib.read_virtualfile(vfname=voutimg, kind="image").contents
28+
... # The image header
29+
... header = image.header.contents
30+
... # Access the header properties
31+
... print(image.type, header.n_bands, header.n_rows, header.n_columns)
32+
... print(header.pad[:])
33+
... # The x and y coordinates
34+
... x = image.x[: header.n_columns]
35+
... y = image.y[: header.n_rows]
36+
... # The data array (with paddings)
37+
... data = np.reshape(
38+
... image.data[: header.n_bands * header.mx * header.my],
39+
... (header.my, header.mx, header.n_bands),
40+
... )
41+
... # The data array (without paddings)
42+
... pad = header.pad[:]
43+
... data = data[pad[2] : header.my - pad[3], pad[0] : header.mx - pad[1], :]
44+
... print(data.shape)
45+
1 3 180 360
46+
[2, 2, 2, 2]
47+
(180, 360, 3)
48+
"""
49+
50+
_fields_: ClassVar = [
51+
# Data type, e.g. GMT_FLOAT
52+
("type", ctp.c_int),
53+
# Array with color lookup values
54+
("colormap", ctp.POINTER(ctp.c_int)),
55+
# Number of colors in a paletted image
56+
("n_indexed_colors", ctp.c_int),
57+
# Pointer to full GMT header for the image
58+
("header", ctp.POINTER(_GMT_GRID_HEADER)),
59+
# Pointer to actual image
60+
("data", ctp.POINTER(ctp.c_ubyte)),
61+
# Pointer to an optional transparency layer stored in a separate variable
62+
("alpha", ctp.POINTER(ctp.c_ubyte)),
63+
# Color interpolation
64+
("color_interp", ctp.c_char_p),
65+
# Pointer to the x-coordinate vector
66+
("x", ctp.POINTER(ctp.c_double)),
67+
# Pointer to the y-coordinate vector
68+
("y", ctp.POINTER(ctp.c_double)),
69+
# Book-keeping variables "hidden" from the API
70+
("hidden", ctp.c_void_p),
71+
]
72+
73+
def to_dataarray(self) -> xr.DataArray:
74+
"""
75+
Convert a _GMT_IMAGE object to an :class:`xarray.DataArray` object.
76+
77+
Returns
78+
-------
79+
dataarray
80+
A :class:`xarray.DataArray` object.
81+
82+
Examples
83+
--------
84+
>>> from pygmt.clib import Session
85+
>>> with Session() as lib:
86+
... with lib.virtualfile_out(kind="image") as voutimg:
87+
... lib.call_module("read", ["@earth_day_01d", voutimg, "-Ti"])
88+
... # Read the image from the virtual file
89+
... image = lib.read_virtualfile(voutimg, kind="image")
90+
... # Convert to xarray.DataArray and use it later
91+
... da = image.contents.to_dataarray()
92+
>>> da # doctest: +NORMALIZE_WHITESPACE, +ELLIPSIS
93+
<xarray.DataArray 'z' (band: 3, y: 180, x: 360)> Size: 2MB
94+
array([[[ 10, 10, 10, ..., 10, 10, 10],
95+
[ 10, 10, 10, ..., 10, 10, 10],
96+
[ 10, 10, 10, ..., 10, 10, 10],
97+
...,
98+
[192, 193, 193, ..., 193, 192, 191],
99+
[204, 206, 206, ..., 205, 206, 204],
100+
[208, 210, 210, ..., 210, 210, 208]],
101+
<BLANKLINE>
102+
[[ 10, 10, 10, ..., 10, 10, 10],
103+
[ 10, 10, 10, ..., 10, 10, 10],
104+
[ 10, 10, 10, ..., 10, 10, 10],
105+
...,
106+
[186, 187, 188, ..., 187, 186, 185],
107+
[196, 198, 198, ..., 197, 197, 196],
108+
[199, 201, 201, ..., 201, 202, 199]],
109+
<BLANKLINE>
110+
[[ 51, 51, 51, ..., 51, 51, 51],
111+
[ 51, 51, 51, ..., 51, 51, 51],
112+
[ 51, 51, 51, ..., 51, 51, 51],
113+
...,
114+
[177, 179, 179, ..., 178, 177, 177],
115+
[185, 187, 187, ..., 187, 186, 185],
116+
[189, 191, 191, ..., 191, 191, 189]]])
117+
Coordinates:
118+
* x (x) float64 3kB -179.5 -178.5 -177.5 -176.5 ... 177.5 178.5 179.5
119+
* y (y) float64 1kB 89.5 88.5 87.5 86.5 ... -86.5 -87.5 -88.5 -89.5
120+
* band (band) uint8 3B 0 1 2
121+
Attributes:
122+
title:
123+
history:
124+
description:
125+
long_name: z
126+
actual_range: [ 1.79769313e+308 -1.79769313e+308]
127+
128+
>>> da.coords["x"] # doctest: +NORMALIZE_WHITESPACE, +ELLIPSIS
129+
<xarray.DataArray 'x' (x: 360)> Size: 3kB
130+
array([-179.5, -178.5, -177.5, ..., 177.5, 178.5, 179.5])
131+
Coordinates:
132+
* x (x) float64 3kB -179.5 -178.5 -177.5 -176.5 ... 177.5 178.5 179.5
133+
134+
>>> da.coords["y"] # doctest: +NORMALIZE_WHITESPACE, +ELLIPSIS
135+
<xarray.DataArray 'y' (y: 180)> Size: 1kB
136+
array([ 89.5, 88.5, 87.5, 86.5, 85.5, 84.5, 83.5, 82.5, 81.5, 80.5,
137+
79.5, 78.5, 77.5, 76.5, 75.5, 74.5, 73.5, 72.5, 71.5, 70.5,
138+
69.5, 68.5, 67.5, 66.5, 65.5, 64.5, 63.5, 62.5, 61.5, 60.5,
139+
...
140+
-0.5, -1.5, -2.5, -3.5, -4.5, -5.5, -6.5, -7.5, -8.5, -9.5,
141+
...
142+
-60.5, -61.5, -62.5, -63.5, -64.5, -65.5, -66.5, -67.5, -68.5, -69.5,
143+
-70.5, -71.5, -72.5, -73.5, -74.5, -75.5, -76.5, -77.5, -78.5, -79.5,
144+
-80.5, -81.5, -82.5, -83.5, -84.5, -85.5, -86.5, -87.5, -88.5, -89.5])
145+
Coordinates:
146+
* y (y) float64 1kB 89.5 88.5 87.5 86.5 ... -86.5 -87.5 -88.5 -89.5
147+
148+
>>> da.gmt.registration, da.gmt.gtype
149+
(1, 0)
150+
"""
151+
152+
# Get image header
153+
header: _GMT_GRID_HEADER = self.header.contents
154+
155+
# Get DataArray without padding
156+
pad = header.pad[:]
157+
data: np.ndarray = np.reshape(
158+
a=self.data[: header.n_bands * header.mx * header.my],
159+
newshape=(header.my, header.mx, header.n_bands),
160+
)[pad[2] : header.my - pad[3], pad[0] : header.mx - pad[1], :]
161+
162+
# Get x and y coordinates
163+
coords: dict[str, list | np.ndarray] = {
164+
"x": self.x[: header.n_columns],
165+
"y": self.y[: header.n_rows],
166+
"band": np.array([0, 1, 2], dtype=np.uint8),
167+
}
168+
169+
# Create the xarray.DataArray object
170+
image = xr.DataArray(
171+
data=data,
172+
coords=coords,
173+
dims=("y", "x", "band"),
174+
name=header.name,
175+
attrs=header.data_attrs,
176+
).transpose("band", "y", "x")
177+
178+
# Set GMT accessors.
179+
# Must put at the end, otherwise info gets lost after certain image operations.
180+
image.gmt.registration = header.registration
181+
image.gmt.gtype = header.gtype
182+
return image

0 commit comments

Comments
 (0)