Source code for specutils.fitting.continuum

import astropy.units as u
from astropy.modeling.polynomial import Chebyshev1D
from astropy.modeling.fitting import LevMarLSQFitter

from ..fitting import fit_lines
from ..manipulation.smoothing import median_smooth
from ..spectra import SpectralRegion


__all__ = ['fit_continuum', 'fit_generic_continuum']




def fit_generic_continuum(spectrum, median_window=3, model=Chebyshev1D(3),
                          fitter=LevMarLSQFitter(),
                          exclude_regions=None, weights=None):
    """
    Basic fitting of the continuum of an input spectrum. The input
    spectrum is smoothed using a median filter to remove the spikes.

    Parameters
    ----------
    spectrum : Spectrum1D
        The spectrum object overwhich the equivalent width will be calculated.
    model : list of `~astropy.modeling.Model`
        The list of models that contain the initial guess.
    median_window : float
        The width of the median smoothing kernel used to filter the data before
        fitting the continuum. See the ``kernel_size`` parameter of
        `~scipy.signal.medfilt` for more information.
    fitter : `~astropy.fitting._FitterMeta`
        The astropy fitter to use for fitting the model.
        Default: `~astropy.modeling.fitting.LevMarLSQFitter`
    exclude_regions : list of 2-tuples
        List of regions to exclude in the fitting. Passed through
        to the fitmodels routine.
    weights : list  (NOT IMPLEMENTED YET)
        List of weights to define importance of fitting regions.

    Returns
    -------
    continuum_model
        Fitted continuum as a model of whatever class ``model`` provides.

    Notes
    -----
       * Could add functionality to set the bounds in
         ``model`` if they are not set.
       * The models in the list of ``model`` are added together and passed as a
         compound model to the `~astropy.modeling.fitting.Fitter` class instance.

    """
    # Simple median smooth to remove spikes and peaks
    spectrum_smoothed = median_smooth(spectrum, median_window)

    return fit_continuum(spectrum_smoothed, model=model, fitter=fitter,
                         exclude_regions=exclude_regions, weights=weights)





def fit_continuum(spectrum, model=Chebyshev1D(3), fitter=LevMarLSQFitter(),
                  exclude_regions=None, window=None, weights=None):
    """
    Entry point for fitting using the `~astropy.modeling.fitting`
    machinery.

    Parameters
    ----------
    spectrum : Spectrum1D
        The spectrum object overwhich the equivalent width will be calculated.
    model: list of `~astropy.modeling.Model`
        The list of models that contain the initial guess.
    fitter : `~astropy.fitting._FitterMeta`
        The astropy fitter to use for fitting the model.
        Default: `~astropy.modeling.fitting.LevMarLSQFitter`
    exclude_regions : list of 2-tuples
        List of regions to exclude in the fitting. Passed through
        to the fitmodels routine.
    window : tuple of wavelengths
        Start and end wavelengths used for fitting.
    weights : list  (NOT IMPLEMENTED YET)
        List of weights to define importance of fitting regions.

    Returns
    -------
    models : list of `~astropy.modeling.Model`
        The list of models that contain the fitted model parameters.
    """
    if weights is not None:
        raise NotImplementedError("Weights support is not yet implemented.")

    w = window
    if type(w) in [list, tuple] and all([_is_valid_sequence(x) for x in w]):
        w = SpectralRegion(w)

    # Fit the flux to the model
    continuum_spectrum = fit_lines(spectrum, model, fitter, exclude_regions,
                                   weights, w)

    return continuum_spectrum



# Checks for sequences of of 2-tuples with Quantities
def _is_valid_sequence(value):
    if type(value) in [list, tuple]:
        return len(value) == 2 and \
               isinstance(value[0], u.Quantity) and \
               isinstance(value[1], u.Quantity)

    return False