| from __future__ import annotations |
|
|
| from collections.abc import Mapping |
|
|
| import numpy as np |
| import torch |
|
|
| from model.sa_pinn import Burgers2D, Equation, Helmholtz2D, Laplace1D |
|
|
|
|
| CASES = ("laplace", "helmholtz", "burgers") |
|
|
|
|
| def exact_solution(case: str, coordinates: np.ndarray) -> np.ndarray | None: |
| if case == "laplace": |
| return np.sin(np.pi * coordinates[:, 0:1]) |
| if case == "helmholtz": |
| return np.sin(np.pi * coordinates[:, 0:1]) * np.sin( |
| 4.0 * np.pi * coordinates[:, 1:2] |
| ) |
| if case == "burgers": |
| return None |
| raise ValueError(f"unsupported case: {case}") |
|
|
|
|
| def build_equation(case: str, data_config: Mapping) -> Equation: |
| if case == "laplace": |
| return Laplace1D() |
| if case == "helmholtz": |
| return Helmholtz2D(float(data_config["wave_number"])) |
| if case == "burgers": |
| return Burgers2D(float(data_config["viscosity"])) |
| raise ValueError(f"unsupported case: {case}") |
|
|
|
|
| def generate_data(case: str, config: Mapping, seed: int) -> dict[str, np.ndarray | None]: |
| generator = np.random.default_rng(seed) |
| n_pde = int(config["n_pde"]) |
| n_solution = int(config["n_sol"]) |
| n_boundary = int(config["n_bnd"]) |
| test_resolution = int(config["test_res"]) |
| if min(n_pde, n_solution, n_boundary, test_resolution) <= 0: |
| raise ValueError("all data counts must be positive") |
|
|
| if case == "laplace": |
| lower, upper = map(float, config["domain"]) |
| x_pde = generator.uniform(lower, upper, (n_pde, 1)) |
| x_boundary = np.array([[lower], [upper]], dtype=np.float64) |
| u_boundary = exact_solution(case, x_boundary) |
| x_data = np.linspace(lower, upper, n_solution)[:, None] |
| u_data = exact_solution(case, x_data) |
| noise_std = float(config.get("noise_std", 0.0)) |
| if noise_std: |
| u_data = u_data + noise_std * generator.standard_normal(u_data.shape) |
| x_test = np.linspace(lower, upper, test_resolution)[:, None] |
| test_shape = (test_resolution,) |
| elif case == "helmholtz": |
| x_lower, x_upper = map(float, config["domain_x"]) |
| y_lower, y_upper = map(float, config["domain_y"]) |
| x_pde = np.column_stack( |
| ( |
| generator.uniform(x_lower, x_upper, n_pde), |
| generator.uniform(y_lower, y_upper, n_pde), |
| ) |
| ) |
| if n_boundary < 4 or n_boundary % 4: |
| raise ValueError("Helmholtz n_bnd must be divisible by 4 and at least 4") |
| per_edge = n_boundary // 4 |
| x_axis = np.linspace(x_lower, x_upper, per_edge) |
| y_axis = np.linspace(y_lower, y_upper, per_edge) |
| x_boundary = np.vstack( |
| ( |
| np.column_stack((x_axis, np.full(per_edge, y_lower))), |
| np.column_stack((x_axis, np.full(per_edge, y_upper))), |
| np.column_stack((np.full(per_edge, x_lower), y_axis)), |
| np.column_stack((np.full(per_edge, x_upper), y_axis)), |
| ) |
| ) |
| u_boundary = exact_solution(case, x_boundary) |
| x_data = np.column_stack( |
| ( |
| generator.uniform(x_lower, x_upper, n_solution), |
| generator.uniform(y_lower, y_upper, n_solution), |
| ) |
| ) |
| u_data = exact_solution(case, x_data) |
| grid_x = np.linspace(x_lower, x_upper, test_resolution) |
| grid_y = np.linspace(y_lower, y_upper, test_resolution) |
| mesh_x, mesh_y = np.meshgrid(grid_x, grid_y, indexing="xy") |
| x_test = np.column_stack((mesh_x.ravel(), mesh_y.ravel())) |
| test_shape = mesh_x.shape |
| elif case == "burgers": |
| x_lower, x_upper = map(float, config["domain_x"]) |
| t_lower, t_upper = map(float, config["domain_t"]) |
| x_pde = np.column_stack( |
| ( |
| generator.uniform(x_lower, x_upper, n_pde), |
| generator.uniform(t_lower, t_upper, n_pde), |
| ) |
| ) |
| if n_boundary < 2: |
| raise ValueError("Burgers n_bnd must be at least 2") |
| left_count = n_boundary // 2 |
| right_count = n_boundary - left_count |
| left_time = generator.uniform(t_lower, t_upper, left_count) |
| right_time = generator.uniform(t_lower, t_upper, right_count) |
| x_boundary = np.vstack( |
| ( |
| np.column_stack((np.full(left_count, x_lower), left_time)), |
| np.column_stack((np.full(right_count, x_upper), right_time)), |
| ) |
| ) |
| u_boundary = np.zeros((n_boundary, 1)) |
| initial_x = generator.uniform(x_lower, x_upper, n_solution) |
| x_data = np.column_stack((initial_x, np.full(n_solution, t_lower))) |
| u_data = -np.sin(np.pi * initial_x)[:, None] |
| grid_x = np.linspace(x_lower, x_upper, test_resolution) |
| grid_t = np.linspace(t_lower, t_upper, test_resolution) |
| mesh_x, mesh_t = np.meshgrid(grid_x, grid_t, indexing="xy") |
| x_test = np.column_stack((mesh_x.ravel(), mesh_t.ravel())) |
| test_shape = mesh_x.shape |
| else: |
| raise ValueError(f"unsupported case: {case}") |
|
|
| return { |
| "x_pde": np.asarray(x_pde, dtype=np.float64), |
| "x_boundary": np.asarray(x_boundary, dtype=np.float64), |
| "u_boundary": np.asarray(u_boundary, dtype=np.float64), |
| "x_data": np.asarray(x_data, dtype=np.float64), |
| "u_data": np.asarray(u_data, dtype=np.float64), |
| "x_test": np.asarray(x_test, dtype=np.float64), |
| "u_exact": exact_solution(case, np.asarray(x_test, dtype=np.float64)), |
| "test_shape": test_shape, |
| } |
|
|
|
|
| def to_tensors( |
| data: Mapping[str, np.ndarray | None], |
| device: torch.device, |
| dtype: torch.dtype, |
| ) -> dict[str, torch.Tensor | None]: |
| return { |
| key: ( |
| torch.as_tensor(value, dtype=dtype, device=device) |
| if isinstance(value, np.ndarray) |
| else None |
| ) |
| for key, value in data.items() |
| if key != "test_shape" |
| } |
|
|
|
|
| def point_counts(data: Mapping[str, np.ndarray | None]) -> dict[str, int]: |
| return { |
| "pde": int(data["x_pde"].shape[0]), |
| "boundary": int(data["x_boundary"].shape[0]), |
| "data": int(data["x_data"].shape[0]) if data.get("x_data") is not None else 0, |
| } |
|
|
|
|
| def relative_l2(prediction: np.ndarray, reference: np.ndarray | None) -> float | None: |
| if reference is None: |
| return None |
| return float( |
| np.linalg.norm(prediction.reshape(-1) - reference.reshape(-1)) |
| / (np.linalg.norm(reference.reshape(-1)) + 1.0e-12) |
| ) |
|
|
