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