Repository: fatiando/verde
Pull request #236 (2020-03-11 15:44:24)

Add function for rolling windows on irregular data
The `rolling_window` function returns the indices of points falling on
rolling/moving windows. It works for irregularly sampled data and is a
nice complement to `xarray.DataArray.rolling`  (which only works on
grids). Uses a KDTree to select points but is not compatible with
pykdtree (it's missing the method we need).

<!--
Please describe changes proposed and WHY you made them. If fixing an issue,
include the text "Fixes #XXX" (replace XXX by the issue number. GitHub will
automatically close it when this gets merged.
-->





**Reminders**:

- [ ] Run `make format` and `make check` to make sure the code follows the style guide.
- [ ] Add tests for new features or tests that would have caught the bug that you're fixing.
- [ ] Add new public functions/methods/classes to `doc/api/index.rst` and the base `__init__.py` file for the package.
- [ ] Write detailed docstrings for all functions/classes/methods. It often helps to design better code if you write the docstrings first.
- [ ] If adding new functionality, add an example to the docstring, gallery, and/or tutorials.
- [ ] Add your full name, affiliation, and [ORCID](https://orcid.org) (optional) to the `AUTHORS.md` file (if you haven't already) in case you'd like to be listed as an author on the [Zenodo](https://zenodo.org/communities/fatiando) archive of the next release.


leouieda: They idea for this is to do rolling operations on irregular data. In gravmag, the main usage is for Euler deconvolution (which is what I plan on using this for). I thought of implementing something like `DataArray.rolling` which is much more sophisticated, but there is no way to make it general enough for our purposes. So for now it just returns the indices of points for each window. 

Any thoughts?
leouieda: @santisoler that's spot on, thanks! I'll add the warnings (exceptions would be a bit too strict since there is nothing particularly wrong about non-overlapping windows) and some tests to make sure they are raised. Thank you for the review :+1: 
leouieda: @santisoler implemented your suggestions and this is ready for another round of reviews.

The only thing I left out is checking if `size == spacing`. The output of `block_split` is actually a bit different so they 2 functions aren't entirely the interchangeable. `block_split` returns a label for each block while `rolling_split` returns the indices for each window. So even if `size==spacing`, `block_split` returns a single array that has be used in a `groupby` operation while `rolling_split` returns multiple arrays used for indexing.


santisoler: You're welcome @leouieda! 

> I'll add the warnings (exceptions would be a bit too strict since there is nothing particularly wrong about non-overlapping windows) and some tests to make sure they are raised.

That's ok for me. I cannot think of a use case where we want to use non-overlapping windows (and thus leaving _unwindowed_ data), but you're right it's not a problem to have them implemented.

> The only thing I left out is checking if size == spacing. The output of block_split is actually a bit different so they 2 functions aren't entirely the interchangeable. block_split returns a label for each block while rolling_split returns the indices for each window. So even if size==spacing, block_split returns a single array that has be used in a groupby operation while rolling_split returns multiple arrays used for indexing.

You're right, they are not exactly making the same thing. I was misled by the similar names.
leouieda: OK, a TODO for this:

- [ ] Rename to `rolling_window` to avoid confusion with `block_split` since they are fundamentally different
- [ ] See if it's possible to preserve the shape of the window center coordinates by reshaping the list of indices.
leouieda: @santisoler done with the changes. What do you think? I was able to preserve the shape of the windows arrays.
leouieda: Merging away :rocket: 

Patch:
diff --git a/doc/api/index.rst b/doc/api/index.rst
index 4f1e070..384d975 100644
--- a/doc/api/index.rst
+++ b/doc/api/index.rst
@@ -60,6 +60,7 @@ Coordinate Manipulation
     project_region
     inside
     block_split
+    rolling_window
 
 Utilities
 ---------
diff --git a/verde/__init__.py b/verde/__init__.py
index d1fd2ee..7f56591 100644
--- a/verde/__init__.py
+++ b/verde/__init__.py
@@ -7,6 +7,7 @@ from .coordinates import (
     grid_coordinates,
     inside,
     block_split,
+    rolling_window,
     profile_coordinates,
     get_region,
     pad_region,
diff --git a/verde/coordinates.py b/verde/coordinates.py
index b7e912d..02532dd 100644
--- a/verde/coordinates.py
+++ b/verde/coordinates.py
@@ -1,6 +1,8 @@
 """
 Functions for generating and manipulating coordinates.
 """
+import warnings
+
 import numpy as np
 from sklearn.utils import check_random_state
 
@@ -735,6 +737,7 @@ def block_split(coordinates, spacing=None, adjust="spacing", region=None, shape=
     See also
     --------
     BlockReduce : Apply a reduction operation to the data in blocks (windows).
+    rolling_window : Select points on a rolling (moving) window.
 
     Examples
     --------
@@ -770,15 +773,236 @@ def block_split(coordinates, spacing=None, adjust="spacing", region=None, shape=
     """
     if region is None:
         region = get_region(coordinates)
-    block_coords = tuple(
-        i.ravel()
-        for i in grid_coordinates(
-            region, spacing=spacing, shape=shape, adjust=adjust, pixel_register=True
-        )
+    block_coords = grid_coordinates(
+        region, spacing=spacing, shape=shape, adjust=adjust, pixel_register=True
     )
     tree = kdtree(block_coords)
     labels = tree.query(np.transpose(n_1d_arrays(coordinates, 2)))[1]
-    return block_coords, labels
+    return n_1d_arrays(block_coords, len(block_coords)), labels
+
+
+def rolling_window(
+    coordinates, size, spacing=None, shape=None, region=None, adjust="spacing"
+):
+    """
+    Select points on a rolling (moving) window.
+
+    A window of the given size is moved across the region at a given step
+    (specified by *spacing* or *shape*). Returns the indices of points falling
+    inside each window step. You can use the indices to select points falling
+    inside a given window.
+
+    The size of the step when moving the windows can be specified by the
+    *spacing* parameter. Alternatively, the number of windows in the
+    South-North and West-East directions can be specified using the *shape*
+    parameter. **One of the two must be given.**
+
+    Parameters
+    ----------
+    coordinates : tuple of arrays
+        Arrays with the coordinates of each data point. Should be in the
+        following order: (easting, northing, vertical, ...).
+    size : float
+        The size of the windows. Units should match the units of *coordinates*.
+    spacing : float, tuple = (s_north, s_east), or None
+        The window size in the South-North and West-East directions,
+        respectively. A single value means that the size is equal in both
+        directions.
+    shape : tuple = (n_north, n_east) or None
+        The number of blocks in the South-North and West-East directions,
+        respectively.
+    region : list = [W, E, S, N]
+        The boundaries of a given region in Cartesian or geographic
+        coordinates. If not region is given, will use the bounding region of
+        the given points.
+    adjust : {'spacing', 'region'}
+        Whether to adjust the spacing or the region if required. Ignored if
+        *shape* is given instead of *spacing*. Defaults to adjusting the
+        spacing.
+
+    Returns
+    -------
+    window_coordinates : tuple of arrays
+        Coordinate arrays for the center of each window.
+    indices : array
+        Each element of the array corresponds the indices of points falling
+        inside a window. The array will have the same shape as the
+        *window_coordinates*. Use the array elements to index the coordinates
+        for each window. The indices will depend on the number of dimensions in
+        the input coordinates. For example, if the coordinates are 2D arrays,
+        each window will contain indices for 2 dimensions (row, column).
+
+    See also
+    --------
+    block_split : Split a region into blocks and label points accordingly.
+
+    Examples
+    --------
+
+    Generate a set of sample coordinates on a grid and determine the indices
+    of points for each rolling window:
+
+    >>> from verde import grid_coordinates
+    >>> coords = grid_coordinates((-5, -1, 6, 10), spacing=1)
+    >>> print(coords[0])
+    [[-5. -4. -3. -2. -1.]
+     [-5. -4. -3. -2. -1.]
+     [-5. -4. -3. -2. -1.]
+     [-5. -4. -3. -2. -1.]
+     [-5. -4. -3. -2. -1.]]
+    >>> print(coords[1])
+    [[ 6.  6.  6.  6.  6.]
+     [ 7.  7.  7.  7.  7.]
+     [ 8.  8.  8.  8.  8.]
+     [ 9.  9.  9.  9.  9.]
+     [10. 10. 10. 10. 10.]]
+    >>> # Get the rolling window indices
+    >>> window_coords, indices = rolling_window(coords, size=2, spacing=2)
+    >>> # Window coordinates will be 2D arrays. Their shape is the number of
+    >>> # windows in each dimension
+    >>> print(window_coords[0].shape, window_coords[1].shape)
+    (2, 2) (2, 2)
+    >>> # The there are the easting and northing coordinates for the center of
+    >>> # each rolling window
+    >>> for coord in window_coords:
+    ...     print(coord)
+    [[-4. -2.]
+     [-4. -2.]]
+    [[7. 7.]
+     [9. 9.]]
+    >>> # The indices of points falling on each window will have the same shape
+    >>> # as the window center coordinates
+    >>> print(indices.shape)
+    (2, 2)
+    >>> # The points in the first window. Indices are 2D positions because the
+    >>> # coordinate arrays are 2D.
+    >>> print(len(indices[0, 0]))
+    2
+    >>> for dimension in indices[0, 0]:
+    ...     print(dimension)
+    [0 0 0 1 1 1 2 2 2]
+    [0 1 2 0 1 2 0 1 2]
+    >>> for dimension in indices[0, 1]:
+    ...     print(dimension)
+    [0 0 0 1 1 1 2 2 2]
+    [2 3 4 2 3 4 2 3 4]
+    >>> for dimension in indices[1, 0]:
+    ...     print(dimension)
+    [2 2 2 3 3 3 4 4 4]
+    [0 1 2 0 1 2 0 1 2]
+    >>> for dimension in indices[1, 1]:
+    ...     print(dimension)
+    [2 2 2 3 3 3 4 4 4]
+    [2 3 4 2 3 4 2 3 4]
+    >>> # To get the coordinates for each window, use indexing
+    >>> print(coords[0][indices[0, 0]])
+    [-5. -4. -3. -5. -4. -3. -5. -4. -3.]
+    >>> print(coords[1][indices[0, 0]])
+    [6. 6. 6. 7. 7. 7. 8. 8. 8.]
+
+    If the coordinates are 1D, the indices will also be 1D:
+
+    >>> coords1d = [coord.ravel() for coord in coords]
+    >>> window_coords, indices = rolling_window(coords1d, size=2, spacing=2)
+    >>> print(len(indices[0, 0]))
+    1
+    >>> print(indices[0, 0][0])
+    [ 0  1  2  5  6  7 10 11 12]
+    >>> print(indices[0, 1][0])
+    [ 2  3  4  7  8  9 12 13 14]
+    >>> print(indices[1, 0][0])
+    [10 11 12 15 16 17 20 21 22]
+    >>> print(indices[1, 1][0])
+    [12 13 14 17 18 19 22 23 24]
+    >>> # The returned indices can be used in the same way as before
+    >>> print(coords1d[0][indices[0, 0]])
+    [-5. -4. -3. -5. -4. -3. -5. -4. -3.]
+    >>> print(coords1d[1][indices[0, 0]])
+    [6. 6. 6. 7. 7. 7. 8. 8. 8.]
+
+    By default, the windows will span the entire data region. You can also
+    control the specific region you'd like the windows to cover:
+
+    >>> # Coordinates on a larger region but with the same spacing as before
+    >>> coords = grid_coordinates((-10, 5, 0, 20), spacing=1)
+    >>> # Get the rolling window indices but limited to the region from before
+    >>> window_coords, indices = rolling_window(
+    ...     coords, size=2, spacing=2, region=(-5, -1, 6, 10),
+    ... )
+    >>> # The windows should still be in the same place as before
+    >>> for coord in window_coords:
+    ...     print(coord)
+    [[-4. -2.]
+     [-4. -2.]]
+    [[7. 7.]
+     [9. 9.]]
+    >>> # And indexing the coordinates should also provide the same result
+    >>> print(coords[0][indices[0, 0]])
+    [-5. -4. -3. -5. -4. -3. -5. -4. -3.]
+    >>> print(coords[1][indices[0, 0]])
+    [6. 6. 6. 7. 7. 7. 8. 8. 8.]
+
+    """
+    shapes = [coord.shape for coord in coordinates]
+    if not all(shape == shapes[0] for shape in shapes):
+        raise ValueError(
+            "Coordinate arrays must have the same shape. Given shapes: {}".format(
+                shapes
+            )
+        )
+    if region is None:
+        region = get_region(coordinates)
+    # Calculate the region spanning the centers of the rolling windows
+    window_region = [
+        dimension + (-1) ** (i % 2) * size / 2 for i, dimension in enumerate(region)
+    ]
+    _check_rolling_window_overlap(window_region, size, shape, spacing)
+    centers = grid_coordinates(
+        window_region, spacing=spacing, shape=shape, adjust=adjust
+    )
+    # pykdtree doesn't support query_ball_point yet and we need that
+    tree = kdtree(coordinates, use_pykdtree=False)
+    # Coordinates must be transposed because the kd-tree wants them as columns
+    # of a matrix
+    # Use p=inf (infinity norm) to get square windows instead of circular ones
+    indices1d = tree.query_ball_point(
+        np.transpose(n_1d_arrays(centers, 2)), r=size / 2, p=np.inf
+    )
+    # Make the indices array the same shape as the center coordinates array.
+    # That preserves the information of the number of windows in each
+    # dimension. Need to first create an empty array of object type because
+    # otherwise numpy tries to use the index tuples as dimensions (even if
+    # given ndim=1 explicitly). Can't make it 1D and then reshape because the
+    # reshape is ignored for some reason. The workaround is to create the array
+    # with the correct shape and assign the values to a raveled view of the
+    # array.
+    indices = np.empty(centers[0].shape, dtype="object")
+    # Need to convert the indices to int arrays because unravel_index doesn't
+    # like empty lists but can handle empty integer arrays in case a window has
+    # no points inside it.
+    indices.ravel()[:] = [
+        np.unravel_index(np.array(i, dtype="int"), shape=shapes[0]) for i in indices1d
+    ]
+    return centers, indices
+
+
+def _check_rolling_window_overlap(region, size, shape, spacing):
+    """
+    Warn the user if there is no overlap between neighboring windows.
+    """
+    if shape is not None:
+        ndims = len(shape)
+        dimensions = [region[i * ndims + 1] - region[i * ndims] for i in range(ndims)]
+        # The - 1 is because we need to divide by the number of intervals, not
+        # the number of nodes.
+        spacing = tuple(dim / (n - 1) for dim, n in zip(dimensions, shape))
+    spacing = np.atleast_1d(spacing)
+    if np.any(spacing > size):
+        warnings.warn(
+            f"Rolling windows do not overlap (size '{size}' and spacing '{spacing}'). "
+            "Some data points may not be included in any window. "
+            "Increase size or decrease spacing to avoid this."
+        )
 
 
 def longitude_continuity(coordinates, region):


Test patch:
diff --git a/verde/tests/test_coordinates.py b/verde/tests/test_coordinates.py
index 7012732..ec38c46 100644
--- a/verde/tests/test_coordinates.py
+++ b/verde/tests/test_coordinates.py
@@ -1,6 +1,8 @@
 """
 Test the coordinate generation functions
 """
+import warnings
+
 import numpy as np
 import numpy.testing as npt
 import pytest
@@ -11,9 +13,51 @@ from ..coordinates import (
     profile_coordinates,
     grid_coordinates,
     longitude_continuity,
+    rolling_window,
 )
 
 
+def test_rolling_window_invalid_coordinate_shapes():
+    "Shapes of input coordinates must all be the same"
+    coordinates = [np.arange(10), np.arange(10).reshape((5, 2))]
+    with pytest.raises(ValueError):
+        rolling_window(coordinates, size=2, spacing=1)
+
+
+def test_rolling_window_empty():
+    "Make sure empty windows return an empty index"
+    coords = grid_coordinates((-5, -1, 6, 10), spacing=1)
+    # Use a larger region to make sure the first window is empty
+    # Doing this will raise a warning for non-overlapping windows. Capture it
+    # so it doesn't pollute the test log.
+    with warnings.catch_warnings(record=True):
+        windows = rolling_window(coords, size=0.001, spacing=1, region=(-7, 1, 4, 12))[
+            1
+        ]
+    assert windows[0, 0][0].size == 0 and windows[0, 0][1].size == 0
+    # Make sure we can still index with an empty array
+    assert coords[0][windows[0, 0]].size == 0
+
+
+def test_rolling_window_warnings():
+    "Should warn users if the windows don't overlap"
+    coords = grid_coordinates((-5, -1, 6, 10), spacing=1)
+    # For exact same size there will be 1 point overlapping so should not warn
+    with warnings.catch_warnings(record=True) as warn:
+        rolling_window(coords, size=2, spacing=2)
+        assert not any(issubclass(w.category, UserWarning) for w in warn)
+    args = [dict(spacing=3), dict(spacing=(4, 1)), dict(shape=(1, 2))]
+    for arg in args:
+        with warnings.catch_warnings(record=True) as warn:
+            rolling_window(coords, size=2, **arg)
+            # Filter out the user warnings from some deprecation warnings that
+            # might be thrown by other packages.
+            userwarnings = [w for w in warn if issubclass(w.category, UserWarning)]
+            assert len(userwarnings) == 1
+            assert issubclass(userwarnings[-1].category, UserWarning)
+            assert str(userwarnings[-1].message).split()[0] == "Rolling"
+
+
 def test_spacing_to_shape():
     "Check that correct spacing and region are returned"
     region = (-10, 0, 0, 5)