| """Generate sparse structured observations and collocation coordinates lazily.""" |
|
|
| import argparse |
| import importlib.util |
| from pathlib import Path |
|
|
| import numpy as np |
| import yaml |
|
|
|
|
| ROOT = Path(__file__).resolve().parents[1] |
|
|
|
|
| def load_model_module(): |
| spec = importlib.util.spec_from_file_location("pinn_tc_model", ROOT / "model/pinn-tc.py") |
| module = importlib.util.module_from_spec(spec) |
| spec.loader.exec_module(module) |
| return module |
|
|
|
|
| def structured_observations(config, mode): |
| data = config["data"] |
| extent = float(data["horizontal_extent_m"]) |
| levels = np.linspace(float(data["pressure_min_pa"]), float(data["pressure_max_pa"]), int(data["pressure_levels"])) |
| times = np.asarray(data["observation_times_s"], dtype=np.float64) |
| if times.shape != (int(data["time_steps"]),) or not np.array_equal(times, [-10800.0, 0.0, 10800.0]): |
| raise ValueError("paper observation times must be [-3h, 0h, +3h]") |
| line = np.linspace(-extent, extent, int(data["observation_line_points"])) |
| coordinates = [] |
| for time_index, time in enumerate(times): |
| active_mode = ("Cross" if time_index % 2 == 0 else "Plus") if mode == "Switch" else mode |
| for pressure in levels[:: int(data["observation_pressure_stride"])]: |
| if active_mode == "Plus": |
| coordinates.extend((y, x, time, pressure) for y, x in zip(np.zeros_like(line), line)) |
| coordinates.extend((y, x, time, pressure) for y, x in zip(line, np.zeros_like(line))) |
| elif active_mode == "Cross": |
| coordinates.extend((y, x, time, pressure) for y, x in zip(line, line)) |
| coordinates.extend((y, x, time, pressure) for y, x in zip(line, -line)) |
| else: |
| raise ValueError("observation mode must be Cross, Plus, or Switch") |
| edge = np.linspace(-extent, extent, int(data["boundary_points_per_edge"])) |
| coordinates.extend((-extent, value, time, pressure) for value in edge) |
| coordinates.extend((extent, value, time, pressure) for value in edge) |
| coordinates.extend((value, -extent, time, pressure) for value in edge[1:-1]) |
| coordinates.extend((value, extent, time, pressure) for value in edge[1:-1]) |
| return np.unique(np.asarray(coordinates, dtype=np.float32), axis=0) |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--mode", choices=("Cross", "Plus", "Switch")) |
| parser.add_argument("--force", action="store_true") |
| args = parser.parse_args() |
| config = yaml.safe_load((ROOT / "conf/config.yaml").read_text()) |
| mode = args.mode or config["data"]["observation_mode"] |
| output = ROOT / config["data"]["root"] |
| output.mkdir(parents=True, exist_ok=True) |
| path = output / "training_points.npz" |
| if path.exists() and not args.force: |
| print(f"exists={path.relative_to(ROOT)} (use --force to regenerate)") |
| return |
| module = load_model_module() |
| observations = structured_observations(config, mode) |
| targets = module.analytic_vortex_numpy(observations)[:, :3] |
| rng = np.random.default_rng(int(config["seed"])) |
| count = int(config["data"]["collocation_points"]) |
| collocation = np.column_stack(( |
| rng.uniform(-config["data"]["horizontal_extent_m"], config["data"]["horizontal_extent_m"], (count, 2)), |
| rng.uniform(-config["data"]["time_extent_s"], config["data"]["time_extent_s"], count), |
| rng.uniform(config["data"]["pressure_min_pa"], config["data"]["pressure_max_pa"], count), |
| )).astype(np.float32) |
| |
| np.savez(path, observation_coordinates=observations, observation_targets=targets, |
| collocation_coordinates=collocation, input_order=np.asarray(module.INPUT_ORDER), |
| supervised_order=np.asarray(module.OUTPUT_ORDER[:3]), observation_mode=np.asarray(mode), |
| dense_shape=np.asarray([data_dim := int(config["data"]["grid_points"]), data_dim, |
| int(config["data"]["pressure_levels"]), int(config["data"]["time_steps"])])) |
| print(f"generated={path.relative_to(ROOT)} observations={len(observations)} collocation={len(collocation)} mode={mode}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
