File size: 3,457 Bytes
e9b4f6f | 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 | from __future__ import annotations
import os
import sys
from pathlib import Path
import numpy as np
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
os.chdir(ROOT)
from onescience.utils.YParams import YParams
def make_edges(num_nodes: int) -> np.ndarray:
src = np.arange(num_nodes, dtype=np.int64)
dst = (src + 1) % num_nodes
edges = np.stack(
[np.concatenate([src, dst]), np.concatenate([dst, src])],
axis=0,
)
return edges
def write_placeholder_training_files(param_dir: Path, rng: np.random.Generator) -> None:
param_dir.mkdir(parents=True, exist_ok=True)
np.save(param_dir / "Cd.npy", rng.normal(size=(1,)).astype(np.float32))
np.save(param_dir / "I1.npy", rng.normal(size=(8, 3)).astype(np.float32))
np.save(param_dir / "I2.npy", rng.normal(size=(8, 3)).astype(np.float32))
np.save(param_dir / "Press.npy", rng.normal(size=(8, 1)).astype(np.float32))
np.save(param_dir / "Velo.npy", rng.normal(size=(8, 3)).astype(np.float32))
def write_sample(sample_dir: Path, num_nodes: int, rng: np.random.Generator) -> tuple[np.ndarray, np.ndarray]:
sample_dir.mkdir(parents=True, exist_ok=True)
pos = rng.uniform(low=-1.0, high=1.0, size=(num_nodes, 3)).astype(np.float32)
sdf = rng.uniform(low=0.0, high=0.2, size=(num_nodes, 1)).astype(np.float32)
normals = rng.normal(size=(num_nodes, 3)).astype(np.float32)
normals /= np.linalg.norm(normals, axis=1, keepdims=True) + 1e-6
x = np.concatenate([pos, sdf, normals], axis=1).astype(np.float32)
y = np.concatenate(
[
0.1 * pos + rng.normal(scale=0.01, size=(num_nodes, 3)),
rng.normal(scale=0.05, size=(num_nodes, 1)),
],
axis=1,
).astype(np.float32)
surf = np.zeros((num_nodes,), dtype=np.bool_)
surf[num_nodes // 2 :] = True
edge_index = make_edges(num_nodes)
np.save(sample_dir / "x.npy", x)
np.save(sample_dir / "y.npy", y)
np.save(sample_dir / "pos.npy", pos)
np.save(sample_dir / "surf.npy", surf)
np.save(sample_dir / "edge_index.npy", edge_index)
return x, y
def main() -> None:
cfg = YParams(str(ROOT / "conf/config.yaml"), "datapipe")
data_dir = ROOT / cfg.source.data_dir
preprocessed_dir = ROOT / cfg.source.preprocessed_save_dir
stats_dir = ROOT / cfg.source.stats_dir
stats_dir.mkdir(parents=True, exist_ok=True)
rng = np.random.default_rng(42)
all_x: list[np.ndarray] = []
all_y: list[np.ndarray] = []
for fold_id in range(9):
param_name = f"param{fold_id}"
param_dir = data_dir / param_name
write_placeholder_training_files(param_dir, rng)
sample_name = "sample_000"
(param_dir / sample_name).mkdir(parents=True, exist_ok=True)
x, y = write_sample(preprocessed_dir / param_name / sample_name, num_nodes=8, rng=rng)
all_x.append(x)
all_y.append(y)
x_all = np.concatenate(all_x, axis=0)
y_all = np.concatenate(all_y, axis=0)
np.save(stats_dir / "mean_in.npy", x_all.mean(axis=0).astype(np.float32))
np.save(stats_dir / "std_in.npy", (x_all.std(axis=0) + 1e-6).astype(np.float32))
np.save(stats_dir / "mean_out.npy", y_all.mean(axis=0).astype(np.float32))
np.save(stats_dir / "std_out.npy", (y_all.std(axis=0) + 1e-6).astype(np.float32))
print(f"Fake ShapeNetCar data generated under {ROOT / 'data/mlcfd_data'}")
if __name__ == "__main__":
main()
|