From 9909f90b4781be89e3f3ff7c87893928b3e3be6e Mon Sep 17 00:00:00 2001 From: mgunyho <20118130+mgunyho@users.noreply.github.com> Date: Wed, 31 May 2023 15:43:06 +0300 Subject: [PATCH] Implement multidimensional initial guess and bounds for `curvefit` (#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 <14371165+Illviljan@users.noreply.github.com> * Add type hints to test Co-authored-by: Illviljan <14371165+Illviljan@users.noreply.github.com> --------- Co-authored-by: Illviljan <14371165+Illviljan@users.noreply.github.com> --- doc/whats-new.rst | 2 + xarray/core/dataarray.py | 99 +++++++++++++++++++++++-- xarray/core/dataset.py | 103 +++++++++++++++++++------- xarray/tests/test_dataarray.py | 131 ++++++++++++++++++++++++++++++++- 4 files changed, 298 insertions(+), 37 deletions(-) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 2c9d72ebfba..55fbda9b096 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -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ó `_. Breaking changes ~~~~~~~~~~~~~~~~ diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py index ee8573ebec7..ea6c1684f0d 100644 --- a/xarray/core/dataarray.py +++ b/xarray/core/dataarray.py @@ -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 @@ -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: @@ -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` @@ -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 + + 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") + + array([1.05692036, 1.73549639, 2.94215771]) + Coordinates: + * x (x) int64 0 1 2 + param >> fit_result["curvefit_coefficients"].sel(param="amplitude") + + array([0.1005489 , 0.19631423, 0.30003579]) + Coordinates: + * x (x) int64 0 1 2 + param >> 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") + + array([1.0569213 , 1.73550052, 2.94215733]) + Coordinates: + * x (x) int64 0 1 2 + param >> fit_result["curvefit_coefficients"].sel(param="amplitude") + + array([0.10054889, 0.1963141 , 0.3000358 ]) + Coordinates: + * x (x) int64 0 1 2 + param 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 @@ -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: @@ -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` @@ -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) @@ -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)() @@ -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": { diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py index a479b1a765d..e3138a32716 100644 --- a/xarray/tests/test_dataarray.py +++ b/xarray/tests/test_dataarray.py @@ -4415,7 +4415,7 @@ def exp_decay(t, n0, tau=1): da = da.chunk({"x": 1}) fit = da.curvefit( - coords=[da.t], func=exp_decay, p0={"n0": 4}, bounds={"tau": [2, 6]} + coords=[da.t], func=exp_decay, p0={"n0": 4}, bounds={"tau": (2, 6)} ) assert_allclose(fit.curvefit_coefficients, expected, rtol=1e-3) @@ -4436,12 +4436,141 @@ def exp_decay(t, n0, tau=1): assert param_defaults == {"n0": 4, "tau": 6} assert bounds_defaults == {"n0": (-np.inf, np.inf), "tau": (5, np.inf)} + # DataArray as bound + param_defaults, bounds_defaults = xr.core.dataset._initialize_curvefit_params( + params=params, + p0={"n0": 4}, + bounds={"tau": [DataArray([3, 4], coords=[("x", [1, 2])]), np.inf]}, + func_args=func_args, + ) + assert param_defaults["n0"] == 4 + assert ( + param_defaults["tau"] == xr.DataArray([4, 5], coords=[("x", [1, 2])]) + ).all() + assert bounds_defaults["n0"] == (-np.inf, np.inf) + assert ( + bounds_defaults["tau"][0] == DataArray([3, 4], coords=[("x", [1, 2])]) + ).all() + assert bounds_defaults["tau"][1] == np.inf + param_names = ["a"] params, func_args = xr.core.dataset._get_func_args(np.power, param_names) assert params == param_names with pytest.raises(ValueError): xr.core.dataset._get_func_args(np.power, []) + @requires_scipy + @pytest.mark.parametrize("use_dask", [True, False]) + def test_curvefit_multidimensional_guess(self, use_dask: bool) -> None: + if use_dask and not has_dask: + pytest.skip("requires dask") + + def sine(t, a, f, p): + return a * np.sin(2 * np.pi * (f * t + p)) + + t = np.arange(0, 2, 0.02) + da = DataArray( + np.stack([sine(t, 1.0, 2, 0), sine(t, 1.0, 2, 0)]), + coords={"x": [0, 1], "t": t}, + ) + + # Fitting to a sine curve produces a different result depending on the + # initial guess: either the phase is zero and the amplitude is positive + # or the phase is 0.5 * 2pi and the amplitude is negative. + + expected = DataArray( + [[1, 2, 0], [-1, 2, 0.5]], + coords={"x": [0, 1], "param": ["a", "f", "p"]}, + ) + + # Different initial guesses for different values of x + a_guess = DataArray([1, -1], coords=[da.x]) + p_guess = DataArray([0, 0.5], coords=[da.x]) + + if use_dask: + da = da.chunk({"x": 1}) + + fit = da.curvefit( + coords=[da.t], + func=sine, + p0={"a": a_guess, "p": p_guess, "f": 2}, + ) + assert_allclose(fit.curvefit_coefficients, expected) + + with pytest.raises( + ValueError, + match=r"Initial guess for 'a' has unexpected dimensions .* should only have " + "dimensions that are in data dimensions", + ): + # initial guess with additional dimensions should be an error + da.curvefit( + coords=[da.t], + func=sine, + p0={"a": DataArray([1, 2], coords={"foo": [1, 2]})}, + ) + + @requires_scipy + @pytest.mark.parametrize("use_dask", [True, False]) + def test_curvefit_multidimensional_bounds(self, use_dask: bool) -> None: + if use_dask and not has_dask: + pytest.skip("requires dask") + + def sine(t, a, f, p): + return a * np.sin(2 * np.pi * (f * t + p)) + + t = np.arange(0, 2, 0.02) + da = xr.DataArray( + np.stack([sine(t, 1.0, 2, 0), sine(t, 1.0, 2, 0)]), + coords={"x": [0, 1], "t": t}, + ) + + # Fit a sine with different bounds: positive amplitude should result in a fit with + # phase 0 and negative amplitude should result in phase 0.5 * 2pi. + + expected = DataArray( + [[1, 2, 0], [-1, 2, 0.5]], + coords={"x": [0, 1], "param": ["a", "f", "p"]}, + ) + + if use_dask: + da = da.chunk({"x": 1}) + + fit = da.curvefit( + coords=[da.t], + func=sine, + p0={"f": 2, "p": 0.25}, # this guess is needed to get the expected result + bounds={ + "a": ( + DataArray([0, -2], coords=[da.x]), + DataArray([2, 0], coords=[da.x]), + ), + }, + ) + assert_allclose(fit.curvefit_coefficients, expected) + + # Scalar lower bound with array upper bound + fit2 = da.curvefit( + coords=[da.t], + func=sine, + p0={"f": 2, "p": 0.25}, # this guess is needed to get the expected result + bounds={ + "a": (-2, DataArray([2, 0], coords=[da.x])), + }, + ) + assert_allclose(fit2.curvefit_coefficients, expected) + + with pytest.raises( + ValueError, + match=r"Upper bound for 'a' has unexpected dimensions .* should only have " + "dimensions that are in data dimensions", + ): + # bounds with additional dimensions should be an error + da.curvefit( + coords=[da.t], + func=sine, + bounds={"a": (0, DataArray([1], coords={"foo": [1]}))}, + ) + class TestReduce: @pytest.fixture(autouse=True)