1import numpy as np
 2
 3META = {'name': 'kpz_field_forecast', 'domain': 'pde', 'description': 'One-step forecasting of a periodic 3-D KPZ field.'}
 4
 5def _step(h, nu=0.22, lam=1.8, dt=0.025):
 6    lap = np.zeros_like(h)
 7    grad2 = np.zeros_like(h)
 8    for ax in (2, 3, 4):
 9        hp, hm = np.roll(h, -1, ax), np.roll(h, 1, ax)
10        lap += hp - 2*h + hm
11        grad2 += ((hp-hm)/2.0)**2
12    return h + dt*(nu*lap + 0.5*lam*grad2)
13
14def get_dataset(seed, n_train, n_test):
15    rng = np.random.default_rng(seed)
16    n = n_train + n_test
17    h = rng.normal(size=(n,1,8,8,8)).astype(np.float32)
18    for ax in (2,3,4):
19        h = 0.55*h + 0.15*(np.roll(h,1,ax)+np.roll(h,-1,ax))
20    h /= h.std(axis=(2,3,4), keepdims=True) + 1e-6
21    y = _step(h).astype(np.float32)
22    return {'xtr':h[:n_train], 'ytr':y[:n_train], 'xte':h[n_train:], 'yte':y[n_train:], 'task':'regression', 'metric':'mse', 'out_dim':1}