1"""Custom manifold-valued one-step dynamics benchmark."""
 2import numpy as np
 3
 4META = {
 5    "name": "sphere_one_step_dynamics",
 6    "domain": "geometric_dynamics",
 7    "description": "Predict one geodesic step of a controlled smooth vector field on S2."
 8}
 9
10
11def _exp(x, v):
12    r = np.linalg.norm(v, axis=1, keepdims=True)
13    return np.cos(r) * x + np.sinc(r / np.pi) * v
14
15
16def get_dataset(seed, n_train, n_test):
17    rng = np.random.default_rng(seed)
18
19    def make(n):
20        z = rng.uniform(-0.85, 0.85, n)
21        angle = rng.uniform(0, 2 * np.pi, n)
22        q = np.sqrt(1 - z * z)
23        x = np.column_stack([q * np.cos(angle), q * np.sin(angle), z]).astype(np.float32)
24        c = np.array([0.35, -0.42, 0.28], dtype=np.float32)
25        b = c[None, :] - x * (x * c[None, :]).sum(1, keepdims=True)
26        y = _exp(x, 0.12 * b).astype(np.float32)
27        return x, y
28
29    xtr, ytr = make(n_train)
30    xte, yte = make(n_test)
31    return {
32        "xtr": xtr, "ytr": ytr, "xte": xte, "yte": yte,
33        "task": "regression", "metric": "mse", "out_dim": 3,
34    }