| |
| """Generate the FAST-DeepONet Darcy inclusion benchmark from a fixed seed.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| from pathlib import Path |
|
|
| import numpy as np |
| from scipy import sparse |
| from scipy.sparse.linalg import spsolve |
|
|
|
|
| GRID = 64 |
| SEED = 33 |
| SPLITS = (("train", 60), ("tune", 10), ("calib", 20), ("test", 10)) |
| OFFICIAL_COMMIT = "9541c112176c42470cd39138dacbab8af6892232" |
|
|
|
|
| def make_basis(x: np.ndarray, y: np.ndarray) -> np.ndarray: |
| modes = ( |
| (1, 0, 0, 0), (0, 1, 0, 0), (1, 1, 0, 0), (1, 0, 1, 0), |
| (0, 1, 0, 1), (1, 1, 1, 1), (2, 0, 0, 0), (0, 2, 0, 0), |
| ) |
| fields = [] |
| for kx, ky, px, py in modes: |
| x_wave = np.cos(2 * np.pi * kx * x + px * np.pi / 2) if kx else 1.0 |
| y_wave = np.cos(2 * np.pi * ky * y + py * np.pi / 2) if ky else 1.0 |
| field = x_wave * y_wave |
| fields.append(field / np.max(np.abs(field))) |
| return np.stack(fields) |
|
|
|
|
| def site_locations(rng: np.random.Generator) -> np.ndarray: |
| axis = np.linspace(0.2, 0.8, 4) |
| candidates = np.array([(a, b) for a in axis for b in axis]) |
| sites = candidates[rng.permutation(len(candidates))[:10]] |
| return sites + rng.uniform(-0.04, 0.04, size=sites.shape) |
|
|
|
|
| def solve_pressure(log_coefficient: np.ndarray) -> np.ndarray: |
| coefficient = np.exp(log_coefficient) |
| inner = coefficient[1:-1, 1:-1] |
| harmonic = lambda neighbor: 2 * inner * neighbor / (inner + neighbor) |
| east = harmonic(coefficient[2:, 1:-1]) |
| west = harmonic(coefficient[:-2, 1:-1]) |
| north = harmonic(coefficient[1:-1, 2:]) |
| south = harmonic(coefficient[1:-1, :-2]) |
| side = GRID - 2 |
| nodes = np.arange(side * side).reshape(side, side) |
| spacing_squared = (1.0 / GRID) ** 2 |
|
|
| connections = ( |
| (nodes[:-1, :], nodes[1:, :], east[:-1, :]), |
| (nodes[1:, :], nodes[:-1, :], west[1:, :]), |
| (nodes[:, :-1], nodes[:, 1:], north[:, :-1]), |
| (nodes[:, 1:], nodes[:, :-1], south[:, 1:]), |
| ) |
| row = np.concatenate([nodes.ravel()] + [src.ravel() for src, _, _ in connections]) |
| col = np.concatenate([nodes.ravel()] + [dst.ravel() for _, dst, _ in connections]) |
| diagonal = (east + west + north + south).ravel() / spacing_squared |
| values = np.concatenate( |
| [diagonal] + [-edge.ravel() / spacing_squared for _, _, edge in connections] |
| ) |
| operator = sparse.csr_matrix( |
| (values, (row, col)), shape=(side * side, side * side) |
| ) |
| pressure = np.zeros((GRID, GRID)) |
| pressure[1:-1, 1:-1] = spsolve( |
| operator, np.ones(side * side) |
| ).reshape(side, side) |
| return pressure |
|
|
|
|
| def generate_sample( |
| rng: np.random.Generator, |
| basis: np.ndarray, |
| sites: np.ndarray, |
| x: np.ndarray, |
| y: np.ndarray, |
| ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: |
| factors = rng.normal(size=len(basis)) |
| weights = np.sqrt(np.exp(-0.5 * np.arange(len(basis)))) |
| log_coefficient = np.tensordot(factors * weights, basis, axes=1) * 0.8 |
|
|
| count = rng.integers(3, 7) |
| chosen = sites[rng.permutation(len(sites))[:count]] |
| mask = np.zeros((GRID, GRID)) |
| for center in chosen: |
| sign = rng.choice([-1, 1]) |
| magnitude = rng.uniform(1.2, 1.8) |
| inclusion = ( |
| (x - center[0]) ** 2 + (y - center[1]) ** 2 <= 0.045 ** 2 |
| ).astype(np.float64) |
| log_coefficient += sign * magnitude * inclusion |
| mask = np.maximum(mask, inclusion) |
|
|
| log_coefficient = np.clip(log_coefficient, -3, 3) |
| pressure = solve_pressure(log_coefficient) |
| return tuple( |
| array.astype(np.float32).ravel() |
| for array in (log_coefficient, pressure, mask) |
| ) |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument( |
| "--output-dir", |
| type=Path, |
| default=Path(__file__).resolve().parents[1] / "data", |
| ) |
| parser.add_argument( |
| "--overwrite", action="store_true", help="Replace existing NPZ samples" |
| ) |
| args = parser.parse_args() |
| root = args.output_dir.resolve() |
| existing = [path for split, _ in SPLITS for path in (root / split).glob("*.npz")] |
| if existing and not args.overwrite: |
| parser.error(f"{root} already contains NPZ samples; use --overwrite to replace") |
|
|
| axis = (np.arange(GRID) + 0.5) / GRID |
| x, y = np.meshgrid(axis, axis, indexing="ij") |
| rng = np.random.default_rng(SEED) |
| basis = make_basis(x, y) |
| sites = site_locations(rng) |
|
|
| for split, count in SPLITS: |
| directory = root / split |
| directory.mkdir(parents=True, exist_ok=True) |
| for index in range(count): |
| permeability, pressure, mask = generate_sample(rng, basis, sites, x, y) |
| np.savez( |
| directory / f"darcy_inclusion_{split}_{index:03d}.npz", |
| nx=np.int64(GRID), |
| ny=np.int64(GRID), |
| pde_name="darcy_inclusion", |
| init_value=permeability, |
| last_value=pressure, |
| sparse_mask=mask, |
| ) |
| print(f"{split}: generated {count} samples in {directory}", flush=True) |
|
|
| manifest = { |
| "source": "https://github.com/kwonji39/FAST-DeepONet", |
| "source_commit": OFFICIAL_COMMIT, |
| "experiment": "Table 2 Darcy inclusion", |
| "seed": SEED, |
| "grid_size": [GRID, GRID], |
| "splits": dict(SPLITS), |
| "arrays": ["init_value", "last_value", "sparse_mask"], |
| } |
| root.mkdir(parents=True, exist_ok=True) |
| (root / "DATA_MANIFEST.json").write_text( |
| json.dumps(manifest, indent=2) + "\n", encoding="utf-8" |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
