added filter_non_finite flag to fit_lines

specutils/fitting/fitmodels.py

@@ -259,7 +259,8 @@ def _generate_line_list_table(spectrum, emission_inds, absorption_inds):
 
 
 def fit_lines(spectrum, model, fitter=fitting.LevMarLSQFitter(calc_uncertainties=True),
-              exclude_regions=None, weights=None, window=None, get_fit_info=False,
+              exclude_regions=None, weights=None, window=None, filter_non_finite=False,
+              get_fit_info=False,
               **kwargs):
     """
     Fit the input models to the spectrum. The parameter values of the
@@ -288,6 +289,9 @@ def fit_lines(spectrum, model, fitter=fitting.LevMarLSQFitter(calc_uncertainties
         Flag to return the ``fit_info`` from the underlying scipy optimizer used
         in the fitting. If True, the returned model will have a ``fit_info``
         key populated in its ``meta`` dictionary.
+    filter_non_finite : bool, optional
+        Flag to filter non-finite value in flux or weight arrrays. Can
+        only be used if fitter is a `LevMarLSQFitter`.
     Additional keyword arguments are passed directly into the call to the
     ``fitter``.
 
@@ -361,6 +365,7 @@ def fit_lines(spectrum, model, fitter=fitting.LevMarLSQFitter(calc_uncertainties
             window=model_window,
             get_fit_info=get_fit_info,
             ignore_units=ignore_units,
+            filter_non_finite=filter_non_finite,
             **kwargs
         )
 
@@ -377,7 +382,7 @@ def fit_lines(spectrum, model, fitter=fitting.LevMarLSQFitter(calc_uncertainties
 
 def _fit_lines(spectrum, model, fitter=fitting.LevMarLSQFitter(calc_uncertainties=True),
                exclude_regions=None, weights=None, window=None, get_fit_info=False,
-               ignore_units=False, **kwargs):
+               ignore_units=False, filter_non_finite=False, **kwargs):
     """
     Fit the input model (initial conditions) to the spectrum.  Output will be
     the same model with the parameters set based on the fitting.
@@ -503,8 +508,19 @@ def _fit_lines(spectrum, model, fitter=fitting.LevMarLSQFitter(calc_uncertaintie
         if weights is not None:
             weights = weights[nmask]
 
+    if filter_non_finite is True:
+        if isinstance(fitter, fitting.LevMarLSQFitter):  # fine, just checking for clarity
+            pass 
+        else:
+            warnings.warn('`filter_non_finite` can only be True when fitter is LevMarLSQFitter. Setting to False.')
+            filter_non_finite = False
+
+    print('FILTER NON FINITE', filter_non_finite)
+    print(flux)
+    print(weights)
+
     fit_model = fitter(model, dispersion,
-                       flux, weights=weights, **kwargs)
+                       flux, weights=weights, filter_non_finite=filter_non_finite, **kwargs)
 
     if hasattr(fitter, 'fit_info') and get_fit_info:
         fit_model.meta['fit_info'] = fitter.fit_info

specutils/tests/test_fitting.py

@@ -2,6 +2,7 @@
 import numpy as np
 import pytest
 from astropy.modeling import models
+from astropy.modeling.fitting import NonFiniteValueError
 from astropy.nddata import StdDevUncertainty
 from astropy.tests.helper import assert_quantity_allclose
 from astropy.utils.exceptions import AstropyUserWarning
@@ -705,3 +706,30 @@ def test_fit_subspectrum():
     with pytest.warns(AstropyUserWarning, match='The fit may be unsuccessful'):
         g_fit = fit_lines(sub_spectrum, g_init)
     yfit = g_fit(sub_spectrum.spectral_axis)  # noqa
+
+def test_fitting_filter_nonfinite():
+
+    # Test that by default, fit_lines has filter_non_finite=False
+    # and a spectrum with nonfinite values in weights will fail
+    # fit. Then test that if filter_non_finite=False, the fit will
+    # work.
+
+    # Create a simple spectrum with a Gaussian.
+    np.random.seed(0)
+    x = np.linspace(0., 10., 20)
+    y = 3 * np.exp(-0.5 * (x - 6.3) ** 2 / 0.8 ** 2)
+
+    uncerts = np.ones_like(x)
+    uncerts[0:5] = np.nan
+    uncerts = StdDevUncertainty(uncerts)
+
+    spectrum = Spectrum1D(flux=y*u.Jy, spectral_axis=x*u.um, uncertainty=uncerts)
+
+    with pytest.raises(NonFiniteValueError):  # test that it fails first
+        g_fit = fit_lines(spectrum, models.Gaussian1D(), weights='unc')
+        y_fit = g_fit(x*u.um)
+
+    # now test that the fit works when filter_non_finite=True
+    with pytest.warns(AstropyUserWarning, match='Non-Finite input data has been removed by the fitter.'):
+        g_fit = fit_lines(spectrum, models.Gaussian1D(), weights='unc', filter_non_finite=True)
+    y_fit = g_fit(x*u.um)