Units Implimentation

saba/main.py

@@ -1,7 +1,9 @@
 from __future__ import (absolute_import, unicode_literals, division,
                         print_function)
 import numpy as np
-from collections import OrderedDict
+import warnings
+
+from collections import OrderedDict, defaultdict
 from sherpa.fit import Fit
 from sherpa.data import Data1D, Data1DInt, Data2D, Data2DInt, DataSimulFit
 from sherpa.data import BaseData
@@ -12,13 +14,13 @@
 from sherpa.optmethods import GridSearch, LevMar, MonCar, NelderMead
 from sherpa.estmethods import Confidence, Covariance, Projection
 from sherpa.sim import MCMC
-import warnings
-
 from astropy.extern.six.moves import range
 from astropy.utils import format_doc
 from astropy.utils.exceptions import AstropyUserWarning
 from astropy.tests.helper import catch_warnings
 
+from astropy.modeling.utils import _combine_equivalency_dict
+from astropy.units import Quantity
 
 with catch_warnings(AstropyUserWarning) as warns:
     """this is to stop the import warning
@@ -29,7 +31,6 @@
     if "SherpaFitter" not in w.message.message:
         warnings.warn(w)
 
-# from astropy.modeling
 
 __all__ = ('SherpaFitter', 'SherpaMCMC')
 
@@ -96,8 +97,6 @@ class EstMethod(SherpaWrapper):
                       'projection': Projection}
 
 
-
-
 class SherpaMCMC(object):
     """
     An interface which makes use of sherpa's MCMC(pyBLoCXS) functionality.
@@ -325,8 +324,7 @@ def __init__(self, optimizer="levmar", statistic="leastsq", estmethod="covarianc
         # sherpa doesn't currently have a docstring for est_method but maybe the future
         setattr(self.__class__, 'est_config', property(lambda s: s._est_config, doc=self._est_method.__doc__))
 
-
-    def __call__(self, models, x, y, z=None, xbinsize=None, ybinsize=None, err=None, bkg=None, bkg_scale=1, **kwargs):
+    def __call__(self, models, x, y, z=None, xbinsize=None, ybinsize=None, err=None, bkg=None, bkg_scale=1, equivalencies=None, **kwargs):
         """
         Fit the astropy model with a the sherpa fit routines.
 
@@ -362,6 +360,11 @@ def __call__(self, models, x, y, z=None, xbinsize=None, ybinsize=None, err=None,
         """
 
         tie_list = []
+
+        _models, _x, _y, _z, _xbinsize, _ybinsize, _err, _bkg = self.remove_units(models, x, y, z, xbinsize, ybinsize, err, bkg, equivalencies)
+
+
+
         try:
             n_inputs = models[0].n_inputs
         except TypeError:
@@ -389,7 +392,7 @@ def __call__(self, models, x, y, z=None, xbinsize=None, ybinsize=None, err=None,
         self._fitter = Fit(self._data.data, self._fitmodel.sherpa_model, self._stat_method, self._opt_method, self._est_method, **kwargs)
         self.fit_info = self._fitter.fit()
 
-        return self._fitmodel.get_astropy_model()
+        return self.restore_units(self._fitmodel.get_astropy_model())
 
     def est_errors(self, sigma=None, maxiters=None, numcores=1, methoddict=None, parlist=None):
         """
@@ -439,6 +442,188 @@ def get_sampler(self, *args, **kwargs):
         """
         return SherpaMCMC(self, *args, **kwargs)
 
+    def remove_units(self, models, x, y, z=None, xbinsize=None, ybinsize=None, err=None, bkg=None, equivalencies=None):
+        """
+        This stripts data of there units so they can be fit.
+        """
+        #try to get a model, x, y, z, xbin, ybin, err, bkg tuple
+        #see how many models
+        try:
+            models._supports_unit_fitting
+            n_models = 1
+            models = [models]
+        except AttributeError:
+            n_models = len(models)
+
+        #try to figure out the is multiple datasets
+        #assume all data has same type?
+        try:
+            n_data = x.ndim
+            if z  is None:
+                assert x.ndim == y.ndim, AstropyUserWarning("x and y dimensions don't match")
+            else:
+                assert x.ndim == y.ndim == z.ndim, AstropyUserWarning("x, y and z dimensions don't match")
+        except AttributeError:
+            n_data = len(x)
+            if z  is None:
+                assert len(x) == len(y), AstropyUserWarning("x and y dimensions don't match")
+            else:
+                assert len(x) == len(y) == len(z), AstropyUserWarning("x, y and z dimensions don't match")
+
+        if n_data == 1 and n_models > 1:
+            n_data = n_models
+            x = np.asanyarray(x)
+            y = np.asanyarray(y)
+
+            x = [x.copy() for _ in range(n_data)]
+            y = [y.copy() for _ in range(n_data)]
+
+            if z is not None:
+                z = np.asanyarray(z)
+                z = [z.copy() for _ in range(n_data)]
+
+        elif n_data > 1 and  n_models == 1 :
+                n_models  = n_data
+                models = [models[0].copy() for _ in range(n_models)]
+
+        elif n_data is not 1 and n_models is not 1 and not n_data ==  n_models:
+            raise Exception("Don't know how to handle multiple models "
+                            "unless there is one foreach dataset")
+
+        else:
+            x = [np.asanyarray(x)]
+            y = [np.asanyarray(y)]
+            if z is not None:
+                z = [np.asanyarray(z)]
+
+        if z is None:
+            z = n_data * [None]
+
+        if xbinsize is None:
+            xbinsize = n_data * [None]
+
+        if ybinsize is None:
+            ybinsize = n_data * [None]
+
+        if err is None:
+            err = n_data * [None]
+
+        if bkg is None:
+            bkg = n_data * [None]
+
+        #iterate over all the things
+        self._units_sets = defaultdict(list)
+        _x = []
+        _y = []
+        _z = []
+        _xbinsize = []
+        _ybinsize = []
+        _err = []
+        _bkg = []
+        _models = []
+
+        for  model, xx, yy, zz, xxbin, yybin, eerr, bbkg in zip(models, x, y, z, xbinsize, ybinsize, err, bkg):
+            if model._supports_unit_fitting:
+                input_units_equivalencies = _combine_equivalency_dict(model.inputs,
+                                                                      equivalencies,
+                                                                      model.input_units_equivalencies)
+                if model.input_units is not None:
+                    if isinstance(xx, Quantity):
+                        self._units_sets['x'].append(1 * xx.unit)
+                        xx = xx.to(model.input_units['x'], equivalencies=input_units_equivalencies['x'])
+                    else:
+                        self._units_sets['x'].append(None)
+
+
+
+                    if isinstance(yy, Quantity) and zz is not None:
+                        self._units_sets['y'].append(1 * yy.unit)
+                        yy = yy.to(model.input_units['y'], equivalencies=input_units_equivalencies['y'])
+                    elif  isinstance(yy, Quantity):
+                        self._units_sets['y'].append(1 * yy.unit)
+                    else:
+                        self._units_sets['y'].append(None)
+
+                    if  xxbin is not None and isinstance(xxbin, Quantity):
+                        xxbin = xxbin.to(model.input_units['x'], equivalencies=input_units_equivalencies['x'])
+
+                    if  yybin is not None and isinstance(yybin, Quantity) and z is not None:
+                        yybin = yybin.to(model.input_units['y'], equivalencies=input_units_equivalencies['y'])
+
+                    if isinstance(zz, Quantity):
+                        self._units_sets['z'].append(1 * zz.unit)
+                    else:
+                        self._units_sets['z'].append(None)
+
+                    if eerr is not None and isinstance(eerr, Quantity):
+                        if z is not None:
+                            eerr = eerr.to(self._units_sets['z'])
+                        else:
+                            eerr = eerr.to(self._units_sets['y'])
+
+                    if bbkg is not None and isinstance(bbkg, Quantity):
+                        if z is not None:
+                            bbkg = bbkg.to(self._units_sets['z'])
+                        else:
+                            bbkg = bbkg.to(self._units_sets['y'])
+
+
+                    _x.append(xx)
+                    _y.append(yy)
+                    _z.append(zz)
+                    _xbinsize.append(xxbin)
+                    _ybinsize.append(yybin)
+                    _err.append(eerr)
+                    _bkg.append(bbkg)
+                    _models.append(model.without_units_for_data(x=xx, y=yy, z=zz))
+                else:
+                    _x.append(xx)
+                    _y.append(yy)
+                    _z.append(zz)
+                    _xbinsize.append(xxbin)
+                    _ybinsize.append(yybin)
+                    _err.append(eerr)
+                    _bkg.append(bbkg)
+                    _models.append(model)
+                    self._units_sets['x'].append(None)
+                    self._units_sets['y'].append(None)
+                    self._units_sets['z'].append(None)
+            else:
+                if isinstance(xx, Quantity) or isinstance(yy, Quantity) or isinstance(zz, Quantity):
+                    warnings.warn(AstropyUserWarning("This model{0} does not support being fit to data "
+                                                     "with units the units will be ignored this may "
+                                                     "produce erroneous results".format(len(self._units_sets['x']))))
+                _x.append(xx)
+                _y.append(yy)
+                _z.append(zz)
+                _xbinsize.append(xxbin)
+                _ybinsize.append(yybin)
+                _err.append(eerr)
+                _bkg.append(bbkg)
+                _models.append(model)
+                self._units_sets['x'].append(None)
+                self._units_sets['y'].append(None)
+                self._units_sets['z'].append(None)
+
+        return _models, _x, _y, _z, _xbinsize, _ybinsize, _err, _bkg
+
+    def restore_units(self,models):
+        """
+        This retores the units to data .
+        """
+        _models = []
+        try:
+            if models._supports_unit_fitting and self._units_sets['x'][0] is not None and self._units_sets['y'][0] is not None:
+                return models.with_units_from_data(x=self._units_sets['x'][0], y=self._units_sets['y'][0], z=self._units_sets['z'][0])
+        except AttributeError:
+            for n, model in enumerate(models):
+                if self._units_sets['x'][n] is not None and self._units_sets['y'][n] is not None:
+                    _models.append(model.with_units_from_data(x=1 * self._units_sets['x'][n], y=1 * self._units_sets['y'][n], z=1 * self._units_sets['z'][n]))
+                else:
+                    _models.append(model)
+                return _models
+        return models
+
 
 class Dataset(SherpaWrapper):
 
@@ -712,6 +897,8 @@ def get_astropy_model(self):
         else:
             return return_models[0]
 
+    def are_model_units_sane(models):
+        pass #  TODO
 
 class Data1DIntBkg(Data1DInt):
     """

saba/tests/test_main.py

@@ -21,6 +21,7 @@
 
 from saba import SherpaFitter, Dataset, ConvertedModel
 from astropy.modeling.models import Gaussian1D, Gaussian2D, Polynomial1D
+from astropy import units as u
 
 _RANDOM_SEED = 0x1337
 np.random.seed(_RANDOM_SEED)
@@ -343,3 +344,25 @@ def test_check_fit(self, nbins=100):
             x_med = (xx[med_ind] + xx[med_ind + 1]) / 2.0
 
             assert_allclose(self.params[nn], x_med, atol=0.1)
+
+
+
+class TestSherpaQuantiesFitter(object):
+
+    def setup_class(self):
+        self.x = np.linspace(1, 5, 30) * u.micron
+        self.y = np.exp(-0.5 * (self.x - 2.5 * u.micron)**2 / (200 * u.nm)**2) * u.mJy
+        self.model_micron = Gaussian1D(mean=3 * u.micron, stddev=1 * u.micron, amplitude=1 * u.Jy)
+        self.model_thz =  Gaussian1D(mean=110 * u.THz, stddev=10 * u.THz, amplitude=1 * u.Jy)
+        self.fitter = SherpaFitter()
+
+    def test_basic(self):
+        self.fitter(self.model_micron, self.x, self.y)
+
+
+    def test_user_equivalencies(self):
+        self.fitter(self.model_thz, self.x, self.y, equivalencies={'x': u.spectral()})
+
+    def test_model_equivalencies(self):
+        self.model_thz.input_units_equivalencies = {'x': u.spectral()}
+        self.fitter(self.model_thz, self.x, self.y)