Skip to content

Commit

Permalink
Simplify the tests by converting GMT data containers into dataframe o…
Browse files Browse the repository at this point in the history
…r dataarray
seisman committed Jul 10, 2024
1 parent 76a511e commit 591172b
Showing 1 changed file with 46 additions and 49 deletions.
95 changes: 46 additions & 49 deletions pygmt/tests/test_clib_read_data.py
Original file line number Diff line number Diff line change
@@ -4,9 +4,21 @@

from pathlib import Path

import numpy as np
import pandas as pd
import pytest
import xarray as xr
from pygmt.clib import Session
from pygmt.helpers import GMTTempFile
from pygmt.io import load_dataarray
from pygmt.src import which


@pytest.fixture(scope="module", name="expected_xrgrid")
def fixture_expected_xrgrid():
"""
The expected xr.DataArray object for the static_earth_relief.nc file.
"""
return load_dataarray(which("@static_earth_relief.nc"))


def test_clib_read_data_dataset():
@@ -27,49 +39,42 @@ def test_clib_read_data_dataset():
print("10.0 11.0 12.0 TEXT123 TEXT456789", file=fp)

with Session() as lib:
data_ptr = lib.read_data(tmpfile.name, kind="dataset")
ds = data_ptr.contents

assert ds.n_tables == 1
assert ds.n_segments == 2
assert ds.n_columns == 3

tbl = ds.table[0].contents
assert tbl.min[: tbl.n_columns] == [1.0, 2.0, 3.0]
assert tbl.max[: tbl.n_columns] == [10.0, 11.0, 12.0]


def test_clib_read_data_grid():
ds = lib.read_data(tmpfile.name, kind="dataset").contents
df = ds.to_dataframe(header=0)
expected_df = pd.DataFrame(
data={
"x": [1.0, 4.0, 7.0, 10.0],
"y": [2.0, 5.0, 8.0, 11.0],
"z": [3.0, 6.0, 9.0, 12.0],
"name": pd.Series(
[
"TEXT1 TEXT23",
"TEXT4 TEXT567",
"TEXT8 TEXT90",
"TEXT123 TEXT456789",
],
dtype=pd.StringDtype(),
),
}
)
pd.testing.assert_frame_equal(df, expected_df)


def test_clib_read_data_grid(expected_xrgrid):
"""
Test the Session.read_data method for grids.
The test is adapted from the doctest in the _GMT_GRID class.
"""
with Session() as lib:
data_ptr = lib.read_data(
"@static_earth_relief.nc",
kind="grid",
mode="GMT_CONTAINER_AND_DATA",
)
grid = data_ptr.contents
header = grid.header.contents
assert header.n_rows == 14
assert header.n_columns == 8
assert header.n_bands == 1
assert header.wesn[:] == [-55.0, -47.0, -24.0, -10.0]
assert header.z_min == 190.0
assert header.z_max == 981.0
grid = lib.read_data("@static_earth_relief.nc", kind="grid").contents
xrgrid = grid.to_dataarray()
xr.testing.assert_equal(xrgrid, expected_xrgrid)
# Explicitely check n_bands
assert grid.header.contents.n_bands == 1

assert grid.data # The data is read
data = np.reshape(grid.data[: header.mx * header.my], (header.my, header.mx))
data = data[
header.pad[2] : header.my - header.pad[3],
header.pad[0] : header.mx - header.pad[1],
]
assert data[3][4] == 250.0


def test_clib_read_data_grid_two_steps():
def test_clib_read_data_grid_two_steps(expected_xrgrid):
"""
Test the Session.read_data method for grids in two steps, first reading the header
and then the data.
@@ -92,19 +97,11 @@ def test_clib_read_data_grid_two_steps():
assert not grid.data # The data is not read yet

# Read the data
data_ptr = lib.read_data(
infile, kind="grid", mode="GMT_DATA_ONLY", data=data_ptr
)
grid = data_ptr.contents

assert grid.data # The data is read
header = grid.header.contents
data = np.reshape(grid.data[: header.mx * header.my], (header.my, header.mx))
data = data[
header.pad[2] : header.my - header.pad[3],
header.pad[0] : header.mx - header.pad[1],
]
assert data[3][4] == 250.0
lib.read_data(infile, kind="grid", mode="GMT_DATA_ONLY", data=data_ptr)
xrgrid = data_ptr.contents.to_dataarray()
xr.testing.assert_equal(xrgrid, expected_xrgrid)
# Explicitely check n_bands
assert grid.header.contents.n_bands == 1


def test_clib_read_data_image_as_grid():

0 comments on commit 591172b

Please sign in to comment.