1import numpy as np
 2
 3META = {
 4    "name": "poisson_field_operator",
 5    "domain": "pde",
 6    "description": "Small manufactured 2-D Poisson/BVP operator: Fourier forcing coefficients to a spatial solution field.",
 7}
 8
 9
10def get_dataset(seed, n_train, n_test):
11    rng = np.random.RandomState(seed)
12    side = 8
13    xx, yy = np.meshgrid(np.linspace(0, 1, side), np.linspace(0, 1, side), indexing="ij")
14    modes = [(1, 1), (1, 2), (2, 1), (2, 2), (1, 3), (3, 1)]
15
16    def sample(n, rs):
17        z = rs.normal(size=(n, len(modes))).astype(np.float32)
18        fields = []
19        for i, (kx, ky) in enumerate(modes):
20            basis = np.sin(np.pi * kx * xx) * np.sin(np.pi * ky * yy)
21            # Inverse-Laplacian scaling is the manufactured Poisson structure.
22            fields.append(z[:, i, None] * basis.reshape(1, -1) / (kx * kx + ky * ky))
23        u = np.sum(fields, axis=0).astype(np.float32)
24        # Heteroscedastic observation/model discrepancy is localized near a boundary.
25        noise_scale = (0.012 + 0.030 * np.exp(-((xx - .15) ** 2 + (yy - .15) ** 2) / .025)).reshape(1, -1)
26        u += rs.normal(size=u.shape).astype(np.float32) * noise_scale
27        return z, u
28
29    xtr, ytr = sample(n_train, rng)
30    xte, yte = sample(n_test, np.random.RandomState(seed + 5000))
31    return {"xtr": xtr, "ytr": ytr, "xte": xte, "yte": yte,
32            "task": "regression", "metric": "mse", "out_dim": side * side}