Repository: fatiando/verde
Pull request #48 (2018-05-20 22:55:08)

Base data component names on number of components
Get rid of the `data_type` optional attribute and guess the names based
solely on the number of components returned by `predict`. Makes the
base class more flexible.


Patch:
diff --git a/verde/base.py b/verde/base.py
index f935ab8..2ce9e8d 100644
--- a/verde/base.py
+++ b/verde/base.py
@@ -33,14 +33,11 @@ class BaseGridder(BaseEstimator):
       ``_``.
 
     The child class can define the following attributes to control the names of
-    coordinates and data values in the output ``xarray.Dataset`` and
-    ``pandas.DataFrame`` and how distances are calculated:
+    coordinates and how distances are calculated:
 
     * ``coordinate_system``: either ``'cartesian'`` or ``'geographic'``. Will
       influence dimension names and distance calculations. Defaults to
       ``'cartesian'``.
-    * ``data_type``: one of ``'scalar'``, ``'vector2d'``, or ``'vector3d'``.
-      Defaults to ``'scalar'``.
 
     Examples
     --------
@@ -181,7 +178,6 @@ class BaseGridder(BaseEstimator):
         if shape is None and spacing is None:
             shape = (101, 101)
         dims = get_dims(self, dims)
-        data_names = get_data_names(self, data_names)
         region = get_region(self, region)
         coordinates = grid_coordinates(region, shape=shape, spacing=spacing,
                                        adjust=adjust)
@@ -189,6 +185,7 @@ class BaseGridder(BaseEstimator):
             data = check_data(self.predict(coordinates))
         else:
             data = check_data(self.predict(projection(*coordinates)))
+        data_names = get_data_names(data, data_names)
         coords = {dims[1]: coordinates[0][0, :], dims[0]: coordinates[1][:, 0]}
         attrs = {'metadata': 'Generated by {}'.format(repr(self))}
         data_vars = {name: (dims, value, attrs)
@@ -238,13 +235,13 @@ class BaseGridder(BaseEstimator):
 
         """
         dims = get_dims(self, dims)
-        data_names = get_data_names(self, data_names)
         region = get_region(self, region)
         coordinates = scatter_points(region, size, random_state)
         if projection is None:
             data = check_data(self.predict(coordinates))
         else:
             data = check_data(self.predict(projection(*coordinates)))
+        data_names = get_data_names(data, data_names)
         columns = [(dims[0], coordinates[1]), (dims[1], coordinates[0])]
         columns.extend(zip(data_names, data))
         return pd.DataFrame(dict(columns), columns=[c[0] for c in columns])
@@ -294,13 +291,13 @@ class BaseGridder(BaseEstimator):
         """
         coordsys = getattr(self, 'coordinate_system', 'cartesian')
         dims = get_dims(self, dims)
-        data_names = get_data_names(self, data_names)
         east, north, distances = profile_coordinates(
             point1, point2, size, coordinate_system=coordsys)
         if projection is None:
             data = check_data(self.predict((east, north)))
         else:
             data = check_data(self.predict(projection(east, north)))
+        data_names = get_data_names(data, data_names)
         columns = [(dims[0], north), (dims[1], east), ('distance', distances)]
         columns.extend(zip(data_names, data))
         return pd.DataFrame(dict(columns), columns=[c[0] for c in columns])
@@ -337,39 +334,39 @@ def get_dims(instance, dims):
     return ('northing', 'easting')
 
 
-def get_data_names(instance, data_names):
+def get_data_names(data, data_names):
     """
-    Get default names for data fields for an instance if not given as arguments
+    Get default names for data fields if none are given based on the data.
 
     Examples
     --------
 
-    >>> grd = BaseGridder()
-    >>> get_data_names(grd, data_names=None)
+    >>> import numpy as np
+    >>> east, north, up = [np.arange(10)]*3
+    >>> get_data_names((east,), data_names=None)
     ('scalars',)
-    >>> grd.data_type = 'vector2d'
-    >>> get_data_names(grd, data_names=None)
+    >>> get_data_names((east, north), data_names=None)
     ('east_component', 'north_component')
-    >>> grd.data_type = 'vector3d'
-    >>> get_data_names(grd, data_names=None)
+    >>> get_data_names((east, north, up), data_names=None)
     ('east_component', 'north_component', 'vertical_component')
-    >>> get_data_names(grd, data_names=('ringo', 'george'))
+    >>> get_data_names((up, north), data_names=('ringo', 'george'))
     ('ringo', 'george')
 
     """
     if data_names is not None:
+        if len(data) != len(data_names):
+            raise ValueError(
+                "Data has {} components but only {} names provided: {}"
+                .format(len(data), len(data_names), str(data_names)))
         return data_names
-    valid_types = ['scalar', 'vector2d', 'vector3d']
-    data_type = getattr(instance, 'data_type', 'scalar')
-    if data_type not in valid_types:
-        raise ValueError(
-            "Invalid data type for {}: '{}'. Must be one of {}."
-            .format(instance.__class__.__name__, data_type, str(valid_types)))
-    if data_type == 'vector2d':
-        return ('east_component', 'north_component')
-    elif data_type == 'vector3d':
-        return ('east_component', 'north_component', 'vertical_component')
-    return ('scalars',)
+    data_types = [('scalars',),
+                  ('east_component', 'north_component'),
+                  ('east_component', 'north_component', 'vertical_component')]
+    if len(data) > len(data_types):
+        raise ValueError(" ".join([
+            "Default data names only available for up to 3 components.",
+            "Must provide custom names through the 'data_names' argument."]))
+    return data_types[len(data) - 1]
 
 
 def get_region(instance, region):


Test patch:
diff --git a/verde/tests/test_base.py b/verde/tests/test_base.py
index ab4953e..b7e80aa 100644
--- a/verde/tests/test_base.py
+++ b/verde/tests/test_base.py
@@ -46,42 +46,28 @@ def test_get_dims_fails():
 
 def test_get_data_names():
     "Tests that get_data_names returns the expected results"
-    grd = BaseGridder()
-    assert get_data_names(grd, data_names=None) == ('scalars',)
-    assert get_data_names(grd, data_names=('a', 'b')) == ('a', 'b')
-
-    grd.data_type = 'scalar'
-    assert get_data_names(grd, data_names=None) == ('scalars',)
-    assert get_data_names(grd, data_names=('a', 'b')) == ('a', 'b')
-
-    grd.data_type = 'vector2d'
-    assert get_data_names(grd, data_names=None) == ('east_component',
-                                                    'north_component')
-    assert get_data_names(grd, data_names=('a', 'b')) == ('a', 'b')
-
-    grd.data_type = 'vector3d'
-    assert get_data_names(grd, data_names=None) == ('east_component',
-                                                    'north_component',
-                                                    'vertical_component')
-    assert get_data_names(grd, data_names=('a', 'b')) == ('a', 'b')
-
-    # Make sure the given dims is returned no matter what
-    grd.data_type = 'an invalid type'
-    assert get_data_names(grd, data_names=('a', 'b')) == ('a', 'b')
+    data1 = (np.arange(10), )
+    data2 = tuple([np.arange(10)]*2)
+    data3 = tuple([np.arange(10)]*3)
+    # Test the default names
+    assert get_data_names(data1, data_names=None) == ('scalars',)
+    assert get_data_names(data2, data_names=None) == ('east_component',
+                                                      'north_component')
+    assert get_data_names(data3, data_names=None) == ('east_component',
+                                                      'north_component',
+                                                      'vertical_component')
+    # Test custom names
+    assert get_data_names(data1, data_names=('a',)) == ('a',)
+    assert get_data_names(data2, data_names=('a', 'b')) == ('a', 'b')
+    assert get_data_names(data3, data_names=('a', 'b', 'c')) == ('a', 'b', 'c')
 
 
 def test_get_data_names_fails():
     "Check if fails for invalid data types"
-    grd = BaseGridder()
-    with pytest.raises(ValueError):
-        grd.data_type = 'Scalars'
-        get_data_names(grd, data_names=None)
     with pytest.raises(ValueError):
-        grd.data_type = 'Vector3d'
-        get_data_names(grd, data_names=None)
+        get_data_names(tuple([np.arange(5)]*4), data_names=None)
     with pytest.raises(ValueError):
-        grd.data_type = 'some totally not valid name'
-        get_data_names(grd, data_names=None)
+        get_data_names(tuple([np.arange(5)]*2), data_names=("meh",))
 
 
 def test_get_region():