TD-to-PDE Continuation Training / poisson_track.py

Failed on benchmark

Raw ⬇ ZIP
 1import numpy as np
 2
 3META = {
 4    'name': 'local_poisson_boundary',
 5    'domain': 'pde',
 6    'description': 'Unit-disc Poisson Dirichlet regression with harmonic cubic solution; inputs are interior points and targets are observed values.'
 7}
 8
 9def get_dataset(seed, n_train=400, n_test=400):
10    rng = np.random.RandomState(seed)
11    def sample(n):
12        r = np.sqrt(rng.uniform(0.0, 1.0, n))
13        th = rng.uniform(0.0, 2.0 * np.pi, n)
14        x = np.stack([r * np.cos(th), r * np.sin(th)], axis=1).astype(np.float32)
15        y = (x[:, 0] ** 3 - 3.0 * x[:, 0] * x[:, 1] ** 2).astype(np.float32)[:, None]
16        return x, y
17    xtr, ytr = sample(n_train)
18    xte, yte = sample(n_test)
19    return {'xtr': xtr, 'ytr': ytr, 'xte': xte, 'yte': yte,
20            'task': 'regression', 'metric': 'mse', 'out_dim': 1,
21            'input_shape': (2,)}