1import math
 2import numpy as np
 3
 4META = {
 5    'name': 'two_view_gauge_localization',
 6    'domain': 'geometry/localization',
 7    'description': 'Planar two-view global and relay-frame observations for target localization with an unknown yaw and translation gauge.'
 8}
 9
10def _rot(a):
11    c, s = np.cos(a), np.sin(a)
12    return np.array([[c, -s], [s, c]], dtype=np.float32)
13
14def _make(seed, n):
15    rng = np.random.default_rng(seed)
16    x = rng.uniform(-2, 2, (n, 2)).astype(np.float32)
17    p = rng.uniform(-2, 2, (n, 2)).astype(np.float32)
18    psi = rng.uniform(-math.pi, math.pi, n)
19    q1 = rng.uniform(-1, 1, (n, 2)).astype(np.float32)
20    dq = rng.uniform(-2, 2, (n, 2)).astype(np.float32)
21    dq += (np.linalg.norm(dq, axis=1) < .25)[:, None] * np.array([.4, .1], np.float32)
22    q2 = q1 + dq
23    ev1 = np.empty_like(q1); ev2 = np.empty_like(q1); et = np.empty_like(q1)
24    for i, a in enumerate(psi):
25        R = _rot(-a)
26        ev1[i] = R @ (q1[i] - x[i]); ev2[i] = R @ (q2[i] - x[i]); et[i] = R @ (p[i] - x[i])
27    noise = .035
28    for z in (q1, q2, ev1, ev2, et):
29        z += rng.normal(0, noise, z.shape).astype(np.float32)
30    return np.concatenate([q1, q2, ev1, ev2, et], axis=1).astype(np.float32), p.astype(np.float32)
31
32def get_dataset(seed, n_train, n_test):
33    xtr, ytr = _make(seed, n_train)
34    xte, yte = _make(seed + 5000, n_test)
35    return {'xtr': xtr, 'ytr': ytr, 'xte': xte, 'yte': yte,
36            'task': 'regression', 'metric': 'mse', 'out_dim': 2}