1import numpy as np
 2
 3META = {"name": "deformation_registration", "domain": "geometry", "description": "3D coordinate deformation regression with a smooth twist and radial compression; exposes spatial coordinates and deformation gradients."}
 4
 5def _map(x, amp=1.0):
 6    z=x[:,2]; th=amp*1.8*z
 7    c=np.cos(th); s=np.sin(th)
 8    xx=.68*x[:,0]; yy=.68*x[:,1]
 9    return np.stack([c*xx-s*yy, s*xx+c*yy, z], axis=1).astype(np.float32)
10
11def get_dataset(seed, n_train=400, n_test=200):
12    rng=np.random.default_rng(seed)
13    xtr=rng.uniform(-1,1,(n_train,3)).astype(np.float32)
14    xte=np.random.default_rng(seed+10000).uniform(-1,1,(n_test,3)).astype(np.float32)
15    # Fixed moderately large deformation makes the structural regularizer relevant.
16    return {"xtr":xtr,"ytr":_map(xtr,2.0),"xte":xte,"yte":_map(xte,2.0),
17            "task":"regression","metric":"mse","input_shape":(3,),"out_dim":3}