1"""Local custom track: supervised harmonic Dirichlet/Poisson-style collocation."""
 2import numpy as np
 3
 4META = {
 5    "name": "poisson_harmonic_square",
 6    "domain": "pde",
 7    "description": "Unit-square collocation for the harmonic Dirichlet solution u(x,y)=x^2-y^2; interior PDE residual is Delta u=0 and values provide boundary/data supervision."
 8}
 9
10def get_dataset(seed, n_train, n_test):
11    rng = np.random.default_rng(int(seed))
12    # Keep the train/test distributions identical while retaining a genuine
13    # collocation domain; the PDE loss is evaluated on every training point.
14    xtr = rng.uniform(-1.0, 1.0, size=(int(n_train), 2)).astype(np.float32)
15    xte = rng.uniform(-1.0, 1.0, size=(int(n_test), 2)).astype(np.float32)
16    def solution(x):
17        return (x[:, 0] ** 2 - x[:, 1] ** 2)[:, None].astype(np.float32)
18    return {"xtr": xtr, "ytr": solution(xtr), "xte": xte, "yte": solution(xte),
19            "task": "regression", "metric": "mse", "input_shape": (2,), "out_dim": 1}