Repository: astropy/astropy

Issue:
Can't evaluate model sets over common co-ordinate array for model_set_axis > 0
Related to last week's #7142 & #7143, I'm still not 100% clear on the intended usage of `model_set_axis` myself -- but shouldn't one be able to do `some_model(x, model_set_axis=False)` when `some_model` was created with a `model_set_axis` other than 0? This is not working in the cases I have tried.

```Python
import numpy as np
from astropy.modeling.models import *
from astropy.modeling.fitting import *

p = Polynomial1D(1, n_models=2, model_set_axis=1)

x = np.arange(0,4)

# Generate data for fitting 2 models and re-stack them along the last axis:
y = np.array([2*x+1, x+4])
y = np.moveaxis(y, 0, -1)

f = LinearLSQFitter()

# This seems to fit the model_set correctly:
fit = f(p, x, y)
print(fit.parameters)

# But this fails to evaluate the model, with a broadcasting error (the same
# thing happens with Polynomial1D, Polynomial2D, Chebyshev1D & Legendre1D):
model_y = fit(x, model_set_axis=False)
```

```
[ 1.  4.  2.  1.]
Traceback (most recent call last):
  File "test_model_set_axis_1.py", line 21, in <module>
    model_y = fit(x, model_set_axis=False)
  File "/home/jturner/repos/astropy/astropy/modeling/core.py", line 382, in __call__
    ('equivalencies', None)])
  File "/home/jturner/repos/astropy/astropy/modeling/core.py", line 363, in __call__
    return super(cls, self).__call__(*inputs, **kwargs)
  File "/home/jturner/repos/astropy/astropy/modeling/core.py", line 761, in __call__
    outputs = self.evaluate(*chain(inputs, parameters))
  File "/home/jturner/repos/astropy/astropy/modeling/polynomial.py", line 806, in evaluate
    return self.horner(x, coeffs)
  File "/home/jturner/repos/astropy/astropy/modeling/polynomial.py", line 840, in horner
    c0 = coeffs[-i] + c0 * x
ValueError: operands could not be broadcast together with shapes (1,2) (1,4) 
```

The reason I care about this is that I was trying to finish fixing #6819, but `FittingWithOutlierRemoval` has to evaluate the model set at each iteration, over the single `x` (/ `y`) grid supplied by the user -- ie. using `model_set_axis=False`. Since FittingWithOutlierRemoval can't just assume that `model_set_axis=0` for its input, I was checking the behaviour with other axes and found it broken in the models.

Thanks,

James.


Discussion:
astropy-bot[bot]: Hi there @nden :wave: - thanks for the pull request! I'm just a friendly :robot: that checks for issues related to the changelog and making sure that this pull request is milestoned and labeled correctly. This is mainly intended for the maintainers, so if you are not a maintainer you can ignore this, and a maintainer will let you know if any action is required on your part :smiley:.

I noticed the following issues with this pull request:

* The milestone has not been set (this can only be set by a maintainer)
* Changelog entry not present (or pull request number missing) and neither the **Affects-dev** nor the **no-changelog-entry-needed** label are set

Would it be possible to fix these? Thanks!

*If there are any issues with this message, please report them [here](https://github.com/astropy/astropy-bot/issues).*
nden: @pllim The error comes from Cython. Is this known/understood? Do I need to rebase?
pllim: Huh, no segfault this time... @astrofrog , will this be fixed by your fix too?

```
  File "Cython/Compiler/Code.py", line 521, in Cython.Compiler.Code.sub_tempita
ModuleNotFoundError: No module named 'scipy.sparse.Tempita'
```
nden: Rebased.
pllim: It fails with `/home/travis/.travis/job_stages: line 57:  5286 Segmentation fault` now, which is a known and unrelated issue.
jehturner: Great, thanks! As painful as this may have been to troubleshoot, I'm glad the changes don't look unduly extensive. That's better than I'd hoped :slightly_smiling_face:.

I'm having a bit of trouble following the changes to `Parameter.shape` at first glance though (not being familiar in detail with what Parameters do or how their shapes are defined)... I'm also not sure I understand exactly what pivots are. If only there were some developer documentation on the internals! (OK, there's the doc string, as far as it goes.)

nden: @jehturner I pushed another commit which may be cleaner and provides some explanation of why I think `Parameter.shape` is a problem. My understanding of the intention or how it actually works is not great either. My theory is that the docstring which says 

```
If we are dealing with a model *set* the shape is the shape of
the parameter within a single model in the set
```
is correct and I assume that the code returns incorrectly a shape of (1,) instead of () for each coefficient in the model:

`p = Polynomial1D(1, n_models=2, model_set_axis=1)`

The shape of a parameter is important because it's used in `prepare_inputs_model_set` as `max_param_shape` and determines the how the input should be reshaped and what the pivot is.

I *think* `pivot` is the dimension which should store results for each model. 

The complexity in the code comes partially from trying to work out higher dimensions, e.g. parameters are arrays which broadcast together but have different shapes. 
Anyway, that's the extent to which I understand what's happening. 


Patch:
diff --git a/CHANGES.rst b/CHANGES.rst
index a29607457c..2e3ae7062c 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -299,6 +299,8 @@ astropy.io.votable
 astropy.modeling
 ^^^^^^^^^^^^^^^^
 
+- Fix model set evaluation over common input when model_set_axis > 0. [#7222]
+
 astropy.nddata
 ^^^^^^^^^^^^^^
 
diff --git a/astropy/modeling/core.py b/astropy/modeling/core.py
index 5f69f1621a..133fe07ebb 100644
--- a/astropy/modeling/core.py
+++ b/astropy/modeling/core.py
@@ -1585,11 +1585,17 @@ class Model(metaclass=_ModelMeta):
 
         return outputs
 
-    def prepare_outputs(self, format_info, *outputs, **kwargs):
+    def prepare_outputs(self, format_info, *outputs, model_set_axis=None, **kwargs):
+        if model_set_axis is None:
+            # By default the model_set_axis for the input is assumed to be the
+            # same as that for the parameters the model was defined with
+            # TODO: Ensure that negative model_set_axis arguments are respected
+            model_set_axis = self.model_set_axis
+
         if len(self) == 1:
             return _prepare_outputs_single_model(self, outputs, format_info)
         else:
-            return _prepare_outputs_model_set(self, outputs, format_info)
+            return _prepare_outputs_model_set(self, outputs, format_info, model_set_axis)
 
     def copy(self):
         """
@@ -3313,7 +3319,6 @@ def _prepare_inputs_model_set(model, params, inputs, n_models, model_set_axis,
                            _input.shape[model_set_axis + 1:])
         else:
             input_shape = _input.shape
-
         for param in params:
             try:
                 check_broadcast(input_shape, param.shape)
@@ -3364,16 +3369,15 @@ def _prepare_inputs_model_set(model, params, inputs, n_models, model_set_axis,
     return reshaped, (pivots,)
 
 
-def _prepare_outputs_model_set(model, outputs, format_info):
+def _prepare_outputs_model_set(model, outputs, format_info, model_set_axis):
     pivots = format_info[0]
 
     outputs = list(outputs)
-
     for idx, output in enumerate(outputs):
         pivot = pivots[idx]
-        if pivot < output.ndim and pivot != model.model_set_axis:
+        if pivot < output.ndim and pivot != model_set_axis:
             outputs[idx] = np.rollaxis(output, pivot,
-                                       model.model_set_axis)
+                                       model_set_axis)
 
     return tuple(outputs)
 
@@ -3396,8 +3400,8 @@ def _validate_input_shapes(inputs, argnames, n_models, model_set_axis,
         return
 
     all_shapes = []
-
     for idx, _input in enumerate(inputs):
+
         input_shape = np.shape(_input)
         # Ensure that the input's model_set_axis matches the model's
         # n_models
@@ -3409,10 +3413,15 @@ def _validate_input_shapes(inputs, argnames, n_models, model_set_axis,
                     "least {1}-dimensional.".format(
                         model_set_axis, model_set_axis + 1))
             elif input_shape[model_set_axis] != n_models:
+                try:
+                    argname = argnames[idx]
+                except IndexError:
+                    # the case of model.inputs = ()
+                    argname = str(idx)
                 raise ValueError(
                     "Input argument {0!r} does not have the correct "
                     "dimensions in model_set_axis={1} for a model set with "
-                    "n_models={2}.".format(argnames[idx], model_set_axis,
+                    "n_models={2}.".format(argname, model_set_axis,
                                            n_models))
         all_shapes.append(input_shape)
 
diff --git a/astropy/modeling/parameters.py b/astropy/modeling/parameters.py
index b98e3a9b71..d5089e18ce 100644
--- a/astropy/modeling/parameters.py
+++ b/astropy/modeling/parameters.py
@@ -505,8 +505,16 @@ class Parameter(OrderedDescriptor):
 
             if model_axis < 0:
                 model_axis = len(shape) + model_axis
-
-            shape = shape[:model_axis] + shape[model_axis + 1:]
+                shape = shape[:model_axis] + shape[model_axis + 1:]
+            else:
+                # When a model set is initialized, the dimension of the parameters
+                # is increased by model_set_axis+1. To find the shape of a parameter
+                # within a single model the extra dimensions need to be removed first.
+                # The following dimension shows the number of models.
+                # The rest of the shape tuple represents the shape of the parameter
+                # in a single model.
+
+                shape = shape[model_axis + 1:]
 
         return shape
 


Test patch:
diff --git a/astropy/modeling/tests/test_input.py b/astropy/modeling/tests/test_input.py
index 6c8f895413..ba1088ea23 100644
--- a/astropy/modeling/tests/test_input.py
+++ b/astropy/modeling/tests/test_input.py
@@ -258,7 +258,7 @@ class TestEvaluation:
 
         p1 = models.Polynomial1D(4, n_models=2)
         y1 = p1(np.array([self.x1, self.x1]).T, model_set_axis=-1)
-        assert y1.shape == (2, 20)
+        assert y1.shape == (20, 2)
         assert_allclose(y1[0, :], y1[1, :], atol=10 ** (-12))
 
     def test_p2_1set_1pset(self):
@@ -279,21 +279,26 @@ class TestEvaluation:
 
     def test_nset_domain(self):
         """
-        Polynomial evaluation of multiple data sets with different domain
+        Test model set with negative model_set_axis.
+
+        In this case model_set_axis=-1 is identical to model_set_axis=1.
         """
 
         xx = np.array([self.x1, self.x1]).T
         xx[0, 0] = 100
         xx[1, 0] = 100
         xx[2, 0] = 99
-        p1 = models.Polynomial1D(5, n_models=2)
+        p1 = models.Polynomial1D(5, c0=[1, 2], c1=[3, 4], n_models=2)
         yy = p1(xx, model_set_axis=-1)
-        x1 = xx[:, 0]
-        x2 = xx[:, 1]
-        p1 = models.Polynomial1D(5)
-        assert_allclose(p1(x1), yy[0, :], atol=10 ** (-12))
-        p1 = models.Polynomial1D(5)
-        assert_allclose(p1(x2), yy[1, :], atol=10 ** (-12))
+        assert_allclose(xx.shape, yy.shape)
+        yy1 = p1(xx, model_set_axis=1)
+        assert_allclose(yy, yy1)
+        #x1 = xx[:, 0]
+        #x2 = xx[:, 1]
+        #p1 = models.Polynomial1D(5)
+        #assert_allclose(p1(x1), yy[0, :], atol=10 ** (-12))
+        #p1 = models.Polynomial1D(5)
+        #assert_allclose(p1(x2), yy[1, :], atol=10 ** (-12))
 
     def test_evaluate_gauss2d(self):
         cov = np.array([[1., 0.8], [0.8, 3]])
diff --git a/astropy/modeling/tests/test_model_sets.py b/astropy/modeling/tests/test_model_sets.py
new file mode 100644
index 0000000000..b5805236a6
--- /dev/null
+++ b/astropy/modeling/tests/test_model_sets.py
@@ -0,0 +1,224 @@
+# Licensed under a 3-clause BSD style license - see LICENSE.rst
+"""
+This module tests model set evaluation for some common use cases.
+"""
+import pytest
+import numpy as np
+from numpy.testing.utils import assert_allclose
+
+from ..models import Polynomial1D, Polynomial2D
+from ..fitting import LinearLSQFitter
+from ..core import Model
+from ..parameters import Parameter
+
+
+x = np.arange(4)
+xx = np.array([x, x + 10])
+xxx = np.arange(24).reshape((3, 4, 2))
+
+
+class TParModel(Model):
+    """
+    A toy model to test parameters machinery
+    """
+    #standard_broadasting = False
+    inputs = ('x',)
+    outputs = ('x',)
+    coeff = Parameter()
+    e = Parameter()
+
+    def __init__(self, coeff, e, **kwargs):
+        super().__init__(coeff=coeff, e=e, **kwargs)
+
+    @staticmethod
+    def evaluate(x, coeff, e):
+        return x*coeff + e
+
+
+def test_model_axis_1():
+    """
+    Test that a model initialized with model_set_axis=1
+    can be evaluated with model_set_axis=False.
+    """
+
+    p1 = Polynomial1D(1, n_models=2, model_set_axis=1)
+    p1.c0 = [2, 3]
+    p1.c1 = [1, 2]
+    t1 = Polynomial1D(1, c0=2, c1=1)
+    t2 = Polynomial1D(1, c0=3, c1=2)
+
+    with pytest.raises(ValueError):
+        p1(x)
+
+    with pytest.raises(ValueError):
+        p1(xx)
+
+    y = p1(x, model_set_axis=False)
+    assert len(y) == 2
+    assert_allclose(y[0], t1(x))
+    assert_allclose(y[1], t2(x))
+
+    y = p1(xx, model_set_axis=False)
+    assert len(y) == 2
+    assert_allclose(y[0], t1(xx))
+    assert_allclose(y[1], t2(xx))
+
+    y = p1(xxx, model_set_axis=False)
+    assert_allclose(y[0], t1(xxx))
+    assert_allclose(y[1], t2(xxx))
+    assert len(y) == 2
+
+
+def test_model_axis_2():
+    """
+    Test that a model initialized with model_set_axis=2
+    can be evaluated with model_set_axis=False.
+    """
+    p1 = Polynomial1D(1, c0=[[[1, 2,3 ]]], c1=[[[10, 20, 30]]],
+                      n_models=3, model_set_axis=2)
+    t1 = Polynomial1D(1, c0=1, c1=10)
+    t2 = Polynomial1D(1, c0=2, c1=20)
+    t3 = Polynomial1D(1, c0=3, c1=30)
+
+    with pytest.raises(ValueError):
+        p1(x)
+
+    with pytest.raises(ValueError):
+        p1(xx)
+
+    y = p1(x, model_set_axis=False)
+    assert y.shape == (4, 1, 3)
+    assert_allclose(y[:, :, 0].flatten(), t1(x))
+    assert_allclose(y[:, :, 1].flatten(), t2(x))
+    assert_allclose(y[:, :, 2].flatten(), t3(x))
+
+    p2 = Polynomial2D(1, c0_0=[[[0,1,2]]], c0_1=[[[3,4,5]]],
+                      c1_0=[[[5,6,7]]], n_models=3, model_set_axis=2)
+    t1 = Polynomial2D(1, c0_0=0, c0_1=3, c1_0=5)
+    t2 = Polynomial2D(1, c0_0=1, c0_1=4, c1_0=6)
+    t3 = Polynomial2D(1, c0_0=2, c0_1=5, c1_0=7)
+
+    assert p2.c0_0.shape == ()
+    y = p2(x, x, model_set_axis=False)
+    assert y.shape == (4, 1, 3)
+    # These are columns along the 2nd axis.
+    assert_allclose(y[:, :, 0].flatten(), t1(x, x))
+    assert_allclose(y[:, :, 1].flatten(), t2(x, x))
+    assert_allclose(y[:, :, 2].flatten(), t3(x, x))
+
+
+def test_axis_0():
+    """
+    Test that a model initialized with model_set_axis=0
+    can be evaluated with model_set_axis=False.
+    """
+    x = np.arange(4)
+    xx = np.array([x, x + 10])
+    xxx = np.arange(24).reshape((3, 4, 2))
+
+    p1 = Polynomial1D(1, n_models=2, model_set_axis=0)
+    p1.c0 = [2, 3]
+    p1.c1 = [1, 2]
+    t1 = Polynomial1D(1, c0=2, c1=1)
+    t2 = Polynomial1D(1, c0=3, c1=2)
+
+    with pytest.raises(ValueError):
+        p1(x)
+
+    y = p1(xx)
+    assert len(y) == 2
+    assert_allclose(y[0], t1(xx[0]))
+    assert_allclose(y[1], t2(xx[1]))
+
+    y = p1(x, model_set_axis=False)
+    assert len(y) == 2
+    assert_allclose(y[0], t1(x))
+    assert_allclose(y[1], t2(x))
+
+    y = p1(xx, model_set_axis=False)
+    assert len(y) == 2
+    assert_allclose(y[0], t1(xx))
+    assert_allclose(y[1], t2(xx))
+    y = p1(xxx, model_set_axis=False)
+    assert_allclose(y[0], t1(xxx))
+    assert_allclose(y[1], t2(xxx))
+    assert len(y) == 2
+
+
+def test_negative_axis():
+    p1 = Polynomial1D(1, c0=[1, 2], c1=[3, 4], n_models=2, model_set_axis=-1)
+    x = np.arange(4)
+    xx = np.array([x, x + 10])
+    t1 = Polynomial1D(1, c0=1,c1=3)
+    t2 = Polynomial1D(1, c0=2,c1=4)
+
+    with pytest.raises(ValueError):
+        p1(x)
+
+    with pytest.raises(ValueError):
+        p1(xx)
+
+    xxt = xx.T
+    y = p1(xxt)
+    assert_allclose(y[: ,0], t1(xxt[: ,0]))
+    assert_allclose(y[: ,1], t2(xxt[: ,1]))
+
+
+def test_shapes():
+    p2 = Polynomial1D(1, n_models=3, model_set_axis=2)
+    assert p2.c0.shape == ()
+    assert p2.c1.shape == ()
+
+    p1 = Polynomial1D(1, n_models=2, model_set_axis=1)
+    assert p1.c0.shape == ()
+    assert p1.c1.shape == ()
+
+    p1 = Polynomial1D(1, c0=[1, 2], c1=[3, 4], n_models=2, model_set_axis=-1)
+    assert p1.c0.shape == ()
+    assert p1.c1.shape == ()
+
+    e1 = [1, 2]
+    e2 = [3, 4]
+
+    a1 = np.array([[10, 20], [30, 40]])
+    a2 = np.array([[50, 60], [70, 80]])
+
+    t = TParModel([a1, a2], [e1, e2], n_models=2, model_set_axis=-1)
+    assert t.coeff.shape == (2, 2)
+    assert t.e.shape == (2,)
+
+
+
+    t = TParModel([[a1, a2]], [[e1, e2]], n_models=2, model_set_axis=1)
+    assert t.coeff.shape == (2, 2)
+    assert t.e.shape == (2,)
+
+    t = TParModel([a1, a2], [e1, e2], n_models=2, model_set_axis=0)
+    assert t.coeff.shape == (2, 2)
+    assert t.e.shape == (2,)
+
+    t = TParModel([a1, a2], e=[1, 2], n_models=2, model_set_axis=0)
+    assert t.coeff.shape == (2, 2)
+    assert t.e.shape == ()
+
+
+def test_linearlsqfitter():
+    """
+    Issue #7159
+    """
+    x = np.arange(0,4)
+    p = Polynomial1D(1, n_models=2, model_set_axis=1)
+
+    # Generate data for fitting 2 models and re-stack them along the last axis:
+    y = np.array([2*x+1, x+4])
+    y = np.rollaxis(y, 0, -1).T
+
+    f = LinearLSQFitter()
+    # This seems to fit the model_set correctly:
+    fit = f(p, x, y)
+    model_y = fit(x, model_set_axis=False)
+
+    m1 = Polynomial1D(1, c0=fit.c0[0][0], c1=fit.c1[0][0])
+    m2 = Polynomial1D(1, c0=fit.c0[0][1], c1=fit.c1[0][1])
+    assert_allclose(model_y[0], m1(x))
+    assert_allclose(model_y[1], m2(x))
diff --git a/astropy/modeling/tests/test_parameters.py b/astropy/modeling/tests/test_parameters.py
index b28ae5f403..2b38cb13c1 100644
--- a/astropy/modeling/tests/test_parameters.py
+++ b/astropy/modeling/tests/test_parameters.py
@@ -534,22 +534,22 @@ class TestParameterInitialization:
     def test_two_model_nonzero_model_set_axis(self):
         # An example where the model set axis is the *last* axis of the
         # parameter arrays
-        coeff = np.array([[[10, 20], [30, 40]], [[50, 60], [70, 80]]])
+        coeff = np.array([[[10, 20, 30], [30, 40, 50]], [[50, 60, 70], [70, 80, 90]]])
         coeff = np.rollaxis(coeff, 0, 3)
-        e = np.array([[1, 2], [3, 4]])
+        e = np.array([[1, 2, 3], [3, 4, 5]])
         e = np.rollaxis(e, 0, 2)
-        t = TParModel(coeff, e, model_set_axis=-1)
+        t = TParModel(coeff, e, n_models=2, model_set_axis=-1)
         assert len(t) == 2
         assert t.model_set_axis == -1
         assert len(t.param_sets) == 2
         assert np.issubdtype(t.param_sets.dtype, np.object_)
-        assert np.all(t.param_sets[0] == [[[10, 50], [20, 60]],
-                                          [[30, 70], [40, 80]]])
-        assert np.all(t.param_sets[1] == [[[1, 3], [2, 4]]])
-        assert np.all(t.parameters == [10, 50, 20, 60, 30, 70, 40, 80,
-                                       1, 3, 2, 4])
-        assert t.coeff.shape == (2, 2)
-        assert t.e.shape == (2,)
+        assert np.all(t.param_sets[0] == [[[10, 50], [20, 60], [30, 70]],
+                                          [[30, 70], [40, 80], [50, 90]]])
+        assert np.all(t.param_sets[1] == [[[1, 3], [2, 4], [3, 5]]])
+        assert np.all(t.parameters == [10, 50, 20, 60, 30, 70, 30, 70, 40, 80,
+                                       50, 90, 1, 3, 2, 4, 3, 5])
+        assert t.coeff.shape == (2, 3)
+        assert t.e.shape == (3,)
 
     def test_wrong_number_of_params(self):
         with pytest.raises(InputParameterError):