1import numpy as np
 2
 3META = {
 4    "name": "rough_energy_regression",
 5    "domain": "optimization",
 6    "description": "Regression on a smooth baseline plus a finite Weierstrass rough component; used to test dilation-matched updates in a neural training system."
 7}
 8
 9A, B, N, K = 0.8, 3.0, 12, 0.25
10
11def phi(x):
12    return np.sin(2*np.pi*x)
13
14def rough(x):
15    x = np.asarray(x, dtype=np.float64)
16    z = np.zeros(x.shape[0])
17    for n in range(N+1):
18        z += A**n * phi(B**n*x[:, 0])
19    return z
20
21def get_dataset(seed, n_train=400, n_test=400):
22    rng = np.random.RandomState(int(seed))
23    xtr = rng.uniform(0.2, 3.8, size=(n_train, 4)).astype(np.float32)
24    xte = rng.uniform(0.2, 3.8, size=(n_test, 4)).astype(np.float32)
25    def target(x):
26        smooth = 0.35*(x[:,0]-2.0)**2 + 0.15*x[:,1] - 0.10*x[:,2] + 0.08*x[:,3]
27        return (smooth + K*rough(x)).astype(np.float32)
28    return {"xtr":xtr, "ytr":target(xtr)[:,None], "xte":xte,
29            "yte":target(xte)[:,None], "task":"regression", "metric":"mse",
30            "input_shape":(4,), "out_dim":1}