File size: 3,485 Bytes
600bbdc | 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 | import json
from pathlib import Path
import numpy as np
from common import load_config
def bytes_feature(value: bytes):
import tensorflow.compat.v1 as tf
return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))
def int64_feature(value: int):
import tensorflow.compat.v1 as tf
return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))
def make_positions(rng, num_steps: int, num_particles: int, dim: int, bounds):
low, high = float(bounds[0]), float(bounds[1])
position = rng.uniform(low + 0.2, high - 0.2, size=(num_particles, dim))
frames = []
for _ in range(num_steps):
velocity = rng.normal(0.0, 0.015, size=(num_particles, dim))
position = np.clip(position + velocity, low + 0.05, high - 0.05)
frames.append(position.astype(np.float32))
return np.stack(frames, axis=0)
def make_particle_types(num_particles: int, num_node_types: int):
particle_type = np.full((num_particles,), 5, dtype=np.int64)
particle_type[::5] = 3
particle_type = np.clip(particle_type, 0, num_node_types - 1)
return particle_type
def write_split(path: Path, split: str, num_sequences: int, cfg, rng) -> None:
import tensorflow.compat.v1 as tf
num_steps = int(cfg.data.fake.num_steps)
num_particles = int(cfg.data.fake.num_particles)
dim = int(cfg.dim)
bounds = cfg.data.fake.bounds
writer_path = path / f"{split}.tfrecord"
with tf.io.TFRecordWriter(str(writer_path)) as writer:
for index in range(num_sequences):
positions = make_positions(rng, num_steps, num_particles, dim, bounds)
particle_type = make_particle_types(num_particles, int(cfg.data.num_node_types))
context = tf.train.Features(
feature={
"key": int64_feature(index),
"particle_type": bytes_feature(particle_type.tobytes()),
}
)
feature_list = tf.train.FeatureList(
feature=[bytes_feature(frame.tobytes()) for frame in positions]
)
example = tf.train.SequenceExample(
context=context,
feature_lists=tf.train.FeatureLists(
feature_list={"position": feature_list}
),
)
writer.write(example.SerializeToString())
def write_metadata(path: Path, cfg) -> None:
dim = int(cfg.dim)
bounds = [float(v) for v in cfg.data.fake.bounds]
metadata = {
"sequence_length": int(cfg.data.fake.num_steps) - 1,
"dt": 1.0,
"default_connectivity_radius": float(cfg.data.fake.connectivity_radius),
"bounds": [bounds for _ in range(dim)],
"dim": dim,
"vel_mean": [0.0] * dim,
"vel_std": [1.0] * dim,
"acc_mean": [0.0] * dim,
"acc_std": [1.0] * dim,
}
with open(path / "metadata.json", "w", encoding="utf-8") as f:
json.dump(metadata, f, indent=2)
def main():
cfg = load_config()
data_dir = Path(cfg.data.data_dir)
data_dir.mkdir(parents=True, exist_ok=True)
write_metadata(data_dir, cfg)
rng = np.random.default_rng(int(cfg.data.fake.seed))
for split in ("train", "valid", "test"):
split_cfg = getattr(cfg.data, split)
write_split(data_dir, split_cfg.split, int(split_cfg.num_sequences), cfg, rng)
print(f"Fake Lagrangian TFRecord data written to {data_dir}")
if __name__ == "__main__":
main()
|