# 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 import assert_allclose from astropy.modeling.models import Polynomial1D, Polynomial2D from astropy.modeling.fitting import LinearLSQFitter from astropy.modeling.core import Model from astropy.modeling.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. """ model_axis = 1 n_models = 2 p1 = Polynomial1D(1, n_models=n_models, model_set_axis=model_axis) 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 y.shape[model_axis] == n_models assert_allclose(y[:, 0], t1(x)) assert_allclose(y[:, 1], t2(x)) y = p1(xx, model_set_axis=False) assert y.shape[model_axis] == n_models assert_allclose(y[:, 0, :], t1(xx)) assert_allclose(y[:, 1, :], t2(xx)) y = p1(xxx, model_set_axis=False) assert y.shape[model_axis] == n_models assert_allclose(y[:, 0, :, :], t1(xxx)) assert_allclose(y[:, 1, :, :], t2(xxx)) 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 == (1, 4, 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 == (1, 4, 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. """ 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) 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 """ 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)) def test_model_set_axis_outputs(): fitter = LinearLSQFitter() model_set = Polynomial2D(1, n_models=2, model_set_axis=2) y2, x2 = np.mgrid[: 5, : 5] # z = np.moveaxis([x2 + y2, 1 - 0.1 * x2 + 0.2 * y2]), 0, 2) z = np.rollaxis(np.array([x2 + y2, 1 - 0.1 * x2 + 0.2 * y2]), 0, 3) model = fitter(model_set, x2, y2, z) res = model(x2, y2, model_set_axis=False) assert z.shape == res.shape # Test initializing with integer model_set_axis # and evaluating with a different model_set_axis model_set = Polynomial1D(1, c0=[1, 2], c1=[2, 3], n_models=2, model_set_axis=0) y0 = model_set(xx) y1 = model_set(xx.T, model_set_axis=1) assert_allclose(y0[0], y1[:, 0]) assert_allclose(y0[1], y1[:, 1]) model_set = Polynomial1D(1, c0=[[1, 2]], c1=[[2, 3]], n_models=2, model_set_axis=1) y0 = model_set(xx.T) y1 = model_set(xx, model_set_axis=0) assert_allclose(y0[:, 0], y1[0]) assert_allclose(y0[:, 1], y1[1]) with pytest.raises(ValueError): model_set(x)