| """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: |
| 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", |
| ] |
|
|
