datasets.py

import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader

# =========================
# Helpers
# =========================
def rand_fourier_series_1d(N:int, kmax:int=10, decay:float=2.0, rng=None):
    """Random smooth 1D signal via truncated Fourier series."""
    rng = np.random.default_rng() if rng is None else rng
    x = np.linspace(0, 2*np.pi, N, endpoint=False)
    a = np.zeros(N, dtype=np.float32)
    for k in range(1, kmax+1):
        c = rng.normal(0, 1.0)/(k**decay)
        s = rng.normal(0, 1.0)/(k**decay)
        a += c*np.cos(k*x) + s*np.sin(k*x)
    return a.astype(np.float32)

def smooth_random_2d(N:int, kmax:int=6, decay:float=2.5, rng=None):
    """Random smooth 2D field from a low-frequency spectral stencil."""
    rng = np.random.default_rng() if rng is None else rng
    A = np.zeros((N,N), dtype=np.complex64)
    for kx in range(-kmax, kmax+1):
        for ky in range(-kmax, kmax+1):
            if kx==0 and ky==0:
                continue
            amp = rng.normal(0,1.0) / ((1+abs(kx)+abs(ky))**decay)
            phase = 2*np.pi*rng.random()
            A[kx % N, ky % N] = amp*np.exp(1j*phase)
    f = np.fft.ifft2(A).real.astype(np.float32)
    f -= f.mean()
    return f

# Frequency helpers on a unit-length domain (L=1)
def rfftfreq_1d(N):   # 0..N/2 cycles per unit
    return np.fft.rfftfreq(N, d=1.0/N)

def fftfreq_1d(N):    # negative..positive cycles per unit
    return np.fft.fftfreq(N, d=1.0/N)

# =========================
# Existing Datasets
# =========================
class Heat1D(Dataset):
    def __init__(self, N=64, T=0.1, nu=0.01, nsamples=3000, kmax=12, seed=1337):
        self.N, self.T, self.nu = N, T, nu
        self.nsamples, self.kmax = nsamples, kmax
        self.rng = np.random.default_rng(seed)
    def __len__(self): return self.nsamples
    def __getitem__(self, idx):
        a = rand_fourier_series_1d(self.N, self.kmax, 2.0, self.rng)
        A = np.fft.rfft(a)
        k = rfftfreq_1d(self.N)
        lam = (2*np.pi*k)**2
        decay = np.exp(-self.nu * lam * self.T)
        uT = np.fft.irfft(decay * A, n=self.N).astype(np.float32)
        return torch.from_numpy(a)[None], torch.from_numpy(uT)[None]

class Heat2D(Dataset):
    def __init__(self, N=64, T=0.1, nu=0.05, nsamples=1500, kmax=6, seed=42):
        self.N, self.T, self.nu = N, T, nu
        self.nsamples, self.kmax = nsamples, kmax
        self.rng = np.random.default_rng(seed)
    def __len__(self): return self.nsamples
    def __getitem__(self, idx):
        a = smooth_random_2d(self.N, self.kmax, 2.5, self.rng)
        A = np.fft.fft2(a)
        kx = fftfreq_1d(self.N); ky = fftfreq_1d(self.N)
        KX, KY = np.meshgrid(kx, ky, indexing='ij')
        lam = (2*np.pi)**2 * (KX**2 + KY**2)
        decay = np.exp(-self.nu * lam * self.T)
        U = np.fft.ifft2(A * decay).real.astype(np.float32)
        return torch.from_numpy(a)[None], torch.from_numpy(U)[None]

class Poisson2D(Dataset):
    def __init__(self, N=64, nsamples=1500, kmax=6, decay=2.5, seed=7):
        self.N, self.nsamples = N, nsamples
        self.kmax, self.decay = kmax, decay
        self.rng = np.random.default_rng(seed)
    def __len__(self): return self.nsamples
    def __getitem__(self, idx):
        f = smooth_random_2d(self.N, self.kmax, self.decay, self.rng)
        F = np.fft.fft2(f)
        kx = fftfreq_1d(self.N); ky = fftfreq_1d(self.N)
        KX, KY = np.meshgrid(kx, ky, indexing='ij')
        denom = (2*np.pi)**2 * (KX**2 + KY**2)
        denom[0,0] = np.inf  # set mean mode to zero
        U = np.fft.ifft2(F / denom).real.astype(np.float32)
        return torch.from_numpy(f)[None], torch.from_numpy(U)[None]

class Helmholtz2D(Dataset):
    def __init__(self, N=64, lam=1.0, nsamples=1500, kmax=6, decay=2.5, seed=9):
        self.N, self.lam, self.nsamples = N, lam, nsamples
        self.kmax, self.decay = kmax, decay
        self.rng = np.random.default_rng(seed)
    def __len__(self): return self.nsamples
    def __getitem__(self, idx):
        f = smooth_random_2d(self.N, self.kmax, self.decay, self.rng)
        F = np.fft.fft2(f)
        kx = fftfreq_1d(self.N); ky = fftfreq_1d(self.N)
        KX, KY = np.meshgrid(kx, ky, indexing='ij')
        denom = self.lam + (2*np.pi)**2 * (KX**2 + KY**2)
        U = np.fft.ifft2(F / denom).real.astype(np.float32)
        return torch.from_numpy(f)[None], torch.from_numpy(U)[None]

# =========================
# New gradient-sensitive PDEs
# =========================
class Wave1D(Dataset):
    """
    u_tt = c^2 u_xx, periodic, u_t(.,0)=0  =>  û_k(T) = cos(c |k| 2π T) û_k(0)
    """
    def __init__(self, N=64, T=0.25, c=1.0, nsamples=3000, kmax=20, decay=1.2, seed=2025):
        self.N, self.T, self.c = N, T, c
        self.nsamples, self.kmax, self.decay = nsamples, kmax, decay
        self.rng = np.random.default_rng(seed)
    def __len__(self): return self.nsamples
    def __getitem__(self, idx):
        a = rand_fourier_series_1d(self.N, self.kmax, self.decay, self.rng)
        A = np.fft.rfft(a)
        k = rfftfreq_1d(self.N)
        omega = 2*np.pi*np.abs(k)*self.c
        U = np.fft.irfft(np.cos(omega*self.T) * A, n=self.N).astype(np.float32)
        return torch.from_numpy(a)[None], torch.from_numpy(U)[None]

class Advection1D(Dataset):
    """
    u_t + c u_x = 0 (periodic). Solution: u(x,T) = a(x - cT) => phase shift in Fourier space.
    """
    def __init__(self, N=64, T=0.25, c=0.5, nsamples=3000, kmax=24, decay=1.0, seed=17):
        self.N, self.T, self.c = N, T, c
        self.nsamples, self.kmax, self.decay = nsamples, kmax, decay
        self.rng = np.random.default_rng(seed)
    def __len__(self): return self.nsamples
    def __getitem__(self, idx):
        a = rand_fourier_series_1d(self.N, self.kmax, self.decay, self.rng)
        A = np.fft.rfft(a)
        k = rfftfreq_1d(self.N)
        phase = np.exp(-1j * 2*np.pi * k * (self.c * self.T))
        U = np.fft.irfft(A * phase, n=self.N).astype(np.float32)
        return torch.from_numpy(a)[None], torch.from_numpy(U)[None]

class Advection2D(Dataset):
    """
    u_t + v_x u_x + v_y u_y = 0 (periodic). u(x,y,T) = a(x - v_x T, y - v_y T).
    """
    def __init__(self, N=64, T=0.25, vx=0.5, vy=0.3, nsamples=1500, kmax=8, decay=1.2, seed=23):
        self.N, self.T, self.vx, self.vy = N, T, vx, vy
        self.nsamples, self.kmax, self.decay = nsamples, kmax, decay
        self.rng = np.random.default_rng(seed)
    def __len__(self): return self.nsamples
    def __getitem__(self, idx):
        a = smooth_random_2d(self.N, self.kmax, self.decay, self.rng)
        A = np.fft.fft2(a)
        kx = fftfreq_1d(self.N); ky = fftfreq_1d(self.N)
        KX, KY = np.meshgrid(kx, ky, indexing='ij')
        phase = np.exp(-1j * 2*np.pi * (KX*self.vx + KY*self.vy) * self.T)
        U = np.fft.ifft2(A * phase).real.astype(np.float32)
        return torch.from_numpy(a)[None], torch.from_numpy(U)[None]

# =========================
# Loaders
# =========================
def make_loaders(task:str, res:int, batch_size:int=32, **kw):
    """
    Returns (train, val, test) DataLoaders with an 80/10/10 split.
    task in {
      'heat1d','heat2d','poisson2d','helmholtz2d',
      'wave1d','advection1d','advection2d'
    }
    """
    if task=='heat1d':
        ds = Heat1D(N=res, T=kw.get('T',0.1), nu=kw.get('nu',0.01),
                    nsamples=kw.get('nsamples',3000))
    elif task=='heat2d':
        ds = Heat2D(N=res, T=kw.get('T',0.1), nu=kw.get('nu',0.05),
                    nsamples=kw.get('nsamples',1500))
    elif task=='poisson2d':
        ds = Poisson2D(N=res, nsamples=kw.get('nsamples',1500))
    elif task=='helmholtz2d':
        ds = Helmholtz2D(N=res, lam=kw.get('lam',1.0), nsamples=kw.get('nsamples',1500))
    elif task=='wave1d':
        ds = Wave1D(N=res, T=kw.get('T',0.25), c=kw.get('c',1.0),
                    nsamples=kw.get('nsamples',3000))
    elif task=='advection1d':
        ds = Advection1D(N=res, T=kw.get('T',0.25), c=kw.get('c',0.5),
                         nsamples=kw.get('nsamples',3000))
    elif task=='advection2d':
        ds = Advection2D(N=res, T=kw.get('T',0.25), vx=kw.get('vx',0.5), vy=kw.get('vy',0.3),
                         nsamples=kw.get('nsamples',1500))
    else:
        raise ValueError(f"Unknown task {task}")

    # deterministic split (avoid Windows worker issues)
    n=len(ds); a=int(0.8*n); b=int(0.1*n)
    torch.manual_seed(1234)
    indices = torch.randperm(n).tolist()
    train_idx = indices[:a]
    val_idx   = indices[a:a+b]
    test_idx  = indices[a+b:]

    train = torch.utils.data.Subset(ds, train_idx)
    val   = torch.utils.data.Subset(ds, val_idx)
    test  = torch.utils.data.Subset(ds, test_idx)

    # num_workers=0 keeps it portable across OSes/terminals
    return (
        DataLoader(train, batch_size=batch_size, shuffle=True,  num_workers=0, pin_memory=False),
        DataLoader(val,   batch_size=batch_size, shuffle=False, num_workers=0, pin_memory=False),
        DataLoader(test,  batch_size=batch_size, shuffle=False, num_workers=0, pin_memory=False),
    )