diff --git a/pygmt/clib/session.py b/pygmt/clib/session.py index 8d7ee8d71ed..77925597e3f 100644 --- a/pygmt/clib/session.py +++ b/pygmt/clib/session.py @@ -1375,6 +1375,8 @@ def virtualfile_from_data( extra_arrays : list of 1d arrays Optional. A list of numpy arrays in addition to x, y and z. All of these arrays must be of the same size as the x/y/z arrays. + required_z : bool + State whether the 'z' column is required. Returns ------- diff --git a/pygmt/helpers/utils.py b/pygmt/helpers/utils.py index 19b6e2ec2f7..2f3ce181164 100644 --- a/pygmt/helpers/utils.py +++ b/pygmt/helpers/utils.py @@ -39,8 +39,9 @@ def data_kind(data, x=None, y=None, z=None, required_z=False): x/y : 1d arrays or None x and y columns as numpy arrays. z : 1d array or None - z column as numpy array. To be used optionally when x and y - are given. + z column as numpy array. To be used optionally when x and y are given. + required_z : bool + State whether the 'z' column is required. Returns ------- @@ -80,7 +81,10 @@ def data_kind(data, x=None, y=None, z=None, required_z=False): elif hasattr(data, "__geo_interface__"): kind = "geojson" elif data is not None: - if required_z and data.shape[1] < 3: + if required_z and ( + getattr(data, "shape", (3, 3))[1] < 3 # np.array, pd.DataFrame + or len(getattr(data, "data_vars", (0, 1, 2))) < 3 # xr.Dataset + ): raise GMTInvalidInput("data must provide x, y, and z columns.") kind = "matrix" else: diff --git a/pygmt/tests/test_blockm.py b/pygmt/tests/test_blockm.py index 39a0e315f7f..d71b372b6b2 100644 --- a/pygmt/tests/test_blockm.py +++ b/pygmt/tests/test_blockm.py @@ -3,9 +3,11 @@ """ import os +import numpy as np import numpy.testing as npt import pandas as pd import pytest +import xarray as xr from pygmt import blockmean, blockmode from pygmt.datasets import load_sample_bathymetry from pygmt.exceptions import GMTInvalidInput @@ -31,12 +33,13 @@ def test_blockmean_input_dataframe(dataframe): npt.assert_allclose(output.iloc[0], [245.888877, 29.978707, -384.0]) -def test_blockmean_input_table_matrix(dataframe): +@pytest.mark.parametrize("array_func", [np.array, xr.Dataset]) +def test_blockmean_input_table_matrix(array_func, dataframe): """ Run blockmean using table input that is not a pandas.DataFrame but still a matrix. """ - table = dataframe.values + table = array_func(dataframe) output = blockmean(table=table, spacing="5m", region=[245, 255, 20, 30]) assert isinstance(output, pd.DataFrame) assert output.shape == (5849, 3) diff --git a/pygmt/tests/test_clib.py b/pygmt/tests/test_clib.py index 2be988ec487..22fd4d6d7b8 100644 --- a/pygmt/tests/test_clib.py +++ b/pygmt/tests/test_clib.py @@ -419,7 +419,36 @@ def test_virtual_file_bad_direction(): print("This should have failed") -def test_virtualfile_from_data_required_z_matrix(): +@pytest.mark.parametrize( + "array_func,kind", + [(np.array, "matrix"), (pd.DataFrame, "vector"), (xr.Dataset, "vector")], +) +def test_virtualfile_from_data_required_z_matrix(array_func, kind): + """ + Test that function works when third z column in a matrix is needed and + provided. + """ + shape = (5, 3) + dataframe = pd.DataFrame( + data=np.arange(shape[0] * shape[1]).reshape(shape), columns=["x", "y", "z"] + ) + data = array_func(dataframe) + with clib.Session() as lib: + with lib.virtualfile_from_data(data=data, required_z=True) as vfile: + with GMTTempFile() as outfile: + lib.call_module("info", f"{vfile} ->{outfile.name}") + output = outfile.read(keep_tabs=True) + bounds = "\t".join( + [ + f"<{i.min():.0f}/{i.max():.0f}>" + for i in (dataframe.x, dataframe.y, dataframe.z) + ] + ) + expected = f"<{kind} memory>: N = {shape[0]}\t{bounds}\n" + assert output == expected + + +def test_virtualfile_from_data_required_z_matrix_missing(): """ Test that function fails when third z column in a matrix is needed but not provided.