1from builtins import range
2from ._tools import *
3
4def hardcoded_expansion(x, degree):
5    nvars = x.shape[1]
6    exp_dim = mdp.nodes._expanded_dim(degree, nvars)
7    exp = numx.zeros((x.shape[0], exp_dim), 'd')
8    # degree 1
9    exp[:,:nvars] = x.copy()
10    # degree 2
11    k = nvars
12    if degree>=2:
13        for i in range(nvars):
14            for j in range(i,nvars):
15                exp[:,k] = x[:,i]*x[:,j]
16                k += 1
17    # degree 3
18    if degree>=3:
19        for i in range(nvars):
20            for j in range(i,nvars):
21                for l in range(j,nvars):
22                    exp[:,k] = x[:,i]*x[:,j]*x[:,l]
23                    k += 1
24    # degree 4
25    if degree>=4:
26        for i in range(nvars):
27            for j in range(i,nvars):
28                for l in range(j,nvars):
29                    for m in range(l,nvars):
30                        exp[:,k] = x[:,i]*x[:,j]*x[:,l]*x[:,m]
31                        k += 1
32    # degree 5
33    if degree>=5:
34        for i in range(nvars):
35            for j in range(i,nvars):
36                for l in range(j,nvars):
37                    for m in range(l,nvars):
38                        for n in range(m,nvars):
39                            exp[:,k] = \
40                                     x[:,i]*x[:,j]*x[:,l]*x[:,m]*x[:,n]
41                            k += 1
42    return exp
43
44def test_expansion():
45    for degree in range(1,6):
46        for dim in range(1,5):
47            expand = mdp.nodes.PolynomialExpansionNode(degree=degree)
48            mat,mix,inp = get_random_mix((10,dim))
49            des = hardcoded_expansion(inp, degree)
50            exp = expand.execute(inp)
51            assert_array_almost_equal(exp, des, decimal)
52