1import numpy as np
 2
 3META = {
 4    "name": "local_shock_characteristic",
 5    "domain": "pde",
 6    "description": "1-D five-variable shock/acoustic local prediction task with equilibrium and oscillatory modes.",
 7}
 8
 9
10def get_dataset(seed, n_train, n_test):
11    def make(n, stream):
12        rng = np.random.default_rng(int(seed) + stream)
13        L = 9
14        x = np.linspace(-1.0, 1.0, L)
15        xs, ys = [], []
16        for _ in range(int(n)):
17            center = rng.uniform(-0.35, 0.35)
18            width = rng.uniform(0.035, 0.11)
19            shock = 0.5 * (1.0 + np.tanh((x - center) / width))
20            h = 1.0 + rng.uniform(0.05, 0.35) * shock + 0.012 * rng.normal(size=L)
21            q = rng.uniform(-0.18, 0.18) + 0.05 * shock + 0.009 * rng.normal(size=L)
22            theta = 1.0 + rng.uniform(-0.15, 0.15) + 0.03 * shock + 0.009 * rng.normal(size=L)
23            envelope = np.exp(-((x - center) / rng.uniform(0.18, 0.32)) ** 2)
24            amp = rng.uniform(0.03, 0.22)
25            ring = ((-1.0) ** np.arange(L)) * envelope
26            a4 = amp * ring + 0.003 * rng.normal(size=L)
27            a5 = -0.8 * amp * ring + 0.003 * rng.normal(size=L)
28            state = np.stack([h, q, theta, a4, a5], axis=1)
29            # Stable local one-step target: center value plus a small physical-like flux.
30            d = state[1:] - state[:-1]
31            target = state[4] - 0.18 * (state[5] - state[3]) + 0.08 * (d[3] - d[2])
32            xs.append(state.astype(np.float32).reshape(-1))
33            ys.append(np.float32(target[0]))
34        return np.asarray(xs, dtype=np.float32), np.asarray(ys, dtype=np.float32)[:, None]
35
36    xtr, ytr = make(n_train, 0)
37    xte, yte = make(n_test, 5000)
38    return {"xtr": xtr, "ytr": ytr, "xte": xte, "yte": yte,
39            "task": "regression", "metric": "mse", "out_dim": 1}