File size: 5,624 Bytes
b887635 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 | #!/usr/bin/env python3
"""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()
|