1import math
 2import numpy as np
 3
 4META = {
 5    'name': 'symmetric_polynomial_regression',
 6    'domain': 'symmetric_interactions',
 7    'description': 'Permutation-invariant degree-four polynomial regression benchmark.'
 8}
 9
10def _compositions(d, m):
11    out = []
12    def rec(i, left, cur):
13        if i == d - 1:
14            out.append(tuple(cur + [left]))
15        else:
16            for k in range(left + 1):
17                rec(i + 1, left - k, cur + [k])
18    rec(0, m, [])
19    return out
20
21def get_dataset(seed, n_train, n_test):
22    rng = np.random.RandomState(int(seed))
23    d, m = 8, 4
24    alphas = _compositions(d, m)
25    mult = np.array([math.factorial(m) / np.prod([math.factorial(a) for a in al]) for al in alphas])
26    coeff = rng.normal(0, 0.20, len(alphas)) / np.sqrt(mult)
27    def make(n):
28        x = rng.normal(size=(n, d)).astype(np.float32)
29        z = np.ones((n, len(alphas)), dtype=np.float64)
30        for k, alpha in enumerate(alphas):
31            for j, p in enumerate(alpha):
32                if p:
33                    z[:, k] *= x[:, j].astype(np.float64) ** p
34        y = z.dot(mult * coeff) + rng.normal(0, 0.05, n)
35        return x, y.astype(np.float32)[:, None]
36    xtr, ytr = make(n_train)
37    xte, yte = make(n_test)
38    return {'xtr': xtr, 'ytr': ytr, 'xte': xte, 'yte': yte,
39            'input_shape': (d,), 'out_dim': 1, 'task': 'regression', 'metric': 'mse'}