Repository: fatiando/verde
Pull request #49 (2018-05-20 23:49:40)

Fit trends to vector data
The `VectorTrend` class fits a separate `Trend` to each component of the
data. Components can be accessed separately using
`VectorTrend.component_`. Operations of the `VectorTrend` will produce
multi-component data.

Had to generalize the `check_fit_input` function to allow for
multi-component data and weights.


Patch:
diff --git a/examples/vector_trend.py b/examples/vector_trend.py
new file mode 100644
index 0000000..c4e0110
--- /dev/null
+++ b/examples/vector_trend.py
@@ -0,0 +1,83 @@
+"""
+Trends in vector data
+=====================
+
+Verde provides the :class:`verde.VectorTrend` class to estimate a polynomial
+trend on each component of vector data, like GPS velocities. You can access
+each trend as a separate :class:`verde.Trend` or operate on all vector
+components directly (for gridding, etc).
+"""
+import matplotlib.pyplot as plt
+import cartopy.crs as ccrs
+import cartopy.feature as cfeature
+# We need these two classes to set proper ticklabels for Cartopy maps
+from cartopy.mpl.ticker import LongitudeFormatter, LatitudeFormatter
+import numpy as np
+import verde as vd
+
+
+# Fetch the GPS data from the U.S. West coast. The data has a strong trend
+# toward the North-West because of the relative movement along the San Andreas
+# Fault System.
+data = vd.datasets.fetch_california_gps()
+
+# We'll fit a degree 4 trend on both the East and North components and weight
+# the data using the variance of each component.
+trend = vd.VectorTrend(degree=4)
+trend.fit(coordinates=(data.longitude, data.latitude),
+          data=(data.velocity_east, data.velocity_north),
+          weights=((data.std_east.min()/data.std_east)**2,
+                   (data.std_north.min()/data.std_north)**2))
+print("Vector trend estimator:", trend)
+
+# The separate Trend objects for each component can be accessed through the
+# 'component_' attribute. You could grid them individually if you wanted.
+print("East component trend:", trend.component_[0])
+print("East trend coefficients:", trend.component_[0].coef_)
+print("North component trend:", trend.component_[1])
+print("North trend coefficients:", trend.component_[1].coef_)
+
+# We can make a grid with both trend components as data variables
+grid = trend.grid(spacing=0.1, dims=['latitude', 'longitude'])
+print("\nGridded 2-component trend:")
+print(grid)
+
+
+# Now we can map both trends along with the original data for comparison
+def plot_trend(ax, component, title):
+    "Make a map of the given trend component on the given axes"
+    crs = ccrs.PlateCarree()
+    ax.set_title(title)
+    # Plot the trend in pseudo color
+    maxabs = np.abs([component.min(), component.max()]).max()
+    tmp = ax.pcolormesh(component.longitude, component.latitude,
+                        component.values, vmin=-maxabs, vmax=maxabs,
+                        cmap='seismic', transform=crs)
+    cb = plt.colorbar(tmp, ax=ax, orientation='horizontal', pad=0.05)
+    cb.set_label('meters/year')
+    # Plot the original data
+    ax.quiver(data.longitude.values, data.latitude.values,
+              data.velocity_east.values, data.velocity_north.values,
+              scale=0.2, transform=crs, color='y', label="Original data")
+    # and the residuals
+    ax.quiver(data.longitude.values, data.latitude.values,
+              trend.residual_[0].values, trend.residual_[1].values,
+              scale=0.2, transform=crs, color='k', label='Residuals')
+    # Setup the map ticks
+    ax.set_xticks(np.arange(-124, -115, 4), crs=crs)
+    ax.set_yticks(np.arange(33, 42, 2), crs=crs)
+    ax.xaxis.set_major_formatter(LongitudeFormatter())
+    ax.yaxis.set_major_formatter(LatitudeFormatter())
+    # ax.coastlines(color='white')
+    ax.set_extent(vd.get_region(data.longitude, data.latitude), crs=crs)
+
+
+# Make a plot of the data using Cartopy to handle projections and coastlines
+fig, axes = plt.subplots(1, 2, figsize=(9, 7),
+                         subplot_kw=dict(projection=ccrs.Mercator()))
+# Plot the two trend components
+plot_trend(axes[0], grid.east_component, "East component trend")
+plot_trend(axes[1], grid.north_component, "North component trend")
+axes[0].legend(loc="lower left")
+plt.tight_layout()
+plt.show()
diff --git a/verde/__init__.py b/verde/__init__.py
index ac19549..6a590a3 100644
--- a/verde/__init__.py
+++ b/verde/__init__.py
@@ -10,7 +10,7 @@ from .coordinates import scatter_points, grid_coordinates, inside, \
     profile_coordinates, block_region, get_region
 from .grid_math import block_reduce, distance_mask
 from .scipy_bridge import ScipyGridder
-from .trend import Trend, trend_jacobian
+from .trend import Trend, trend_jacobian, VectorTrend
 from .chain import Chain
 from .spline import Spline, spline_jacobian
 
diff --git a/verde/base.py b/verde/base.py
index 2ce9e8d..a41eeff 100644
--- a/verde/base.py
+++ b/verde/base.py
@@ -411,11 +411,14 @@ def check_fit_input(coordinates, data, weights):
     coordinates : tuple of arrays
         Arrays with the coordinates of each data point. Should be in the
         following order: (easting, northing, vertical, ...).
-    data : array
-        The data values of each data point.
+    data : array or tuple of arrays
+        The data values of each data point. Data can have more than one
+        component. In such cases, data should be a tuple of arrays.
     weights : None or array
         If not None, then the weights assigned to each data point.
         Typically, this should be 1 over the data uncertainty squared.
+        If the data has multiple components, the weights have the same number
+        of components.
 
     Returns
     -------
@@ -424,12 +427,24 @@ def check_fit_input(coordinates, data, weights):
         ravel the array before returning.
 
     """
-    if any(coord.shape != data.shape for coord in coordinates):
+    data = check_data(data)
+    if any(i.shape != j.shape for i in coordinates for j in data):
         raise ValueError(
             "Coordinate and data arrays must have the same shape.")
     if weights is not None:
-        if weights.shape != data.shape:
+        weights = check_data(weights)
+        if len(weights) != len(data):
             raise ValueError(
-                "Weights must have the same shape as the data array.")
-        weights = weights.ravel()
+                "Number of data '{}' and weights '{}' must be equal."
+                .format(len(data), len(weights)))
+        if any(i.size != j.size for i in weights for j in data):
+            raise ValueError(
+                "Weights must have the same size as the data array.")
+        weights = tuple(i.ravel() for i in weights)
+    else:
+        weights = tuple([None]*len(data))
+    if len(weights) == 1:
+        weights = weights[0]
+    if len(data) == 1:
+        data = data[0]
     return coordinates, data, weights
diff --git a/verde/trend.py b/verde/trend.py
index aacf16d..f9f4897 100644
--- a/verde/trend.py
+++ b/verde/trend.py
@@ -65,7 +65,8 @@ class Trend(BaseGridder):
 
     See also
     --------
-    verde.trend_jacobian ; Make the Jacobian matrix for 2D polynomial
+    verde.trend_jacobian : Make the Jacobian matrix for 2D polynomial
+    verde.VectorTrend : Polynomial trend estimator for multi-component data
 
     """
 
@@ -222,7 +223,8 @@ def trend_jacobian(easting, northing, degree, dtype='float64'):
 
     See also
     --------
-    verde.Trend ; Polynomial trend estimator
+    verde.Trend : Polynomial trend estimator
+    verde.VectorTrend : Polynomial trend estimator for multi-component data
 
     """
     if easting.shape != northing.shape:
@@ -234,3 +236,114 @@ def trend_jacobian(easting, northing, degree, dtype='float64'):
     for col, (i, j) in enumerate(combinations):
         out[:, col] = (easting.ravel()**i)*(northing.ravel()**j)
     return out
+
+
+class VectorTrend(BaseGridder):
+    """
+    Fit 2D polynomial trends to multi-component vector data.
+
+    Each data component is fitted to a separated :class:`verde.Trend`. The
+    trends can be evaluated separately through the ``VectorTrend.component_``
+    attribute or jointly using the methods of this class.
+
+    Parameters
+    ----------
+    degree : int
+        The degree of the polynomials. Must be >= 1.
+
+    Attributes
+    ----------
+    components_ : list
+        List of the fitted :class:`verde.Trend` on each component of the data.
+    region_ : tuple
+        The boundaries (``[W, E, S, N]``) of the data used to fit the
+        interpolator. Used as the default region for the
+        :meth:`~verde.Trend.grid` and :meth:`~verde.Trend.scatter` methods.
+    residual_ : list of array
+        The difference between each component of the input data and the
+        predicted polynomial trend.
+
+    See also
+    --------
+    verde.Trend : Polynomial trend estimator
+    verde.trend_jacobian : Make the Jacobian matrix for 2D polynomial
+
+    """
+
+    def __init__(self, degree):
+        self.degree = degree
+
+    def fit(self, coordinates, data, weights=None):
+        """
+        Fit the trend to the given multi-component data.
+
+        The data region is captured and used as default for the
+        :meth:`~verde.VectorTrend.grid` and :meth:`~verde.VectorTrend.scatter`
+        methods.
+
+        All input arrays must have the same shape. If weights are given, there
+        must be a separate array for each component of the data.
+
+        Parameters
+        ----------
+        coordinates : tuple of arrays
+            Arrays with the coordinates of each data point. Should be in the
+            following order: (easting, northing, vertical, ...). Only easting
+            and northing will be used, all subsequent coordinates will be
+            ignored.
+        data : tuple of array
+            The data values of each component at each data point. Must be a
+            tuple.
+        weights : None or tuple of array
+            If not None, then the weights assigned to each data point of each
+            data component. Typically, this should be 1 over the data
+            uncertainty squared.
+
+        Returns
+        -------
+        self
+            Returns this estimator instance for chaining operations.
+
+        """
+        if not isinstance(data, tuple):
+            raise ValueError(
+                "Data must be a tuple of arrays. {} given.".format(type(data)))
+        if weights is not None and not isinstance(weights, tuple):
+            raise ValueError(
+                "Weights must be a tuple of arrays. {} given."
+                .format(type(weights)))
+        coordinates, data, weights = check_fit_input(coordinates, data,
+                                                     weights)
+        self.region_ = get_region(*coordinates[:2])
+        self.component_ = [
+            Trend(degree=self.degree).fit(coordinates, data_comp, weight_comp)
+            for data_comp, weight_comp in zip(data, weights)]
+        self.residual_ = tuple(
+            data_comp - comp.predict(coordinates).reshape(data_comp.shape)
+            for data_comp, comp in zip(data, self.component_))
+        return self
+
+    def predict(self, coordinates):
+        """
+        Evaluate the polynomial trend of each component on a of points.
+
+        Requires a fitted estimator (see :meth:`~verde.VectorTrend.fit`).
+
+        Parameters
+        ----------
+        coordinates : tuple of arrays
+            Arrays with the coordinates of each data point. Should be in the
+            following order: (easting, northing, vertical, ...). Only easting
+            and northing will be used, all subsequent coordinates will be
+            ignored.
+
+        Returns
+        -------
+        data : tuple of array
+            The trend values for each vector component evaluated on the given
+            points. The order of components will be the same as was provided to
+            :meth:`~verde.VectorTrend.fit`.
+
+        """
+        check_is_fitted(self, ['component_'])
+        return tuple(comp.predict(coordinates) for comp in self.component_)


Test patch:
diff --git a/verde/tests/test_base.py b/verde/tests/test_base.py
index b7e80aa..a21bd18 100644
--- a/verde/tests/test_base.py
+++ b/verde/tests/test_base.py
@@ -186,3 +186,5 @@ def test_check_fit_input_fails_weights():
     weights = np.arange(30)
     with pytest.raises(ValueError):
         check_fit_input(coords, data, weights)
+    with pytest.raises(ValueError):
+        check_fit_input(coords, (data, data), weights)
diff --git a/verde/tests/test_trend.py b/verde/tests/test_trend.py
index fa9a1f1..def8b63 100644
--- a/verde/tests/test_trend.py
+++ b/verde/tests/test_trend.py
@@ -6,7 +6,8 @@ import numpy as np
 import numpy.testing as npt
 import pytest
 
-from ..trend import Trend, polynomial_power_combinations, trend_jacobian
+from ..trend import Trend, polynomial_power_combinations, trend_jacobian, \
+    VectorTrend
 from ..coordinates import grid_coordinates
 
 
@@ -19,6 +20,24 @@ def simple_model():
     return (east, north), coefs, data
 
 
+@pytest.fixture()
+def simple_2d_model():
+    "A single degree=1 2-component model"
+    east, north = grid_coordinates((1000, 5000, -5000, -1000), shape=(50, 50))
+    coefs = ([-20, 0.5, 3], [10, 2, -0.4])
+    data = tuple(c[0] + c[1]*east + c[2]*north for c in coefs)
+    return (east, north), coefs, data
+
+
+@pytest.fixture()
+def simple_3d_model():
+    "A single degree=1 3-component model"
+    east, north = grid_coordinates((1000, 5000, -5000, -1000), shape=(50, 50))
+    coefs = ([-20, 0.5, 3], [10, 2, -0.4], [30, -10, -1.3])
+    data = tuple(c[0] + c[1]*east + c[2]*north for c in coefs)
+    return (east, north), coefs, data
+
+
 def test_trend(simple_model):
     "Test the trend estimation on a simple problem"
     coords, coefs, data = simple_model
@@ -55,3 +74,49 @@ def test_trend_jacobian_fails():
     east, north = np.arange(50), np.arange(30)
     with pytest.raises(ValueError):
         trend_jacobian(east, north, degree=1)
+
+
+def test_vector_trend(simple_2d_model):
+    "Test the vector trend estimation on a simple problem"
+    coords, coefs, data = simple_2d_model
+    trend = VectorTrend(degree=1).fit(coords, data)
+    for i, coef in enumerate(coefs):
+        npt.assert_allclose(trend.component_[i].coef_, coef)
+        npt.assert_allclose(trend.predict(coords)[i], data[i])
+        npt.assert_allclose(trend.residual_[i], 0, atol=1e-5)
+
+
+def test_vector_trend_3d(simple_3d_model):
+    "Test the vector trend estimation on a simple problem"
+    coords, coefs, data = simple_3d_model
+    trend = VectorTrend(degree=1).fit(coords, data)
+    for i, coef in enumerate(coefs):
+        npt.assert_allclose(trend.component_[i].coef_, coef)
+        npt.assert_allclose(trend.predict(coords)[i], data[i])
+        npt.assert_allclose(trend.residual_[i], 0, atol=1e-5)
+
+
+def test_vector_trend_fails(simple_2d_model):
+    "Test the vector trend estimation on a simple problem"
+    coords, _, data = simple_2d_model
+    trend = VectorTrend(degree=1)
+    with pytest.raises(ValueError):
+        trend.fit(coords, list(data))
+    with pytest.raises(ValueError):
+        trend.fit(coords, data, weights=[np.ones_like(data)]*2)
+
+
+def test_vector_trend_weights(simple_2d_model):
+    "Use weights to account for outliers"
+    coords, coefs, data = simple_2d_model
+    outlier = np.abs(data[0]).max()*3
+    data_out = tuple(i.copy() for i in data)
+    weights = tuple(np.ones_like(i) for i in data)
+    for i, coef in enumerate(coefs):
+        data_out[i][20, 20] += outlier
+        weights[i][20, 20] = 1e-10
+    trend = VectorTrend(degree=1).fit(coords, data_out, weights)
+    for i, coef in enumerate(coefs):
+        npt.assert_allclose(trend.component_[i].coef_, coef)
+        npt.assert_allclose(trend.residual_[i][20, 20], outlier)
+        npt.assert_allclose(trend.predict(coords)[i], data[i])