How model sets work

modeling_model_sets.py

from astropy.modeling.models import *
import numpy as np
from astropy.modeling.fitting import *

# How to create a multiset model
# There are two arguments that control this - `n_models` and `model_set_axis`.
# The trick with model_sets is that input arguments and model parameters need
# to broadcast appropriately.

x = np.arange(4)
x1 = np.array([x, x]).T
print(x1)

#[[0 0]
# [1 1]
# [2 2]
# [3 3]]

x0 = np.array([x,x])
print(x0)

# [[0 1 2 3]
# [0 1 2 3]]

# If model_set_axis is False the input arguments are used to evaluate all models
# regardless of the dimension of the array.

pf= Polynomial1D(1, n_models=2, model_set_axis=False)
pf.c0.shape

# ()
pf(x)
# array([[ 0.,  0.,  0.,  0.],
#       [ 0.,  0.,  0.,  0.]])

pf(x0)
# array([[[ 0.,  0.,  0.,  0.],
#        [ 0.,  0.,  0.,  0.]],
#
#       [[ 0.,  0.,  0.,  0.],
#        [ 0.,  0.,  0.,  0.]]])

pf(x1)
# array([[[ 0.,  0.,  0.,  0.],
#        [ 0.,  0.,  0.,  0.]],
#
#       [[ 0.,  0.,  0.,  0.],
#        [ 0.,  0.,  0.,  0.]]])

# As comparison for parameter.shape on a single model
ps = Polynomial1D(1)
ps.c0.shape
# ()

# If model_set_axis is no False, inputs and parameters must be multidimensional.

# if model_set_axis=0 the input for each model is taken along
# the 0-th axis of the input array.
p0 = Polynomial1D(1, n_models=2, model_set_axis=0)
p0.c0.shape

# ()
p0(x0)
# array([[ 0.,  0.,  0.,  0.],
#       [ 0.,  0.,  0.,  0.]])

# p0(x1) and p0(x0) raise ValueError as expected

# For model_set_axis=1
p1= Polynomial1D(1, n_models=2, model_set_axis=1)
p1.c0.shape

# (1,)
p1(x1)

# array([[ 0.,  0.],
#       [ 0.,  0.],
#       [ 0.,  0.],
#       [ 0.,  0.]])

# p1(x) and p1(x0) raise ValueError

# If model_set_axis is None, it behaves as model_set_axis=0
pn = Polynomial1D(1, n_models=2, model_set_axis=None)
pn.c0.shape

# ()

pn(x0)
# array([[ 0.,  0.,  0.,  0.],
#       [ 0.,  0.,  0.,  0.]])


# fitting of model sets
f=LinearLSQFitter()

# model_set_axis=False
fpf = f(pf, x, pf(x))
fpf(x)

# array([[ 0.,  0.,  0.,  0.],
#       [ 0.,  0.,  0.,  0.]])

# The two commands below raise an error
#  440             raise ValueError('{0} gives unsupported >2D derivative matrix for '
# --> 441                              'this x/y'.format(type(model_copy).__name__))
#    442 
#    443         # Subtract any terms fixed by the user from (a copy of) the RHS, in

# ValueError: Polynomial1D gives unsupported >2D derivative matrix for this x/y

fpf = f(pf, x1, pf(x1))
fpf= f(pf, x0, pf(x0))

# model_set_axis = 0
# Raises the above error - anything that passes  2D arrays to the 1D polynomial derivative fails
# because the derivatives do not support 2D input

# Passing 1D array as independent variable and ND array as dependent var works
fp0= f(p0, x, p0(x0))
fp0(x0)

# array([[ 0.,  0.,  0.,  0.],
#       [ 0.,  0.,  0.,  0.]])

fp1= f(p1, x, p1(x1)))
fp1(x1)

# array([[ 0.,  0.],
#       [ 0.,  0.],
#       [ 0.,  0.],
#       [ 0.,  0.]])

# If model parameters, inputs to model evaluate and fitters are consistent, everything works

# When a model is initialized with one value of model_set_axis should it be possible to evaluate it with
# a different model_set_axis value? Is this flexibility necessary?
# it works (I think by accident) when a model is initialized with model_set_axis=0
# and evaluated with model_set_axis=False

# it fails (I think correctly) when a model is initialized  with model_set_axis=1
# and evaluated with model_set_axis=False. 
# I think this is correct because model_set_axis=1 means the coefficients will be applied to two 
# different columns in the input, while with model_set_axis=False each set of coefficients will be applied 
# to all columns.
# The reason is basically making sure there's no ambiguity.
# There's one unfortunate side effect and it is the use case when
# the input is a 1D array. In this case there's no ambiguity and each model can be evaluated 
# on the 1D input array. But this doesn't work. Should it?

# Also, if the answer to the above is 'no', should model_set_axis be allowed as an argument
# to the __call__ function?

#### More on shapes of model parameters ###########

# The shape of a parameter is stored in model._param_metrics.
# It is stored there as the shape of the vaue of the parameter.
# However, Parameter.shape returns the shape of the value as if a single model was initialized.
# For example

pf
# <Polynomial1D(1, c0=[ 0., 0.], c1=[ 0., 0.], n_models=2)>

pf.c0.shape
# ()

pf._param_metrics['c0']['shape'
# (2,)
                       
p0
# <Polynomial1D(1, c0=[ 0., 0.], c1=[ 0., 0.], n_models=2)>
                        
p0.c0
# ()
   
p0._param_metrics['c0']['shape']
# (2,)
                        
p1
# <Polynomial1D(1, c0=[[ 0., 0.]], c1=[[ 0., 0.]], n_models=2)>

p1.c0
# (1,)

p1._param_metrics['c0']['shape']
# (1, 2)
                        
# For a model set of 3 models
p2 = Polynomial1D(1, n_models=3, model_set_axis=2)

p2
# <Polynomial1D(1, c0=[[[ 0., 0., 0.]]], c1=[[[ 0., 0., 0.]]], n_models=3)>

p2.c0.shape
# (1, 1)

p2._param_metrics['c0']['shape']
# (1, 1, 3)

# Also, if the answer to the above is 'no', should model_set_axis be allowed as an argument to the __call__ function?

Indeed, if __call__ is always supposed to be consistent with __init__ then why does it have a model_set_axis argument at all? I don't think this possibility had actually occurred to me... Also, if it's an accident that model_set_axis=False works for axis 0, doesn't that just re-open issue #7142?

I'm still mulling this over (Friday brain) but regarding your "should X work" questions, I'm inclined to say "the simpler the better", except that it's useful to be able to supply a single ordinate array for evaluating multiple models (as when fitting)... I actually found it a bit confusing when getting started with modeling that you have to supply x to fitters but (by default) x0 to models.