Pass in data to Figure.plot (#93)

Use the new clib context managers to pass in data to the plot module.
Made some utils functions to help simplify the code in the method. Add
new test plots to cover all of the method. Add more documentation and
plots to the first steps notebook.

doc/first-steps.ipynb

@@ -207,15 +207,34 @@
     "# get the same numbers\n",
     "np.random.seed(42)\n",
     "\n",
-    "lon = np.random.uniform(150, 240, 30)\n",
-    "lat = np.random.uniform(-30, 60, 30)"
+    "ndata = 30\n",
+    "region = [150, 240, -30, 60]\n",
+    "lon = np.random.uniform(region[0], region[1], ndata)\n",
+    "lat = np.random.uniform(region[2], region[3], ndata)\n",
+    "magnitude = np.random.uniform(1, 9, ndata)\n",
+    "depth = np.random.uniform(0, 1, ndata)"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Save the data to a file in 2 columns:"
+    "There are 3 ways to pass data into `Figure.plot`:\n",
+    "\n",
+    "1. x and y coordinates as 1d numpy arrays using the `x` and `y` arguments.\n",
+    "2. A whole data table as a 2d numpy array using the `data` argument.\n",
+    "3. A file name using the `data` argument.\n",
+    "\n",
+    "Let's explore all of these options."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Using x and y\n",
+    "\n",
+    "Now we can plot the data using `Figure.plot` and passing the x and y coordinate arrays:"
    ]
   },
   {
@@ -224,14 +243,65 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "np.savetxt('first-steps-data.txt', np.transpose([lon, lat]))"
+    "fig = gmt.Figure()\n",
+    "fig.coast(region='g', projection='G200/30/6i', frame='ag', \n",
+    "          resolution='i', area_thresh=5000, land='white', \n",
+    "          water='DarkTurquoise')\n",
+    "# Plot using circles (c) of 0.3 cm\n",
+    "fig.plot(x=lon, y=lat, style='c0.3c', color='black')\n",
+    "fig.show()"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Now we can plot the data using `Figure.plot` and passing the file name as the `data` argument:"
+    "We can make the size of the markers follow a data value by passing in the argument `sizes`:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig = gmt.Figure()\n",
+    "fig.coast(region='g', projection='G200/30/6i', frame='ag', \n",
+    "          resolution='i', area_thresh=5000, land='white', \n",
+    "          water='DarkTurquoise')\n",
+    "fig.plot(x=lon, y=lat, sizes=0.1*magnitude, style='cc', color='black')\n",
+    "fig.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We can also map the colors of the markers to data by passing `color` as an array and providing a colormap name (`cmap`):"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig = gmt.Figure()\n",
+    "fig.coast(region='g', projection='G200/30/6i', frame='ag', \n",
+    "          resolution='i', area_thresh=5000, land='white', \n",
+    "          water='DarkTurquoise')\n",
+    "fig.plot(x=lon, y=lat, style='cc', sizes=0.1*magnitude, \n",
+    "         color=depth, cmap='viridis')\n",
+    "fig.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Using a matrix\n",
+    "\n",
+    "Sometimes, you'll have data that you loaded from a file in the form of a single numpy 2d array (like a table). You can plot these data by passing in `data`."
    ]
   },
   {
@@ -240,23 +310,70 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "fig_globe = gmt.Figure()\n",
-    "fig_globe.coast(\n",
-    "    region='g', projection='G200/30/6i', frame='ag', \n",
-    "    resolution='i', area_thresh=5000, land='white', \n",
-    "    water='DarkTurquoise')\n",
-    "fig_globe.plot(data='first-steps-data.txt', style='s0.3c',\n",
-    "               color='black')\n",
-    "fig_globe.show()"
+    "data = np.transpose([lon, lat, magnitude, depth])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We can use the `columns` argument to specify the columns that correspond to x, y, color, and size, respectively. We'll use it to specify that we want to set the colors using the depths:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig = gmt.Figure()\n",
+    "fig.coast(region=region, projection='M6i', frame='ag', \n",
+    "          resolution='i', area_thresh=5000, land='grey')\n",
+    "fig.plot(data=data, style='c0.8c', cmap='viridis', columns='0,1,3')\n",
+    "fig.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Using a file name\n",
+    "\n",
+    "If you have data in a file and you don't need to do any calculations on it, you can pass the file name to `plot` directly. The syntax for plotting will be the same as the example for a data matrix.\n",
+    "\n",
+    "First, we must save our sample data to a file:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "np.savetxt('first-steps-data.txt', data)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig = gmt.Figure()\n",
+    "fig.coast(region='g', projection='N-165/6i', frame='ag', \n",
+    "          resolution='i', area_thresh=5000, land='chocolate')\n",
+    "fig.plot(data='first-steps-data.txt', style='c0.4c', \n",
+    "         cmap='viridis', columns='0,1,3')\n",
+    "fig.show()"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "GMT ships some sample data that can be accessed by passing `@data_file` as the `data` argument. \n",
+    "### Using sample data from GMT\n",
     "\n",
-    "For example, we can plot the earthquake epicenters from the `tut_quakes.ngdc` sample dataset:"
+    "GMT ships some sample data that can be accessed by passing `@data_file` as the `data` argument. For example, we can plot the earthquake epicenters from the `tut_quakes.ngdc` sample dataset:"
    ]
   },
   {

gmt/base_plotting.py

@@ -3,8 +3,9 @@
 Does not define any special non-GMT methods (savefig, show, etc).
 """
 from .clib import LibGMT
-from .utils import build_arg_string
+from .utils import build_arg_string, dummy_context, data_kind
 from .decorators import fmt_docstring, use_alias, kwargs_to_strings
+from .exceptions import GMTError
 
 
 class BasePlotting():
@@ -102,17 +103,22 @@ def coast(self, **kwargs):
 
     @fmt_docstring
     @use_alias(R='region', J='projection', B='frame', P='portrait', S='style',
-               G='color', W='pen', i='columns')
+               G='color', W='pen', i='columns', C='cmap')
     @kwargs_to_strings(R='sequence', i='sequence_comma')
-    def plot(self, data, **kwargs):
+    def plot(self, x=None, y=None, data=None, sizes=None, **kwargs):
         """
         Plot lines, polygons, and symbols on maps.
 
         Used to be psxy.
 
-        Takes (x,y) pairs as inputs or reads them from a file and plots lines,
+        Takes a matrix, (x,y) pairs, or a file name as input and plots lines,
         polygons, or symbols at those locations on a map.
 
+        Must provide either *data* or *x* and *y*.
+
+        If providing data through *x* and *y*, *color* (G) can be a 1d array
+        that will be mapped to a colormap.
+
         If a symbol is selected and no symbol size given, then psxy will
         interpret the third column of the input data as symbol size. Symbols
         whose size is <= 0 are skipped. If no symbols are specified then the
@@ -129,9 +135,15 @@ def plot(self, data, **kwargs):
 
         Parameters
         ----------
-        data : str or array
-            *Required*. Input data table as an array or a file name.
-            **Only accepts file names for now.**
+        x, y : 1d arrays
+            Arrays of x and y coordinates of the data points.
+        data : str or 2d array
+            Either a data file name or a 2d numpy array with the tabular data.
+            Use option *columns* (i) to choose which columns are x, y, color,
+            and size, respectively.
+        sizes : 1d array
+            The sizes of the data points in units specified in *style* (S).
+            Only valid if using *x* and *y*.
         {J}
         {R}
         A : bool or str
@@ -157,10 +169,34 @@ def plot(self, data, **kwargs):
 
         """
         kwargs = self._preprocess(**kwargs)
-        assert isinstance(data, str), 'Only accepts file names for now.'
-        arg_str = ' '.join([data, build_arg_string(kwargs)])
+
+        kind = data_kind(data, x, y)
+
+        extra_arrays = []
+        if 'G' in kwargs and not isinstance(kwargs['G'], str):
+            if kind != 'vectors':
+                raise GMTError(
+                    "Can't use arrays for color if data is matrix or file.")
+            extra_arrays.append(kwargs['G'])
+            del kwargs['G']
+        if sizes is not None:
+            if kind != 'vectors':
+                raise GMTError(
+                    "Can't use arrays for sizes if data is matrix or file.")
+            extra_arrays.append(sizes)
+
         with LibGMT() as lib:
-            lib.call_module('psxy', arg_str)
+            # Choose how data will be passed in to the module
+            if kind == 'file':
+                file_context = dummy_context(data)
+            elif kind == 'matrix':
+                file_context = lib.matrix_to_vfile(data)
+            elif kind == 'vectors':
+                file_context = lib.vectors_to_vfile(x, y, *extra_arrays)
+
+            with file_context as fname:
+                arg_str = ' '.join([fname, build_arg_string(kwargs)])
+                lib.call_module('plot', arg_str)
 
     @fmt_docstring
     @use_alias(R='region', J='projection', B='frame', P='portrait')