# Copyright (c) 2017 The Verde Developers.
# Distributed under the terms of the BSD 3-Clause License.
# SPDX-License-Identifier: BSD-3-Clause
#
# This code is part of the Fatiando a Terra project (https://www.fatiando.org)
#
"""
General utilities.
"""
import functools
import dask
import numpy as np
import pandas as pd
import xarray as xr
from scipy.spatial import cKDTree
try:
from pykdtree.kdtree import KDTree as pyKDTree
except ImportError:
pyKDTree = None # noqa: N816
try:
import numba
except ImportError:
numba = None
from .base.utils import (
check_coordinates,
check_data,
check_data_names,
check_extra_coords_names,
n_1d_arrays,
)
def dispatch(function, delayed=False, client=None):
"""
Decide how to wrap a function for Dask depending on the options given.
Parameters
----------
function : callable
The function that will be called.
delayed : bool
If True, will wrap the function in :func:`dask.delayed`.
client : None or dask.distributed Client
If *delayed* is False and *client* is not None, will return a partial
execution of the ``client.submit`` with the function as first argument.
Returns
-------
function : callable
The function wrapped in Dask.
"""
if delayed:
return dask.delayed(function)
if client is not None:
return functools.partial(client.submit, function)
return function
def parse_engine(engine):
"""
Choose the best engine available and check if it's valid.
Parameters
----------
engine : str
The name of the engine. If ``"auto"`` will favor numba if it's
available.
Returns
-------
engine : str
The name of the engine that should be used.
"""
engines = {"auto", "numba", "numpy"}
if engine not in engines:
raise ValueError("Invalid engine '{}'. Must be in {}.".format(engine, engines))
if engine == "auto":
if numba is None:
return "numpy"
return "numba"
return engine
def dummy_jit(**kwargs): # noqa: U100
"""
Replace numba.jit if not installed with a function that raises RunTimeError
Use as a decorator.
Parameters
----------
function
A function that you would decorate with :func:`numba.jit`.
Returns
-------
function
A function that raises :class:`RunTimeError` warning that numba isn't
installed.
"""
def dummy_decorator(function):
"The actual decorator"
@functools.wraps(function)
def dummy_function(*args, **kwargs): # noqa: U100
"Just raise an exception."
raise RuntimeError("Could not find numba.")
return dummy_function
return dummy_decorator
[docs]def variance_to_weights(variance, tol=1e-15, dtype="float64"):
"""
Converts data variances to weights for gridding.
Weights are defined as the inverse of the variance, scaled to the range
[0, 1], i.e. ``variance.min()/variance``.
Any variance that is smaller than *tol* will automatically receive a weight
of 1 to avoid zero division or blown up weights.
Parameters
----------
variance : array or tuple of arrays
An array with the variance of each point. If there are multiple arrays
in a tuple, will calculated weights for each of them separately. Can
have NaNs but they will be converted to zeros and therefore receive a
weight of 1.
tol : float
The tolerance, or cutoff threshold, for small variances.
dtype : str or numpy dtype
The type of the output weights array.
Returns
-------
weights : array or tuple of arrays
Data weights in the range [0, 1] with the same shape as *variance*. If
more than one variance array was provided, then this will be a tuple
with the weights corresponding to each variance array.
Examples
--------
>>> print(variance_to_weights([0, 2, 0.2, 1e-16]))
[1. 0.1 1. 1. ]
>>> print(variance_to_weights([0, 0, 0, 0]))
[1. 1. 1. 1.]
>>> for w in variance_to_weights(([0, 1, 10], [2, 4.0, 8])):
... print(w)
[1. 1. 0.1]
[1. 0.5 0.25]
"""
variance = check_data(variance)
weights = []
for var in variance:
var = np.nan_to_num(np.atleast_1d(var), copy=False)
w = np.ones_like(var, dtype=dtype)
nonzero = var > tol
if np.any(nonzero):
nonzero_var = var[nonzero]
w[nonzero] = nonzero_var.min() / nonzero_var
weights.append(w)
if len(weights) == 1:
return weights[0]
return tuple(weights)
[docs]def maxabs(*args, nan=True):
"""
Calculate the maximum absolute value of the given array(s).
Use this to set the limits of your colorbars and center them on zero.
Parameters
----------
args
One or more arrays. If more than one are given, a single maximum will
be calculated across all arrays.
Returns
-------
maxabs : float
The maximum absolute value across all arrays.
Examples
--------
>>> maxabs((1, -10, 25, 2, 3))
25
>>> maxabs((1, -10.5, 25, 2), (0.1, 100, -500), (-200, -300, -0.1, -499))
500.0
If the array contains NaNs, we'll use the ``nan`` version of of the numpy
functions by default. You can turn this off through the *nan* argument.
>>> import numpy as np
>>> maxabs((1, -10, 25, 2, 3, np.nan))
25.0
>>> maxabs((1, -10, 25, 2, 3, np.nan), nan=False)
nan
"""
arrays = [np.atleast_1d(i) for i in args]
if nan:
npmin, npmax = np.nanmin, np.nanmax
else:
npmin, npmax = np.min, np.max
absolute = [npmax(np.abs([npmin(i), npmax(i)])) for i in arrays]
return npmax(absolute)
[docs]def make_xarray_grid(
coordinates,
data,
data_names,
dims=("northing", "easting"),
extra_coords_names=None,
):
"""
Create an :class:`xarray.Dataset` grid from 2D numpy arrays
This function creates an :class:`xarray.Dataset` out of 2d gridded data
including easting and northing coordinates, any extra coordinates (like
upward elevation, time, etc) and data arrays.
Use this to transform the outputs of :func:`verde.grid_coordinates` and
the ``predict`` method of a gridder into an :class:`xarray.Dataset`.
.. important::
This function does not perform any interpolation or gridding. If you
are dealing with 1d arrays and need to interpolate them on a regular
grid (a.k.a _grid_ them) you might want to use one of the available
interpolators in Verde, like :class:`verde.Spline`,
:class:`verde.KNeighbors`, :class:`Linear`, etc.
.. note::
This is a utility function to help create 2D grids (i.e., grids with
two ``dims`` coordinates). For arbitrary N-dimensional arrays, use
:mod:`xarray` directly.
Parameters
----------
coordinates : tuple of arrays
Arrays with coordinates of each point in the grid. Each array must
contain values for a dimension in the order: easting, northing,
vertical, etc. All arrays must be 2d and need to have the same *shape*.
These coordinates can be generated through
:func:`verde.grid_coordinates`.
data : array, tuple of arrays or None
Array or tuple of arrays with data values on each point in the grid.
Each array must contain values for a dimension in the same order as
the coordinates. All arrays need to have the same *shape*.
If None, the :class:`xarray.Dataset` will not have any ``data_var``
array.
data_names : str or list
The name(s) of the data variables in the output grid.
Ignored if ``data`` is None.
dims : list (optional)
The names of the northing and easting data dimensions, respectively,
in the output grid. Must be defined in the following order: northing
dimension, easting dimension.
**NOTE: This is an exception to the "easting" then
"northing" pattern but is required for compatibility with xarray.**
The easting and northing coordinates in the :class:`xarray.Dataset`
will have the same names as the passed dimensions.
extra_coords_names : str or list (optional)
Name or list of names for any additional coordinates besides the
easting and northing ones. Ignored if coordinates has
only two elements. The extra coordinates are non-index coordinates of
the grid array.
Returns
-------
grid : :class:`xarray.Dataset`
A 2D grid with one or more data variables.
Examples
--------
>>> import numpy as np
>>> import verde as vd
>>> # Create the coordinates of the regular grid
>>> coordinates = vd.grid_coordinates((-10, -6, 8, 10), spacing=2)
>>> # And some dummy data for each point of the grid
>>> data = np.ones_like(coordinates[0])
>>> # Create the grid
>>> grid = make_xarray_grid(coordinates, data, data_names="dummy")
>>> print(grid) # doctest: +SKIP
<xarray.Dataset>
Dimensions: (northing: 2, easting: 3)
Coordinates:
* easting (easting) float64 -10.0 -8.0 -6.0
* northing (northing) float64 8.0 10.0
Data variables:
dummy (northing, easting) float64 1.0 1.0 1.0 1.0 1.0 1.0
>>> # Create a grid with an extra coordinate
>>> coordinates = vd.grid_coordinates(
... (-10, -6, 8, 10), spacing=2, extra_coords=5
... )
>>> # And some dummy data for each point of the grid
>>> data = np.ones_like(coordinates[0])
>>> # Create the grid
>>> grid = make_xarray_grid(
... coordinates, data, data_names="dummy", extra_coords_names="upward"
... )
>>> print(grid) # doctest: +SKIP
<xarray.Dataset>
Dimensions: (northing: 2, easting: 3)
Coordinates:
* easting (easting) float64 -10.0 -8.0 -6.0
* northing (northing) float64 8.0 10.0
upward (northing, easting) float64 5.0 5.0 5.0 5.0 5.0 5.0
Data variables:
dummy (northing, easting) float64 1.0 1.0 1.0 1.0 1.0 1.0
>>> # Create a grid containing only coordinates and no data
>>> coordinates = vd.grid_coordinates(
... (-10, -6, 8, 10), spacing=2, extra_coords=-7
... )
>>> grid = make_xarray_grid(
... coordinates,
... data=None,
... data_names=None,
... extra_coords_names="upward",
... )
>>> print(grid) # doctest: +SKIP
<xarray.Dataset>
Dimensions: (easting: 3, northing: 2)
Coordinates:
* easting (easting) float64 -10.0 -8.0 -6.0
* northing (northing) float64 8.0 10.0
upward (northing, easting) float64 -7.0 -7.0 -7.0 -7.0 -7.0 -7.0
Data variables:
*empty*
"""
# Check dimensions of the horizontal coordinates of the regular grid
ndim = get_ndim_horizontal_coords(*coordinates[:2])
# Convert 2d horizontal coordinates to 1d arrays if needed
if ndim == 2:
coordinates = meshgrid_to_1d(coordinates)
# dims is like shape with order (rows, cols) for the array
# so the first element is northing and second is easting
coords = {dims[1]: coordinates[0], dims[0]: coordinates[1]}
# Extra coordinates are handled like 2D data arrays with
# the same dims and the data.
if coordinates[2:]:
extra_coords_names = check_extra_coords_names(coordinates, extra_coords_names)
for name, extra_coord in zip(extra_coords_names, coordinates[2:]):
coords[name] = (dims, extra_coord)
# Initialize data_vars as None. If data is not None, build data_vars as
# a dirctionary to be passed to xr.Dataset constructor.
data_vars = None
if data is not None:
data = check_data(data)
data_names = check_data_names(data, data_names)
data_vars = {name: (dims, value) for name, value in zip(data_names, data)}
return xr.Dataset(data_vars, coords)
def meshgrid_to_1d(coordinates):
"""
Convert horizontal coordinates of 2d grids into 1d-arrays
Parameters
----------
coordinates : tuple of arrays
Arrays with coordinates of each point in the grid. Each array must
contain values for a dimension in the order: easting, northing,
vertical, etc. All arrays must be 2d and need to have the same
*shape*. The horizontal coordinates should be actual meshgrids.
Returns
-------
coordinates : tuple of arrays
Arrays with coordinates of each point in the grid. The horizontal
coordinates have been converted to 1d-arrays, having only a single
coordinate point per its corresponding axis.
All extra coordinates have not been modified.
Examples
--------
>>> import verde as vd
>>> coordinates = vd.grid_coordinates(
... region=(0, 4, -3, 3), spacing=1, extra_coords=2
... )
>>> easting, northing, height = meshgrid_to_1d(coordinates)
>>> print(easting)
[0. 1. 2. 3. 4.]
>>> print(northing)
[-3. -2. -1. 0. 1. 2. 3.]
>>> print(height)
[[2. 2. 2. 2. 2.]
[2. 2. 2. 2. 2.]
[2. 2. 2. 2. 2.]
[2. 2. 2. 2. 2.]
[2. 2. 2. 2. 2.]
[2. 2. 2. 2. 2.]
[2. 2. 2. 2. 2.]]
"""
check_coordinates(coordinates)
check_meshgrid(coordinates)
easting, northing = coordinates[0][0, :], coordinates[1][:, 0]
coordinates = (easting, northing, *coordinates[2:])
return coordinates
def meshgrid_from_1d(coordinates):
"""
Convert horizontal coordinates of 2d grids from 1d-arrays to 2d-arrays
Parameters
----------
coordinates : tuple of arrays
Arrays with coordinates of each point in the grid. Each array must
contain values for a dimension in the order: easting, northing,
vertical, etc. The horizontal coordinates (easting and northing) should
be 1d arrays, any extra coordinate should be an array with a shape of
``(northing.size, easting.size)``.
Returns
-------
coordinates : tuple of arrays
Arrays with coordinates of each point in the grid.
The horizontal coordinates have been converted to 2d-arrays with the
same shape, forming a meshgrid. All extra coordinates have not been
modified.
Examples
--------
>>> easting = np.linspace(0, 4, 5)
>>> northing = np.linspace(-3, 3, 7)
>>> height = 2 * np.ones((7, 5))
>>> coordinates = (easting, northing, height)
>>> easting, northing, height = meshgrid_from_1d(coordinates)
>>> print(easting)
[[0. 1. 2. 3. 4.]
[0. 1. 2. 3. 4.]
[0. 1. 2. 3. 4.]
[0. 1. 2. 3. 4.]
[0. 1. 2. 3. 4.]
[0. 1. 2. 3. 4.]
[0. 1. 2. 3. 4.]]
>>> print(northing)
[[-3. -3. -3. -3. -3.]
[-2. -2. -2. -2. -2.]
[-1. -1. -1. -1. -1.]
[ 0. 0. 0. 0. 0.]
[ 1. 1. 1. 1. 1.]
[ 2. 2. 2. 2. 2.]
[ 3. 3. 3. 3. 3.]]
"""
ndim = get_ndim_horizontal_coords(*coordinates[:2])
if ndim != 1:
raise ValueError(
"Horizontal coordinates must be 1d-arrays. " + f"{ndim}d-arrays provided."
)
easting, northing = np.meshgrid(coordinates[0], coordinates[1])
coordinates = (easting, northing, *coordinates[2:])
coordinates = check_coordinates(coordinates)
return coordinates
def get_ndim_horizontal_coords(easting, northing):
"""
Return the number of dimensions of the horizontal coordinates arrays
Also check if the two horizontal coordinates arrays same dimensions.
Parameters
----------
easting : nd-array
Array for the easting coordinates
northing : nd-array
Array for the northing coordinates
Returns
-------
ndim : int
Number of dimensions of the ``easting`` and ``northing`` arrays.
"""
ndim = np.ndim(easting)
if ndim != np.ndim(northing):
raise ValueError(
"Horizontal coordinates dimensions mismatch. "
+ f"The easting coordinate array has {easting.ndim} dimensions "
+ f"while the northing has {northing.ndim}."
)
return ndim
def check_meshgrid(coordinates):
"""
Check if the given horizontal coordinates define a meshgrid
Check if the rows of the easting 2d-array are identical. Check if the
columns of the northing 2d-array are identical. This function does not
check if the easting and northing coordinates are evenly spaced.
Parameters
----------
coordinates : tuple of arrays
Arrays with coordinates of each point in the grid. Each array must
contain values for a dimension in the order: easting, northing,
vertical, etc. Only easting and northing will be checked, the other
ones will be ignored. All arrays must be 2d and need to have the same
*shape*.
"""
# Get the two first arrays as easting and northing
easting, northing = coordinates[:2]
# Check if all elements of easting along the zeroth axis are equal
msg = (
"Invalid coordinate array. The arrays for the horizontal "
+ "coordinates of a regular grid must be meshgrids."
)
if not np.allclose(easting[0, :], easting):
raise ValueError(msg)
# Check if all elements of northing along the first axis are equal
# (need to make northing[:, 0] a vertical array so numpy can compare)
if not np.allclose(northing[:, 0][:, None], northing):
raise ValueError(msg)
[docs]def grid_to_table(grid):
"""
Convert a grid to a table with the values and coordinates of each point.
Takes a 2D grid as input, extracts the coordinates and runs them through
:func:`numpy.meshgrid` to create a 2D table. Works for 2D grids and any
number of variables. Use cases includes passing gridded data to functions
that expect data in XYZ format, such as :class:`verde.BlockReduce`
Parameters
----------
grid : :class:`xarray.Dataset` or :class:`xarray.DataArray`
A 2D grid with one or more data variables.
Returns
-------
table : :class:`pandas.DataFrame`
Table with coordinates and variable values for each point in the grid.
Column names are taken from the grid. If *grid* is a
:class:`xarray.DataArray` that doesn't have a ``name`` attribute
defined, the column with data values will be called ``"scalars"``.
Examples
--------
>>> import xarray as xr
>>> import numpy as np
>>> # Create a sample grid with a single data variable
>>> temperature = xr.DataArray(
... np.arange(20).reshape((4, 5)),
... coords=(np.arange(4), np.arange(5, 10)),
... dims=['northing', 'easting']
... )
>>> print(temperature.values)
[[ 0 1 2 3 4]
[ 5 6 7 8 9]
[10 11 12 13 14]
[15 16 17 18 19]]
>>> # For DataArrays, the data column will be "scalars" by default
>>> table = grid_to_table(temperature)
>>> list(sorted(table.columns))
['easting', 'northing', 'scalars']
>>> print(table.scalars.values)
[ 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19]
>>> print(table.northing.values)
[0 0 0 0 0 1 1 1 1 1 2 2 2 2 2 3 3 3 3 3]
>>> print(table.easting.values)
[5 6 7 8 9 5 6 7 8 9 5 6 7 8 9 5 6 7 8 9]
>>> # If the DataArray defines a "name", we will use that instead
>>> temperature.name = "temperature_K"
>>> table = grid_to_table(temperature)
>>> list(sorted(table.columns))
['easting', 'northing', 'temperature_K']
>>> # Conversion of Datasets will preserve the data variable names
>>> grid = xr.Dataset({"temperature": temperature})
>>> table = grid_to_table(grid)
>>> list(sorted(table.columns))
['easting', 'northing', 'temperature']
>>> print(table.temperature.values)
[ 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19]
>>> print(table.northing.values)
[0 0 0 0 0 1 1 1 1 1 2 2 2 2 2 3 3 3 3 3]
>>> print(table.easting.values)
[5 6 7 8 9 5 6 7 8 9 5 6 7 8 9 5 6 7 8 9]
>>> # Grids with multiple data variables will have more columns.
>>> wind_speed = xr.DataArray(
... np.arange(20, 40).reshape((4, 5)),
... coords=(np.arange(4), np.arange(5, 10)),
... dims=['northing', 'easting']
... )
>>> grid['wind_speed'] = wind_speed
>>> table = grid_to_table(grid)
>>> list(sorted(table.columns))
['easting', 'northing', 'temperature', 'wind_speed']
>>> print(table.northing.values)
[0 0 0 0 0 1 1 1 1 1 2 2 2 2 2 3 3 3 3 3]
>>> print(table.easting.values)
[5 6 7 8 9 5 6 7 8 9 5 6 7 8 9 5 6 7 8 9]
>>> print(table.temperature.values)
[ 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19]
>>> print(table.wind_speed.values)
[20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39]
"""
if hasattr(grid, "data_vars"):
# It's a Dataset
data_names = list(grid.data_vars.keys())
data_arrays = [grid[name].values.ravel() for name in data_names]
coordinate_names = list(grid[data_names[0]].dims)
else:
# It's a DataArray
data_names = [grid.name if grid.name is not None else "scalars"]
data_arrays = [grid.values.ravel()]
coordinate_names = list(grid.dims)
north = grid.coords[coordinate_names[0]].values
east = grid.coords[coordinate_names[1]].values
# Need to flip the coordinates because the names are in northing and
# easting order
coordinates = [i.ravel() for i in np.meshgrid(east, north)][::-1]
data_dict = dict(zip(coordinate_names, coordinates))
data_dict.update(dict(zip(data_names, data_arrays)))
return pd.DataFrame(data_dict)
def kdtree(coordinates, use_pykdtree=True, **kwargs):
"""
Create a KD-Tree object with the given coordinate arrays.
Automatically transposes and flattens the coordinate arrays into a single
matrix for use in the KD-Tree classes.
All other keyword arguments are passed to the KD-Tree class.
If installed, package ``pykdtree`` will be used instead of
:class:`scipy.spatial.cKDTree` for better performance. Not all features are
available in ``pykdtree`` so if you require the scipy version set
``use_pykdtee=False``.
Parameters
----------
coordinates : tuple of arrays
Arrays with the coordinates of each data point. Should be in the
following order: (easting, northing, vertical, ...). All coordinate
arrays are used.
use_pykdtree : bool
If True, will prefer ``pykdtree`` (if installed) over
:class:`scipy.spatial.cKDTree`. Otherwise, always use the scipy
version.
Returns
-------
tree : :class:`scipy.spatial.cKDTree` or ``pykdtree.kdtree.KDTree``
The tree instance initialized with the given coordinates and arguments.
"""
points = np.transpose(n_1d_arrays(coordinates, len(coordinates)))
if pyKDTree is not None and use_pykdtree:
tree = pyKDTree(points, **kwargs)
else:
tree = cKDTree(points, **kwargs)
return tree
def partition_by_sum(array, parts):
"""
Partition an array into parts of approximately equal sum.
Does not change the order of the array elements.
Produces the partition indices on the array. Use :func:`numpy.split` to
divide the array along these indices.
.. warning::
Depending on the input and number of parts, there might not exist
partition points. In these cases, the function will raise
``ValueError``. This is more likely to happen as the number of parts
approaches the number of elements in the array.
Parameters
----------
array : array or array-like
The 1D array that will be partitioned. The array will be raveled before
computations.
parts : int
Number of parts to split the array. Can be at most the number of
elements in the array.
Returns
-------
indices : array
The indices in which the array should be split.
Notes
-----
Solution from https://stackoverflow.com/a/54024280
Examples
--------
>>> import numpy as np
>>> array = np.arange(10)
>>> split_points = partition_by_sum(array, parts=2)
>>> print(split_points)
[7]
>>> for part in np.split(array, split_points):
... print(part, part.sum())
[0 1 2 3 4 5 6] 21
[7 8 9] 24
>>> split_points = partition_by_sum(array, parts=3)
>>> print(split_points)
[6 8]
>>> for part in np.split(array, split_points):
... print(part, part.sum())
[0 1 2 3 4 5] 15
[6 7] 13
[8 9] 17
>>> split_points = partition_by_sum(array, parts=5)
>>> print(split_points)
[4 6 7 9]
>>> for part in np.split(array, split_points):
... print(part, part.sum())
[0 1 2 3] 6
[4 5] 9
[6] 6
[7 8] 15
[9] 9
>>> # Use an array with a random looking element order
>>> array = [5, 6, 4, 6, 8, 1, 2, 6, 3, 3]
>>> split_points = partition_by_sum(array, parts=2)
>>> print(split_points)
[4]
>>> for part in np.split(array, split_points):
... print(part, part.sum())
[5 6 4 6] 21
[8 1 2 6 3 3] 23
>>> # Splits can have very different sums but this is best that can be done
>>> # without changing the order of the array.
>>> split_points = partition_by_sum(array, parts=5)
>>> print(split_points)
[1 3 4 7]
>>> for part in np.split(array, split_points):
... print(part, part.sum())
[5] 5
[6 4] 10
[6] 6
[8 1 2] 11
[6 3 3] 12
"""
array = np.atleast_1d(array).ravel()
if parts > array.size:
raise ValueError(
"Cannot partition an array of size {} into {} parts of equal sum.".format(
array.size, parts
)
)
cumulative_sum = array.cumsum()
# Ideally, we want each part to have the same number of points (total /
# parts).
ideal_sum = cumulative_sum[-1] // parts
# If the parts are ideal, the cumulative sum of each part will be this
ideal_cumsum = np.arange(1, parts) * ideal_sum
# Find the places in the real cumulative sum where the ideal values would
# be. These are the split points. Between each split point, the sum of
# elements will be approximately the ideal sum. Need to insert to the right
# side so that we find cumsum[i - 1] <= ideal < cumsum[i]. This way, if a
# part has ideal sum, the last element (i - 1) will be included. Otherwise,
# we would never have ideal sums.
indices = np.searchsorted(cumulative_sum, ideal_cumsum, side="right")
# Check for repeated split points, which indicates that there is no way to
# split the array.
if np.unique(indices).size != indices.size:
raise ValueError(
"Could not find partition points to split the array into {} parts "
"of equal sum.".format(parts)
)
return indices
