Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Implementation of polyfit and polyval #3733

Merged
merged 23 commits into from
Mar 25, 2020
Merged
Changes from 1 commit
Commits
Show all changes
23 commits
Select commit Hold shift + click to select a range
80936fd
[WIP] Implementation of polyfit and polyval - minimum testing - no docs
aulemahal Jan 30, 2020
05355af
Formatting with black, flake8
aulemahal Jan 30, 2020
7f26cb1
Fix failing test
aulemahal Jan 30, 2020
fdc1819
More intelligent skipna switching
aulemahal Jan 30, 2020
924f4eb
Add docs | Change coeff order to fit numpy | move polyval
aulemahal Jan 31, 2020
1f6030a
Move doc patching to class
aulemahal Jan 31, 2020
0f922e7
conditional doc patching
aulemahal Jan 31, 2020
bfec9c8
Fix windows fail - more efficient nan skipping
aulemahal Jan 31, 2020
2d3235a
Fix typo in least_squares
aulemahal Jan 31, 2020
31440ae
Move polyfit to dataset
aulemahal Feb 11, 2020
da831a5
Add more tests | fix some edge cases
aulemahal Feb 12, 2020
f13123f
Skip test without dask
aulemahal Feb 12, 2020
896313f
Fix 1D case | add docs
aulemahal Feb 12, 2020
5b70971
skip polyval test without dask
aulemahal Feb 12, 2020
5af3c64
Explicit docs | More restrictive polyval
aulemahal Feb 13, 2020
d731fd9
Merge branch 'master' into implement-polyfit-polyval
aulemahal Mar 13, 2020
3400dea
Small typo in polyfit docstrings
aulemahal Mar 13, 2020
198ff81
Apply suggestions from code review
aulemahal Mar 13, 2020
34726a1
Polyfit : fix style in docstring | add see also section
aulemahal Mar 13, 2020
860a892
Merge branch 'implement-polyfit-polyval' of github.com:aulemahal/xarr…
aulemahal Mar 13, 2020
3f7fe4c
Clean up docstrings and documentation.
aulemahal Mar 13, 2020
31e92bb
Merge master
aulemahal Mar 25, 2020
7eeba59
Move whats new entry to 0.16 | fix PEP8 issue in test_dataarray
aulemahal Mar 25, 2020
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Add docs | Change coeff order to fit numpy | move polyval
aulemahal committed Jan 31, 2020
commit 924f4ebdf5fb8e6af33d6d512b421c4e15302d8f
3 changes: 2 additions & 1 deletion xarray/__init__.py
Original file line number Diff line number Diff line change
@@ -15,7 +15,7 @@
from .core.alignment import align, broadcast
from .core.combine import auto_combine, combine_by_coords, combine_nested
from .core.common import ALL_DIMS, full_like, ones_like, zeros_like
from .core.computation import apply_ufunc, dot, where
from .core.computation import apply_ufunc, dot, where, polyval
from .core.concat import concat
from .core.dataarray import DataArray
from .core.dataset import Dataset
@@ -65,6 +65,7 @@
"open_mfdataset",
"open_rasterio",
"open_zarr",
"polyval",
"register_dataarray_accessor",
"register_dataset_accessor",
"save_mfdataset",
46 changes: 46 additions & 0 deletions xarray/core/computation.py
Original file line number Diff line number Diff line change
@@ -1299,3 +1299,49 @@ def where(cond, x, y):
dataset_join="exact",
dask="allowed",
)


def polyval(coord, coeffs, degree_dim="degree"):
"""Evaluate a polynomial at specific values

Parameters
----------
coord : DataArray or array-like
The 1D coordinate along which to evaluate the polynomial.
coeffs : DataArray or array-like
Coefficients of the polynomials.
degree_dim : str or int, optional
Name of dimension or axis number along which the degree increases in `coeffs`.
aulemahal marked this conversation as resolved.
Show resolved Hide resolved

See also
--------
`xarray.DataArray.polyfit`
aulemahal marked this conversation as resolved.
Show resolved Hide resolved
`numpy.polyval`
aulemahal marked this conversation as resolved.
Show resolved Hide resolved
"""
from .dataarray import DataArray
from .missing import get_clean_interp_index

if not isinstance(coord, DataArray):
coord = DataArray(coord, dims=("x",), name="x")
aulemahal marked this conversation as resolved.
Show resolved Hide resolved
x = get_clean_interp_index(coord, coord.name)

if isinstance(degree_dim, str):
deg_coord = coeffs[degree_dim]
else: # Axis number
degree_axis = degree_dim
if isinstance(coeffs, DataArray):
degree_dim = coeffs.dims[degree_axis]
deg_coord = coeffs[degree_dim]
else:
degree_dim = "degree"
deg_coord = np.arange(coeffs.shape[degree_axis])[::-1]
coeffs = DataArray(coeffs)
coeffs.coords[coeffs.dims[degree_axis]] = deg_coord
coeffs = coeffs.rename({coeffs.dims[degree_axis]: degree_dim})

lhs = DataArray(
np.vander(x, int(deg_coord.max()) + 1),
dims=(coord.name, degree_dim),
coords={coord.name: coord, degree_dim: np.arange(deg_coord.max() + 1)[::-1]},
)
return (lhs * coeffs).sum(degree_dim)
48 changes: 21 additions & 27 deletions xarray/core/dataarray.py
Original file line number Diff line number Diff line change
@@ -3241,12 +3241,28 @@ def polyfit(
self,
dim: Hashable,
deg: int,
skipna: bool = None,
rcond: float = None,
w: Union[Hashable, Any] = None,
full: bool = False,
cov: bool = False,
skipna: bool = None,
):
"""Least squares polynomial fit.

This replicates the behaviour of `numpy.polyfit` but differs by skipping
invalid values when `skipna = True`.

Parameters
----------
dim : hashable
Coordinate along which to fit the polynomials.
skipna : bool, optional
If True, removes all invalid values before fitting each 1D slices of the array
Defaults is True if data is stored in as dask.array or if there is any
invalid values, False otherwise.

Documentation of `numpy.polyfit` follows:
"""
if skipna is None:
if isinstance(self.data, pycompat.dask_array_type):
skipna = True
@@ -3255,7 +3271,7 @@ def polyfit(

x = get_clean_interp_index(self, dim)
order = int(deg) + 1
lhs = np.vander(x, order, increasing=True)
lhs = np.vander(x, order)

dims_to_stack = [dimname for dimname in self.dims if dimname != dim]
stacked_dim = utils.get_temp_dimname(dims_to_stack, "stacked")
@@ -3293,7 +3309,7 @@ def polyfit(
coeffs = DataArray(
coeffs / scale[:, np.newaxis],
dims=(degree_dim, stacked_dim),
coords={degree_dim: np.arange(order), stacked_dim: rhs[stacked_dim]},
coords={degree_dim: np.arange(order)[::-1], stacked_dim: rhs[stacked_dim]},
name=name + "_polyfit_coefficients",
).unstack(stacked_dim)

@@ -3315,7 +3331,7 @@ def polyfit(
sing = DataArray(
sing,
dims=(degree_dim,),
coords={degree_dim: np.arange(order)},
coords={degree_dim: np.arange(order)[::-1]},
name=name + "_polyfit_singular_values",
)
return coeffs, residuals, rank, sing, rcond
@@ -3344,29 +3360,7 @@ def polyfit(
str = property(StringAccessor)


def polyval(
coord: Union["DataArray", Any],
coefficients: Union["DataArray", Any],
degree_dim: Union[str, int] = "degree",
) -> "DataArray":
if not isinstance(coord, DataArray):
coord = DataArray(coord, dims=("x",), name="x")
x = get_clean_interp_index(coord, coord.name)

if isinstance(degree_dim, str):
deg_coord = coefficients[degree_dim]
else: # Axis number
if isinstance(coefficients, DataArray):
deg_coord = coefficients[coefficients.dims[degree_dim]]
else:
deg_coord = np.arange(coefficients.shape[degree_dim])

lhs = DataArray(
np.vander(x, int(deg_coord.max()) + 1, increasing=True),
dims=(coord.name, degree_dim),
coords={coord.name: coord, degree_dim: np.arange(deg_coord.max() + 1)},
)
return (lhs * coefficients).sum(degree_dim)
DataArray.polyfit.__doc__ += np.polyfit.__doc__


# priority most be higher than Variable to properly work with binary ufuncs
12 changes: 12 additions & 0 deletions xarray/tests/test_computation.py
Original file line number Diff line number Diff line change
@@ -1120,3 +1120,15 @@ def test_where():
actual = xr.where(cond, 1, 0)
expected = xr.DataArray([1, 0], dims="x")
assert_identical(expected, actual)


def test_polyval():
x = np.arange(10)
da = xr.DataArray(
np.stack((1.0 + x + 2.0 * x ** 2, 1.0 + 2.0 * x + 3.0 * x ** 2)),
dims=("d", "x"),
coords={"x": x, "d": [0, 1]},
)
coeffs = da.polyfit("x", 2)
da_pv = xr.polyval(da.x, coeffs)
xr.testing.assert_allclose(da, da_pv.T)
27 changes: 5 additions & 22 deletions xarray/tests/test_dataarray.py
Original file line number Diff line number Diff line change
@@ -4155,27 +4155,22 @@ def test_rank(self):
def test_polyfit(self):
x = np.arange(10)
da = DataArray(
np.stack((1.0 + x + 2.0 * x ** 2, 1.0 + 2.0 * x + 3.0 * x ** 2)),
np.stack((1.0 + 1.0 * x + 2.0 * x ** 2, 1.0 + 2.0 * x + 3.0 * x ** 2)),
dims=("d", "x"),
coords={"x": x, "d": [0, 1]},
)
coeffs = da.polyfit("x", 2)
expected = DataArray(
[[1, 1, 2], [1, 2, 3]],
[[2, 1, 1], [3, 2, 1]],
dims=("d", "degree"),
coords={"degree": [0, 1, 2], "d": [0, 1]},
coords={"degree": [2, 1, 0], "d": [0, 1]},
).T
assert_allclose(coeffs, expected)

# With NaN
da[0, 1] = np.nan
coeffs = da.polyfit("x", 2)
expected_na = DataArray(
[[np.nan, np.nan, np.nan], [1, 2, 3]],
dims=("d", "degree"),
coords={"degree": [0, 1, 2], "d": [0, 1]},
).T
assert_allclose(coeffs, expected_na)
coeffs = da.polyfit("x", 2, skipna=True)
assert_allclose(coeffs, expected)

# Skipna + Full output
coeffs, resids, rank, sing, rcond = da.polyfit("x", 2, skipna=True, full=True)
@@ -4184,18 +4179,6 @@ def test_polyfit(self):
np.testing.assert_almost_equal(resids, [0, 0])


def test_polyval():
x = np.arange(10)
da = DataArray(
np.stack((1.0 + x + 2.0 * x ** 2, 1.0 + 2.0 * x + 3.0 * x ** 2)),
dims=("d", "x"),
coords={"x": x, "d": [0, 1]},
)
coeffs = da.polyfit("x", 2)
da_pv = xr.core.dataarray.polyval(da.x, coeffs)
assert_allclose(da, da_pv.T)


@pytest.fixture(params=[1])
def da(request):
if request.param == 1: