File size: 2,665 Bytes
371b59f | 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 | from __future__ import annotations
from collections.abc import Mapping
from pathlib import Path
import numpy as np
import scipy.io
import torch
def load_allen_cahn(path: Path) -> dict[str, np.ndarray | tuple[int, int]]:
if not path.is_file():
raise FileNotFoundError(f"Allen-Cahn data not found: {path}")
data = scipy.io.loadmat(path)
missing = {"tt", "x", "uu"}.difference(data)
if missing:
raise ValueError(f"Allen-Cahn data is missing fields: {sorted(missing)}")
time_axis = np.asarray(data["tt"], dtype=np.float64).reshape(-1)
space_axis = np.asarray(data["x"], dtype=np.float64).reshape(-1)
raw_solution = np.real(np.asarray(data["uu"]))
expected_shape = (space_axis.size, time_axis.size)
if raw_solution.shape != expected_shape:
raise ValueError(
f"uu shape must be {expected_shape} in [x,t] order, got {raw_solution.shape}"
)
exact_grid = raw_solution.T
mesh_x, mesh_t = np.meshgrid(space_axis, time_axis, indexing="xy")
coordinates = np.column_stack((mesh_x.ravel(), mesh_t.ravel()))
exact = exact_grid.reshape(-1, 1)
return {
"time": time_axis,
"space": space_axis,
"coordinates": coordinates,
"exact": exact,
"grid_shape": exact_grid.shape,
}
def sample_training_data(
dataset: Mapping,
n_train: int,
seed: int,
device: torch.device,
dtype: torch.dtype,
) -> tuple[torch.Tensor, torch.Tensor]:
coordinates = np.asarray(dataset["coordinates"])
exact = np.asarray(dataset["exact"])
if n_train <= 0 or n_train > coordinates.shape[0]:
raise ValueError(
f"n_train must be between 1 and {coordinates.shape[0]}, got {n_train}"
)
generator = np.random.default_rng(seed)
indices = generator.choice(coordinates.shape[0], n_train, replace=False)
return (
torch.as_tensor(coordinates[indices], dtype=dtype, device=device),
torch.as_tensor(exact[indices], dtype=dtype, device=device),
)
def batched_predict(
model: torch.nn.Module,
coordinates: np.ndarray,
batch_size: int,
device: torch.device,
dtype: torch.dtype,
) -> np.ndarray:
if batch_size <= 0:
raise ValueError("evaluation batch size must be positive")
predictions = []
model.eval()
with torch.no_grad():
for start in range(0, coordinates.shape[0], batch_size):
batch = torch.as_tensor(
coordinates[start : start + batch_size], dtype=dtype, device=device
)
predictions.append(model.predict_u(batch).cpu().numpy())
return np.concatenate(predictions, axis=0)
|