Repository: scipy/scipy
Pull request: BUG: `optimize.curve_fit`: fix sigma dimension issue with `nan_policy="omit"`
#### Reference issue
Closes #21772.

#### What does this implement/fix?

Fixes the bug in optimize.curve_fit where using nan_policy="omit" did not remove corresponding entries of sigma.
In this PR, I fix this bug by checking for the dimension of sigma and removing the entries where either xdata or ydata is NaN.

#### Additional information
<!--Any additional information you think is important.-->


Problem statement:
BUG: `optimize.curve_fit` with `nan_policy="omit"` fails to handle sigma properly
### Describe your issue.

Hello,
I have encountered an issue with curve_fit. It is bypassable, but I'd say it is unexpected behaviour. Generally speaking, it occurs when both of the following is true:
-- I have xdata or ydata with NaNs and use nan_policy = "omit" to handle them;
-- I have specified sigma.
In this case, I get ValueError "Sigma has incorrect shape" in optimize._minpack_py.py despite sigma being of the same size as xdata and ydata.

I checked what is going in _minpack_py.py and figured out the following.
It occurs because "nan_policy = omit" leads to the removal of NaN indices from xdata and ydata (lines 960-967), but not from sigma. Then, after line 969, the shape of sigma is checked by comparing it with ydata. And even if sigma and ydata were originally of the same size, they end up to be different after ydata truncation, throwing an error.

So, to bypass this problem, one has either to remove indices corresponding NaNs in xdata and ydata from sigma, or not using nan_policy = "omit" at all and remove NaN indices manually. I suggest improving nan_policy = "omit" in a way that it would remove the corresponding indices from sigma, if needed.

Sorry if my bug report is incorrect, I am first time doing this.

### Reproducing Code Example

```python
import numpy as np
from scipy.optimize import curve_fit

def exponential(x, a, b):
	return b * np.exp(a * x)

np.random.seed(0)
N = 90
x = np.linspace(1, 10, N, endpoint = False)
sigma = x * 0.05
y = exponential(x, 0.2, 0.5) + np.random.normal(0, sigma, N)

p0 = [0.1, 1.0]

# without nans, works perfect
popt, pcov = curve_fit(exponential, x, y, p0 = p0, sigma = sigma)
perr = np.diag(np.sqrt(pcov))
print("without nans:")
print(f"popt: {popt}")
print(f"perr: {perr}")

# with nans, have to set nan_policy = "omit"...
x[11] = np.nan
y[22] = np.nan
popt, pcov = curve_fit(exponential, x, y, p0 = p0, nan_policy = "omit")
perr = np.diag(np.sqrt(pcov))
print("with nans:")
print(f"popt: {popt}")
print(f"perr: {perr}")

#...but when using sigma, it says "sigma has incorrect shape"
x[11] = np.nan
y[22] = np.nan
popt, pcov = curve_fit(exponential, x, y, p0 = p0, sigma = sigma, nan_policy = "omit")
```


### Error message

```shell
Traceback (most recent call last):
  File "/root/Science/Spiral_shape_analysis/fit_results/../spiralcutter_advanced/scipy_test.py", line 34, in <module>
    popt, pcov = curve_fit(exponential, x, y, p0 = p0, sigma = sigma, nan_policy = "omit")
  File "/usr/local/lib/python3.10/dist-packages/scipy/optimize/_minpack_py.py", line 985, in curve_fit
    raise ValueError("`sigma` has incorrect shape.")
ValueError: `sigma` has incorrect shape.
```


### SciPy/NumPy/Python version and system information

```shell
>>> print(scipy.__version__, numpy.__version__, sys.version_info)
1.14.1 2.1.2 sys.version_info(major=3, minor=10, micro=12, releaselevel='final', serial=0)
>>> scipy.show_config()
Build Dependencies:
  blas:
    detection method: pkgconfig
    found: true
    include directory: /opt/_internal/cpython-3.10.14/lib/python3.10/site-packages/scipy_openblas32/include
    lib directory: /opt/_internal/cpython-3.10.14/lib/python3.10/site-packages/scipy_openblas32/lib
    name: scipy-openblas
    openblas configuration: OpenBLAS 0.3.27.dev DYNAMIC_ARCH NO_AFFINITY Zen MAX_THREADS=64
    pc file directory: /project
    version: 0.3.27.dev
  lapack:
    detection method: pkgconfig
    found: true
    include directory: /opt/_internal/cpython-3.10.14/lib/python3.10/site-packages/scipy_openblas32/include
    lib directory: /opt/_internal/cpython-3.10.14/lib/python3.10/site-packages/scipy_openblas32/lib
    name: scipy-openblas
    openblas configuration: OpenBLAS 0.3.27.dev DYNAMIC_ARCH NO_AFFINITY Zen MAX_THREADS=64
    pc file directory: /project
    version: 0.3.27.dev
  pybind11:
    detection method: config-tool
    include directory: unknown
    name: pybind11
    version: 2.12.0
Compilers:
  c:
    commands: cc
    linker: ld.bfd
    name: gcc
    version: 10.2.1
  c++:
    commands: c++
    linker: ld.bfd
    name: gcc
    version: 10.2.1
  cython:
    commands: cython
    linker: cython
    name: cython
    version: 3.0.11
  fortran:
    commands: gfortran
    linker: ld.bfd
    name: gcc
    version: 10.2.1
  pythran:
    include directory: ../../tmp/pip-build-env-h_xz8lfs/overlay/lib/python3.10/site-packages/pythran
    version: 0.16.1
Machine Information:
  build:
    cpu: x86_64
    endian: little
    family: x86_64
    system: linux
  cross-compiled: false
  host:
    cpu: x86_64
    endian: little
    family: x86_64
    system: linux
Python Information:
  path: /opt/python/cp310-cp310/bin/python
  version: '3.10'
```

Hints from issue discussion:
Oops, thanks for reporting!

@AtsushiSakai Looks like we missed this. Do you have a moment to submit a PR?

@Vallastro Just to confirm, this would also be a problem when `sigma` is 2D. In that case, we would remove both the rows and columns where there are NaNs in either `x` or `y`?
> Just to confirm, this would also be a problem when `sigma` is 2D. In that case, we would remove both the rows and columns where there are NaNs in either `x` or `y`?

I guess, yes. In this case, if `ydata[i]` is invalidated, then for any `j` covariation between `xdata[i]` and `xdata[j]`, represented in `sigma[i, j]` and `sigma[j, i]` has no sense. And if `N` elements of `ydata` out of original `M` are remaining after removing NaNs, then removing columns and rows from `sigma` will make it `N` x `N` from `M` x `M`, which is expected.
@mdhaber Is this issue accepting PRs? If so, could you please point me to the files that need to be changed? I would love to contribute.

Commit message: Added a regression test for omit nan

Patch to review:
diff --git a/scipy/optimize/_minpack_py.py b/scipy/optimize/_minpack_py.py
index 8c0564c06bbe..5927fe2c45cd 100644
--- a/scipy/optimize/_minpack_py.py
+++ b/scipy/optimize/_minpack_py.py
@@ -961,9 +961,20 @@ def curve_fit(f, xdata, ydata, p0=None, sigma=None, absolute_sigma=False,
             has_nan = has_nan.any(axis=tuple(range(has_nan.ndim-1)))
             has_nan |= np.isnan(ydata)
 
+            # Store the size of ydata before purging the nan values
+            prev_ydata_size = ydata.size
+
             xdata = xdata[..., ~has_nan]
             ydata = ydata[~has_nan]
 
+            # Also omit the corresponding entries form sigma
+            if sigma is not None:
+                sigma = np.asarray(sigma)
+                if sigma.shape == (prev_ydata_size, ):
+                    sigma = sigma[~has_nan]
+                if sigma.shape == (prev_ydata_size, prev_ydata_size):
+                    sigma = sigma[~has_nan, ~has_nan]
+
     # Determine type of sigma
     if sigma is not None:
         sigma = np.asarray(sigma)
diff --git a/scipy/optimize/tests/test_minpack.py b/scipy/optimize/tests/test_minpack.py
index d771308f069f..dc4131e2cf70 100644
--- a/scipy/optimize/tests/test_minpack.py
+++ b/scipy/optimize/tests/test_minpack.py
@@ -804,6 +804,27 @@ def test_maxfev_and_bounds(self):
         assert_allclose(popt1, 2, atol=1e-14)
         assert_allclose(popt2, 2, atol=1e-14)
 
+    def test_curvefit_omitnan(self):
+        def exponential(x, a, b):
+            return b * np.exp(a * x)
+
+        np.random.seed(0)
+        N = 100
+        x = np.linspace(1, 10, N)
+        sigma = x * 0.05
+        y = exponential(x, 0.2, 0.5) + np.random.normal(0, sigma, N)
+
+        p0 = [0.1, 1.0]
+
+        popt1, pcov1 = curve_fit(exponential, x, y, p0 = p0, sigma = sigma)
+
+        x[11] = np.nan
+        y[22] = np.nan
+        popt2, pcov2 = curve_fit(exponential, x, y, p0 = p0, sigma = sigma, nan_policy = "omit")
+
+        assert_allclose(popt1, popt2, atol=1e-2)
+        assert_allclose(pcov1, pcov2, atol=1e-2)
+
     def test_curvefit_simplecovariance(self):
 
         def func(x, a, b):


Code review comments:
diff_hunk: @@ -804,6 +804,27 @@ def test_maxfev_and_bounds(self):
         assert_allclose(popt1, 2, atol=1e-14)
         assert_allclose(popt2, 2, atol=1e-14)
 
+    def test_curvefit_omitnan(self):
+        def exponential(x, a, b):
+            return b * np.exp(a * x)
+
+        np.random.seed(0)
path: scipy/optimize/tests/test_minpack.py
text: @user1:
We should avoid setting the global random state and instead create and use a generator 
```suggestion
        rng = np.random.default_rng(578285731148908)
```

@author:
Thank you for pointing this out. I didn't know this.

diff_hunk: @@ -961,9 +961,20 @@ def curve_fit(f, xdata, ydata, p0=None, sigma=None, absolute_sigma=False,
             has_nan = has_nan.any(axis=tuple(range(has_nan.ndim-1)))
             has_nan |= np.isnan(ydata)
 
+            # Store the size of ydata before purging the nan values
+            prev_ydata_size = ydata.size
+
             xdata = xdata[..., ~has_nan]
             ydata = ydata[~has_nan]
 
+            # Also omit the corresponding entries form sigma
path: scipy/optimize/_minpack_py.py
text: @user1:
```suggestion
            # Also omit the corresponding entries from sigma
```

@author:
Oops!

diff_hunk: @@ -961,9 +961,20 @@ def curve_fit(f, xdata, ydata, p0=None, sigma=None, absolute_sigma=False,
             has_nan = has_nan.any(axis=tuple(range(has_nan.ndim-1)))
             has_nan |= np.isnan(ydata)
 
+            # Store the size of ydata before purging the nan values
+            prev_ydata_size = ydata.size
+
             xdata = xdata[..., ~has_nan]
             ydata = ydata[~has_nan]
 
+            # Also omit the corresponding entries form sigma
+            if sigma is not None:
+                sigma = np.asarray(sigma)
+                if sigma.shape == (prev_ydata_size, ):
path: scipy/optimize/_minpack_py.py
text: @user1:
```suggestion
                if sigma.shape == (prev_ydata_size,):
```

diff_hunk: @@ -804,6 +804,27 @@ def test_maxfev_and_bounds(self):
         assert_allclose(popt1, 2, atol=1e-14)
         assert_allclose(popt2, 2, atol=1e-14)
 
+    def test_curvefit_omitnan(self):
+        def exponential(x, a, b):
+            return b * np.exp(a * x)
+
+        np.random.seed(0)
+        N = 100
+        x = np.linspace(1, 10, N)
+        sigma = x * 0.05
+        y = exponential(x, 0.2, 0.5) + np.random.normal(0, sigma, N)
+
+        p0 = [0.1, 1.0]
+
+        popt1, pcov1 = curve_fit(exponential, x, y, p0 = p0, sigma = sigma)
+
+        x[11] = np.nan
+        y[22] = np.nan
+        popt2, pcov2 = curve_fit(exponential, x, y, p0 = p0, sigma = sigma, nan_policy = "omit")
path: scipy/optimize/tests/test_minpack.py
text: @user1:
```suggestion
        popt1, pcov1 = curve_fit(exponential, x, y, p0=p0, sigma=sigma)

        x[11] = np.nan
        y[22] = np.nan
        popt2, pcov2 = curve_fit(exponential, x, y, p0=p0, sigma=sigma, nan_policy="omit")
```