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