Support various datetime types as input (#464)

Currently supported datetime types: - [numpy.datetime64](https://numpy.org/doc/stable/reference/arrays.datetime.html) - [pandas.DateTimeIndex](https://pandas.pydata.org/docs/user_guide/timeseries.html) - Raw datetime strings in both ISO and non-ISO formats, e.g., `2010-01-01`, `1/1/2018`, `Jul 5, 2019` - Python's [built-in datetime and date](https://docs.python.org/3/library/datetime.html) * Add a function array_to_datetime to convert any legal array/list to pandas.DateTimeIndex and numpy.datetime64 types * Test that xarray.DataArray datetimes can be plotted Co-authored-by: Wei Ji <[email protected]>
GenericMappingTools · Jul 13, 2020 · f401f85 · f401f85
1 parent 0f12972
commit f401f85
Show file tree

Hide file tree

Showing 5 changed files with 147 additions and 13 deletions.
diff --git a/pygmt/clib/conversion.py b/pygmt/clib/conversion.py
@@ -2,6 +2,7 @@
 Functions to convert data types into ctypes friendly formats.
 """
 import numpy as np
+import pandas as pd
 
 from ..exceptions import GMTInvalidInput
 
@@ -237,3 +238,82 @@ def kwargs_to_ctypes_array(argument, kwargs, dtype):
     if argument in kwargs:
         return dtype(*kwargs[argument])
     return None
+
+
+def array_to_datetime(array):
+    """
+    Convert an 1d datetime array from various types into pandas.DatetimeIndex
+    (i.e., numpy.datetime64).
+
+    If the input array is not in legal datetime formats, raise a "ParseError"
+    exception.
+
+    Parameters
+    ----------
+    array : list or 1d array
+        The input datetime array in various formats.
+
+        Supported types:
+
+        - str
+        - numpy.datetime64
+        - pandas.DateTimeIndex
+        - datetime.datetime and datetime.date
+
+    Returns
+    -------
+    array : 1d datetime array in pandas.DatetimeIndex (i.e., numpy.datetime64)
+
+    Examples
+    --------
+    >>> import datetime
+    >>> # numpy.datetime64 array
+    >>> x = np.array(
+    ...     ["2010-06-01", "2011-06-01T12", "2012-01-01T12:34:56"],
+    ...     dtype="datetime64",
+    ... )
+    >>> array_to_datetime(x)
+    DatetimeIndex(['2010-06-01 00:00:00', '2011-06-01 12:00:00',
+                   '2012-01-01 12:34:56'],
+                  dtype='datetime64[ns]', freq=None)
+
+    >>> # pandas.DateTimeIndex array
+    >>> x = pd.date_range("2013", freq="YS", periods=3)
+    >>> array_to_datetime(x)  # doctest: +NORMALIZE_WHITESPACE
+    DatetimeIndex(['2013-01-01', '2014-01-01', '2015-01-01'],
+                  dtype='datetime64[ns]', freq='AS-JAN')
+
+    >>> # Python's built-in date and datetime
+    >>> x = [datetime.date(2018, 1, 1), datetime.datetime(2019, 1, 1)]
+    >>> array_to_datetime(x)  # doctest: +NORMALIZE_WHITESPACE
+    DatetimeIndex(['2018-01-01', '2019-01-01'],
+        dtype='datetime64[ns]', freq=None)
+
+    >>> # Raw datetime strings in various format
+    >>> x = [
+    ...     "2018",
+    ...     "2018-02",
+    ...     "2018-03-01",
+    ...     "2018-04-01T01:02:03",
+    ...     "5/1/2018",
+    ...     "Jun 05, 2018",
+    ...     "2018/07/02",
+    ... ]
+    >>> array_to_datetime(x)
+    DatetimeIndex(['2018-01-01 00:00:00', '2018-02-01 00:00:00',
+                   '2018-03-01 00:00:00', '2018-04-01 01:02:03',
+                   '2018-05-01 00:00:00', '2018-06-05 00:00:00',
+                   '2018-07-02 00:00:00'],
+                  dtype='datetime64[ns]', freq=None)
+
+    >>> # Mixed datetime types
+    >>> x = [
+    ...    "2018-01-01",
+    ...    np.datetime64("2018-01-01"),
+    ...    datetime.datetime(2018, 1, 1),
+    ... ]
+    >>> array_to_datetime(x)  # doctest: +NORMALIZE_WHITESPACE
+    DatetimeIndex(['2018-01-01', '2018-01-01', '2018-01-01'],
+                  dtype='datetime64[ns]', freq=None)
+    """
+    return pd.to_datetime(array)
diff --git a/pygmt/clib/session.py b/pygmt/clib/session.py
@@ -22,6 +22,7 @@
     vectors_to_arrays,
     dataarray_to_matrix,
     as_c_contiguous,
+    array_to_datetime,
 )
 
 FAMILIES = [
@@ -48,12 +49,13 @@
 REGISTRATIONS = ["GMT_GRID_PIXEL_REG", "GMT_GRID_NODE_REG"]
 
 DTYPES = {
-    "float64": "GMT_DOUBLE",
-    "float32": "GMT_FLOAT",
-    "int64": "GMT_LONG",
-    "int32": "GMT_INT",
-    "uint64": "GMT_ULONG",
-    "uint32": "GMT_UINT",
+    np.float64: "GMT_DOUBLE",
+    np.float32: "GMT_FLOAT",
+    np.int64: "GMT_LONG",
+    np.int32: "GMT_INT",
+    np.uint64: "GMT_ULONG",
+    np.uint32: "GMT_UINT",
+    np.datetime64: "GMT_DATETIME",
 }
 
 
@@ -708,15 +710,22 @@ def _check_dtype_and_dim(self, array, ndim):
         True
 
         """
-        if array.dtype.name not in DTYPES:
-            raise GMTInvalidInput(
-                "Unsupported numpy data type '{}'.".format(array.dtype.name)
-            )
+        # check the array has the given dimension
         if array.ndim != ndim:
             raise GMTInvalidInput(
                 "Expected a numpy 1d array, got {}d.".format(array.ndim)
             )
-        return self[DTYPES[array.dtype.name]]
+
+        # check the array has a valid/known data type
+        if array.dtype.type not in DTYPES:
+            try:
+                # Try to convert any unknown numpy data types to np.datetime64
+                array = np.asarray(array, dtype=np.datetime64)
+            except ValueError:
+                raise GMTInvalidInput(
+                    "Unsupported numpy data type '{}'.".format(array.dtype.type)
+                )
+        return self[DTYPES[array.dtype.type]]
 
     def put_vector(self, dataset, column, vector):
         """
@@ -762,7 +771,13 @@ def put_vector(self, dataset, column, vector):
         )
 
         gmt_type = self._check_dtype_and_dim(vector, ndim=1)
-        vector_pointer = vector.ctypes.data_as(ctp.c_void_p)
+        if gmt_type == self["GMT_DATETIME"]:
+            vector_pointer = (ctp.c_char_p * len(vector))()
+            vector_pointer[:] = np.char.encode(
+                np.datetime_as_string(array_to_datetime(vector))
+            )
+        else:
+            vector_pointer = vector.ctypes.data_as(ctp.c_void_p)
         status = c_put_vector(
             self.session_pointer, dataset, column, gmt_type, vector_pointer
         )

diff --git a/pygmt/tests/baseline/test_plot_datetime.png b/pygmt/tests/baseline/test_plot_datetime.png
diff --git a/pygmt/tests/test_clib.py b/pygmt/tests/test_clib.py
@@ -349,7 +349,7 @@ def test_put_vector_invalid_dtype():
             mode="GMT_CONTAINER_ONLY",
             dim=[2, 3, 1, 0],  # columns, rows, layers, dtype
         )
-        data = np.array([37, 12, 556], dtype="complex128")
+        data = np.array([37, 12, 556], dtype="object")
         with pytest.raises(GMTInvalidInput):
             lib.put_vector(dataset, column=1, vector=data)
 

diff --git a/pygmt/tests/test_plot.py b/pygmt/tests/test_plot.py
@@ -2,9 +2,13 @@
 """
 Tests plot.
 """
+import datetime
 import os
 
 import numpy as np
+import pandas as pd
+import xarray as xr
+
 import pytest
 
 from .. import Figure
@@ -284,3 +288,38 @@ def test_plot_scalar_xy():
     fig.plot(x=0, y=0, style="t1c")
     fig.plot(x=1.5, y=-1.5, style="s1c")
     return fig
+
+
+@pytest.mark.mpl_image_compare
+def test_plot_datetime():
+    """Test various datetime input data"""
+    fig = Figure()
+    fig.basemap(projection="X15c/5c", region="2010-01-01/2020-01-01/0/10", frame=True)
+
+    # numpy.datetime64 types
+    x = np.array(
+        ["2010-06-01", "2011-06-01T12", "2012-01-01T12:34:56"], dtype="datetime64"
+    )
+    y = [1.0, 2.0, 3.0]
+    fig.plot(x, y, style="c0.2c", pen="1p")
+
+    # pandas.DatetimeIndex
+    x = pd.date_range("2013", freq="YS", periods=3)
+    y = [4, 5, 6]
+    fig.plot(x, y, style="t0.2c", pen="1p")
+
+    # xarray.DataArray
+    x = xr.DataArray(data=pd.date_range(start="2015-03", freq="QS", periods=3))
+    y = [7.5, 6, 4.5]
+    fig.plot(x, y, style="s0.2c", pen="1p")
+
+    # raw datetime strings
+    x = ["2016-02-01", "2017-03-04T00:00"]
+    y = [7, 8]
+    fig.plot(x, y, style="a0.2c", pen="1p")
+
+    # the Python built-in datetime and date
+    x = [datetime.date(2018, 1, 1), datetime.datetime(2019, 1, 1)]
+    y = [8.5, 9.5]
+    fig.plot(x, y, style="i0.2c", pen="1p")
+    return fig