1import numpy as np
 2META = {"name":"transport_density","domain":"transport","description":"1D nonnegative transported densities: reconstruct the next state from a density input."}
 3
 4def _density(x, c, w):
 5    y=np.exp(-0.5*((x-c)/w)**2)+0.35*np.exp(-0.5*((x-(c-.14))/(.65*w))**2)
 6    y=np.maximum(y,1e-12); y/=np.trapz(y,x); return y
 7
 8def get_dataset(seed, n_train, n_test):
 9    rng=np.random.RandomState(seed); x=np.linspace(0,1,32,dtype=np.float32)
10    def make(n):
11        U=[]; V=[]
12        for _ in range(n):
13            c=rng.uniform(.18,.82); w=rng.uniform(.035,.09); shift=rng.uniform(-.10,.10)
14            u=_density(x,c,w); v=_density(x,np.clip(c+shift,.08,.92),w)
15            U.append(u.astype(np.float32)); V.append(v.astype(np.float32))
16        return np.asarray(U),np.asarray(V)
17    xtr,ytr=make(n_train); xte,yte=make(n_test)
18    return {"xtr":xtr,"ytr":ytr,"xte":xte,"yte":yte,"task":"regression","metric":"mse","out_dim":32}