Skip to content

Commit 7e71225

Browse files
committed
Add tests for reading datasets and grids
1 parent c55dd0c commit 7e71225

File tree

1 file changed

+145
-0
lines changed

1 file changed

+145
-0
lines changed

pygmt/tests/test_clib_read_data.py

+145
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,145 @@
1+
"""
2+
Test the Session.read_data method.
3+
"""
4+
5+
import ctypes as ctp
6+
from pathlib import Path
7+
8+
import numpy as np
9+
from pygmt.clib import Session
10+
from pygmt.datatypes import _GMT_DATASET, _GMT_GRID
11+
from pygmt.helpers import GMTTempFile
12+
13+
14+
def test_clib_read_data_dataset():
15+
"""
16+
Test the Session.read_data method for datasets.
17+
18+
The test is adapted from the doctest in the _GMT_DATASET class.
19+
"""
20+
with GMTTempFile(suffix=".txt") as tmpfile:
21+
# Prepare the sample data file
22+
with Path(tmpfile.name).open(mode="w", encoding="utf-8") as fp:
23+
print("# x y z name", file=fp)
24+
print(">", file=fp)
25+
print("1.0 2.0 3.0 TEXT1 TEXT23", file=fp)
26+
print("4.0 5.0 6.0 TEXT4 TEXT567", file=fp)
27+
print(">", file=fp)
28+
print("7.0 8.0 9.0 TEXT8 TEXT90", file=fp)
29+
print("10.0 11.0 12.0 TEXT123 TEXT456789", file=fp)
30+
31+
with Session() as lib:
32+
data_ptr = lib.read_data(
33+
"GMT_IS_DATASET",
34+
"GMT_IS_PLP",
35+
"GMT_READ_NORMAL",
36+
None,
37+
tmpfile.name,
38+
None,
39+
)
40+
ds = ctp.cast(data_ptr, ctp.POINTER(_GMT_DATASET)).contents
41+
42+
assert ds.n_tables == 1
43+
assert ds.n_segments == 2
44+
assert ds.n_columns == 3
45+
46+
tbl = ds.table[0].contents
47+
assert tbl.min[: tbl.n_columns] == [1.0, 2.0, 3.0]
48+
assert tbl.max[: tbl.n_columns] == [10.0, 11.0, 12.0]
49+
50+
51+
def test_clib_read_data_grid():
52+
"""
53+
Test the Session.read_data method for grids.
54+
55+
The test is adapted from the doctest in the _GMT_GRID class.
56+
"""
57+
with Session() as lib:
58+
data_ptr = lib.read_data(
59+
"GMT_IS_GRID",
60+
"GMT_IS_SURFACE",
61+
"GMT_CONTAINER_AND_DATA",
62+
None,
63+
"@static_earth_relief.nc",
64+
None,
65+
)
66+
grid = ctp.cast(data_ptr, ctp.POINTER(_GMT_GRID)).contents
67+
header = grid.header.contents
68+
assert header.n_rows == 14
69+
assert header.n_columns == 8
70+
assert header.n_bands == 1
71+
assert header.wesn[:] == [-55.0, -47.0, -24.0, -10.0]
72+
assert header.z_min == 190.0
73+
assert header.z_max == 981.0
74+
75+
assert grid.data # The data is read
76+
data = np.reshape(grid.data[: header.mx * header.my], (header.my, header.mx))
77+
data = data[
78+
header.pad[2] : header.my - header.pad[3],
79+
header.pad[0] : header.mx - header.pad[1],
80+
]
81+
assert data[3][4] == 250.0
82+
83+
84+
def test_clib_read_data_grid_two_steps():
85+
"""
86+
Test the Session.read_data method for grids in two steps, first reading the header
87+
and then the data.
88+
89+
The test is adapted from the doctest in the _GMT_GRID class.
90+
"""
91+
family, geometry = "GMT_IS_GRID", "GMT_IS_SURFACE"
92+
infile = "@static_earth_relief.nc"
93+
with Session() as lib:
94+
# Read the header first
95+
data_ptr = lib.read_data(
96+
family, geometry, "GMT_CONTAINER_ONLY", None, infile, None
97+
)
98+
grid = ctp.cast(data_ptr, ctp.POINTER(_GMT_GRID)).contents
99+
header = grid.header.contents
100+
assert header.n_rows == 14
101+
assert header.n_columns == 8
102+
assert header.n_bands == 1
103+
assert header.wesn[:] == [-55.0, -47.0, -24.0, -10.0]
104+
assert header.z_min == 190.0
105+
assert header.z_max == 981.0
106+
107+
assert not grid.data # The data is not read yet
108+
109+
# Read the data
110+
data_ptr = lib.read_data(
111+
family, geometry, "GMT_DATA_ONLY", None, infile, data_ptr
112+
)
113+
grid = ctp.cast(data_ptr, ctp.POINTER(_GMT_GRID)).contents
114+
115+
assert grid.data # The data is read
116+
header = grid.header.contents
117+
data = np.reshape(grid.data[: header.mx * header.my], (header.my, header.mx))
118+
data = data[
119+
header.pad[2] : header.my - header.pad[3],
120+
header.pad[0] : header.mx - header.pad[1],
121+
]
122+
assert data[3][4] == 250.0
123+
124+
125+
def test_clib_read_data_image_as_grid():
126+
"""
127+
Test the Session.read_data method for images.
128+
"""
129+
with Session() as lib:
130+
data_ptr = lib.read_data(
131+
"GMT_IS_GRID",
132+
"GMT_IS_SURFACE",
133+
"GMT_CONTAINER_AND_DATA",
134+
None,
135+
"@earth_day_01d_p",
136+
None,
137+
)
138+
image = ctp.cast(data_ptr, ctp.POINTER(_GMT_GRID)).contents
139+
header = image.header.contents
140+
assert header.n_rows == 180
141+
assert header.n_columns == 360
142+
assert header.n_bands == 1
143+
assert header.wesn[:] == [-180.0, 180.0, -90.0, 90.0]
144+
145+
assert image.data

0 commit comments

Comments
 (0)