1import numpy as np
 2
 3META = {
 4    'name': 'bloch_gauge_pde',
 5    'domain': 'pde',
 6    'description': '1D periodic Bloch-factor field surrogate with exact reciprocal-lattice gauge covariance.'
 7}
 8
 9def get_dataset(seed, n_train, n_test):
10    rng = np.random.default_rng(int(seed))
11    n = int(n_train) + int(n_test)
12    c = rng.uniform(-1.0, 1.0, n).astype('float32')
13    x = rng.uniform(0.0, 1.0, n).astype('float32')
14    q0 = rng.uniform(-np.pi, np.pi, n).astype('float32')
15    m = rng.integers(-2, 3, n).astype('float32')
16    q = q0 + m * (2*np.pi)
17    amp = np.exp(0.30j*c*np.cos(2*np.pi*x) + 0.10j*c*np.sin(4*np.pi*x))
18    f0 = amp * (1 + 0.12*q0*np.cos(2*np.pi*x) + 0.05j*q0*np.sin(2*np.pi*x))
19    y = f0 * np.exp(-1j*m*2*np.pi*x)
20    X = np.stack([c, x, q], axis=1).astype('float32')
21    Y = np.stack([y.real, y.imag], axis=1).astype('float32')
22    return {'xtr': X[:n_train], 'ytr': Y[:n_train], 'xte': X[n_train:], 'yte': Y[n_train:], 'task': 'regression', 'metric': 'mse', 'out_dim': 2}