Added a new quantity_asanyarray function to properly deal with lists of quantities, and use this to fix the use of models with units when n_models > 1

astrofrog · astrofrog · commit 6b4266c4bfc0 · 2017-01-17T15:55:22.000Z
diff --git a/astropy/modeling/core.py b/astropy/modeling/core.py
@@ -35,6 +35,7 @@
 from ..extern.six.moves import copyreg, zip
 from ..table import Table
 from ..units import Quantity, UnitBase, UnitsError, dimensionless_unscaled
+from ..units.utils import quantity_asanyarray
 from ..utils import (sharedmethod, find_current_module,
                      InheritDocstrings, OrderedDescriptorContainer,
                      check_broadcast, IncompatibleShapeError)
@@ -1455,7 +1456,8 @@ def _initialize_parameters(self, args, kwargs):
             if arg is None:
                 # A value of None implies using the default value, if exists
                 continue
-            params[self.param_names[idx]] = np.asanyarray(arg, dtype=np.float)
+            params[self.param_names[idx]] = quantity_asanyarray(arg, dtype=np.float)
+
 
         # At this point the only remaining keyword arguments should be
         # parameter names; any others are in error.
@@ -1469,7 +1471,7 @@ def _initialize_parameters(self, args, kwargs):
                 if value is None:
                     continue
                 else:
-                    params[param_name] = np.asanyarray(value, dtype=np.float)
+                    params[param_name] = quantity_asanyarray(value, dtype=np.float)
 
         if kwargs:
             # If any keyword arguments were left over at this point they are
diff --git a/astropy/units/utils.py b/astropy/units/utils.py
@@ -19,6 +19,7 @@
 from numpy import finfo
 
 from ..extern import six
+from .. import units as u
 
 _float_finfo = finfo(float)
 # take float here to ensure comparison with another float is fast
@@ -219,3 +220,12 @@ def resolve_fractions(a, b):
     elif not a_is_fraction and b_is_fraction:
         a = Fraction(a)
     return a, b
+
+
+def quantity_asanyarray(a, dtype=None):
+    if isinstance(a, np.ndarray):
+        return np.asanyarray(a, dtype=dtype)
+    elif any(isinstance(x, u.Quantity) for x in a):
+        return u.Quantity(a, dtype=dtype)
+    else:
+        raise ValueError("Unexpected object passed to quantity_asanyarray: {0}".format(a))
