import numpy as np

META = {
    'name': 'local_poisson_boundary',
    'domain': 'pde',
    'description': 'Unit-disc Poisson Dirichlet regression with harmonic cubic solution; inputs are interior points and targets are observed values.'
}

def get_dataset(seed, n_train=400, n_test=400):
    rng = np.random.RandomState(seed)
    def sample(n):
        r = np.sqrt(rng.uniform(0.0, 1.0, n))
        th = rng.uniform(0.0, 2.0 * np.pi, n)
        x = np.stack([r * np.cos(th), r * np.sin(th)], axis=1).astype(np.float32)
        y = (x[:, 0] ** 3 - 3.0 * x[:, 0] * x[:, 1] ** 2).astype(np.float32)[:, None]
        return x, y
    xtr, ytr = sample(n_train)
    xte, yte = sample(n_test)
    return {'xtr': xtr, 'ytr': ytr, 'xte': xte, 'yte': yte,
            'task': 'regression', 'metric': 'mse', 'out_dim': 1,
            'input_shape': (2,)}