1import numpy as np
 2
 3META = {
 4    'name': 'pde_patch_consensus',
 5    'domain': 'pde',
 6    'description': 'Noisy observations of an exact 1D advection-diffusion solution.'
 7}
 8V = 0.7
 9D = 0.015
10K = 2*np.pi
11
12def _field(x, t):
13    return np.sin(K*(x - V*t))*np.exp(-D*K*K*t)
14
15def get_dataset(seed, n_train, n_test):
16    rng = np.random.default_rng(seed)
17    def sample(n, noise):
18        x = rng.random(n); t = rng.random(n)
19        y = _field(x, t) + noise*rng.normal(size=n)
20        return np.stack([x, t], axis=1).astype('float32'), y.astype('float32')
21    xtr, ytr = sample(n_train, 0.05)
22    xte, yte = sample(n_test, 0.0)
23    return {'xtr': xtr, 'ytr': ytr, 'xte': xte, 'yte': yte,
24            'task': 'regression', 'metric': 'mse', 'out_dim': 1}