1import numpy as np
 2
 3META = {
 4    'name': 'mixed_junction_pde',
 5    'domain': 'pde',
 6    'description': 'Manufactured mixed Dirichlet-Neumann junction field on a square; the corner at the origin has Kondratjev r^1/2 sin(theta/2) regularity.'
 7}
 8
 9def get_dataset(seed, n_train, n_test):
10    rng = np.random.RandomState(seed)
11    # Positive quadrant is a local model of a straight mixed junction.
12    xtr = rng.uniform(0.0, 1.0, (n_train, 2)).astype(np.float32)
13    xte = rng.uniform(0.0, 1.0, (n_test, 2)).astype(np.float32)
14    def field(x):
15        r = np.sqrt(np.sum(x*x, axis=1) + 1e-14)
16        th = np.arctan2(x[:, 1], x[:, 0])
17        # Singular harmonic mode for a mixed corner, plus smooth harmonic-like terms.
18        return (r**0.5 * np.sin(0.5*th) +
19                0.20*np.sin(np.pi*x[:, 0])*np.sin(np.pi*x[:, 1]) +
20                0.10*x[:, 0]*x[:, 1]).astype(np.float32)
21    return {'xtr': xtr, 'ytr': field(xtr).reshape(-1, 1), 'xte': xte, 'yte': field(xte).reshape(-1, 1),
22            'task': 'regression', 'metric': 'mse', 'out_dim': 1}