"""Synthetic datasets for the four main experiments in arXiv:1910.03193. All equations and split rules follow the paper. The numerical choices that the paper omits are configurable and documented in ``config/config.yaml``. """ from __future__ import annotations import copy import hashlib import json import os import tempfile from pathlib import Path from typing import Any, Dict, Mapping, Sequence, Tuple import numpy as np import torch from torch.utils.data import Dataset try: from scipy.integrate import cumulative_trapezoid, solve_ivp from scipy.interpolate import CubicSpline except ImportError: # Deferred error keeps model-only imports usable. cumulative_trapezoid = None solve_ivp = None CubicSpline = None def deep_update(base: Dict[str, Any], override: Mapping[str, Any]) -> Dict[str, Any]: """Recursively update ``base`` without mutating the caller's mapping.""" for key, value in override.items(): if isinstance(value, Mapping) and isinstance(base.get(key), Mapping): base[key] = deep_update(dict(base[key]), value) else: base[key] = copy.deepcopy(value) return base def resolve_config(config: Mapping[str, Any], smoke_test: bool = False) -> Dict[str, Any]: resolved = copy.deepcopy(dict(config)) smoke_override = resolved.pop("smoke_test", {}) if smoke_test: resolved = deep_update(resolved, smoke_override) resolved.setdefault("project", {})["paper_scale"] = not smoke_test return resolved class OperatorDataset(Dataset): """Triplets with compact storage for repeated PDE branch functions.""" def __init__( self, branch_functions: np.ndarray, trunk: np.ndarray, target: np.ndarray, function_index: np.ndarray | None = None, metadata: Mapping[str, Any] | None = None, ) -> None: branch_functions = np.asarray(branch_functions, dtype=np.float32) trunk = np.asarray(trunk, dtype=np.float32) target = np.asarray(target, dtype=np.float32) if function_index is None: function_index = np.arange(len(trunk), dtype=np.int64) function_index = np.asarray(function_index, dtype=np.int64) if branch_functions.ndim != 2 or trunk.ndim != 2 or target.ndim != 2: raise ValueError("branch, trunk and target arrays must all have rank two") if target.shape[1] != 1 or len(trunk) != len(target) or len(trunk) != len(function_index): raise ValueError("trunk, target and function_index lengths must agree") if len(function_index) and ( function_index.min() < 0 or function_index.max() >= len(branch_functions) ): raise ValueError("function_index refers outside branch_functions") self.branch_functions = branch_functions self.trunk = trunk self.target = target self.function_index = function_index self.metadata = dict(metadata or {}) def __len__(self) -> int: return len(self.trunk) def __getitem__(self, index: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: function_id = self.function_index[index] return ( torch.from_numpy(self.branch_functions[function_id]), torch.from_numpy(self.trunk[index]), torch.from_numpy(self.target[index]), ) def expanded_branch(self) -> np.ndarray: return self.branch_functions[self.function_index] class FunctionSpaceSampler: """GRF or Chebyshev function sampler on a reusable fine grid.""" def __init__(self, config: Mapping[str, Any], domain_end: float) -> None: self.config = dict(config) self.domain_end = float(domain_end) self.grid_size = int(self.config["grf_grid_size"]) self.grid = np.linspace(0.0, self.domain_end, self.grid_size, dtype=np.float64) self.kind = str(self.config.get("type", "grf")).lower() self._cholesky: np.ndarray | None = None if self.kind == "grf": length_scale = float(self.config["length_scale"]) distances = self.grid[:, None] - self.grid[None, :] covariance = np.exp(-(distances**2) / (2.0 * length_scale**2)) jitter = float(self.config.get("jitter", 1.0e-13)) identity = np.eye(self.grid_size, dtype=np.float64) for attempt in range(6): try: self._cholesky = np.linalg.cholesky(covariance + jitter * identity) break except np.linalg.LinAlgError: jitter *= 10.0 if self._cholesky is None: raise np.linalg.LinAlgError("GRF covariance Cholesky failed after jitter fallback") elif self.kind != "chebyshev": raise ValueError(f"Unsupported function space {self.kind!r}") def sample(self, count: int, rng: np.random.Generator) -> np.ndarray: if self.kind == "grf": standard_normal = rng.standard_normal((self.grid_size, count)) return (self._cholesky @ standard_normal).T cheb = self.config.get("chebyshev", {}) degree = int(cheb.get("degree", 10)) bound = float(cheb.get("coefficient_bound", 1.0)) coefficients = rng.uniform(-bound, bound, size=(count, degree + 1)) mapped_grid = 2.0 * self.grid / self.domain_end - 1.0 return np.stack( [np.polynomial.chebyshev.chebval(mapped_grid, row) for row in coefficients], axis=0 ) def interpolate(self, values: np.ndarray, points: np.ndarray) -> np.ndarray: values = np.asarray(values, dtype=np.float64) points = np.asarray(points, dtype=np.float64) method = str(self.config.get("interpolation", "cubic")).lower() if method == "cubic": _require_scipy("cubic GRF interpolation") return np.asarray(CubicSpline(self.grid, values, axis=-1)(points)) if method != "linear": raise ValueError(f"Unsupported interpolation method {method!r}") if values.ndim == 1: return np.interp(points, self.grid, values) return np.stack([np.interp(points, self.grid, row) for row in values], axis=0) def _require_scipy(operation: str) -> None: if solve_ivp is None or CubicSpline is None or cumulative_trapezoid is None: raise ImportError(f"SciPy is required for {operation}; install it in the execution environment") def _rowwise_linear_interpolation( grid: np.ndarray, values: np.ndarray, points: np.ndarray ) -> np.ndarray: points = np.clip(np.asarray(points), grid[0], grid[-1]) right = np.searchsorted(grid, points, side="right") right = np.clip(right, 1, len(grid) - 1) left = right - 1 fraction = (points - grid[left]) / (grid[right] - grid[left]) rows = np.arange(len(points)) return values[rows, left] * (1.0 - fraction) + values[rows, right] * fraction def solve_antiderivative( input_grid: np.ndarray, input_values: np.ndarray, query_points: np.ndarray ) -> np.ndarray: _require_scipy("antiderivative reference generation") integral = cumulative_trapezoid(input_values, input_grid, axis=-1, initial=0.0) return _rowwise_linear_interpolation(input_grid, integral, query_points) def solve_nonlinear_ode( input_grid: np.ndarray, input_values: np.ndarray, query_points: np.ndarray, solver_config: Mapping[str, Any], ) -> np.ndarray: _require_scipy("nonlinear ODE reference generation") query_points = np.asarray(query_points, dtype=np.float64) if query_points.ndim != 1: raise ValueError("query_points must be one-dimensional") if not len(query_points): return np.empty(0, dtype=np.float64) interpolant = CubicSpline(input_grid, input_values) maximum = float(np.max(query_points)) if maximum == 0.0: return np.zeros_like(query_points) solution = solve_ivp( lambda x, state: -state**2 + interpolant(x), (0.0, maximum), np.zeros(1, dtype=np.float64), method=str(solver_config.get("method", "RK45")), rtol=float(solver_config.get("rtol", 1.0e-7)), atol=float(solver_config.get("atol", 1.0e-9)), dense_output=True, ) if not solution.success: raise RuntimeError(f"nonlinear ODE solve failed: {solution.message}") return np.asarray(solution.sol(query_points)[0]) def solve_pendulum( input_grid: np.ndarray, input_values: np.ndarray, query_points: np.ndarray, k: float, solver_config: Mapping[str, Any], ) -> np.ndarray: _require_scipy("pendulum reference generation") query_points = np.asarray(query_points, dtype=np.float64) if not len(query_points): return np.empty(0, dtype=np.float64) interpolant = CubicSpline(input_grid, input_values) maximum = float(np.max(query_points)) if maximum == 0.0: return np.zeros_like(query_points) def right_hand_side(time: float, state: np.ndarray) -> np.ndarray: return np.asarray((state[1], -k * np.sin(state[0]) + interpolant(time))) solution = solve_ivp( right_hand_side, (0.0, maximum), np.zeros(2, dtype=np.float64), method=str(solver_config.get("method", "RK45")), rtol=float(solver_config.get("rtol", 1.0e-7)), atol=float(solver_config.get("atol", 1.0e-9)), dense_output=True, ) if not solution.success: raise RuntimeError(f"pendulum solve failed: {solution.message}") return np.asarray(solution.sol(query_points)[0]) def _solve_tridiagonal( lower: np.ndarray, diagonal: np.ndarray, upper: np.ndarray, rhs: np.ndarray ) -> np.ndarray: """Thomas algorithm for a nonsingular tridiagonal system.""" lower = np.asarray(lower, dtype=np.float64).copy() diagonal = np.asarray(diagonal, dtype=np.float64).copy() upper = np.asarray(upper, dtype=np.float64).copy() rhs = np.asarray(rhs, dtype=np.float64).copy() for index in range(1, len(diagonal)): if abs(diagonal[index - 1]) < np.finfo(np.float64).eps: raise np.linalg.LinAlgError("zero pivot in tridiagonal solve") multiplier = lower[index - 1] / diagonal[index - 1] diagonal[index] -= multiplier * upper[index - 1] rhs[index] -= multiplier * rhs[index - 1] output = np.empty_like(rhs) output[-1] = rhs[-1] / diagonal[-1] for index in range(len(diagonal) - 2, -1, -1): output[index] = (rhs[index] - upper[index] * output[index + 1]) / diagonal[index] return output def solve_diffusion_reaction( spatial_input: np.ndarray, *, diffusion: float, reaction: float, space_points: int, time_points: int, solver_config: Mapping[str, Any], ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: """Fully implicit time stepping with second-order centered spatial differences.""" x_grid = np.linspace(0.0, 1.0, int(space_points), dtype=np.float64) t_grid = np.linspace(0.0, 1.0, int(time_points), dtype=np.float64) source = np.asarray(spatial_input, dtype=np.float64) if source.shape != x_grid.shape: raise ValueError(f"spatial_input must have shape {(len(x_grid),)}, got {source.shape}") field = np.zeros((len(t_grid), len(x_grid)), dtype=np.float64) if len(x_grid) < 3 or len(t_grid) < 2: raise ValueError("PDE grid requires at least 3 spatial and 2 temporal points") dx = x_grid[1] - x_grid[0] dt = t_grid[1] - t_grid[0] ratio = float(diffusion) * dt / (dx * dx) tolerance = float(solver_config.get("pde_newton_tolerance", 1.0e-10)) max_iterations = int(solver_config.get("pde_newton_max_iterations", 20)) interior_source = source[1:-1] interior_size = len(interior_source) off_diagonal = np.full(interior_size - 1, -ratio, dtype=np.float64) for time_index in range(1, len(t_grid)): old = field[time_index - 1, 1:-1] estimate = old.copy() for _ in range(max_iterations): padded = np.pad(estimate, (1, 1), mode="constant") laplacian_term = padded[:-2] - 2.0 * estimate + padded[2:] residual = ( estimate - old - ratio * laplacian_term - dt * float(reaction) * estimate**2 - dt * interior_source ) diagonal = 1.0 + 2.0 * ratio - 2.0 * dt * float(reaction) * estimate update = _solve_tridiagonal( off_diagonal, diagonal, off_diagonal, -residual ) estimate += update if np.max(np.abs(update)) <= tolerance: break else: raise RuntimeError( f"PDE Newton solve did not converge at time index {time_index}" ) field[time_index, 1:-1] = estimate if not np.isfinite(field).all(): raise FloatingPointError("PDE solver produced NaN or infinity") return x_grid, t_grid, field def _generate_ode_like( config: Mapping[str, Any], experiment: str, count: int, seed: int, ) -> OperatorDataset: experiment_config = config["experiments"][experiment] function_config = config["function_space"] solver_config = config["solver_defaults"] rng = np.random.default_rng(seed) sampler = FunctionSpaceSampler(function_config, float(experiment_config["domain_end"])) sensors = np.linspace( 0.0, float(experiment_config["domain_end"]), int(experiment_config["sensor_points"]), dtype=np.float64, ) branch = np.empty((count, len(sensors)), dtype=np.float32) trunk = rng.uniform(0.0, float(experiment_config["domain_end"]), size=(count, 1)) target = np.empty((count, 1), dtype=np.float32) chunk_size = int(function_config.get("generation_chunk_size", 128)) for start in range(0, count, chunk_size): stop = min(count, start + chunk_size) fine_values = sampler.sample(stop - start, rng) branch[start:stop] = sampler.interpolate(fine_values, sensors).astype(np.float32) local_queries = trunk[start:stop, 0] if experiment == "antiderivative": target[start:stop, 0] = solve_antiderivative( sampler.grid, fine_values, local_queries ).astype(np.float32) continue for local_index, values in enumerate(fine_values): query = np.asarray([local_queries[local_index]]) if experiment == "nonlinear_ode": answer = solve_nonlinear_ode(sampler.grid, values, query, solver_config) elif experiment == "pendulum": answer = solve_pendulum( sampler.grid, values, query, float(experiment_config["k"]), solver_config, ) else: raise ValueError(f"Unsupported ODE-like experiment {experiment!r}") target[start + local_index, 0] = answer[0] return OperatorDataset( branch, trunk.astype(np.float32), target, metadata={"experiment": experiment, "seed": seed, "function_count": count}, ) def _generate_pde( config: Mapping[str, Any], split: str, seed: int, ) -> OperatorDataset: experiment_config = config["experiments"]["diffusion_reaction"] if split == "train": function_count = int(experiment_config["train_functions"]) points_per_function = int(experiment_config["points_per_function"]) else: function_count = int(experiment_config["test_functions"]) points_per_function = int(experiment_config["test_points_per_function"]) rng = np.random.default_rng(seed) sampler = FunctionSpaceSampler(config["function_space"], 1.0) sensors = np.linspace(0.0, 1.0, int(experiment_config["sensor_points"])) branch = np.empty((function_count, len(sensors)), dtype=np.float32) total_points = function_count * points_per_function trunk = np.empty((total_points, 2), dtype=np.float32) target = np.empty((total_points, 1), dtype=np.float32) function_index = np.repeat(np.arange(function_count, dtype=np.int64), points_per_function) chunk_size = int(config["function_space"].get("generation_chunk_size", 128)) cursor = 0 for start in range(0, function_count, chunk_size): stop = min(function_count, start + chunk_size) fine_batch = sampler.sample(stop - start, rng) branch[start:stop] = sampler.interpolate(fine_batch, sensors).astype(np.float32) for values in fine_batch: spatial_grid = np.linspace(0.0, 1.0, int(experiment_config["space_points"])) spatial_input = sampler.interpolate(values, spatial_grid) x_grid, t_grid, field = solve_diffusion_reaction( spatial_input, diffusion=float(experiment_config["diffusion"]), reaction=float(experiment_config["reaction"]), space_points=int(experiment_config["space_points"]), time_points=int(experiment_config["time_points"]), solver_config=config["solver_defaults"], ) grid_size = len(x_grid) * len(t_grid) flat_indices = rng.choice( grid_size, size=points_per_function, replace=points_per_function > grid_size, ) time_indices, space_indices = np.divmod(flat_indices, len(x_grid)) next_cursor = cursor + points_per_function trunk[cursor:next_cursor, 0] = x_grid[space_indices] trunk[cursor:next_cursor, 1] = t_grid[time_indices] target[cursor:next_cursor, 0] = field[time_indices, space_indices] cursor = next_cursor return OperatorDataset( branch, trunk, target, function_index, metadata={ "experiment": "diffusion_reaction", "split": split, "seed": seed, "function_count": function_count, "points_per_function": points_per_function, "group_isolated": True, }, ) def _fingerprint(config: Mapping[str, Any], experiment: str, split: str, seed: int) -> str: payload = json.dumps( {"config": config, "experiment": experiment, "split": split, "seed": seed}, sort_keys=True, separators=(",", ":"), ).encode("utf-8") return hashlib.sha256(payload).hexdigest() def _cache_path( config: Mapping[str, Any], project_root: Path, experiment: str, split: str, seed: int ) -> Tuple[Path, str]: fingerprint = _fingerprint(config, experiment, split, seed) relative_root = Path(config["paths"]["cache"]) root = relative_root if relative_root.is_absolute() else project_root / relative_root return root / f"{experiment}_{split}_{fingerprint[:16]}.npz", fingerprint def _save_cache(path: Path, dataset: OperatorDataset, fingerprint: str) -> None: path.parent.mkdir(parents=True, exist_ok=True) metadata = dict(dataset.metadata) metadata["fingerprint"] = fingerprint with tempfile.NamedTemporaryFile(dir=path.parent, suffix=".npz", delete=False) as handle: temporary_path = Path(handle.name) try: np.savez_compressed( temporary_path, branch_functions=dataset.branch_functions, trunk=dataset.trunk, target=dataset.target, function_index=dataset.function_index, metadata=np.asarray(json.dumps(metadata, sort_keys=True)), ) os.replace(temporary_path, path) finally: if temporary_path.exists(): temporary_path.unlink() def _load_cache(path: Path, fingerprint: str) -> OperatorDataset: with np.load(path, allow_pickle=False) as payload: metadata = json.loads(str(payload["metadata"].item())) if metadata.get("fingerprint") != fingerprint: raise ValueError(f"cache fingerprint mismatch for {path}") return OperatorDataset( payload["branch_functions"], payload["trunk"], payload["target"], payload["function_index"], metadata, ) def build_split( config: Mapping[str, Any], experiment: str, split: str, project_root: str | Path, *, use_cache: bool = True, ) -> OperatorDataset: """Build or load one independent train/test split.""" if split not in {"train", "test"}: raise ValueError("split must be 'train' or 'test'") if experiment not in config.get("experiments", {}): raise KeyError(f"Unknown experiment {experiment!r}") root = Path(project_root).resolve() base_seed = int(config["runtime"]["seed"]) seed = base_seed + (0 if split == "train" else 100_000) path, fingerprint = _cache_path(config, root, experiment, split, seed) if use_cache and path.exists(): return _load_cache(path, fingerprint) if experiment == "diffusion_reaction": dataset = _generate_pde(config, split, seed) else: size_key = "train_size" if split == "train" else "test_size" dataset = _generate_ode_like( config, experiment, int(config["experiments"][experiment][size_key]), seed ) dataset.metadata.update( { "split": split, "paper_scale": bool(config["project"]["paper_scale"]), "fingerprint": fingerprint, } ) if use_cache: _save_cache(path, dataset, fingerprint) return dataset def build_datasets( config: Mapping[str, Any], experiment: str, project_root: str | Path ) -> Tuple[OperatorDataset, OperatorDataset]: train = build_split(config, experiment, "train", project_root) test = build_split(config, experiment, "test", project_root) return train, test def analytic_input(name: str, coordinates: np.ndarray) -> np.ndarray: if name == "linear": return coordinates if name == "sin_pi": return np.sin(np.pi * coordinates) if name == "sin_2pi": return np.sin(2.0 * np.pi * coordinates) if name == "x_sin_2pi": return coordinates * np.sin(2.0 * np.pi * coordinates) raise KeyError(f"Unknown analytic input {name!r}") def generate_ood_data( config: Mapping[str, Any], experiment: str, query_points: int | None = None ) -> Dict[str, np.ndarray]: if experiment == "diffusion_reaction": raise ValueError("Use generate_pde_grid_case for the PDE") experiment_config = config["experiments"][experiment] count = int(query_points or config["inference"]["ood_query_points"]) domain_end = float(experiment_config["domain_end"]) sensors = np.linspace(0.0, domain_end, int(experiment_config["sensor_points"])) queries = np.linspace(0.0, domain_end, count) fine_grid = np.linspace(0.0, domain_end, max(1000, count)) names = experiment_config.get("ood_functions", ["linear", "sin_pi", "sin_2pi"]) all_branch, all_trunk, all_target, all_labels = [], [], [], [] for name in names: fine_values = analytic_input(str(name), fine_grid) branch = analytic_input(str(name), sensors) if experiment == "antiderivative": _require_scipy("antiderivative OOD reference") integral = cumulative_trapezoid(fine_values, fine_grid, initial=0.0) target = np.interp(queries, fine_grid, integral) elif experiment == "nonlinear_ode": target = solve_nonlinear_ode( fine_grid, fine_values, queries, config["solver_defaults"] ) elif experiment == "pendulum": target = solve_pendulum( fine_grid, fine_values, queries, float(experiment_config["k"]), config["solver_defaults"], ) else: raise ValueError(f"Unsupported experiment {experiment!r}") all_branch.append(np.repeat(branch[None, :], count, axis=0)) all_trunk.append(queries[:, None]) all_target.append(target[:, None]) all_labels.extend([str(name)] * count) return { "branch": np.concatenate(all_branch).astype(np.float32), "trunk": np.concatenate(all_trunk).astype(np.float32), "target": np.concatenate(all_target).astype(np.float32), "labels": np.asarray(all_labels), } def generate_pde_grid_case(config: Mapping[str, Any], seed: int) -> Dict[str, np.ndarray]: experiment_config = config["experiments"]["diffusion_reaction"] sampler = FunctionSpaceSampler(config["function_space"], 1.0) rng = np.random.default_rng(seed) fine_values = sampler.sample(1, rng)[0] sensors = np.linspace(0.0, 1.0, int(experiment_config["sensor_points"])) branch_vector = sampler.interpolate(fine_values, sensors).astype(np.float32) spatial_grid = np.linspace(0.0, 1.0, int(experiment_config["space_points"])) source = sampler.interpolate(fine_values, spatial_grid) x_grid, t_grid, field = solve_diffusion_reaction( source, diffusion=float(experiment_config["diffusion"]), reaction=float(experiment_config["reaction"]), space_points=int(experiment_config["space_points"]), time_points=int(experiment_config["time_points"]), solver_config=config["solver_defaults"], ) x_mesh, t_mesh = np.meshgrid(x_grid, t_grid) trunk = np.column_stack((x_mesh.ravel(), t_mesh.ravel())).astype(np.float32) return { "branch": np.repeat(branch_vector[None, :], len(trunk), axis=0), "trunk": trunk, "target": field.reshape(-1, 1).astype(np.float32), "source": source.astype(np.float32), "x": x_grid.astype(np.float32), "t": t_grid.astype(np.float32), "grid_shape": np.asarray(field.shape, dtype=np.int64), } __all__ = [ "OperatorDataset", "FunctionSpaceSampler", "resolve_config", "build_split", "build_datasets", "generate_ood_data", "generate_pde_grid_case", "solve_antiderivative", "solve_nonlinear_ode", "solve_pendulum", "solve_diffusion_reaction", ]