Repository: scipy/scipy
Pull request: ENH: optimize: Add `nan_policy` optional argument for `curve_fit`.
<!-- 
Thanks for contributing a pull request! Please ensure that
your PR satisfies the checklist before submitting:
http://scipy.github.io/devdocs/dev/contributor/development_workflow.html#checklist-before-submitting-a-pr

Also, please name and describe your PR as you would write a
commit message:
http://scipy.github.io/devdocs/dev/contributor/development_workflow.html#writing-the-commit-message

Note that we are a team of volunteers; we appreciate your
patience during the review process.

Again, thanks for contributing!
-->

#### Reference issue
<!--Example: Closes gh-WXYZ.-->
Fix #11841
Ref #11846 (A staled PR to fix #11841)

#### What does this implement/fix?
<!--Please explain your changes.-->
As reported in #11841 and #11846, some users would like to add a nan handling functionality to ignore nans in input data for `optimize.curve_fit`.
This PR adds `nan_policy` optional argument for [`curve_fit`](http://scipy.github.io/devdocs/reference/generated/scipy.optimize.curve_fit.html?highlight=curve_fit#scipy.optimize.curve_fit) based on our [nan_policy guideline](http://scipy.github.io/devdocs/dev/api-dev/nan_policy.html).
If `omit` nan_policy is used, the nans in input data will be ignored.

#### Additional information
<!--Any additional information you think is important.-->


Problem statement:
Ignore NaN with scipy.optimize.curve_fit
Occasionally I need to fit data that's incomplete (missing values) or perform a fitting task over a grid of 3-D data, i.e. fitting M points for every NxN element, where one of the elements needs to be masked out or ignored.

It may be a good feature to have a flag (or flags) to curve_fit that tell it to ignore NaNs in an array being fit and return NaN when there are no floats/data to fit (i.e. all nans).

I've seen documented that when you drop the check for nan/infs in curve_fit you can silently get nonsensical results. It seems a flag here would provide a way to define the behavior you'd like when nans/inf are present.

Hints from issue discussion:
Makes sense, there have been issues with this before IIRC. Also there's gh-4777 which asks for nan/inf handling in `leastsq`.
@rgommers I was specifically thinking about NaNs. I only mentioned infs because they’re called out in the documentation I read about the check for flag.

Commit message: ENH: optimize: remove for loop code based on review.

Patch to review:
diff --git a/scipy/optimize/_minpack_py.py b/scipy/optimize/_minpack_py.py
index 205acce8833f..49fc40948782 100644
--- a/scipy/optimize/_minpack_py.py
+++ b/scipy/optimize/_minpack_py.py
@@ -8,7 +8,7 @@
                    finfo, inexact, issubdtype, dtype)
 from scipy import linalg
 from scipy.linalg import svd, cholesky, solve_triangular, LinAlgError, inv
-from scipy._lib._util import _asarray_validated, _lazywhere
+from scipy._lib._util import _asarray_validated, _lazywhere, _contains_nan
 from scipy._lib._util import getfullargspec_no_self as _getfullargspec
 from ._optimize import OptimizeResult, _check_unknown_options, OptimizeWarning
 from ._lsq import least_squares
@@ -567,8 +567,9 @@ def _initialize_feasible(lb, ub):
 
 
 def curve_fit(f, xdata, ydata, p0=None, sigma=None, absolute_sigma=False,
-              check_finite=True, bounds=(-np.inf, np.inf), method=None,
-              jac=None, *, full_output=False, **kwargs):
+              check_finite=None, bounds=(-np.inf, np.inf), method=None,
+              jac=None, *, full_output=False, nan_policy=None,
+              **kwargs):
     """
     Use non-linear least squares to fit a function, f, to data.
 
@@ -624,7 +625,8 @@ def curve_fit(f, xdata, ydata, p0=None, sigma=None, absolute_sigma=False,
         If True, check that the input arrays do not contain nans of infs,
         and raise a ValueError if they do. Setting this parameter to
         False may silently produce nonsensical results if the input arrays
-        do contain nans. Default is True.
+        do contain nans. Default is True. Note that if `nan_policy` is
+        specified explicitly (not None), this value will be ignored.
     bounds : 2-tuple of array_like or `Bounds`, optional
         Lower and upper bounds on parameters. Defaults to no bounds.
         There are two ways to specify the bounds:
@@ -659,6 +661,20 @@ def curve_fit(f, xdata, ydata, p0=None, sigma=None, absolute_sigma=False,
         `mesg`, and `ier`.
 
         .. versionadded:: 1.9
+    nan_policy : {'raise', 'omit', None}, optional
+        Defines how to handle when input contains nan.
+        The following options are available (default is None):
+
+          * 'raise': throws an error
+          * 'omit': performs the calculations ignoring nan values
+          * None: no special handling of NaNs is performed
+            (except what is done by check_finite); the behavior when NaNs
+            are present is implementation-dependent and may change.
+
+        Note that if this value is specified explicitly (not None),
+        `check_finite` will be set as False.
+
+        .. versionadded:: 1.11
     **kwargs
         Keyword arguments passed to `leastsq` for ``method='lm'`` or
         `least_squares` otherwise.
@@ -851,9 +867,10 @@ def curve_fit(f, xdata, ydata, p0=None, sigma=None, absolute_sigma=False,
         raise ValueError("Method 'lm' only works for unconstrained problems. "
                          "Use 'trf' or 'dogbox' instead.")
 
-    # optimization may produce garbage for float32 inputs, cast them to float64
+    if check_finite is None:
+        check_finite = True if nan_policy is None else False
 
-    # NaNs cannot be handled
+    # optimization may produce garbage for float32 inputs, cast them to float64
     if check_finite:
         ydata = np.asarray_chkfinite(ydata, float)
     else:
@@ -870,6 +887,25 @@ def curve_fit(f, xdata, ydata, p0=None, sigma=None, absolute_sigma=False,
     if ydata.size == 0:
         raise ValueError("`ydata` must not be empty!")
 
+    # nan handling is needed only if check_finite is False because if True,
+    # the x-y data are already checked, and they don't contain nans.
+    if not check_finite and nan_policy is not None:
+        if nan_policy == "propagate":
+            raise ValueError("`nan_policy='propagate'` is not supported "
+                             "by this function.")
+
+        x_contains_nan, nan_policy = _contains_nan(xdata, nan_policy)
+        y_contains_nan, nan_policy = _contains_nan(ydata, nan_policy)
+
+        if (x_contains_nan or y_contains_nan) and nan_policy == 'omit':
+            # ignore NaNs for N dimensional arrays
+            has_nan = np.isnan(xdata)
+            has_nan = has_nan.any(axis=tuple(range(has_nan.ndim-1)))
+            has_nan |= np.isnan(ydata)
+
+            xdata = xdata[..., ~has_nan]
+            ydata = ydata[~has_nan]
+
     # Determine type of sigma
     if sigma is not None:
         sigma = np.asarray(sigma)
@@ -941,7 +977,7 @@ def curve_fit(f, xdata, ydata, p0=None, sigma=None, absolute_sigma=False,
         pcov = np.dot(VT.T / s**2, VT)
 
     warn_cov = False
-    if pcov is None:
+    if pcov is None or np.isnan(pcov).any():
         # indeterminate covariance
         pcov = zeros((len(popt), len(popt)), dtype=float)
         pcov.fill(inf)
diff --git a/scipy/optimize/tests/test_minpack.py b/scipy/optimize/tests/test_minpack.py
index c6155f3fee9b..cf025a139e44 100644
--- a/scipy/optimize/tests/test_minpack.py
+++ b/scipy/optimize/tests/test_minpack.py
@@ -551,6 +551,73 @@ def test_NaN_handling(self):
         assert_raises(ValueError, curve_fit, lambda x, a, b: a*x + b,
                       xdata, ydata, **{"check_finite": True})
 
+    @staticmethod
+    def _check_nan_policy(f, xdata_with_nan, xdata_without_nan,
+                          ydata_with_nan, ydata_without_nan, method):
+        kwargs = {'f': f, 'xdata': xdata_with_nan, 'ydata': ydata_with_nan,
+                  'method': method, 'check_finite': False}
+        # propagate test
+        error_msg = ("`nan_policy='propagate'` is not supported "
+                     "by this function.")
+        with assert_raises(ValueError, match=error_msg):
+            curve_fit(**kwargs, nan_policy="propagate", maxfev=2000)
+
+        # raise test
+        with assert_raises(ValueError, match="The input contains nan"):
+            curve_fit(**kwargs, nan_policy="raise")
+
+        # omit test
+        result_with_nan, _ = curve_fit(**kwargs, nan_policy="omit")
+        result_without_nan, _ = curve_fit(**kwargs, nan_policy="omit")
+        assert_allclose(result_with_nan, result_without_nan)
+
+    @pytest.mark.parametrize('method', ["lm", "trf", "dogbox"])
+    def test_nan_policy_1d(self, method):
+        def f(x, a, b):
+            return a*x + b
+
+        xdata_with_nan = np.array([2, 3, np.nan, 4, 4, np.nan])
+        ydata_with_nan = np.array([1, 2, 5, 3, np.nan, 7])
+        xdata_without_nan = np.array([2, 3, 4])
+        ydata_without_nan = np.array([1, 2, 3])
+
+        self._check_nan_policy(f, xdata_with_nan, xdata_without_nan,
+                               ydata_with_nan, ydata_without_nan, method)
+
+    @pytest.mark.parametrize('method', ["lm", "trf", "dogbox"])
+    def test_nan_policy_2d(self, method):
+        def f(x, a, b):
+            x1 = x[0, :]
+            x2 = x[1, :]
+            return a*x1 + b + x2
+
+        xdata_with_nan = np.array([[2, 3, np.nan, 4, 4, np.nan, 5],
+                                   [2, 3, np.nan, np.nan, 4, np.nan, 7]])
+        ydata_with_nan = np.array([1, 2, 5, 3, np.nan, 7, 10])
+        xdata_without_nan = np.array([[2, 3, 5], [2, 3, 7]])
+        ydata_without_nan = np.array([1, 2, 10])
+
+        self._check_nan_policy(f, xdata_with_nan, xdata_without_nan,
+                               ydata_with_nan, ydata_without_nan, method)
+
+    @pytest.mark.parametrize('n', [2, 3])
+    @pytest.mark.parametrize('method', ["lm", "trf", "dogbox"])
+    def test_nan_policy_2_3d(self, n, method):
+        def f(x, a, b):
+            x1 = x[..., 0, :].squeeze()
+            x2 = x[..., 1, :].squeeze()
+            return a*x1 + b + x2
+
+        xdata_with_nan = np.array([[[2, 3, np.nan, 4, 4, np.nan, 5],
+                                   [2, 3, np.nan, np.nan, 4, np.nan, 7]]])
+        xdata_with_nan = xdata_with_nan.squeeze() if n == 2 else xdata_with_nan
+        ydata_with_nan = np.array([1, 2, 5, 3, np.nan, 7, 10])
+        xdata_without_nan = np.array([[[2, 3, 5], [2, 3, 7]]])
+        ydata_without_nan = np.array([1, 2, 10])
+
+        self._check_nan_policy(f, xdata_with_nan, xdata_without_nan,
+                               ydata_with_nan, ydata_without_nan, method)
+
     def test_empty_inputs(self):
         # Test both with and without bounds (regression test for gh-9864)
         assert_raises(ValueError, curve_fit, lambda x, a: a*x, [], [])


Code review comments:
diff_hunk: @@ -551,6 +551,73 @@ def test_NaN_handling(self):
         assert_raises(ValueError, curve_fit, lambda x, a, b: a*x + b,
                       xdata, ydata, **{"check_finite": True})
 
+    @staticmethod
+    def _check_nan_policy(f, xdata_with_nan, xdata_without_nan,
+                          ydata_with_nan, ydata_without_nan, method):
+        kwargs = {'f': f, 'xdata': xdata_with_nan, 'ydata': ydata_with_nan,
+                  'method': method, 'check_finite': False}
+        # propagate test
+        error_msg = ("`nan_policy='propagate'` is not supported "
+                     "by this function.")
+        with assert_raises(ValueError, match=error_msg):
+            curve_fit(**kwargs, nan_policy="propagate", maxfev=2000)
+
+        # raise test
+        with assert_raises(ValueError, match="The input contains nan"):
+            curve_fit(**kwargs, nan_policy="raise")
+
+        # omit test
+        result_with_nan, _ = curve_fit(**kwargs, nan_policy="omit")
+        result_without_nan, _ = curve_fit(**kwargs, nan_policy="omit")
path: scipy/optimize/tests/test_minpack.py
text: @user1:
I think you forgot to use `xdata_without_nan` and ``ydata_without_nan`` here. Shouldn't it be:

```suggestion
        result_with_nan, _ = curve_fit(**kwargs, nan_policy="omit")
        kwargs['xdata'] = xdata_without_nan
        kwargs['ydata'] = ydata_without_nan
        result_without_nan, _ = curve_fit(**kwargs)
```

? Otherwise, the test doesn't really test anything...

@user1:
Also, we can avoid passing `nan_policy='omit'` argument in the second call to `curve_fit`