1import numpy as np
 2
 3META = {'name':'tangent_surface_vector_regression','domain':'geometry/point-cloud','description':'Noisy rotated sphere point clouds with a tangent-field magnitude target; the input contains ambient vector features with normal corruption.'}
 4
 5def _rot(rng):
 6    q,r=np.linalg.qr(rng.normal(size=(3,3)))
 7    q=q@np.diag(np.sign(np.diag(r)))
 8    if np.linalg.det(q)<0:q[:,0]*=-1
 9    return q.astype(np.float32)
10
11def _make(seed,n):
12    rng=np.random.default_rng(seed); m=20; X=np.empty((n,m,7),np.float32); Y=np.empty((n,1),np.float32)
13    a=np.array([.8,-.35,.55],np.float32)
14    for b in range(n):
15        z=rng.normal(size=(m,3)).astype(np.float32); z/=np.linalg.norm(z,axis=1,keepdims=True)
16        R=_rot(rng); p=z@R.T; p[1:]=p[0]+.55*(z[1:]-z[0])@R.T; p[1:]/=np.linalg.norm(p[1:],axis=1,keepdims=True)
17        s=p@a; v=a[None,:]-s[:,None]*p
18        v=v+.9*rng.normal(size=(m,1)).astype(np.float32)*p
19        X[b,:,:3]=p; X[b,:,3]=s; X[b,:,4:]=v; Y[b,0]=np.linalg.norm(v[0]-((v[0]@p[0])*p[0]))
20    return X,Y
21
22def get_dataset(seed,n_train,n_test):
23    xtr,ytr=_make(seed,n_train); xte,yte=_make(seed+5000,n_test)
24    return {'xtr':xtr,'ytr':ytr,'xte':xte,'yte':yte,'task':'regression','metric':'mse','out_dim':1}