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
Formatting with black, flake8
aulemahal committed Jan 30, 2020
commit 05355af618ccfc68c061c20ec3318e974f71e497
11 changes: 10 additions & 1 deletion xarray/core/dask_array_ops.py
Original file line number Diff line number Diff line change
@@ -99,9 +99,18 @@ def func(x, window, axis=-1):

def least_squares(lhs, rhs, rcond=None, skipna=False):
import dask.array as da

lhs_da = da.from_array(lhs, chunks=(rhs.chunks[0], lhs.shape[1]))
if skipna:
results = da.apply_along_axis(nputils._nanpolyfit_1d, 0, rhs.data, lhs_da.data, dtype=float, shape=(lhs.shape[1] + 1,), rcond=rcond)
results = da.apply_along_axis(
nputils._nanpolyfit_1d,
0,
rhs.data,
lhs_da.data,
dtype=float,
shape=(lhs.shape[1] + 1,),
rcond=rcond,
)
coeffs = results[:-1, ...]
residuals = results[-1, ...]
else:
92 changes: 60 additions & 32 deletions xarray/core/dataarray.py
Original file line number Diff line number Diff line change
@@ -3238,20 +3238,20 @@ def map_blocks(

def polyfit(
self,
dim : Hashable,
deg : int,
rcond : float = None,
w : Union[Hashable, Any] = None,
full : bool = False,
cov : bool = False,
skipna : bool = False
dim: Hashable,
deg: int,
rcond: float = None,
w: Union[Hashable, Any] = None,
full: bool = False,
cov: bool = False,
skipna: bool = False,
):
x = get_clean_interp_index(self, dim)
order = int(deg) + 1
lhs = np.vander(x, order, increasing=True)

dims_to_stack = [dimname for dimname in self.dims if dimname != dim]
stacked_dim = utils.get_temp_dimname(dims_to_stack, 'stacked')
stacked_dim = utils.get_temp_dimname(dims_to_stack, "stacked")
rhs = self.transpose(dim, *dims_to_stack).stack({stacked_dim: dims_to_stack})

if rcond is None:
@@ -3273,38 +3273,64 @@ def polyfit(
scale = np.sqrt((lhs * lhs).sum(axis=0))
lhs /= scale

coeffs, residuals = duck_array_ops.least_squares(lhs, rhs, rcond=rcond, skipna=skipna)
coeffs, residuals = duck_array_ops.least_squares(
lhs, rhs, rcond=rcond, skipna=skipna
)

if self.name:
name = "{}_".format(self.name)
dcherian marked this conversation as resolved.
Show resolved Hide resolved
else:
name = ""

degree_dim = utils.get_temp_dimname(dims_to_stack, 'degree')
coeffs = DataArray(coeffs / scale[:, np.newaxis], dims=(degree_dim, stacked_dim),
coords={degree_dim: np.arange(order),
stacked_dim: rhs[stacked_dim]},
name=(self.name or '') + '_polyfit_coefficients'
).unstack(stacked_dim)
degree_dim = utils.get_temp_dimname(dims_to_stack, "degree")
coeffs = DataArray(
coeffs / scale[:, np.newaxis],
dims=(degree_dim, stacked_dim),
coords={degree_dim: np.arange(order), stacked_dim: rhs[stacked_dim]},
name=name + "_polyfit_coefficients",
).unstack(stacked_dim)

rank = np.linalg.matrix_rank(lhs)
if rank != order and not full:
warnings.warn("Polyfit may be poorly conditioned", np.RankWarning, stacklevel=4)
warnings.warn(
"Polyfit may be poorly conditioned", np.RankWarning, stacklevel=4
)

if full or (cov is True):
residuals = DataArray(residuals, dims=(stacked_dim), coords={stacked_dim: rhs[stacked_dim]},
name=(self.name or '') + '_polyfit_residuals'
).unstack(stacked_dim)
residuals = DataArray(
residuals,
dims=(stacked_dim),
coords={stacked_dim: rhs[stacked_dim]},
name=name + "_polyfit_residuals",
).unstack(stacked_dim)
if full:
sing = np.linalg.svd(lhs, compute_uv=False)
sing = DataArray(sing, dims=(degree_dim,), coords={degree_dim: np.arange(order)},
name=(self.name or '') + '_polyfit_singular_values')
sing = DataArray(
sing,
dims=(degree_dim,),
coords={degree_dim: np.arange(order)},
name=name + "_polyfit_singular_values",
)
return coeffs, residuals, rank, sing, rcond
if cov:
Vbase = np.linalg.inv(np.dot(lhs.T, lhs))
Vbase /= np.outer(scale, scale)
if cov == 'unscaled':
if cov == "unscaled":
fac = 1
else:
if x.shape[0] <= order:
raise ValueError("The number of data points must exceed order to scale the covariance matrix.")
raise ValueError(
"The number of data points must exceed order to scale the covariance matrix."
)
fac = residuals / (x.shape[0] - order)
covariance = DataArray(Vbase, dims=('cov_i', 'cov_j'), name=(self.name or '') + '_polyfit_covariance') * fac
covariance = (
DataArray(
Vbase,
dims=("cov_i", "cov_j"),
name=name + "_polyfit_covariance",
)
* fac
)
return coeffs, covariance
return coeffs

@@ -3314,12 +3340,12 @@ def polyfit(


def polyval(
coord : Union["DataArray", Any],
coefficients : Union["DataArray", Any],
degree_dim : Union[str, int] = 'degree'
) -> "DataArray" :
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')
coord = DataArray(coord, dims=("x",), name="x")
x = get_clean_interp_index(coord, coord.name)

if isinstance(degree_dim, str):
@@ -3330,9 +3356,11 @@ def polyval(
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)})
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)


4 changes: 3 additions & 1 deletion xarray/core/nputils.py
Original file line number Diff line number Diff line change
@@ -233,7 +233,9 @@ def least_squares(lhs, rhs, rcond=None, skipna=False):
coeffs, residuals, _, _ = np.linalg.lstsq(lhs, rhs.data, rcond=rcond)
if skipna:
nan_cols = np.nonzero(np.isnan(coeffs[0, :]))[0]
out = np.apply_along_axis(_nanpolyfit_1d, 0, rhs.data[:, nan_cols], lhs, rcond=rcond)
out = np.apply_along_axis(
_nanpolyfit_1d, 0, rhs.data[:, nan_cols], lhs, rcond=rcond
)
coeffs[:, nan_cols] = out[:-1, :]
residuals[nan_cols] = out[-1, :]
return coeffs, residuals
36 changes: 24 additions & 12 deletions xarray/tests/test_dataarray.py
Original file line number Diff line number Diff line change
@@ -4154,32 +4154,44 @@ def test_rank(self):

def test_polyfit(self):
x = np.arange(10)
da = DataArray(np.stack((1. + x + 2. * x**2, 1. + 2. * x + 3. * x**2)),
dims=('d', 'x'), coords={'x': x, 'd': [0, 1]})
coeffs = da.polyfit('x', 2)
expected = DataArray([[1, 1, 2], [1, 2, 3]], dims=('d', 'degree'),
coords={'degree': [0, 1, 2], 'd': [0, 1]}).T
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)
expected = DataArray(
[[1, 1, 2], [1, 2, 3]],
dims=("d", "degree"),
coords={"degree": [0, 1, 2], "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
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)

# Skipna + Full output
coeffs, resids, rank, sing, rcond = da.polyfit('x', 2, skipna=True, full=True)
coeffs, resids, rank, sing, rcond = da.polyfit("x", 2, skipna=True, full=True)
assert_allclose(coeffs, expected)
assert rank == np.linalg.matrix_rank(np.vander(x, 3))
np.testing.assert_almost_equal(resids, [0, 0])


def test_polyval():
x = np.arange(10)
da = DataArray(np.stack((1. + x + 2. * x**2, 1. + 2. * x + 3. * x**2)),
dims=('d', 'x'), coords={'x': x, 'd': [0, 1]})
coeffs = da.polyfit('x', 2)
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)