Skip to content

Commit

Permalink
Implement multidimensional initial guess and bounds for curvefit (p…
Browse files Browse the repository at this point in the history 
…ydata#7821)

* Add test for multidimensional initial guess to curvefit

* Pass initial guess with *args

* Update curvefit docstrings

* Add examples to curvefit

* Add test for error on invalid p0 coords

* Raise exception on invalid coordinates in initial guess

* Add link from polyfit to curvefit

* Update doc so it matches CI

* Formatting

* Add basic test for multidimensional bounds

* Add tests for curvefit_helpers with array-valued bounds

* First attempt at fixing initialize_curvefit_params, issues warnings

* Alternative implementation of bounds initialization using xr.where(), avoids warnings

* Pass also bounds as *args to _wrapper

* Raise exception on unexpected dimensions in bounds

* Update docstring of bounds

* Update bounds docstring in dataarray also

* Update type hints for curvefit p0 and bounds

* Change list to tuple to pass mypy

* Update whats-new

* Use tuples in error message

Co-authored-by: Illviljan <[email protected]>

* Add type hints to test

Co-authored-by: Illviljan <[email protected]>

---------

Co-authored-by: Illviljan <[email protected]>
  • Loading branch information
@mgunyho @Illviljan @dstansby
2 people authored and dstansby committed Jun 28, 2023
1 parent 628f0a6 commit 28efba0
Show file tree
Hide file tree
Showing 4 changed files with 298 additions and 37 deletions.
2 changes: 2 additions & 0 deletions doc/whats-new.rst
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,8 @@ v2023.05.1 (unreleased)
New Features
~~~~~~~~~~~~

- Added support for multidimensional initial guess and bounds in :py:meth:`DataArray.curvefit` (:issue:`7768`, :pull:`7821`).
By `András Gunyhó <https://github.com/mgunyho>`_.

Breaking changes
~~~~~~~~~~~~~~~~
Expand Down
99 changes: 91 additions & 8 deletions xarray/core/dataarray.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5504,6 +5504,7 @@ def polyfit(
numpy.polyfit
numpy.polyval
xarray.polyval
DataArray.curvefit
"""
return self._to_temp_dataset().polyfit(
dim, deg, skipna=skipna, rcond=rcond, w=w, full=full, cov=cov
Expand Down Expand Up @@ -6158,8 +6159,8 @@ def curvefit(
func: Callable[..., Any],
reduce_dims: Dims = None,
skipna: bool = True,
p0: dict[str, Any] | None = None,
bounds: dict[str, Any] | None = None,
p0: dict[str, float | DataArray] | None = None,
bounds: dict[str, tuple[float | DataArray, float | DataArray]] | None = None,
param_names: Sequence[str] | None = None,
kwargs: dict[str, Any] | None = None,
) -> Dataset:
Expand Down Expand Up @@ -6190,12 +6191,16 @@ def curvefit(
Whether to skip missing values when fitting. Default is True.
p0 : dict-like or None, optional
Optional dictionary of parameter names to initial guesses passed to the
`curve_fit` `p0` arg. If none or only some parameters are passed, the rest will
be assigned initial values following the default scipy behavior.
bounds : dict-like or None, optional
Optional dictionary of parameter names to bounding values passed to the
`curve_fit` `bounds` arg. If none or only some parameters are passed, the rest
will be unbounded following the default scipy behavior.
`curve_fit` `p0` arg. If the values are DataArrays, they will be appropriately
broadcast to the coordinates of the array. If none or only some parameters are
passed, the rest will be assigned initial values following the default scipy
behavior.
bounds : dict-like, optional
Optional dictionary of parameter names to tuples of bounding values passed to the
`curve_fit` `bounds` arg. If any of the bounds are DataArrays, they will be
appropriately broadcast to the coordinates of the array. If none or only some
parameters are passed, the rest will be unbounded following the default scipy
behavior.
param_names : sequence of Hashable or None, optional
Sequence of names for the fittable parameters of `func`. If not supplied,
this will be automatically determined by arguments of `func`. `param_names`
Expand All @@ -6214,6 +6219,84 @@ def curvefit(
[var]_curvefit_covariance
The covariance matrix of the coefficient estimates.
Examples
--------
Generate some exponentially decaying data, where the decay constant and amplitude are
different for different values of the coordinate ``x``:
>>> rng = np.random.default_rng(seed=0)
>>> def exp_decay(t, time_constant, amplitude):
... return np.exp(-t / time_constant) * amplitude
...
>>> t = np.linspace(0, 10, 11)
>>> da = xr.DataArray(
... np.stack(
... [
... exp_decay(t, 1, 0.1),
... exp_decay(t, 2, 0.2),
... exp_decay(t, 3, 0.3),
... ]
... )
... + rng.normal(size=(3, t.size)) * 0.01,
... coords={"x": [0, 1, 2], "time": t},
... )
>>> da
<xarray.DataArray (x: 3, time: 11)>
array([[ 0.1012573 , 0.0354669 , 0.01993775, 0.00602771, -0.00352513,
0.00428975, 0.01328788, 0.009562 , -0.00700381, -0.01264187,
-0.0062282 ],
[ 0.20041326, 0.09805582, 0.07138797, 0.03216692, 0.01974438,
0.01097441, 0.00679441, 0.01015578, 0.01408826, 0.00093645,
0.01501222],
[ 0.29334805, 0.21847449, 0.16305984, 0.11130396, 0.07164415,
0.04744543, 0.03602333, 0.03129354, 0.01074885, 0.01284436,
0.00910995]])
Coordinates:
* x (x) int64 0 1 2
* time (time) float64 0.0 1.0 2.0 3.0 4.0 5.0 6.0 7.0 8.0 9.0 10.0
Fit the exponential decay function to the data along the ``time`` dimension:
>>> fit_result = da.curvefit("time", exp_decay)
>>> fit_result["curvefit_coefficients"].sel(param="time_constant")
<xarray.DataArray 'curvefit_coefficients' (x: 3)>
array([1.05692036, 1.73549639, 2.94215771])
Coordinates:
* x (x) int64 0 1 2
param <U13 'time_constant'
>>> fit_result["curvefit_coefficients"].sel(param="amplitude")
<xarray.DataArray 'curvefit_coefficients' (x: 3)>
array([0.1005489 , 0.19631423, 0.30003579])
Coordinates:
* x (x) int64 0 1 2
param <U13 'amplitude'
An initial guess can also be given with the ``p0`` arg (although it does not make much
of a difference in this simple example). To have a different guess for different
coordinate points, the guess can be a DataArray. Here we use the same initial guess
for the amplitude but different guesses for the time constant:
>>> fit_result = da.curvefit(
... "time",
... exp_decay,
... p0={
... "amplitude": 0.2,
... "time_constant": xr.DataArray([1, 2, 3], coords=[da.x]),
... },
... )
>>> fit_result["curvefit_coefficients"].sel(param="time_constant")
<xarray.DataArray 'curvefit_coefficients' (x: 3)>
array([1.0569213 , 1.73550052, 2.94215733])
Coordinates:
* x (x) int64 0 1 2
param <U13 'time_constant'
>>> fit_result["curvefit_coefficients"].sel(param="amplitude")
<xarray.DataArray 'curvefit_coefficients' (x: 3)>
array([0.10054889, 0.1963141 , 0.3000358 ])
Coordinates:
* x (x) int64 0 1 2
param <U13 'amplitude'
See Also
--------
DataArray.polyfit
Expand Down
103 changes: 75 additions & 28 deletions xarray/core/dataset.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -361,17 +361,24 @@ def _initialize_curvefit_params(params, p0, bounds, func_args):
"""Set initial guess and bounds for curvefit.
Priority: 1) passed args 2) func signature 3) scipy defaults
"""
from xarray.core.computation import where

def _initialize_feasible(lb, ub):
# Mimics functionality of scipy.optimize.minpack._initialize_feasible
lb_finite = np.isfinite(lb)
ub_finite = np.isfinite(ub)
p0 = np.nansum(
[
0.5 * (lb + ub) * int(lb_finite & ub_finite),
(lb + 1) * int(lb_finite & ~ub_finite),
(ub - 1) * int(~lb_finite & ub_finite),
]
p0 = where(
lb_finite,
where(
ub_finite,
0.5 * (lb + ub), # both bounds finite
lb + 1, # lower bound finite, upper infinite
),
where(
ub_finite,
ub - 1, # lower bound infinite, upper finite
0, # both bounds infinite
),
)
return p0

Expand All @@ -381,9 +388,13 @@ def _initialize_feasible(lb, ub):
if p in func_args and func_args[p].default is not func_args[p].empty:
param_defaults[p] = func_args[p].default
if p in bounds:
bounds_defaults[p] = tuple(bounds[p])
if param_defaults[p] < bounds[p][0] or param_defaults[p] > bounds[p][1]:
param_defaults[p] = _initialize_feasible(bounds[p][0], bounds[p][1])
lb, ub = bounds[p]
bounds_defaults[p] = (lb, ub)
param_defaults[p] = where(
(param_defaults[p] < lb) | (param_defaults[p] > ub),
_initialize_feasible(lb, ub),
param_defaults[p],
)
if p in p0:
param_defaults[p] = p0[p]
return param_defaults, bounds_defaults
Expand Down Expand Up @@ -8617,8 +8628,8 @@ def curvefit(
func: Callable[..., Any],
reduce_dims: Dims = None,
skipna: bool = True,
p0: dict[str, Any] | None = None,
bounds: dict[str, Any] | None = None,
p0: dict[str, float | DataArray] | None = None,
bounds: dict[str, tuple[float | DataArray, float | DataArray]] | None = None,
param_names: Sequence[str] | None = None,
kwargs: dict[str, Any] | None = None,
) -> T_Dataset:
Expand Down Expand Up @@ -8649,12 +8660,16 @@ def curvefit(
Whether to skip missing values when fitting. Default is True.
p0 : dict-like, optional
Optional dictionary of parameter names to initial guesses passed to the
`curve_fit` `p0` arg. If none or only some parameters are passed, the rest will
be assigned initial values following the default scipy behavior.
`curve_fit` `p0` arg. If the values are DataArrays, they will be appropriately
broadcast to the coordinates of the array. If none or only some parameters are
passed, the rest will be assigned initial values following the default scipy
behavior.
bounds : dict-like, optional
Optional dictionary of parameter names to bounding values passed to the
`curve_fit` `bounds` arg. If none or only some parameters are passed, the rest
will be unbounded following the default scipy behavior.
Optional dictionary of parameter names to tuples of bounding values passed to the
`curve_fit` `bounds` arg. If any of the bounds are DataArrays, they will be
appropriately broadcast to the coordinates of the array. If none or only some
parameters are passed, the rest will be unbounded following the default scipy
behavior.
param_names : sequence of hashable, optional
Sequence of names for the fittable parameters of `func`. If not supplied,
this will be automatically determined by arguments of `func`. `param_names`
Expand Down Expand Up @@ -8721,29 +8736,53 @@ def curvefit(
"in fitting on scalar data."
)

# Check that initial guess and bounds only contain coordinates that are in preserved_dims
for param, guess in p0.items():
if isinstance(guess, DataArray):
unexpected = set(guess.dims) - set(preserved_dims)
if unexpected:
raise ValueError(
f"Initial guess for '{param}' has unexpected dimensions "
f"{tuple(unexpected)}. It should only have dimensions that are in data "
f"dimensions {preserved_dims}."
)
for param, (lb, ub) in bounds.items():
for label, bound in zip(("Lower", "Upper"), (lb, ub)):
if isinstance(bound, DataArray):
unexpected = set(bound.dims) - set(preserved_dims)
if unexpected:
raise ValueError(
f"{label} bound for '{param}' has unexpected dimensions "
f"{tuple(unexpected)}. It should only have dimensions that are in data "
f"dimensions {preserved_dims}."
)

# Broadcast all coords with each other
coords_ = broadcast(*coords_)
coords_ = [
coord.broadcast_like(self, exclude=preserved_dims) for coord in coords_
]
n_coords = len(coords_)

params, func_args = _get_func_args(func, param_names)
param_defaults, bounds_defaults = _initialize_curvefit_params(
params, p0, bounds, func_args
)
n_params = len(params)
kwargs.setdefault("p0", [param_defaults[p] for p in params])
kwargs.setdefault(
"bounds",
[
[bounds_defaults[p][0] for p in params],
[bounds_defaults[p][1] for p in params],
],
)

def _wrapper(Y, *coords_, **kwargs):
def _wrapper(Y, *args, **kwargs):
# Wrap curve_fit with raveled coordinates and pointwise NaN handling
x = np.vstack([c.ravel() for c in coords_])
# *args contains:
# - the coordinates
# - initial guess
# - lower bounds
# - upper bounds
coords__ = args[:n_coords]
p0_ = args[n_coords + 0 * n_params : n_coords + 1 * n_params]
lb = args[n_coords + 1 * n_params : n_coords + 2 * n_params]
ub = args[n_coords + 2 * n_params :]

x = np.vstack([c.ravel() for c in coords__])
y = Y.ravel()
if skipna:
mask = np.all([np.any(~np.isnan(x), axis=0), ~np.isnan(y)], axis=0)
Expand All @@ -8754,7 +8793,7 @@ def _wrapper(Y, *coords_, **kwargs):
pcov = np.full([n_params, n_params], np.nan)
return popt, pcov
x = np.squeeze(x)
popt, pcov = curve_fit(func, x, y, **kwargs)
popt, pcov = curve_fit(func, x, y, p0=p0_, bounds=(lb, ub), **kwargs)
return popt, pcov

result = type(self)()
Expand All @@ -8764,13 +8803,21 @@ def _wrapper(Y, *coords_, **kwargs):
else:
name = f"{str(name)}_"

input_core_dims = [reduce_dims_ for _ in range(n_coords + 1)]
input_core_dims.extend(
[[] for _ in range(3 * n_params)]
) # core_dims for p0 and bounds

popt, pcov = apply_ufunc(
_wrapper,
da,
*coords_,
*param_defaults.values(),
*[b[0] for b in bounds_defaults.values()],
*[b[1] for b in bounds_defaults.values()],
vectorize=True,
dask="parallelized",
input_core_dims=[reduce_dims_ for d in range(len(coords_) + 1)],
input_core_dims=input_core_dims,
output_core_dims=[["param"], ["cov_i", "cov_j"]],
dask_gufunc_kwargs={
"output_sizes": {
Expand Down
Loading

0 comments on commit 28efba0

Please sign in to comment.