Quantity Support in Modeling [continued]

astropy/modeling/fitting.py

@@ -38,13 +38,11 @@
 from ..extern import six
 from .optimizers import (SLSQP, Simplex)
 from .statistic import (leastsquare)
-
+from .. import units as u
 
 __all__ = ['LinearLSQFitter', 'LevMarLSQFitter', 'SLSQPLSQFitter',
            'SimplexLSQFitter', 'JointFitter', 'Fitter']
 
-
-
 # Statistic functions implemented in `astropy.modeling.statistic.py
 STATISTICS = [leastsquare]
 
@@ -84,6 +82,24 @@ def __new__(mcls, name, bases, members):
         return cls
 
 
+def unit_support(func):
+
+    print("ICI")
+
+    def wrapper(self, model, x, y, z=None, *args, **kwargs):
+        print("HERE")
+        if isinstance(x, u.Quantity) or isinstance(y, u.Quantity):
+            if z is None:
+                model_new = func(self, model, x.value, y.value, *args, **kwargs)
+                return model_new.with_units_from_data(x, y)
+            else:
+                model_new = func(self, model, x.value, y.value, z=z.value, *args, **kwargs)
+                return model_new.with_units_from_data(x, y, z)
+        else:
+            func(self, model, x.value, y, *args, **kwargs)
+
+    return wrapper
+
 @six.add_metaclass(_FitterMeta)
 class Fitter(object):
     """
@@ -212,6 +228,7 @@ def _map_domain_window(self, model, x, y=None):
             ynew = poly_map_domain(y, model.y_domain, model.y_window)
             return xnew, ynew
 
+    @unit_support
     def __call__(self, model, x, y, z=None, weights=None, rcond=None):
         """
         Fit data to this model.
@@ -239,6 +256,8 @@ def __call__(self, model, x, y, z=None, weights=None, rcond=None):
             a copy of the input model with parameters set by the fitter
         """
 
+        print("FINALLY IN FITTER")
+
         if not model.fittable:
             raise ValueError("Model must be a subclass of FittableModel")
 
@@ -407,6 +426,7 @@ def objective_function(self, fps, *args):
         else:
             return np.ravel(weights * (model(*args[2 : -1]) - meas))
 
+    @unit_support
     def __call__(self, model, x, y, z=None, weights=None,
                  maxiter=DEFAULT_MAXITER, acc=DEFAULT_ACC,
                  epsilon=DEFAULT_EPS, estimate_jacobian=False):
@@ -539,6 +559,7 @@ def __init__(self):
         super(SLSQPLSQFitter, self).__init__(optimizer=SLSQP, statistic=leastsquare)
         self.fit_info = {}
 
+    @unit_support
     def __call__(self, model, x, y, z=None, weights=None, **kwargs):
         """
         Fit data to this model.

astropy/modeling/functional_models.py

@@ -24,6 +24,13 @@
            'TrapezoidDisk2D', 'Ring2D', 'custom_model_1d', 'Voigt1D']
 
 
+def quantity_with_unit(parameter, unit):
+    if parameter.unit is None:
+        return parameter.value * unit
+    else:
+        return parameter.quantity.to(unit)
+
+
 class Gaussian1D(Fittable1DModel):
     """
     One dimensional Gaussian model.
@@ -110,6 +117,21 @@ class Gaussian1D(Fittable1DModel):
     mean = Parameter(default=0)
     stddev = Parameter(default=1)
 
+    # input_units = 'mean'  # Input must have same units as mean
+    # output_units = 'amplitude'  # Output must have same units as amplitude
+
+    def with_units_from_data(self, x, y):
+        """
+        Return an instance of the model which has units for which the parameter
+        values are compatible with the data units.
+        """
+
+        amplitude = quantity_with_unit(self.amplitude, y.unit)
+        mean = quantity_with_unit(self.mean, y.unit)
+        stddev = quantity_with_unit(self.stddev, y.unit)
+
+        return Gaussian1D(amplitude=amplitude, mean=mean, stddev=stddev)
+
     def bounding_box(self, factor=5.5):
         """
         Tuple defining the default ``bounding_box`` limits,
@@ -209,6 +231,9 @@ class GaussianAbsorption1D(Fittable1DModel):
     mean = Parameter(default=0)
     stddev = Parameter(default=1)
 
+    input_units = 'mean'
+    output_units = 'amplitude'
+
     def bounding_box(self, factor=5.5):
         """
         Tuple defining the default ``bounding_box`` limits,
@@ -341,6 +366,9 @@ class Gaussian2D(Fittable2DModel):
     y_stddev = Parameter(default=1)
     theta = Parameter(default=0)
 
+    input_units = ('x_mean', 'y_mean')
+    output_units = 'amplitude'
+
     def __init__(self, amplitude=amplitude.default, x_mean=x_mean.default,
                  y_mean=y_mean.default, x_stddev=None, y_stddev=None,
                  theta=0.0, cov_matrix=None, **kwargs):
@@ -497,6 +525,11 @@ class Shift(Model):
     inputs = ('x',)
     outputs = ('x',)
 
+    # Input must have compatible units with offset, but output is kept in the
+    # input units
+    input_units = 'offset'
+    output_units = 'x'
+
     offset = Parameter(default=0)
     fittable = True
 
@@ -524,6 +557,8 @@ class Scale(Model):
     inputs = ('x',)
     outputs = ('x',)
 
+    output_units = lambda factor, x: factor.units * x.units
+
     factor = Parameter(default=1)
     linear = True
     fittable = True
@@ -557,6 +592,8 @@ class RedshiftScaleFactor(Fittable1DModel):
 
     z = Parameter(description='redshift', default=0)
 
+    output_units = 'x'
+
     @staticmethod
     def evaluate(x, z):
         """One dimensional RedshiftScaleFactor model function"""
@@ -722,6 +759,9 @@ class Sine1D(Fittable1DModel):
     frequency = Parameter(default=1)
     phase = Parameter(default=0)
 
+    input_units = lambda frequency: 1 / frequency.unit
+    output_units = 'amplitude'
+
     @staticmethod
     def evaluate(x, amplitude, frequency, phase):
         """One dimensional Sine model function"""
@@ -765,6 +805,10 @@ class Linear1D(Fittable1DModel):
 
     slope = Parameter(default=1)
     intercept = Parameter(default=0)
+
+    input_units = lambda slope, intercept: intercept.unit / slope.unit
+    output_units = 'intercept'
+
     linear = True
 
     @staticmethod
@@ -839,6 +883,9 @@ class Lorentz1D(Fittable1DModel):
     x_0 = Parameter(default=0)
     fwhm = Parameter(default=1)
 
+    input_units = 'x_0'
+    output_units = 'amplitude'
+
     @staticmethod
     def evaluate(x, amplitude, x_0, fwhm):
         """One dimensional Lorentzian model function"""
@@ -991,6 +1038,9 @@ class Const1D(Fittable1DModel):
     """
 
     amplitude = Parameter(default=1)
+
+    output_units = 'amplitude'
+
     linear = True
 
     @staticmethod
@@ -1037,6 +1087,9 @@ class Const2D(Fittable2DModel):
     """
 
     amplitude = Parameter(default=1)
+
+    output_units = 'amplitude'
+
     linear = True
 
     @staticmethod
@@ -1132,6 +1185,9 @@ class Ellipse2D(Fittable2DModel):
     b = Parameter(default=1)
     theta = Parameter(default=0)
 
+    input_units = ('x_0', 'y_0')
+    output_units = 'amplitude'
+
     @staticmethod
     def evaluate(x, y, amplitude, x_0, y_0, a, b, theta):
         """Two dimensional Ellipse model function."""
@@ -1200,6 +1256,9 @@ class Disk2D(Fittable2DModel):
     y_0 = Parameter(default=0)
     R_0 = Parameter(default=1)
 
+    input_units = ('x_0', 'y_0')
+    output_units = 'amplitude'
+
     @staticmethod
     def evaluate(x, y, amplitude, x_0, y_0, R_0):
         """Two dimensional Disk model function"""
@@ -1264,6 +1323,9 @@ class Ring2D(Fittable2DModel):
     r_in = Parameter(default=1)
     width = Parameter(default=1)
 
+    input_units = ('x_0', 'y_0')
+    output_units = 'amplitude'
+
     def __init__(self, amplitude=amplitude.default, x_0=x_0.default,
                  y_0=y_0.default, r_in=r_in.default, width=width.default,
                  r_out=None, **kwargs):
@@ -1372,6 +1434,8 @@ class Box1D(Fittable1DModel):
     x_0 = Parameter(default=0)
     width = Parameter(default=1)
 
+    output_units = 'amplitude'
+
     @staticmethod
     def evaluate(x, amplitude, x_0, width):
         """One dimensional Box model function"""
@@ -1445,6 +1509,9 @@ class Box2D(Fittable2DModel):
     x_width = Parameter(default=1)
     y_width = Parameter(default=1)
 
+    input_units = ('x_0', 'y_0')
+    output_units = 'amplitude'
+
     @staticmethod
     def evaluate(x, y, amplitude, x_0, y_0, x_width, y_width):
         """Two dimensional Box model function"""
@@ -1517,6 +1584,9 @@ class Trapezoid1D(Fittable1DModel):
     width = Parameter(default=1)
     slope = Parameter(default=1)
 
+    input_units = 'x_0'
+    output_units = 'amplitude'
+
     @staticmethod
     def evaluate(x, amplitude, x_0, width, slope):
         """One dimensional Trapezoid model function"""
@@ -1578,6 +1648,9 @@ class TrapezoidDisk2D(Fittable2DModel):
     R_0 = Parameter(default=1)
     slope = Parameter(default=1)
 
+    input_units = ('x_0', 'y_0')
+    output_units = 'amplitude'
+
     @staticmethod
     def evaluate(x, y, amplitude, x_0, y_0, R_0, slope):
         """Two dimensional Trapezoid Disk model function"""
@@ -1656,6 +1729,9 @@ class MexicanHat1D(Fittable1DModel):
     x_0 = Parameter(default=0)
     sigma = Parameter(default=1)
 
+    input_units = 'x_0'
+    output_units = 'amplitude'
+
     @staticmethod
     def evaluate(x, amplitude, x_0, sigma):
         """One dimensional Mexican Hat model function"""
@@ -1700,6 +1776,9 @@ class MexicanHat2D(Fittable2DModel):
     y_0 = Parameter(default=0)
     sigma = Parameter(default=1)
 
+    input_units = ('x_0', 'y_0')
+    output_units = 'amplitude'
+
     @staticmethod
     def evaluate(x, y, amplitude, x_0, y_0, sigma):
         """Two dimensional Mexican Hat model function"""
@@ -1755,6 +1834,8 @@ class AiryDisk2D(Fittable2DModel):
     x_0 = Parameter(default=0)
     y_0 = Parameter(default=0)
     radius = Parameter(default=1)
+    input_units = ('x_0', 'y_0')
+    output_units = 'amplitude'
     _rz = None
     _j1 = None
 
@@ -1835,6 +1916,9 @@ class Moffat1D(Fittable1DModel):
     gamma = Parameter(default=1)
     alpha = Parameter(default=1)
 
+    input_units = 'x_0'
+    output_units = 'amplitude'
+
     @staticmethod
     def evaluate(x, amplitude, x_0, gamma, alpha):
         """One dimensional Moffat model function"""
@@ -1891,6 +1975,9 @@ class Moffat2D(Fittable2DModel):
     gamma = Parameter(default=1)
     alpha = Parameter(default=1)
 
+    input_units = ('x_0', 'y_0')
+    output_units = 'amplitude'
+
     @staticmethod
     def evaluate(x, y, amplitude, x_0, y_0, gamma, alpha):
         """Two dimensional Moffat model function"""