1import numpy as np
 2META = {"name":"symmetric_periodic_poisson","domain":"pde","description":"2-D periodic Poisson inverse on an 8x8 square grid; source fields and exact zero-mean solutions."}
 3
 4def _lap_eigs(n):
 5    k=np.arange(n)
 6    return (4-2*np.cos(2*np.pi*k/n)[:,None]-2*np.cos(2*np.pi*k/n)[None,:]).astype(np.float32)
 7
 8def get_dataset(seed, n_train, n_test):
 9    def sample(n, s):
10        rng=np.random.RandomState(s); N=8
11        # Random Fourier sources with no constant mode, ensuring solvability.
12        f=rng.normal(size=(n,N,N)).astype(np.float32)
13        fh=np.fft.fft2(f,axes=(-2,-1)); fh[:,0,0]=0
14        lam=_lap_eigs(N); uh=fh/(lam[None]+1e-5); uh[:,0,0]=0
15        u=np.fft.ifft2(uh,axes=(-2,-1)).real.astype(np.float32)
16        # Normalize each sample to avoid scale-driven triviality.
17        f=f/(f.std(axis=(1,2),keepdims=True)+1e-4); u=u/(u.std(axis=(1,2),keepdims=True)+1e-4)
18        return f.reshape(n,-1),u.reshape(n,-1)
19    xtr,ytr=sample(n_train,seed); xte,yte=sample(n_test,seed+5000)
20    return {"xtr":xtr,"ytr":ytr,"xte":xte,"yte":yte,"task":"regression","metric":"mse","out_dim":64}