Wrap GMT_Put_Vector and GMT_Write_Data

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.

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)