diff --git a/gmt/clib/core.py b/gmt/clib/core.py index f6f00281ecf..0a522c1ba4f 100644 --- a/gmt/clib/core.py +++ b/gmt/clib/core.py @@ -96,6 +96,7 @@ def __init__(self, libname='libgmt'): self._c_create_data = None self._c_handle_messages = None self._c_put_vector = None + self._c_put_matrix = None self._c_write_data = None self._bind_clib_functions(libname) @@ -176,6 +177,11 @@ def _bind_clib_functions(self, libname): ctypes.c_void_p] self._c_put_vector.restype = ctypes.c_int + self._c_put_matrix = self._libgmt.GMT_Put_Matrix + self._c_put_matrix.argtypes = [ctypes.c_void_p, ctypes.c_void_p, + ctypes.c_uint, ctypes.c_void_p] + self._c_put_matrix.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, @@ -447,15 +453,18 @@ def create_data(self, family, geometry, mode, **kwargs): if mode not in self.data_modes: raise GMTCLibError("Invalid data creation mode '{}'.".format(mode)) geometry_int = self._parse_data_geometry(geometry) + # dim is required (get a segmentation fault if passing it as None + if 'dim' not in kwargs: + kwargs['dim'] = [0]*4 # 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) inc = kwargs_to_ctypes_array('inc', kwargs, ctypes.c_double*2) + pad = self._parse_pad(family, kwargs) # Use the GMT defaults if no value is given registration = kwargs.get('registration', self.get_constant('GMT_GRID_NODE_REG')) - pad = kwargs.get('pad', self.get_constant('GMT_PAD_DEFAULT')) data_ptr = self._c_create_data( self.current_session, @@ -475,6 +484,22 @@ def create_data(self, family, geometry, mode, **kwargs): return data_ptr + def _parse_pad(self, family, kwargs): + """ + Parse and return an appropriate value for pad if none is given. + + Pad is a bit tricky because, for matrix types, pad control the matrix + ordering (row or column major). Using the default pad will set it to + column major and mess things up with the numpy arrays. + """ + pad = kwargs.get('pad', None) + if pad is None: + if 'MATRIX' in family: + pad = 0 + else: + pad = self.get_constant('GMT_PAD_DEFAULT') + return pad + def _parse_data_geometry(self, geometry): """ Parse the geometry argument for GMT data manipulation functions. @@ -548,6 +573,50 @@ def _parse_data_family(self, family): via_value = 0 return family_value + via_value + def _check_dtype_and_dim(self, array, ndim): + """ + Check that a numpy array has the given dimensions and is a valid data + type. + + Parameters + ---------- + array : numpy array + The array to be tested. + ndim : int + The desired dimension of the array. + + Returns + ------- + gmt_type : int + The GMT constant value representing this data type. + + Raises + ------ + GMTCLibError + If the array has the wrong dimensions or is an unsupported data + type. + + Examples + -------- + + >>> import numpy as np + >>> data = np.array([1, 2, 3], dtype='float64') + >>> with LibGMT() as lib: + ... gmttype = lib._check_dtype_and_dim(data, ndim=1) + ... gmttype == lib.get_constant('GMT_DOUBLE') + True + + """ + if array.dtype.name not in self._dtypes: + raise GMTCLibError( + "Unsupported numpy data type '{}'.".format(array.dtype.name) + ) + if array.ndim != ndim: + raise GMTCLibError( + "Expected a numpy 1d array, got {}d.".format(array.ndim) + ) + return self.get_constant(self._dtypes[array.dtype.name]) + def put_vector(self, dataset, column, vector): """ Attach a numpy 1D array as a column on a GMT dataset. @@ -585,15 +654,7 @@ def put_vector(self, dataset, column, vector): 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]) + gmt_type = self._check_dtype_and_dim(vector, ndim=1) vector_pointer = vector.ctypes.data_as(ctypes.c_void_p) status = self._c_put_vector(self.current_session, dataset, @@ -608,6 +669,50 @@ def put_vector(self, dataset, column, vector): ]) ) + def put_matrix(self, dataset, matrix): + """ + Attach a numpy 2D array to a GMT dataset. + + Use this functions to attach numpy array data to a GMT dataset and pass + it to GMT modules. Wraps ``GMT_Put_Matrix``. + + The dataset must be created by :meth:`~gmt.clib.LibGMT.create_data` + first. Use ``|GMT_VIA_MATRIX'`` in the family. + + 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. 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`. + matrix : numpy 2d-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_Matrix`` exits with status != + 0. + + """ + gmt_type = self._check_dtype_and_dim(matrix, ndim=2) + matrix_pointer = matrix.ctypes.data_as(ctypes.c_void_p) + status = self._c_put_matrix(self.current_session, + dataset, + gmt_type, + matrix_pointer) + if status != 0: + raise GMTCLibError( + "Failed to put matrix of type {}.".format(matrix.dtype)) + # pylint: disable=too-many-arguments def write_data(self, family, geometry, mode, wesn, output, data): """ diff --git a/gmt/tests/test_clib.py b/gmt/tests/test_clib.py index 07ffdc76c45..0ff1c082aa2 100644 --- a/gmt/tests/test_clib.py +++ b/gmt/tests/test_clib.py @@ -211,7 +211,7 @@ def test_create_data_dataset(): def test_create_data_grid_dim(): - "Run the function to make sure it doesn't fail badly." + "Create a grid ignoring range and inc." with LibGMT() as lib: # Grids from matrices using dim lib.create_data( @@ -223,14 +223,13 @@ def test_create_data_grid_dim(): def test_create_data_grid_range(): - "Run the function to make sure it doesn't fail badly." + "Create a grid specifying range and inc instead of dim." with LibGMT() as lib: # Grids from matrices using range and int lib.create_data( family='GMT_IS_GRID|GMT_VIA_MATRIX', geometry='GMT_IS_SURFACE', mode='GMT_CONTAINER_ONLY', - dim=[0, 0, 1, 0], ranges=[150., 250., -20., 20.], inc=[0.1, 0.2], ) @@ -279,9 +278,11 @@ def test_put_vector(): 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] + # Turns out wesn doesn't matter for Datasets + wesn = [0]*6 # Save the data to a file to see if it's being accessed correctly - tmp_file = NamedTemporaryFile(delete=False) + tmp_file = NamedTemporaryFile(prefix='gmt-python-', suffix='.txt', + delete=False) tmp_file.close() lib.write_data('GMT_IS_VECTOR', 'GMT_IS_POINT', 'GMT_WRITE_SET', wesn, tmp_file.name, dataset) @@ -337,3 +338,63 @@ def test_put_vector_2d_fails(): data = np.array([[37, 12, 556], [37, 12, 556]], dtype='int32') with pytest.raises(GMTCLibError): lib.put_vector(dataset, column=0, vector=data) + + +def test_put_matrix(): + "Check that assigning a numpy 2d array to a dataset works" + dtypes = 'float32 float64 int32 int64 uint32 uint64'.split() + shape = (3, 4) + for dtype in dtypes: + with LibGMT() as lib: + # Dataset from vectors + dataset = lib.create_data( + family='GMT_IS_DATASET|GMT_VIA_MATRIX', + geometry='GMT_IS_POINT', + mode='GMT_CONTAINER_ONLY', + dim=[shape[1], shape[0], 1, 0], # columns, rows, layers, dtype + ) + data = np.arange(shape[0]*shape[1], dtype=dtype).reshape(shape) + lib.put_matrix(dataset, matrix=data) + # wesn doesn't matter for Datasets + wesn = [0]*6 + # Save the data to a file to see if it's being accessed correctly + tmp_file = NamedTemporaryFile(prefix='gmt-python-', suffix='.txt', + delete=False) + tmp_file.close() + lib.write_data('GMT_IS_MATRIX', 'GMT_IS_POINT', + 'GMT_WRITE_SET', wesn, tmp_file.name, dataset) + # Load the data and check that it's correct + newdata = np.loadtxt(tmp_file.name, dtype=dtype) + os.remove(tmp_file.name) + npt.assert_allclose(newdata, data) + + +def test_put_matrix_grid(): + "Check that assigning a numpy 2d array to a grid works" + dtypes = 'float32 float64 int32 int64 uint32 uint64'.split() + wesn = [10, 15, 30, 40, 0, 0] + inc = [1, 1] + shape = ((wesn[3] - wesn[2])//inc[1] + 1, (wesn[1] - wesn[0])//inc[0] + 1) + for dtype in dtypes: + with LibGMT() as lib: + # Dataset from vectors + grid = lib.create_data( + family='GMT_IS_GRID|GMT_VIA_MATRIX', + geometry='GMT_IS_SURFACE', + mode='GMT_CONTAINER_ONLY', + ranges=wesn[:4], + inc=[1, 1], + registration=lib.get_constant('GMT_GRID_NODE_REG'), + ) + data = np.arange(shape[0]*shape[1], dtype=dtype).reshape(shape) + lib.put_matrix(grid, matrix=data) + # Save the data to a file to see if it's being accessed correctly + tmp_file = NamedTemporaryFile(prefix='gmt-python-', suffix='.txt', + delete=False) + tmp_file.close() + lib.write_data('GMT_IS_MATRIX', 'GMT_IS_SURFACE', + 'GMT_CONTAINER_AND_DATA', wesn, tmp_file.name, grid) + # Load the data and check that it's correct + newdata = np.loadtxt(tmp_file.name, dtype=dtype) + os.remove(tmp_file.name) + npt.assert_allclose(newdata, data)