1import numpy as np
 2
 3META = {
 4    "name": "manifold_pendulum",
 5    "domain": "dynamics_and_embedded_manifolds",
 6    "description": "Controlled pendulum windows with angular state embedded on S1; predict next embedded state and angular velocity.",
 7}
 8
 9def get_dataset(seed, n_train, n_test):
10    def make(n, s):
11        rng = np.random.RandomState(s)
12        X = np.empty((n, 24), dtype=np.float32)
13        Y = np.empty((n, 3), dtype=np.float32)
14        dt = 0.05
15        for i in range(n):
16            theta = rng.uniform(-np.pi, np.pi)
17            omega = rng.uniform(-1.5, 1.5)
18            damp = rng.uniform(0.05, 0.25)
19            g = rng.uniform(8.5, 11.0)
20            rows = []
21            for _ in range(9):
22                u = rng.uniform(-1.0, 1.0)
23                rows.append((theta, omega, u))
24                omega = omega + dt * (-g * np.sin(theta) - damp * omega + u)
25                theta = theta + dt * omega
26            X[i] = np.asarray(rows[:8], dtype=np.float32).reshape(-1)
27            th, om, _ = rows[8]
28            Y[i] = (np.cos(th), np.sin(th), om)
29        return X, Y
30    xtr, ytr = make(n_train, seed)
31    xte, yte = make(n_test, seed + 5000)
32    return {"xtr": xtr, "ytr": ytr, "xte": xte, "yte": yte,
33            "task": "regression", "metric": "mse"}