File size: 4,207 Bytes
7087ec2 | 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 | import math
import os
import sys
from pathlib import Path
import numpy as np
PROJECT_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(PROJECT_ROOT))
from onescience.utils.YParams import YParams
def set_default_data_env():
os.environ.setdefault("ONESCIENCE_BENO_DATA_DIR", str(PROJECT_ROOT / "data"))
def resolve_path(path_value):
path = Path(path_value)
return path if path.is_absolute() else PROJECT_ROOT / path
def load_config():
set_default_data_env()
cfg = YParams(str(PROJECT_ROOT / "conf" / "config.yaml"), "root")
cfg.datapipe.source.data_dir = str(resolve_path(cfg.datapipe.source.data_dir))
cfg.datapipe.source.cache_dir = str(resolve_path(cfg.datapipe.source.cache_dir))
cfg.fake_data.root_dir = str(resolve_path(cfg.fake_data.root_dir))
return cfg
def square_boundary(points, boundary_points):
bottom = [(x, 0.0) for x in points]
top = [(x, 1.0) for x in points]
left = [(0.0, y) for y in points[1:-1]]
right = [(1.0, y) for y in points[1:-1]]
boundary = np.array(bottom + top + left + right, dtype=np.float64)
if boundary.shape[0] == 0:
raise ValueError("resolution is too small to build boundary points")
repeats = int(math.ceil(boundary_points / boundary.shape[0]))
return np.tile(boundary, (repeats, 1))[:boundary_points]
def make_beno_arrays(cfg):
data_cfg = cfg.datapipe.data
fake_cfg = cfg.fake_data
total_samples = int(data_cfg.ntrain) + int(data_cfg.ntest)
resolution = int(data_cfg.resolution)
boundary_points = int(fake_cfg.boundary_points)
rng = np.random.default_rng(int(fake_cfg.seed))
axis = np.linspace(0.0, 1.0, resolution, dtype=np.float64)
yy, xx = np.meshgrid(axis, axis, indexing="ij")
coords = np.stack([xx.reshape(-1), yy.reshape(-1)], axis=-1)
cell_state = np.zeros((resolution, resolution), dtype=np.float64)
cell_state[0, :] = 1.0
cell_state[-1, :] = 1.0
cell_state[:, 0] = 1.0
cell_state[:, -1] = 1.0
cell_state = cell_state.reshape(-1)
boundary_xy = square_boundary(axis, boundary_points)
rhs = np.zeros((total_samples, resolution * resolution, 4), dtype=np.float64)
sol = np.zeros((total_samples, resolution * resolution, 1), dtype=np.float64)
bc = np.zeros((total_samples, boundary_points, 4), dtype=np.float64)
for sample_idx in range(total_samples):
phase = 0.25 * sample_idx
amplitude = 1.0 + 0.1 * sample_idx
forcing = (
amplitude * np.sin(np.pi * xx + phase) * np.sin(np.pi * yy)
+ 0.15 * np.cos(2.0 * np.pi * xx) * np.sin(np.pi * yy + phase)
)
solution = 0.25 * forcing + 0.05 * (xx + yy)
if fake_cfg.noise_std:
solution += float(fake_cfg.noise_std) * rng.standard_normal(solution.shape)
boundary_value = (
0.05 * sample_idx
+ 0.1 * boundary_xy[:, 0]
- 0.08 * boundary_xy[:, 1]
)
rhs[sample_idx, :, 0:2] = coords
rhs[sample_idx, :, 2] = forcing.reshape(-1)
rhs[sample_idx, :, 3] = cell_state
sol[sample_idx, :, 0] = solution.reshape(-1)
bc[sample_idx, :, 0:2] = boundary_xy
bc[sample_idx, :, 2] = boundary_value
return rhs, sol, bc
def clear_matching_cache(cfg):
cache_dir = Path(cfg.datapipe.source.cache_dir)
prefix = cfg.datapipe.source.file_prefix
for split, count in (("train", cfg.datapipe.data.ntrain), ("test", cfg.datapipe.data.ntest)):
cache_file = cache_dir / f"cached_{prefix}_{split}_{count}.pt"
if cache_file.exists():
cache_file.unlink()
def main():
cfg = load_config()
data_dir = Path(cfg.datapipe.source.data_dir)
data_dir.mkdir(parents=True, exist_ok=True)
rhs, sol, bc = make_beno_arrays(cfg)
prefix = cfg.datapipe.source.file_prefix
np.save(data_dir / f"RHS_{prefix}_all.npy", rhs)
np.save(data_dir / f"SOL_{prefix}_all.npy", sol)
np.save(data_dir / f"BC_{prefix}_all.npy", bc)
clear_matching_cache(cfg)
print(f"Fake BENO data written to {data_dir}")
print(f"RHS shape={rhs.shape}, SOL shape={sol.shape}, BC shape={bc.shape}")
if __name__ == "__main__":
main()
|