Repository: pydata/xarray

Issue:
General curve fitting method
Xarray should have a general curve-fitting function as part of its main API.

## Motivation

Yesterday I wanted to fit a simple decaying exponential function to the data in a DataArray and realised there currently isn't an immediate way to do this in xarray. You have to either pull out the `.values` (losing the power of dask), or use `apply_ufunc` (complicated).

This is an incredibly common, domain-agnostic task, so although I don't think we should support various kinds of unusual optimisation procedures (which could always go in an extension package instead), I think a basic fitting method is within scope for the main library. There are [SO questions](https://stackoverflow.com/questions/62987617/using-scipy-curve-fit-with-dask-xarray) asking how to achieve this.

We already have [`.polyfit` and `polyval` anyway](https://github.com/pydata/xarray/pull/3733/files#), which are more specific. (@AndrewWilliams3142 and @aulemahal I expect you will have thoughts on how implement this generally.)

## Proposed syntax

I want something like this to work:

```python
def exponential_decay(xdata, A=10, L=5):
    return A*np.exp(-xdata/L)

# returns a dataset containing the optimised values of each parameter
fitted_params = da.fit(exponential_decay)

fitted_line = exponential_decay(da.x, A=fitted_params['A'], L=fitted_params['L'])

# Compare
da.plot(ax)
fitted_line.plot(ax)
```

It would also be nice to be able to fit in multiple dimensions. That means both for example fitting a 2D function to 2D data:

```python
def hat(xdata, ydata, h=2, r0=1):
    r = xdata**2 + ydata**2
    return h*np.exp(-r/r0)

fitted_params = da.fit(hat)

fitted_hat = hat(da.x, da.y, h=fitted_params['h'], r0=fitted_params['r0'])
```

but also repeatedly fitting a 1D function to 2D data:

```python
# da now has a y dimension too
fitted_params = da.fit(exponential_decay, fit_along=['x'])

# As fitted_params now has y-dependence, broadcasting means fitted_lines does too
fitted_lines = exponential_decay(da.x, A=fitted_params.A, L=fitted_params.L)
```
The latter would be useful for fitting the same curve to multiple model runs, but means we need some kind of `fit_along` or `dim` argument, which would default to all dims.

So the method docstring would end up like
```python
def fit(self, f, fit_along=None, skipna=None, full=False, cov=False):
    """
    Fits the function f to the DataArray.

    Expects the function f to have a signature like
    `result = f(*coords, **params)`
    for example
    `result_da = f(da.xcoord, da.ycoord, da.zcoord, A=5, B=None)`
    The names of the `**params` kwargs will be used to name the output variables.

    Returns
    -------
    fit_results - A single dataset which contains the variables (for each parameter in the fitting function):
    `param1`
        The optimised fit coefficients for parameter one.
    `param1_residuals`
        The residuals of the fit for parameter one.
    ...
    """

```

## Questions

1) Should it wrap `scipy.optimise.curve_fit`, or reimplement it? 

    Wrapping it is simpler, but as it just calls `least_squares` [under the hood](https://github.com/scipy/scipy/blob/v1.5.2/scipy/optimize/minpack.py#L532-L834) then reimplementing it would mean we could use the dask-powered version of `least_squares` (like [`da.polyfit does`](https://github.com/pydata/xarray/blob/9058114f70d07ef04654d1d60718442d0555b84b/xarray/core/dataset.py#L5987)).

2) What form should we expect the curve-defining function to come in?

    `scipy.optimize.curve_fit` expects the curve to act as `ydata = f(xdata, *params) + eps`, but in xarray then `xdata` could be one or multiple coords or dims, not necessarily a single array. Might it work to require a signature like `result_da = f(da.xcoord, da.ycoord, da.zcoord, ..., **params)`? Then the `.fit` method would be work out how many coords to pass to `f` based on the dimension of the `da` and the `fit_along` argument. But then the order of coord arguments in the signature of `f` would matter, which doesn't seem very xarray-like.

3) Is it okay to inspect parameters of the curve-defining function?

    If we tell the user the curve-defining function has to have a signature like `da = func(*coords, **params)`, then we could read the names of the parameters by inspecting the function kwargs. Is that a good idea or might it end up being unreliable? Is the `inspect` standard library module the right thing to use for that? This could also be used to provide default guesses for the fitting parameters.

Discussion:
pep8speaks: Hello @slevang! Thanks for opening this PR. We checked the lines you've touched for [PEP 8](https://www.python.org/dev/peps/pep-0008) issues, and found:

* In the file [`xarray/core/dataarray.py`](https://github.com/pydata/xarray/blob/e06027ae71134e2ae98c502d2104a8e233feacc6/xarray/core/dataarray.py):

> [Line 4292:14](https://github.com/pydata/xarray/blob/e06027ae71134e2ae98c502d2104a8e233feacc6/xarray/core/dataarray.py#L4292): [E241](https://duckduckgo.com/?q=pep8%20E241) multiple spaces after ':'
> [Line 4305:86](https://github.com/pydata/xarray/blob/e06027ae71134e2ae98c502d2104a8e233feacc6/xarray/core/dataarray.py#L4305): [W291](https://duckduckgo.com/?q=pep8%20W291) trailing whitespace
> [Line 4308:71](https://github.com/pydata/xarray/blob/e06027ae71134e2ae98c502d2104a8e233feacc6/xarray/core/dataarray.py#L4308): [W291](https://duckduckgo.com/?q=pep8%20W291) trailing whitespace
> [Line 4310:86](https://github.com/pydata/xarray/blob/e06027ae71134e2ae98c502d2104a8e233feacc6/xarray/core/dataarray.py#L4310): [W291](https://duckduckgo.com/?q=pep8%20W291) trailing whitespace
> [Line 4311:88](https://github.com/pydata/xarray/blob/e06027ae71134e2ae98c502d2104a8e233feacc6/xarray/core/dataarray.py#L4311): [W291](https://duckduckgo.com/?q=pep8%20W291) trailing whitespace
> [Line 4328:85](https://github.com/pydata/xarray/blob/e06027ae71134e2ae98c502d2104a8e233feacc6/xarray/core/dataarray.py#L4328): [W291](https://duckduckgo.com/?q=pep8%20W291) trailing whitespace

* In the file [`xarray/core/dataset.py`](https://github.com/pydata/xarray/blob/e06027ae71134e2ae98c502d2104a8e233feacc6/xarray/core/dataset.py):

> [Line 7013:14](https://github.com/pydata/xarray/blob/e06027ae71134e2ae98c502d2104a8e233feacc6/xarray/core/dataset.py#L7013): [E241](https://duckduckgo.com/?q=pep8%20E241) multiple spaces after ':'
> [Line 7026:86](https://github.com/pydata/xarray/blob/e06027ae71134e2ae98c502d2104a8e233feacc6/xarray/core/dataset.py#L7026): [W291](https://duckduckgo.com/?q=pep8%20W291) trailing whitespace
> [Line 7029:71](https://github.com/pydata/xarray/blob/e06027ae71134e2ae98c502d2104a8e233feacc6/xarray/core/dataset.py#L7029): [W291](https://duckduckgo.com/?q=pep8%20W291) trailing whitespace
> [Line 7031:86](https://github.com/pydata/xarray/blob/e06027ae71134e2ae98c502d2104a8e233feacc6/xarray/core/dataset.py#L7031): [W291](https://duckduckgo.com/?q=pep8%20W291) trailing whitespace
> [Line 7032:88](https://github.com/pydata/xarray/blob/e06027ae71134e2ae98c502d2104a8e233feacc6/xarray/core/dataset.py#L7032): [W291](https://duckduckgo.com/?q=pep8%20W291) trailing whitespace
> [Line 7049:85](https://github.com/pydata/xarray/blob/e06027ae71134e2ae98c502d2104a8e233feacc6/xarray/core/dataset.py#L7049): [W291](https://duckduckgo.com/?q=pep8%20W291) trailing whitespace
> [Line 7079:9](https://github.com/pydata/xarray/blob/e06027ae71134e2ae98c502d2104a8e233feacc6/xarray/core/dataset.py#L7079): [F401](https://duckduckgo.com/?q=pep8%20F401) '.computation.apply_ufunc' imported but unused
> [Line 7103:38](https://github.com/pydata/xarray/blob/e06027ae71134e2ae98c502d2104a8e233feacc6/xarray/core/dataset.py#L7103): [W291](https://duckduckgo.com/?q=pep8%20W291) trailing whitespace
> [Line 7117:17](https://github.com/pydata/xarray/blob/e06027ae71134e2ae98c502d2104a8e233feacc6/xarray/core/dataset.py#L7117): [E128](https://duckduckgo.com/?q=pep8%20E128) continuation line under-indented for visual indent
> [Line 7118:17](https://github.com/pydata/xarray/blob/e06027ae71134e2ae98c502d2104a8e233feacc6/xarray/core/dataset.py#L7118): [E128](https://duckduckgo.com/?q=pep8%20E128) continuation line under-indented for visual indent
> [Line 7119:17](https://github.com/pydata/xarray/blob/e06027ae71134e2ae98c502d2104a8e233feacc6/xarray/core/dataset.py#L7119): [E128](https://duckduckgo.com/?q=pep8%20E128) continuation line under-indented for visual indent
> [Line 7122:35](https://github.com/pydata/xarray/blob/e06027ae71134e2ae98c502d2104a8e233feacc6/xarray/core/dataset.py#L7122): [E231](https://duckduckgo.com/?q=pep8%20E231) missing whitespace after ':'
> [Line 7122:44](https://github.com/pydata/xarray/blob/e06027ae71134e2ae98c502d2104a8e233feacc6/xarray/core/dataset.py#L7122): [E231](https://duckduckgo.com/?q=pep8%20E231) missing whitespace after ':'
> [Line 7122:54](https://github.com/pydata/xarray/blob/e06027ae71134e2ae98c502d2104a8e233feacc6/xarray/core/dataset.py#L7122): [W291](https://duckduckgo.com/?q=pep8%20W291) trailing whitespace
> [Line 7123:44](https://github.com/pydata/xarray/blob/e06027ae71134e2ae98c502d2104a8e233feacc6/xarray/core/dataset.py#L7123): [E231](https://duckduckgo.com/?q=pep8%20E231) missing whitespace after ':'
> [Line 7123:54](https://github.com/pydata/xarray/blob/e06027ae71134e2ae98c502d2104a8e233feacc6/xarray/core/dataset.py#L7123): [W291](https://duckduckgo.com/?q=pep8%20W291) trailing whitespace
> [Line 7124:44](https://github.com/pydata/xarray/blob/e06027ae71134e2ae98c502d2104a8e233feacc6/xarray/core/dataset.py#L7124): [E231](https://duckduckgo.com/?q=pep8%20E231) missing whitespace after ':'
> [Line 7134:47](https://github.com/pydata/xarray/blob/e06027ae71134e2ae98c502d2104a8e233feacc6/xarray/core/dataset.py#L7134): [E231](https://duckduckgo.com/?q=pep8%20E231) missing whitespace after ':'
> [Line 7134:63](https://github.com/pydata/xarray/blob/e06027ae71134e2ae98c502d2104a8e233feacc6/xarray/core/dataset.py#L7134): [E231](https://duckduckgo.com/?q=pep8%20E231) missing whitespace after ':'
> [Line 7134:79](https://github.com/pydata/xarray/blob/e06027ae71134e2ae98c502d2104a8e233feacc6/xarray/core/dataset.py#L7134): [E231](https://duckduckgo.com/?q=pep8%20E231) missing whitespace after ':'
> [Line 7137:20](https://github.com/pydata/xarray/blob/e06027ae71134e2ae98c502d2104a8e233feacc6/xarray/core/dataset.py#L7137): [E225](https://duckduckgo.com/?q=pep8%20E225) missing whitespace around operator
> [Line 7145:1](https://github.com/pydata/xarray/blob/e06027ae71134e2ae98c502d2104a8e233feacc6/xarray/core/dataset.py#L7145): [E305](https://duckduckgo.com/?q=pep8%20E305) expected 2 blank lines after class or function definition, found 1


TomNicholas: This is great, thanks for submitting this!

I just had a go with it, and it worked nicely. I have a couple of suggestions for improving it though:

1) Different fit coefficients as differently-named variables in the output, rather than indexed with a coordinate. This would then be consistent with `Dataset.polyfit`, which returns a set of different variables `[var]_polyfit_coefficients`. We could also get the names from the names of the keyword args when inspecting the function, and if it fails to get names just call them like `param1_fit_coefficients`, `param2_fit_coefficients` etc.

    What you have now works nicely though, so perhaps you could just reorganise the result before returning, like `Dataset.polyfit` does?

2) Initial guesses for each fit parameter. At the moment the user has to pass an ordered array of initial guesses through like

    ```python
    da.curvefit(x=da.x, dim='x', func=linear, kwargs={'p0': [m_guess, c_guess]})
    ```

    but it would be nicer to just pass them as a dictionary like

    ```python
    da.curvefit(x=da.x, dim='x', func=linear, initial_guess={'m': m_guess, 'c': c_guess})
    ```
    or even have the guesses read from the function definition maybe? i.e.
    ```python
    def linear(x, m=m_guess, c=c_guess):
        return m*x + c
    ```

3) (Stretch goal) Ability to fit >1D functions, e.g. fit a 2D gaussian to find a peak in a 2D image. But if we get the API right then this could be left to a later PR.

Also, the whole argument inspection thing probably deserves a few dedicated tests, in addition to testing the fitting functionality.
slevang: > 1. Different fit coefficients as differently-named variables in the output, rather than indexed with a coordinate. 

I think the way I configured things now does replicate the polyfit results. For example:
```python
ds = xr.tutorial.open_dataset('air_temperature')
ds['air2'] = ds.air.copy()
ds.polyfit(dim='time', deg=2)

<xarray.Dataset>
Dimensions:                    (degree: 3, lat: 25, lon: 53)
Coordinates:
  * degree                     (degree) int64 2 1 0
  * lat                        (lat) float64 75.0 72.5 70.0 ... 20.0 17.5 15.0
  * lon                        (lon) float64 200.0 202.5 205.0 ... 327.5 330.0
Data variables:
    air_polyfit_coefficients   (degree, lat, lon) float64 -1.162e-32 ... 1.13...
    air2_polyfit_coefficients  (degree, lat, lon) float64 -1.162e-32 ... 1.14...
```
Compared to this:
```python
def square(x, a, b ,c):
    return a*np.power(x, 2) + b*x + c

ds.curvefit(x=ds.time, dim='time', func=square)

<xarray.Dataset>
Dimensions:                     (lat: 25, lon: 53, param: 3)
Coordinates:
  * lat                         (lat) float32 75.0 72.5 70.0 ... 20.0 17.5 15.0
  * lon                         (lon) float32 200.0 202.5 205.0 ... 327.5 330.0
  * param                       (param) <U1 'a' 'b' 'c'
Data variables:
    air_curvefit_coefficients   (param, lat, lon) float64 -1.162e-32 ... 1.13...
    air2_curvefit_coefficients  (param, lat, lon) float64 -1.162e-32 ... 1.14...
```
In both cases, each variable in the dataset returns a separate coefficients variable, and all fittable coefficients are stacked along a dimension, `degree` for the more specific polyfit case and a generic `param` for curvefit.

> 2. Initial guesses for each fit parameter.

Yeah this would be good. Should be easy to look for default values in the function itself using `inspect`.

> 3. (Stretch goal) Ability to fit >1D functions

Looks like this could be possible with a call to `ravel` like [this](https://scipython.com/blog/non-linear-least-squares-fitting-of-a-two-dimensional-data/). I'll do some experimenting.
TomNicholas: > I think the way I configured things now does replicate the polyfit results.

You're right! My bad. The consistency with polyfit looks good.

> Looks like this could be possible with a call to ravel

Oh nice! That looks like it would allow for ND functions fit to ND data. It looks like there is a [dask version of ravel](https://docs.dask.org/en/latest/array-api.html#dask.array.ravel) which might be useful. (And judging by the comments on that blog post I think @StanczakDominik would appreciate this feature too!)
slevang: Some more progress here.

1. I liked the idea of being able to supply `p0` and `bounds` as a dictionary mapping between parameter names and values, which is much more in the style of xarray. So I've included the logic to translate between this and the scipy args which are just ordered lists. The logic first prioritizes any of these values supplied directly in keyword args, then looks for defaults in the function definition, and otherwise falls back to the scipy defaults.
2. There's now a working implementation of multi-dimensional fits using `ravel` to flatten the coordinate arrays. One limitation is that `func` must be specified like `f((x, y), z)` as scipy takes a single `x` argument for the fitting coordinate. Maybe there is a clever way around this, but I think it's ok.
3. Some changes to the API. 

The best way to specify the fitting coordinates is a bit tricky to figure out. My original use case for this was needing to fit a relationship between two time/lat/lon dataarrays with the fit done over all time. But probably a more common use would be to just fit a curve over one or two dimensions that already exist in your data. So it would be great to handle these possibilities seamlessly. 

What I've settled on for now is a `coords` argument that takes a list of coordinates that should be the same length as the input coordinates of your fitting function. Additionally, there is a `reduce_dim` arg, so that, say if you want to fit a function in time but aggregate over all lat and lon, you can supply those dimensions here. So with a 3D x/y/time array, any of the following should work:
```python
# Fit a 1d function in time, returns parameters with dims (x, y)
da.curvefit(coords='time', ...)
# Fit a 2d function in space, returns parameters with dims (t)
da.curvefit(coords=['x', 'y'], ...)
# Fit a 1d function with another 3d dataarray and aggregate over time, returns parameters with dims (x, y)
da.curvefit(coords=da1, reduce_dim='time', ...)
```
The logic to make this work got a bit complicated, since we need to supply the right `input_core_dims` to `apply_ufunc`, and also to explicitly broadcast the coordinates to ensure the `ravel` operation works. Any thoughts on cleanup here appreciated.

Will eventually need to add tests and improve docs and examples. Tests especially I could use some help on.
slevang: Added a couple usage examples in the docs, including one that replicates the [scipy](https://scipython.com/blog/non-linear-least-squares-fitting-of-a-two-dimensional-data/) example of fitting multiple peaks. Because of the wrapper function and variable args, this requires supplying `param_names` manually. Works nicely though.
slevang: Added some checks that will raise errors if `inspect.signature` is not able to determine the function arguments, either for functions with variable-length args or things like numpy ufuncs that don't seem to work with `signature`. In these cases, `param_names` can be passed manually.

Also added minimal tests, but these should probably be expanded.
slevang: I've been playing around with this some more, and found the performance to be much better using a process-heavy dask scheduler. For example:
```python
import xarray as xr
import numpy as np
import time
import dask
​
def exponential(x, a, xc):
    return np.exp((x - xc) / a)
​
x = np.arange(-5, 5, 0.001)
t = np.arange(-5, 5, 0.01)
X, T = np.meshgrid(x, t)
Z1 = np.random.uniform(low=-5, high=5, size=X.shape)
Z2 = exponential(Z1, 3, X) + np.random.normal(scale=0.1, size=Z1.shape)
​
ds = xr.Dataset(
    data_vars=dict(var1=(["t", "x"], Z1), var2=(["t", "x"], Z2)),
    coords={"t": t, "x": x},
)
​
ds = ds.chunk({'x':10})
​
def test_fit():
    start = time.time()
    fit = ds.var2.curvefit(
        coords=ds.var1,
        func=exponential,
        reduce_dim="t",
    ).compute()
    print(f'Fitting time: {time.time() - start:.2f}s')

with dask.config.set(scheduler='threads'):
    test_fit()
with dask.config.set(scheduler='processes'):
    test_fit()
with dask.distributed.Client() as client:
    test_fit()
with dask.distributed.Client(n_workers=8, threads_per_worker=1) as client:
    test_fit()
```
On my 8-core machine, takes:
```
Fitting time: 8.32s
Fitting time: 2.71s
Fitting time: 4.40s
Fitting time: 3.43s
```
According to [this](https://github.com/scipy/scipy/pull/8740) the underlying scipy routines should be thread safe.
slevang: Thanks for the review @dcherian! The latest commit has been refactored with a couple helper functions and associated tests, and any steps that served no purpose other than consistency with `polyfit` have been removed.

If you can think of any more specific test cases that should be included, happy to add them.
slevang: I think so. I pushed a merge commit to get this up to date with the current release.

Patch:
diff --git a/doc/api.rst b/doc/api.rst
index a140d9e2..6288ce01 100644
--- a/doc/api.rst
+++ b/doc/api.rst
@@ -179,6 +179,7 @@ Computation
    Dataset.integrate
    Dataset.map_blocks
    Dataset.polyfit
+   Dataset.curvefit
 
 **Aggregation**:
 :py:attr:`~Dataset.all`
@@ -375,7 +376,7 @@ Computation
    DataArray.integrate
    DataArray.polyfit
    DataArray.map_blocks
-
+   DataArray.curvefit
 
 **Aggregation**:
 :py:attr:`~DataArray.all`
diff --git a/doc/user-guide/computation.rst b/doc/user-guide/computation.rst
index 28804bcb..2a894f72 100644
--- a/doc/user-guide/computation.rst
+++ b/doc/user-guide/computation.rst
@@ -444,6 +444,89 @@ The inverse operation is done with :py:meth:`~xarray.polyval`,
 .. note::
     These methods replicate the behaviour of :py:func:`numpy.polyfit` and :py:func:`numpy.polyval`.
 
+
+.. _compute.curvefit:
+
+Fitting arbitrary functions
+===========================
+
+Xarray objects also provide an interface for fitting more complex functions using
+:py:meth:`scipy.optimize.curve_fit`. :py:meth:`~xarray.DataArray.curvefit` accepts
+user-defined functions and can fit along multiple coordinates.
+
+For example, we can fit a relationship between two ``DataArray`` objects, maintaining
+a unique fit at each spatial coordinate but aggregating over the time dimension:
+
+.. ipython:: python
+
+    def exponential(x, a, xc):
+        return np.exp((x - xc) / a)
+
+
+    x = np.arange(-5, 5, 0.1)
+    t = np.arange(-5, 5, 0.1)
+    X, T = np.meshgrid(x, t)
+    Z1 = np.random.uniform(low=-5, high=5, size=X.shape)
+    Z2 = exponential(Z1, 3, X)
+    Z3 = exponential(Z1, 1, -X)
+
+    ds = xr.Dataset(
+        data_vars=dict(
+            var1=(["t", "x"], Z1), var2=(["t", "x"], Z2), var3=(["t", "x"], Z3)
+        ),
+        coords={"t": t, "x": x},
+    )
+    ds[["var2", "var3"]].curvefit(
+        coords=ds.var1,
+        func=exponential,
+        reduce_dims="t",
+        bounds={"a": (0.5, 5), "xc": (-5, 5)},
+    )
+
+We can also fit multi-dimensional functions, and even use a wrapper function to
+simultaneously fit a summation of several functions, such as this field containing
+two gaussian peaks:
+
+.. ipython:: python
+
+    def gaussian_2d(coords, a, xc, yc, xalpha, yalpha):
+        x, y = coords
+        z = a * np.exp(
+            -np.square(x - xc) / 2 / np.square(xalpha)
+            - np.square(y - yc) / 2 / np.square(yalpha)
+        )
+        return z
+
+
+    def multi_peak(coords, *args):
+        z = np.zeros(coords[0].shape)
+        for i in range(len(args) // 5):
+            z += gaussian_2d(coords, *args[i * 5 : i * 5 + 5])
+        return z
+
+
+    x = np.arange(-5, 5, 0.1)
+    y = np.arange(-5, 5, 0.1)
+    X, Y = np.meshgrid(x, y)
+
+    n_peaks = 2
+    names = ["a", "xc", "yc", "xalpha", "yalpha"]
+    names = [f"{name}{i}" for i in range(n_peaks) for name in names]
+    Z = gaussian_2d((X, Y), 3, 1, 1, 2, 1) + gaussian_2d((X, Y), 2, -1, -2, 1, 1)
+    Z += np.random.normal(scale=0.1, size=Z.shape)
+
+    da = xr.DataArray(Z, dims=["y", "x"], coords={"y": y, "x": x})
+    da.curvefit(
+        coords=["x", "y"],
+        func=multi_peak,
+        param_names=names,
+        kwargs={"maxfev": 10000},
+    )
+
+.. note::
+    This method replicates the behavior of :py:func:`scipy.optimize.curve_fit`.
+
+
 .. _compute.broadcasting:
 
 Broadcasting by dimension name
diff --git a/doc/whats-new.rst b/doc/whats-new.rst
index 408b59d3..5e394236 100644
--- a/doc/whats-new.rst
+++ b/doc/whats-new.rst
@@ -65,6 +65,8 @@ New Features
   :py:meth:`~pandas.core.groupby.GroupBy.get_group`.
   By `Deepak Cherian <https://github.com/dcherian>`_.
 - Disable the `cfgrib` backend if the `eccodes` library is not installed (:pull:`5083`). By `Baudouin Raoult <https://github.com/b8raoult>`_.
+- Added :py:meth:`DataArray.curvefit` and :py:meth:`Dataset.curvefit` for general curve fitting applications. (:issue:`4300`, :pull:`4849`)
+  By `Sam Levang <https://github.com/slevang>`_.
 
 Breaking changes
 ~~~~~~~~~~~~~~~~
diff --git a/xarray/core/combine.py b/xarray/core/combine.py
index 57324793..375931e1 100644
--- a/xarray/core/combine.py
+++ b/xarray/core/combine.py
@@ -44,6 +44,18 @@ def _infer_tile_ids_from_nested_list(entry, current_pos):
         yield current_pos, entry
 
 
+def _ensure_same_types(series, dim):
+
+    if series.dtype == object:
+        types = set(series.map(type))
+        if len(types) > 1:
+            types = ", ".join(t.__name__ for t in types)
+            raise TypeError(
+                f"Cannot combine along dimension '{dim}' with mixed types."
+                f" Found: {types}."
+            )
+
+
 def _infer_concat_order_from_coords(datasets):
 
     concat_dims = []
@@ -88,11 +100,15 @@ def _infer_concat_order_from_coords(datasets):
                     raise ValueError("Cannot handle size zero dimensions")
                 first_items = pd.Index([index[0] for index in indexes])
 
+                series = first_items.to_series()
+
+                # ensure series does not contain mixed types, e.g. cftime calendars
+                _ensure_same_types(series, dim)
+
                 # Sort datasets along dim
                 # We want rank but with identical elements given identical
                 # position indices - they should be concatenated along another
                 # dimension, not along this one
-                series = first_items.to_series()
                 rank = series.rank(
                     method="dense", ascending=ascending, numeric_only=False
                 )
diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py
index 1f82b48d..3b15144d 100644
--- a/xarray/core/dataarray.py
+++ b/xarray/core/dataarray.py
@@ -4418,6 +4418,84 @@ class DataArray(AbstractArray, DataWithCoords):
         )
         return ds[self.name]
 
+    def curvefit(
+        self,
+        coords: Union[Union[str, "DataArray"], Iterable[Union[str, "DataArray"]]],
+        func: Callable[..., Any],
+        reduce_dims: Union[Hashable, Iterable[Hashable]] = None,
+        skipna: bool = True,
+        p0: Dict[str, Any] = None,
+        bounds: Dict[str, Any] = None,
+        param_names: Sequence[str] = None,
+        kwargs: Dict[str, Any] = None,
+    ):
+        """
+        Curve fitting optimization for arbitrary functions.
+
+        Wraps `scipy.optimize.curve_fit` with `apply_ufunc`.
+
+        Parameters
+        ----------
+        coords : DataArray, str or sequence of DataArray, str
+            Independent coordinate(s) over which to perform the curve fitting. Must share
+            at least one dimension with the calling object. When fitting multi-dimensional
+            functions, supply `coords` as a sequence in the same order as arguments in
+            `func`. To fit along existing dimensions of the calling object, `coords` can
+            also be specified as a str or sequence of strs.
+        func : callable
+            User specified function in the form `f(x, *params)` which returns a numpy
+            array of length `len(x)`. `params` are the fittable parameters which are optimized
+            by scipy curve_fit. `x` can also be specified as a sequence containing multiple
+            coordinates, e.g. `f((x0, x1), *params)`.
+        reduce_dims : str or sequence of str
+            Additional dimension(s) over which to aggregate while fitting. For example,
+            calling `ds.curvefit(coords='time', reduce_dims=['lat', 'lon'], ...)` will
+            aggregate all lat and lon points and fit the specified function along the
+            time dimension.
+        skipna : bool, optional
+            Whether to skip missing values when fitting. Default is True.
+        p0 : dictionary, 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 : dictionary, 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.
+        param_names: seq, optional
+            Sequence of names for the fittable parameters of `func`. If not supplied,
+            this will be automatically determined by arguments of `func`. `param_names`
+            should be manually supplied when fitting a function that takes a variable
+            number of parameters.
+        kwargs : dictionary
+            Additional keyword arguments to passed to scipy curve_fit.
+
+        Returns
+        -------
+        curvefit_results : Dataset
+            A single dataset which contains:
+
+            [var]_curvefit_coefficients
+                The coefficients of the best fit.
+            [var]_curvefit_covariance
+                The covariance matrix of the coefficient estimates.
+
+        See also
+        --------
+        DataArray.polyfit
+        scipy.optimize.curve_fit
+        """
+        return self._to_temp_dataset().curvefit(
+            coords,
+            func,
+            reduce_dims=reduce_dims,
+            skipna=skipna,
+            p0=p0,
+            bounds=bounds,
+            param_names=param_names,
+            kwargs=kwargs,
+        )
+
     # this needs to be at the end, or mypy will confuse with `str`
     # https://mypy.readthedocs.io/en/latest/common_issues.html#dealing-with-conflicting-names
     str = utils.UncachedAccessor(StringAccessor)
diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py
index e959135e..8beacbf8 100644
--- a/xarray/core/dataset.py
+++ b/xarray/core/dataset.py
@@ -1,6 +1,7 @@
 import copy
 import datetime
 import functools
+import inspect
 import sys
 import warnings
 from collections import defaultdict
@@ -456,6 +457,63 @@ def as_dataset(obj: Any) -> "Dataset":
     return obj
 
 
+def _get_func_args(func, param_names):
+    """Use `inspect.signature` to try accessing `func` args. Otherwise, ensure
+    they are provided by user.
+    """
+    try:
+        func_args = inspect.signature(func).parameters
+    except ValueError:
+        func_args = {}
+        if not param_names:
+            raise ValueError(
+                "Unable to inspect `func` signature, and `param_names` was not provided."
+            )
+    if param_names:
+        params = param_names
+    else:
+        params = list(func_args)[1:]
+        if any(
+            [(p.kind in [p.VAR_POSITIONAL, p.VAR_KEYWORD]) for p in func_args.values()]
+        ):
+            raise ValueError(
+                "`param_names` must be provided because `func` takes variable length arguments."
+            )
+    return params, func_args
+
+
+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
+    """
+
+    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),
+            ]
+        )
+        return p0
+
+    param_defaults = {p: 1 for p in params}
+    bounds_defaults = {p: (-np.inf, np.inf) for p in params}
+    for p in params:
+        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])
+        if p in p0:
+            param_defaults[p] = p0[p]
+    return param_defaults, bounds_defaults
+
+
 class DataVariables(Mapping[Hashable, "DataArray"]):
     __slots__ = ("_dataset",)
 
@@ -7074,5 +7132,180 @@ class Dataset(Mapping, ImplementsDatasetReduce, DataWithCoords):
         # apply the selection
         return self.isel(indexers, missing_dims=missing_dims)
 
+    def curvefit(
+        self,
+        coords: Union[Union[str, "DataArray"], Iterable[Union[str, "DataArray"]]],
+        func: Callable[..., Any],
+        reduce_dims: Union[Hashable, Iterable[Hashable]] = None,
+        skipna: bool = True,
+        p0: Dict[str, Any] = None,
+        bounds: Dict[str, Any] = None,
+        param_names: Sequence[str] = None,
+        kwargs: Dict[str, Any] = None,
+    ):
+        """
+        Curve fitting optimization for arbitrary functions.
+
+        Wraps `scipy.optimize.curve_fit` with `apply_ufunc`.
+
+        Parameters
+        ----------
+        coords : DataArray, str or sequence of DataArray, str
+            Independent coordinate(s) over which to perform the curve fitting. Must share
+            at least one dimension with the calling object. When fitting multi-dimensional
+            functions, supply `coords` as a sequence in the same order as arguments in
+            `func`. To fit along existing dimensions of the calling object, `coords` can
+            also be specified as a str or sequence of strs.
+        func : callable
+            User specified function in the form `f(x, *params)` which returns a numpy
+            array of length `len(x)`. `params` are the fittable parameters which are optimized
+            by scipy curve_fit. `x` can also be specified as a sequence containing multiple
+            coordinates, e.g. `f((x0, x1), *params)`.
+        reduce_dims : str or sequence of str
+            Additional dimension(s) over which to aggregate while fitting. For example,
+            calling `ds.curvefit(coords='time', reduce_dims=['lat', 'lon'], ...)` will
+            aggregate all lat and lon points and fit the specified function along the
+            time dimension.
+        skipna : bool, optional
+            Whether to skip missing values when fitting. Default is True.
+        p0 : dictionary, 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 : dictionary, 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.
+        param_names: seq, optional
+            Sequence of names for the fittable parameters of `func`. If not supplied,
+            this will be automatically determined by arguments of `func`. `param_names`
+            should be manually supplied when fitting a function that takes a variable
+            number of parameters.
+        kwargs : dictionary
+            Additional keyword arguments to passed to scipy curve_fit.
+
+        Returns
+        -------
+        curvefit_results : Dataset
+            A single dataset which contains:
+
+            [var]_curvefit_coefficients
+                The coefficients of the best fit.
+            [var]_curvefit_covariance
+                The covariance matrix of the coefficient estimates.
+
+        See also
+        --------
+        Dataset.polyfit
+        scipy.optimize.curve_fit
+        """
+        from scipy.optimize import curve_fit
+
+        if p0 is None:
+            p0 = {}
+        if bounds is None:
+            bounds = {}
+        if kwargs is None:
+            kwargs = {}
+
+        if not reduce_dims:
+            reduce_dims_ = []
+        elif isinstance(reduce_dims, str) or not isinstance(reduce_dims, Iterable):
+            reduce_dims_ = [reduce_dims]
+        else:
+            reduce_dims_ = list(reduce_dims)
+
+        if (
+            isinstance(coords, str)
+            or isinstance(coords, xr.DataArray)
+            or not isinstance(coords, Iterable)
+        ):
+            coords = [coords]
+        coords_ = [self[coord] if isinstance(coord, str) else coord for coord in coords]
+
+        # Determine whether any coords are dims on self
+        for coord in coords_:
+            reduce_dims_ += [c for c in self.dims if coord.equals(self[c])]
+        reduce_dims_ = list(set(reduce_dims_))
+        preserved_dims = list(set(self.dims) - set(reduce_dims_))
+        if not reduce_dims_:
+            raise ValueError(
+                "No arguments to `coords` were identified as a dimension on the calling "
+                "object, and no dims were supplied to `reduce_dims`. This would result "
+                "in fitting on scalar data."
+            )
+
+        # Broadcast all coords with each other
+        coords_ = xr.broadcast(*coords_)
+        coords_ = [
+            coord.broadcast_like(self, exclude=preserved_dims) for coord in 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):
+            # Wrap curve_fit with raveled coordinates and pointwise NaN handling
+            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)
+                x = x[:, mask]
+                y = y[mask]
+                if not len(y):
+                    popt = np.full([n_params], np.nan)
+                    pcov = np.full([n_params, n_params], np.nan)
+                    return popt, pcov
+            x = np.squeeze(x)
+            popt, pcov = curve_fit(func, x, y, **kwargs)
+            return popt, pcov
+
+        result = xr.Dataset()
+        for name, da in self.data_vars.items():
+            if name is xr.core.dataarray._THIS_ARRAY:
+                name = ""
+            else:
+                name = f"{str(name)}_"
+
+            popt, pcov = xr.apply_ufunc(
+                _wrapper,
+                da,
+                *coords_,
+                vectorize=True,
+                dask="parallelized",
+                input_core_dims=[reduce_dims_ for d in range(len(coords_) + 1)],
+                output_core_dims=[["param"], ["cov_i", "cov_j"]],
+                dask_gufunc_kwargs={
+                    "output_sizes": {
+                        "param": n_params,
+                        "cov_i": n_params,
+                        "cov_j": n_params,
+                    },
+                },
+                output_dtypes=(np.float64, np.float64),
+                exclude_dims=set(reduce_dims_),
+                kwargs=kwargs,
+            )
+            result[name + "curvefit_coefficients"] = popt
+            result[name + "curvefit_covariance"] = pcov
+
+        result = result.assign_coords(
+            {"param": params, "cov_i": params, "cov_j": params}
+        )
+        result.attrs = self.attrs.copy()
+
+        return result
+
 
 ops.inject_all_ops_and_reduce_methods(Dataset, array_only=False)


Test patch:
diff --git a/xarray/tests/test_combine.py b/xarray/tests/test_combine.py
index 522b98cf..7664e5e5 100644
--- a/xarray/tests/test_combine.py
+++ b/xarray/tests/test_combine.py
@@ -1,4 +1,5 @@
 from datetime import datetime
+from distutils.version import LooseVersion
 from itertools import product
 
 import numpy as np
@@ -865,5 +866,22 @@ def test_combine_by_coords_raises_for_differing_calendars():
     da_1 = DataArray([0], dims=["time"], coords=[time_1], name="a").to_dataset()
     da_2 = DataArray([1], dims=["time"], coords=[time_2], name="a").to_dataset()
 
-    with raises_regex(TypeError, r"cannot compare .* \(different calendars\)"):
+    if LooseVersion(cftime.__version__) >= LooseVersion("1.5"):
+        error_msg = "Cannot combine along dimension 'time' with mixed types."
+    else:
+        error_msg = r"cannot compare .* \(different calendars\)"
+
+    with raises_regex(TypeError, error_msg):
+        combine_by_coords([da_1, da_2])
+
+
+def test_combine_by_coords_raises_for_differing_types():
+
+    # str and byte cannot be compared
+    da_1 = DataArray([0], dims=["time"], coords=[["a"]], name="a").to_dataset()
+    da_2 = DataArray([1], dims=["time"], coords=[[b"b"]], name="a").to_dataset()
+
+    with raises_regex(
+        TypeError, "Cannot combine along dimension 'time' with mixed types."
+    ):
         combine_by_coords([da_1, da_2])
diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py
index 20e35754..c18f9acf 100644
--- a/xarray/tests/test_dataarray.py
+++ b/xarray/tests/test_dataarray.py
@@ -4690,6 +4690,60 @@ class TestDataArray:
         with pytest.raises(UndefinedVariableError):
             aa.query(x="spam > 50")  # name not present
 
+    @requires_scipy
+    @pytest.mark.parametrize("use_dask", [True, False])
+    def test_curvefit(self, use_dask):
+        if use_dask and not has_dask:
+            pytest.skip("requires dask")
+
+        def exp_decay(t, n0, tau=1):
+            return n0 * np.exp(-t / tau)
+
+        t = np.arange(0, 5, 0.5)
+        da = DataArray(
+            np.stack([exp_decay(t, 3, 3), exp_decay(t, 5, 4), np.nan * t], axis=-1),
+            dims=("t", "x"),
+            coords={"t": t, "x": [0, 1, 2]},
+        )
+        da[0, 0] = np.nan
+
+        expected = DataArray(
+            [[3, 3], [5, 4], [np.nan, np.nan]],
+            dims=("x", "param"),
+            coords={"x": [0, 1, 2], "param": ["n0", "tau"]},
+        )
+
+        if use_dask:
+            da = da.chunk({"x": 1})
+
+        fit = da.curvefit(
+            coords=[da.t], func=exp_decay, p0={"n0": 4}, bounds={"tau": [2, 6]}
+        )
+        assert_allclose(fit.curvefit_coefficients, expected, rtol=1e-3)
+
+        da = da.compute()
+        fit = da.curvefit(coords="t", func=np.power, reduce_dims="x", param_names=["a"])
+        assert "a" in fit.param
+        assert "x" not in fit.dims
+
+    def test_curvefit_helpers(self):
+        def exp_decay(t, n0, tau=1):
+            return n0 * np.exp(-t / tau)
+
+        params, func_args = xr.core.dataset._get_func_args(exp_decay, [])
+        assert params == ["n0", "tau"]
+        param_defaults, bounds_defaults = xr.core.dataset._initialize_curvefit_params(
+            params, {"n0": 4}, {"tau": [5, np.inf]}, func_args
+        )
+        assert param_defaults == {"n0": 4, "tau": 6}
+        assert bounds_defaults == {"n0": (-np.inf, np.inf), "tau": (5, 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, [])
+
 
 class TestReduce:
     @pytest.fixture(autouse=True)