import numpy as np

META = {
    'name': 'pde_patch_consensus',
    'domain': 'pde',
    'description': 'Noisy observations of an exact 1D advection-diffusion solution.'
}
V = 0.7
D = 0.015
K = 2*np.pi

def _field(x, t):
    return np.sin(K*(x - V*t))*np.exp(-D*K*K*t)

def get_dataset(seed, n_train, n_test):
    rng = np.random.default_rng(seed)
    def sample(n, noise):
        x = rng.random(n); t = rng.random(n)
        y = _field(x, t) + noise*rng.normal(size=n)
        return np.stack([x, t], axis=1).astype('float32'), y.astype('float32')
    xtr, ytr = sample(n_train, 0.05)
    xte, yte = sample(n_test, 0.0)
    return {'xtr': xtr, 'ytr': ytr, 'xte': xte, 'yte': yte,
            'task': 'regression', 'metric': 'mse', 'out_dim': 1}