Merge pull request #593 from nmearl/model-initialization

Introduce simple model initializations

specviz/plugins/model_editor/initializers.py

@@ -0,0 +1,302 @@
+"""
+This module is used to initialize spectral models to the data at hand.
+
+This is used by model-fitting code that has to create spectral model
+instances with sensible parameter values such that they can be used as
+first guesses by the fitting algorithms.
+"""
+import numpy as np
+
+__all__ = [
+    'initialize'
+]
+
+AMPLITUDE = 'amplitude'
+POSITION  = 'position'
+WIDTH     = 'width'
+
+
+def _get_model_name(model):
+    class_string = str(model.__class__)
+    return class_string.split('\'>')[0].split(".")[-1]
+
+
+class _Linear1DInitializer(object):
+    """
+    Initialization that is specific to the Linear1D model.
+
+    Notes
+    -----
+    In a way, we need this specialized initializer because
+    the linear 1D model is more like a kind of polynomial.
+    It doesn't mesh well with other non-linear models.
+    """
+    def initialize(self, instance, x, y):
+        """
+        Initialize the model
+
+        Parameters
+        ----------
+        instance: `~astropy.modeling.models`
+            The model to initialize.
+
+        x, y: numpy.ndarray
+            The data to use to initialize from.
+
+        Returns
+        -------
+        instance: `~astropy.modeling.models`
+            The initialized model.
+        """
+
+        # y_range = np.max(y) - np.min(y)
+        # x_range = x[-1] - x[0]
+        # slope = y_range / x_range
+        # y0 = y[0]
+
+        y_mean = np.mean(y)
+
+        instance.slope.value = 0.0
+        instance.intercept.value = y_mean.value
+
+        return instance
+
+
+class _WideBand1DInitializer(object):
+    """
+    Initialization that is applicable to all "wide band"
+    models
+
+    A "wide band" model is one that has an amplitude and
+    a position in wavelength space where this amplitude
+    is defined.
+
+    Parameters
+    ----------
+    factor: float
+        The scale factor to apply to the amplitutde
+    """
+    def __init__(self, factor=1.0):
+        self._factor = factor
+
+    def initialize(self, instance, x, y):
+        """
+        Initialize the model
+
+        Parameters
+        ----------
+        instance: `~astropy.modeling.models`
+            The model to initialize.
+
+        x, y: numpy.ndarray
+            The data to use to initialize from.
+
+        Returns
+        -------
+        instance: `~astropy.modeling.models`
+            The initialized model.
+        """
+        y_mean = np.mean(y)
+        x_range = x[-1] - x[0]
+        position = x_range / 2.0 + x[0]
+
+        name = _get_model_name(instance)
+
+        _setattr(instance, name, AMPLITUDE, y_mean * self._factor)
+        _setattr(instance, name, POSITION, position)
+
+        return instance
+
+
+class _LineProfile1DInitializer(object):
+    """
+    Initialization that is applicable to all "line profile"
+    models.
+
+    A "line profile" model is one that has an amplitude, a width,
+    and a defined position in wavelength space.
+
+    Parameters
+    ----------
+    factor: float
+        The scale factor to apply to the amplitutde
+    """
+    def __init__(self, factor=1.0):
+        self._factor = factor
+
+    def _set_width_attribute(self, instance, name, fwhm):
+        '''
+        Each line profile class has its own way of naming
+        and defining the width parameter. Subclasses should
+        override this method to conform to the specific
+        definitions.
+
+        Parameters
+        ----------
+        name : str
+            The attribute name
+
+        instance: `~astropy.modeling.models`
+            The model to initialize.
+
+        fwhm : float
+            FWHM
+        '''
+        raise NotImplementedError
+
+    def initialize(self, instance, x, y):
+        """
+        Initialize the model
+
+        Parameters
+        ----------
+        instance: `~astropy.modeling.models`
+            The model to initialize.
+
+        x, y: numpy.ndarray
+            The data to use to initialize from.
+
+        Returns
+        -------
+        instance: `~astropy.modeling.models`
+            The initialized model.
+        """
+
+        # X centroid estimates the position
+        centroid = np.sum(x * y) / np.sum(y)
+
+        # width can be estimated by the weighted
+        # 2nd moment of the X coordinate.
+        dx = x - np.mean(x)
+        fwhm = 2 * np.sqrt(np.sum((dx * dx) * y) / np.sum(y))
+
+        # amplitude is derived from area.
+        delta_x = x[1:] - x[:-1]
+        sum_y = np.sum((y[1:] - np.min(y[1:])) * delta_x)
+        height = sum_y / (fwhm / 2.355 * np.sqrt( 2 * np.pi))
+
+        name = _get_model_name(instance)
+
+        _setattr(instance, name, AMPLITUDE, height * self._factor)
+        _setattr(instance, name, POSITION, centroid)
+
+        self._set_width_attribute(instance, name, fwhm)
+
+        return instance
+
+
+class _Width_LineProfile1DInitializer(_LineProfile1DInitializer):
+    def _set_width_attribute(self, instance, name, fwhm):
+        _setattr(instance, name, WIDTH, fwhm)
+
+
+class _Sigma_LineProfile1DInitializer(_LineProfile1DInitializer):
+    def _set_width_attribute(self, instance, name, fwhm):
+        _setattr(instance, name, WIDTH, fwhm / 2.355)
+
+
+def _setattr(instance, mname, pname, value):
+    """
+    Sets parameter value by mapping parameter name to model type.
+
+    Prevents the parameter value setting to be stopped on its tracks
+    by non-existent model names or parameter names.
+
+    Parameters
+    ----------
+    instance: `~astropy.modeling.models`
+        The model to initialize.
+
+    mname: str
+        Model name.
+
+    pname: str
+        Parameter name.
+
+    value: any
+        The value to assign.
+    """
+    # this has to handle both Quantities and plain floats
+    try:
+        setattr(instance, _p_names[mname][pname], value.value)
+    except AttributeError:
+        setattr(instance, _p_names[mname][pname], value)
+    except KeyError:
+        pass
+
+
+# This associates each initializer to its corresponding spectral model.
+# Some models are not really line profiles, but their parameter names
+# and roles are the same as in a typical line profile, so they can be
+# initialized in the same way.
+_initializers = {
+    'Beta1D':                     _WideBand1DInitializer,
+    'Const1D':                    _WideBand1DInitializer,
+    'PowerLaw1D':                 _WideBand1DInitializer,
+    'BrokenPowerLaw1D':           _WideBand1DInitializer,
+    'ExponentialCutoffPowerLaw1D':_WideBand1DInitializer,
+    'LogParabola1D':              _WideBand1DInitializer,
+    'Box1D':                      _Width_LineProfile1DInitializer,
+    'Gaussian1D':                 _Sigma_LineProfile1DInitializer,
+    'Lorentz1D':                  _Width_LineProfile1DInitializer,
+    'Voigt1D':                    _Width_LineProfile1DInitializer,
+    'MexicanHat1D':               _Sigma_LineProfile1DInitializer,
+    'Trapezoid1D':                _Width_LineProfile1DInitializer,
+    'Linear1D':                   _Linear1DInitializer,
+    # 'Spline1D':                   spline.Spline1DInitializer
+}
+
+# Models can have parameter names that are similar amongst them, but not quite the same.
+# This maps the standard names used in the code to the actual names used by astropy.
+_p_names = {
+    'Gaussian1D':                 {AMPLITUDE:'amplitude',  POSITION:'mean', WIDTH:'stddev'},
+    'GaussianAbsorption':         {AMPLITUDE:'amplitude',  POSITION:'mean', WIDTH:'stddev'},
+    'Lorentz1D':                  {AMPLITUDE:'amplitude',  POSITION:'x_0',  WIDTH:'fwhm'},
+    'Voigt1D':                    {AMPLITUDE:'amplitude_L',POSITION:'x_0',  WIDTH:'fwhm_G'},
+    'Box1D':                      {AMPLITUDE:'amplitude',  POSITION:'x_0',  WIDTH:'width'},
+    'MexicanHat1D':               {AMPLITUDE:'amplitude',  POSITION:'x_0',  WIDTH:'sigma'},
+    'Trapezoid1D':                {AMPLITUDE:'amplitude',  POSITION:'x_0',  WIDTH:'width'},
+    'Beta1D':                     {AMPLITUDE:'amplitude',  POSITION:'x_0'},
+    'PowerLaw1D':                 {AMPLITUDE:'amplitude',  POSITION:'x_0'},
+    'ExponentialCutoffPowerLaw1D':{AMPLITUDE:'amplitude',  POSITION:'x_0'},
+    'LogParabola1D':              {AMPLITUDE:'amplitude',  POSITION:'x_0'},
+    'BrokenPowerLaw1D':           {AMPLITUDE:'amplitude',  POSITION:'x_break'},
+    'Const1D':                    {AMPLITUDE:'amplitude'},
+    }
+
+
+def initialize(instance, x, y):
+    """
+    Initialize given model.
+
+    X and Y are for now Quantity arrays with the
+    independent and dependent variables. It's assumed X values
+    are stored in increasing order in the array.
+
+    Parameters
+    ----------
+    instance: `~astropy.modeling.models`
+        The model to initialize.
+
+    x, y: numpy.ndarray
+        The data to use to initialize from.
+
+    Returns
+    -------
+    instance: `~astropy.modeling.models`
+        The initialized model.
+        If there are any errors, the instance is returned
+        uninitialized.
+    """
+    if x is None or y is None:
+        return instance
+
+    name = _get_model_name(instance)
+
+    try:
+        initializer = _initializers[name]()
+
+        return initializer.initialize(instance, x, y)
+
+    except KeyError:
+        return instance

specviz/plugins/model_editor/model_editor.py

@@ -1,20 +1,20 @@
 import os
-import uuid
 import pickle
+import uuid
 
 import numpy as np
-from astropy import units as u
 from astropy.modeling import fitting, models, optimizers
 from qtpy.QtCore import Qt
 from qtpy.QtGui import QIcon
-from qtpy.QtWidgets import (QAction, QDialog, QInputDialog, QMenu, QMessageBox,
-                            QToolButton, QWidget, QFileDialog)
+from qtpy.QtWidgets import (QAction, QDialog, QFileDialog, QInputDialog, QMenu,
+                            QMessageBox, QToolButton, QWidget)
 from qtpy.uic import loadUi
 from specutils.fitting import fit_lines
+from specutils.manipulation import extract_region
 from specutils.spectra import Spectrum1D
-from specutils.spectra.spectral_region import SpectralRegion
 
 from .equation_editor_dialog import ModelEquationEditorDialog
+from .initializers import initialize
 from .items import ModelDataItem
 from .models import ModelFittingModel
 from ...core.plugin import plugin
@@ -24,8 +24,8 @@
     'Linear1D': models.Linear1D,
     'Polynomial1D': models.Polynomial1D,
     'Gaussian1D': models.Gaussian1D,
-    'Voigt': models.Voigt1D,
-    'Lorentzian': models.Lorentz1D,
+    'Voigt1D': models.Voigt1D,
+    'Lorentzian1D': models.Lorentz1D,
 }
 
 FITTERS = {
@@ -181,17 +181,17 @@ def _on_remove_model(self):
 
             # If removing the model resulted in an invalid arithmetic equation,
             # force open the arithmetic editor so the user can fix it.
-            if self.model_tree_view.model().evaluate() is None:
+            if self.model_tree_view.model().equation and self.model_tree_view.model().evaluate() is None:
                 self._on_equation_edit_button_clicked()
 
     def _save_models(self, filename):
         model_editor_model = self.hub.plot_item.data_item.model_editor_model
         models = model_editor_model.fittable_models
+
         with open(filename, 'wb') as handle:
-            pickle.dump(model_editor_model.fittable_models, handle)
+            pickle.dump(models, handle)
 
     def _on_save_model(self, interactive=True):
-
         model_editor_model = self.hub.data_item.model_editor_model
         # There are no models to save
         if not model_editor_model.fittable_models:
@@ -240,7 +240,7 @@ def _add_model(self, model):
         self._redraw_model()
 
     def _add_fittable_model(self, model_type):
-        if issubclass(model_type, MODELS['Polynomial1D']):
+        if issubclass(model_type, models.Polynomial1D):
             text, ok = QInputDialog.getInt(self, 'Polynomial1D',
                                            'Enter Polynomial1D degree:')
             # User decided not to create a model after all
@@ -251,6 +251,17 @@ def _add_fittable_model(self, model_type):
         else:
             model = model_type()
 
+        # Grab any user-defined regions so we may initialize parameters only
+        # for the selected data.
+        inc_regs = self.hub.spectral_regions
+        spec = self._get_selected_plot_data_item().data_item.spectrum
+
+        if inc_regs is not None:
+            spec = extract_region(spec, inc_regs)
+
+        # Initialize the parameters
+        model = initialize(model, spec.spectral_axis, spec.flux)
+
         self._add_model(model)
 
     def _update_model_data_item(self):

specviz/plugins/model_editor/models.py

@@ -126,7 +126,7 @@ def remove_model(self, row):
 
         # Remove the model name from the equation
         self.equation = re.sub(
-            "(\+|-|\*|\/|=|>|<|>=|<=|&|\||%|!|\^|\(|\))\s+?({})".format(
+            "(\+|-|\*|\/|=|>|<|>=|<=|&|\||%|!|\^|\(|\))*\s*?({})".format(
                 model_item.text()),
             "", self._equation)