embray · Jun 26, 2014 · May 30, 2014 · Jun 2, 2014 · Jun 3, 2014 · Jun 5, 2014
diff --git a/CHANGES.rst b/CHANGES.rst
@@ -110,6 +110,11 @@ New Features
     separately from fit statistics, allowing new fitters to be created by
     mixing and matching optimizers and statistic functions. [#1914]
 
+  - Slight overhaul to how inputs to and outputs from models are handled with
+    respect to array-valued parameters and variables, as well as sets of
+    multiple models.  See the associated PR and the modeling section of the
+    v0.4 documentation for more details. [#2634]
+
   - Added a new ``SimplexLSQFitter`` which uses a downhill simplex optimizer
     with a least squares statistic. [#1914]
 

diff --git a/astropy/modeling/core.py b/astropy/modeling/core.py
diff --git a/astropy/modeling/fitting.py b/astropy/modeling/fitting.py
@@ -25,8 +25,10 @@
                         print_function)
 
 import abc
-import warnings
 import inspect
+import operator
+import warnings
+
 from functools import reduce
 
 import numpy as np
@@ -217,30 +219,28 @@ def __call__(self, model, x, y, z=None, weights=None, rcond=None):
         model_copy : `~astropy.modeling.FittableModel`
             a copy of the input model with parameters set by the fitter
         """
+
         if not model.fittable:
             raise ValueError("Model must be a subclass of FittableModel")
+
         if not model.linear:
             raise ModelLinearityError('Model is not linear in parameters, '
                                       'linear fit methods should not be used.')
 
-        if any(model.tied.values()) \
-                or any([tuple(b) != (None, None) for b in model.bounds.values()]) \
-                or model.eqcons or model.ineqcons:
-            raise ValueError("LinearFitter supports only fixed constraints.")
-        multiple = False
+        _validate_constraints(self.supported_constraints, model)
+
         model_copy = model.copy()
         _, fitparam_indices = _model_to_fit_params(model_copy)
+
         if model_copy.n_inputs == 2 and z is None:
             raise ValueError("Expected x, y and z for a 2 dimensional model.")
 
-        farg = _convert_input(x, y, z)
+        farg = _convert_input(x, y, z, n_models=len(model_copy),
+                              model_set_axis=model_copy.model_set_axis)
 
         if len(farg) == 2:
             x, y = farg
-            if y.ndim == 2:
-                if y.shape[1] != model_copy.param_dim:
-                    raise ValueError("Number of data sets (Y array is expected"
-                                     " to equal the number of parameter sets")
+
             # map domain into window
             if hasattr(model_copy, 'domain'):
                 x = self._map_domain_window(model_copy, x)
@@ -252,13 +252,10 @@ def __call__(self, model, x, y, z=None, weights=None, rcond=None):
                 lhs = model_copy.fit_deriv(x, *model_copy.parameters)
             if len(y.shape) == 2:
                 rhs = y
-                multiple = y.shape[1]
             else:
                 rhs = y
         else:
             x, y, z = farg
-            if x.shape[-1] != z.shape[-1]:
-                raise ValueError("x and z should have equal last dimensions")
 
             # map domain into window
             if hasattr(model_copy, 'x_domain'):
@@ -269,14 +266,26 @@ def __call__(self, model, x, y, z=None, weights=None, rcond=None):
                                                    fitparam_indices, x=x, y=y)
             else:
                 lhs = model_copy.fit_deriv(x, y, *model_copy.parameters)
-            if len(z.shape) == 3:
-                rhs = np.array([i.flatten() for i in z]).T
-                multiple = z.shape[0]
+
+            if len(model_copy) > 1:
+                if z.ndim > 2:
+                    # Basically this code here is making the assumption that if
+                    # z has 3 dimensions it represents multiple models where
+                    # the value of z is one plane per model.  It's then
+                    # flattening each plane and transposing so that the model
+                    # axis is *last*.  That's fine, but this could be
+                    # generalized for other dimensionalities of z.
+                    # TODO: See above comment
+                    rhs = np.array([i.flatten() for i in z]).T
+                else:
+                    rhs = z.T
             else:
                 rhs = z.flatten()
+
         # If the derivative is defined along rows (as with non-linear models)
         if model_copy.col_fit_deriv:
             lhs = np.asarray(lhs).T
+
         if weights is not None:
             weights = np.asarray(weights, dtype=np.float)
             if len(x) != len(weights):
@@ -290,9 +299,6 @@ def __call__(self, model, x, y, z=None, weights=None, rcond=None):
                 lhs *= weights[:, np.newaxis]
                 rhs = rhs * weights
 
-        if not multiple and model_copy.param_dim > 1:
-            raise ValueError("Attempting to fit a 1D data set to a model "
-                             "with multiple parameter sets")
         if rcond is None:
             rcond = len(x) * np.finfo(x.dtype).eps
 
@@ -303,20 +309,15 @@ def __call__(self, model, x, y, z=None, weights=None, rcond=None):
         self.fit_info['rank'] = rank
         self.fit_info['singular_values'] = sval
 
-        # If y.n_inputs > model.n_inputs we are doing a simultanious 1D fitting
-        # of several 1D arrays. Otherwise the model is 2D.
-        # if y.n_inputs > self.model.n_inputs:
-        if multiple and model_copy.param_dim != multiple:
-            model_copy.param_dim = multiple
-        # TODO: Changing the model's param_dim needs to be handled more
-        # carefully; for now it's not actually allowed
         lacoef = (lacoef.T / scl).T
         self.fit_info['params'] = lacoef
+
         # TODO: Only Polynomial models currently have an _order attribute;
         # maybe change this to read isinstance(model, PolynomialBase)
         if hasattr(model_copy, '_order') and rank != model_copy._order:
             warnings.warn("The fit may be poorly conditioned\n",
                           AstropyUserWarning)
+
         _fitter_to_model_params(model_copy, lacoef.flatten())
         return model_copy
 
@@ -774,19 +775,42 @@ def __call__(self, *args):
             model.parameters = np.array(mparams)
 
 
-def _convert_input(x, y, z=None):
+def _convert_input(x, y, z=None, n_models=1, model_set_axis=0):
     """Convert inputs to float arrays."""
 
     x = np.asarray(x, dtype=np.float)
     y = np.asarray(y, dtype=np.float)
-    if x.shape[0] != y.shape[0]:
-        raise ValueError("x and y should have the same shape")
+
+    if z is not None:
+        z = np.asarray(z, dtype=np.float)
+
+    # For compatibility with how the linear fitter code currently expects to
+    # work, shift the dependent variable's axes to the expected locations
+    if n_models > 1:
+        if z is None:
+            if y.shape[model_set_axis] != n_models:
+                raise ValueError(
+                    "Number of data sets (y array is expected to equal "
+                    "the number of parameter sets)")
+            # For a 1-D model the y coordinate's model-set-axis is expected to
+            # be last, so that its first dimension is the same length as the x
+            # coordinates.  This is in line with the expectations of
+            # numpy.linalg.lstsq:
+            # http://docs.scipy.org/doc/numpy/reference/generated/numpy.linalg.lstsq.html
+            # That is, each model should be represented by a column.  TODO:
+            # Obviously this is a detail of np.linalg.lstsq and should be
+            # handled specifically by any fitters that use it...
+            y = np.rollaxis(y, model_set_axis, y.ndim)
+        else:
+            # Shape of z excluding model_set_axis
+            z_shape = z.shape[:model_set_axis] + z.shape[model_set_axis + 1:]
+
+            if not (x.shape == y.shape == z_shape):
+                raise ValueError("x, y and z should have the same shape")
+
     if z is None:
         farg = (x, y)
     else:
-        z = np.asarray(z, dtype=np.float)
-        if x.shape != z.shape:
-            raise ValueError("x, y and z should have the same shape")
         farg = (x, y, z)
     return farg
 
@@ -797,31 +821,59 @@ def _convert_input(x, y, z=None):
 # distinction (and the fact that these are not necessarily applicable to any
 # arbitrary fitter--as evidenced for example by the fact that JointFitter has
 # its own versions of these)
+# TODO: Most of this code should be entirely rewritten; it should not be as
+# inefficient as it is.
 def _fitter_to_model_params(model, fps):
     """
     Constructs the full list of model parameters from the fitted and
     constrained parameters.
     """
 
-    _fit_params, _fit_param_indices = _model_to_fit_params(model)
-    if any(model.fixed.values()) or any(model.tied.values()):
-        model.parameters[_fit_param_indices] = fps
+    _, fit_param_indices = _model_to_fit_params(model)
+
+    has_tied = any(model.tied.values())
+    has_fixed = any(model.fixed.values())
+    has_bound = any(b != (None, None) for b in model.bounds.values())
+
+    if not (has_tied or has_fixed or has_bound):
+        # We can just assign directly
+        model.parameters = fps
+        return
+
+    fit_param_indices = set(fit_param_indices)
+    offset = 0
+    for idx, name in enumerate(model.param_names):
+        if idx not in fit_param_indices:
+            continue
+
+        slice_, shape = model._param_metrics[name]
+        # This is determining which range of fps (the fitted parameters) maps
+        # to parameters of the model
+        size = reduce(operator.mul, shape, 1)
+
+        values = fps[offset:offset + size]
+
+        # Check bounds constraints
+        if model.bounds[name] != (None, None):
+            _min, _max = model.bounds[name]
+            if _min is not None:
+                values = np.fmax(values, _min)
+            if _max is not None:
+                values = np.fmin(values, _max)
+
+        model.parameters[slice_] = values
+        offset += size
+
+    # This has to be done in a separate loop due to how tied parameters are
+    # currently evaluated (the fitted parameters need to actually be *set* on
+    # the model first, for use in evaluating the "tied" expression--it might be
+    # better to change this at some point
+    if has_tied:
         for idx, name in enumerate(model.param_names):
             if model.tied[name] != False:
                 value = model.tied[name](model)
                 slice_ = model._param_metrics[name][0]
                 model.parameters[slice_] = value
-    elif any([tuple(b) != (None, None) for b in model.bounds.values()]):
-        for name, par in zip(model.param_names, _fit_params):
-            if model.bounds[name] != (None, None):
-                b = model.bounds[name]
-                if b[0] is not None:
-                    par = max(par, model.bounds[name][0])
-                    if b[1] is not None:
-                        par = min(par, model.bounds[name][1])
-                    setattr(model, name, par)
-    else:
-        model.parameters = fps
 
 
 def _model_to_fit_params(model):
@@ -887,7 +939,7 @@ def _validate_model(model, supported_constraints):
         warnings.warn('Model is linear in parameters; '
                       'consider using linear fitting methods.',
                       AstropyUserWarning)
-    if model.param_dim != 1:
+    elif len(model) != 1:
         # for now only single data sets ca be fitted
         raise ValueError("Non-linear fitters can only fit "
                          "one data set at a time.")