TD-to-PDE Continuation Training / poisson_track.py
Failed on benchmark
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,)}