File size: 4,282 Bytes
2f3c9e4 | 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 | """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)
# The first two random columns already follow the required [y, x] order.
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()
|