diff --git a/gmt/clib/core.py b/gmt/clib/core.py index 838728e9714..f6f00281ecf 100644 --- a/gmt/clib/core.py +++ b/gmt/clib/core.py @@ -75,6 +75,16 @@ class LibGMT(): # pylint: disable=too-many-instance-attributes 'GMT_OUTPUT', ] + # Map numpy dtypes to GMT types + _dtypes = { + 'float64': 'GMT_DOUBLE', + 'float32': 'GMT_FLOAT', + 'int64': 'GMT_LONG', + 'int32': 'GMT_INT', + 'uint64': 'GMT_ULONG', + 'uint32': 'GMT_UINT', + } + def __init__(self, libname='libgmt'): self._logfile = None self._session_id = None @@ -85,6 +95,8 @@ def __init__(self, libname='libgmt'): self._c_call_module = None self._c_create_data = None self._c_handle_messages = None + self._c_put_vector = None + self._c_write_data = None self._bind_clib_functions(libname) @property @@ -158,6 +170,20 @@ def _bind_clib_functions(self, libname): ctypes.c_uint, ctypes.c_char_p] self._c_handle_messages.restype = ctypes.c_int + self._c_put_vector = self._libgmt.GMT_Put_Vector + self._c_put_vector.argtypes = [ctypes.c_void_p, ctypes.c_void_p, + ctypes.c_uint, ctypes.c_uint, + ctypes.c_void_p] + self._c_put_vector.restype = ctypes.c_int + + self._c_write_data = self._libgmt.GMT_Write_Data + self._c_write_data.argtypes = [ctypes.c_void_p, ctypes.c_uint, + ctypes.c_uint, ctypes.c_uint, + ctypes.c_uint, + ctypes.POINTER(ctypes.c_double), + ctypes.c_char_p, ctypes.c_void_p] + self._c_write_data.restype = ctypes.c_int + def __enter__(self): """ Start the GMT session and keep the session argument. @@ -420,9 +446,7 @@ def create_data(self, family, geometry, mode, **kwargs): family_int = self._parse_data_family(family) if mode not in self.data_modes: raise GMTCLibError("Invalid data creation mode '{}'.".format(mode)) - if geometry not in self.data_geometries: - raise GMTCLibError("Invalid data geometry '{}'.".format(geometry)) - + geometry_int = self._parse_data_geometry(geometry) # Convert dim, ranges, and inc to ctypes arrays if given dim = kwargs_to_ctypes_array('dim', kwargs, ctypes.c_uint64*4) ranges = kwargs_to_ctypes_array('ranges', kwargs, ctypes.c_double*4) @@ -436,7 +460,7 @@ def create_data(self, family, geometry, mode, **kwargs): data_ptr = self._c_create_data( self.current_session, family_int, - self.get_constant(geometry), + geometry_int, self.get_constant(mode), dim, ranges, @@ -451,9 +475,36 @@ def create_data(self, family, geometry, mode, **kwargs): return data_ptr + def _parse_data_geometry(self, geometry): + """ + Parse the geometry argument for GMT data manipulation functions. + + Converts the string name to the corresponding integer value. + + Parameters + ---------- + geometry : str + A valid GMT data geometry name (e.g., ``'GMT_IS_POINT'``). See the + ``data_geometries`` attribute for valid names. + + Returns + ------- + geometry_int : int + The converted geometry. + + Raises + ------ + GMTCLibError + If the geometry name is invalid. + + """ + if geometry not in self.data_geometries: + raise GMTCLibError("Invalid data geometry '{}'.".format(geometry)) + return self.get_constant(geometry) + def _parse_data_family(self, family): """ - Parse the data family string into an integer number. + Parse the data family string into an integer number. Valid family names are: GMT_IS_DATASET, GMT_IS_GRID, GMT_IS_PALETTE, GMT_IS_TEXTSET, GMT_IS_MATRIX, and GMT_IS_VECTOR. @@ -496,3 +547,121 @@ def _parse_data_family(self, family): else: via_value = 0 return family_value + via_value + + def put_vector(self, dataset, column, vector): + """ + Attach a numpy 1D array as a column on a GMT dataset. + + Use this functions to attach numpy array data to a GMT dataset and pass + it to GMT modules. Wraps ``GMT_Put_Vector``. + + The dataset must be created by :meth:`~gmt.clib.LibGMT.create_data` + first. Use ``family='GMT_IS_DATASET|GMT_VIA_VECTOR'``. + + Not at all numpy dtypes are supported, only: float64, float32, int64, + int32, uint64, and uint32. + + .. warning:: + The numpy array must be C contiguous in memory. If it comes from a + column slice of a 2d array, for example, you will have to make a + copy. Use :func:`numpy.ascontiguousarray` to make sure your vector + is contiguous (it won't copy if it already is). + + Parameters + ---------- + dataset : ctypes.c_void_p + The ctypes void pointer to a ``GMT_Dataset``. Create it with + :meth:`~gmt.clib.LibGMT.create_data`. + column : int + The column number of this vector in the dataset (starting from 0). + vector : numpy 1d-array + The array that will be attached to the dataset. Must be a 1d C + contiguous array. + + Raises + ------ + GMTCLibError + If given invalid input or ``GMT_Put_Vector`` exits with status != + 0. + + """ + if vector.dtype.name not in self._dtypes: + raise GMTCLibError( + "Unsupported numpy data type '{}'.".format(vector.dtype.name) + ) + if vector.ndim != 1: + raise GMTCLibError( + "Expected a numpy 1d array, got {}d.".format(vector.ndim) + ) + gmt_type = self.get_constant(self._dtypes[vector.dtype.name]) + vector_pointer = vector.ctypes.data_as(ctypes.c_void_p) + status = self._c_put_vector(self.current_session, + dataset, + column, + gmt_type, + vector_pointer) + if status != 0: + raise GMTCLibError( + ' '.join([ + "Failed to put vector of type {}".format(vector.dtype), + "in column {} of dataset.".format(column), + ]) + ) + + # pylint: disable=too-many-arguments + def write_data(self, family, geometry, mode, wesn, output, data): + """ + Write a GMT data container to a file. + + The data container should be created by + :meth:`~gmt.clib.LibGMT.create_data`. + + Wraps ``GMT_Write_Data`` but only allows writing to a file. So the + ``method`` argument is omitted. + + Parameters + ---------- + family : str + A valid GMT data family name (e.g., ``'GMT_IS_DATASET'``). See the + ``data_families`` attribute for valid names. Don't use the + ``GMT_VIA_VECTOR`` or ``GMT_VIA_MATRIX`` constructs for this. Use + ``GMT_IS_VECTOR`` and ``GMT_IS_MATRIX`` instead. + geometry : str + A valid GMT data geometry name (e.g., ``'GMT_IS_POINT'``). See the + ``data_geometries`` attribute for valid names. + mode : str + How the data is to be written to the file. This option varies + depending on the given family. See the GMT API documentation for + details. + wesn : list or numpy array + [xmin, xmax, ymin, ymax, zmin, zmax] of the data. Must have 6 + elements. + output : str + The output file name. + data : ctypes.c_void_p + Pointer to the data container created by + :meth:`~gmt.clib.LibGMT.create_data`. + + Raises + ------ + GMTCLibError + For invalid input arguments or if the GMT API functions returns a + non-zero status code. + + """ + family_int = self._parse_data_family(family) + geometry_int = self._parse_data_geometry(geometry) + status = self._c_write_data(self.current_session, + family_int, + self.get_constant('GMT_IS_FILE'), + geometry_int, + self.get_constant(mode), + (ctypes.c_double*6)(*wesn), + output.encode(), + data) + if status != 0: + raise GMTCLibError( + "Failed to write dataset to '{}'".format(output)) + # Can't test this if by giving a bad file name because if + # output=='', GMT will just write to stdout and spaces are valid + # file names. diff --git a/gmt/tests/test_clib.py b/gmt/tests/test_clib.py index cbfa0a819c8..07ffdc76c45 100644 --- a/gmt/tests/test_clib.py +++ b/gmt/tests/test_clib.py @@ -3,8 +3,11 @@ Test the wrappers for the C API. """ import os +from tempfile import NamedTemporaryFile import pytest +import numpy as np +import numpy.testing as npt from ..clib.core import LibGMT from ..clib.utils import clib_extension, load_libgmt, check_libgmt @@ -47,6 +50,7 @@ def test_constant(): assert lib.get_constant('GMT_SESSION_EXTERNAL') != -99999 assert lib.get_constant('GMT_MODULE_CMD') != -99999 assert lib.get_constant('GMT_PAD_DEFAULT') != -99999 + assert lib.get_constant('GMT_DOUBLE') != -99999 with pytest.raises(GMTCLibError): lib.get_constant('A_WHOLE_LOT_OF_JUNK') @@ -255,3 +259,81 @@ def test_create_data_fails(): ranges=[150., 250., -20., 20.], inc=[0.1, 0.2], ) + + +def test_put_vector(): + "Check that assigning a numpy array to a dataset works" + dtypes = 'float32 float64 int32 int64 uint32 uint64'.split() + for dtype in dtypes: + with LibGMT() as lib: + # Dataset from vectors + dataset = lib.create_data( + family='GMT_IS_DATASET|GMT_VIA_VECTOR', + geometry='GMT_IS_POINT', + mode='GMT_CONTAINER_ONLY', + dim=[3, 5, 1, 0], # columns, rows, layers, dtype + ) + x = np.array([1, 2, 3, 4, 5], dtype=dtype) + y = np.array([6, 7, 8, 9, 10], dtype=dtype) + z = np.array([11, 12, 13, 14, 15], dtype=dtype) + lib.put_vector(dataset, column=lib.get_constant("GMT_X"), vector=x) + lib.put_vector(dataset, column=lib.get_constant("GMT_Y"), vector=y) + lib.put_vector(dataset, column=lib.get_constant("GMT_Z"), vector=z) + wesn = [x.min(), x.max(), y.min(), y.max(), 0, 0] + # Save the data to a file to see if it's being accessed correctly + tmp_file = NamedTemporaryFile(delete=False) + tmp_file.close() + lib.write_data('GMT_IS_VECTOR', 'GMT_IS_POINT', + 'GMT_WRITE_SET', wesn, tmp_file.name, dataset) + # Load the data and check that it's correct + newx, newy, newz = np.loadtxt(tmp_file.name, unpack=True, + dtype=dtype) + os.remove(tmp_file.name) + npt.assert_allclose(newx, x) + npt.assert_allclose(newy, y) + npt.assert_allclose(newz, z) + + +def test_put_vector_invalid_dtype(): + "Check that it fails with an exception for invalid data types" + with LibGMT() as lib: + # Dataset from vectors + dataset = lib.create_data( + family='GMT_IS_DATASET|GMT_VIA_VECTOR', + geometry='GMT_IS_POINT', + mode='GMT_CONTAINER_ONLY', + dim=[2, 3, 1, 0], # columns, rows, layers, dtype + ) + data = np.array([37, 12, 556], dtype='complex128') + with pytest.raises(GMTCLibError): + lib.put_vector(dataset, column=1, vector=data) + + +def test_put_vector_wrong_column(): + "Check that it fails with an exception when giving an invalid column" + with LibGMT() as lib: + # Dataset from vectors + dataset = lib.create_data( + family='GMT_IS_DATASET|GMT_VIA_VECTOR', + geometry='GMT_IS_POINT', + mode='GMT_CONTAINER_ONLY', + dim=[1, 3, 1, 0], # columns, rows, layers, dtype + ) + data = np.array([37, 12, 556], dtype='float32') + with pytest.raises(GMTCLibError): + lib.put_vector(dataset, column=1, vector=data) + + +def test_put_vector_2d_fails(): + "Check that it fails with an exception for multidimensional arrays" + with LibGMT() as lib: + # Dataset from vectors + dataset = lib.create_data( + family='GMT_IS_DATASET|GMT_VIA_VECTOR', + geometry='GMT_IS_POINT', + mode='GMT_CONTAINER_ONLY', + dim=[1, 6, 1, 0], # columns, rows, layers, dtype + ) + data = np.array([[37, 12, 556], [37, 12, 556]], dtype='int32') + with pytest.raises(GMTCLibError): + lib.put_vector(dataset, column=0, vector=data)