File size: 3,300 Bytes
eca4864 | 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 | from __future__ import annotations
import argparse
import sys
from pathlib import Path
SCRIPT_DIR = Path(__file__).resolve().parent
if str(SCRIPT_DIR) not in sys.path:
sys.path.insert(0, str(SCRIPT_DIR))
import h5py
import numpy as np
from common import DEFAULT_CONFIG, load_config, resolve_path
def generate_year(
path: Path,
variables: list[str],
*,
time_steps: int,
height: int,
width: int,
time_step_hours: int,
chunk_time_steps: int,
fill_value: float,
materialize_pattern: bool,
) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
channels = len(variables)
with h5py.File(path, "w") as handle:
fields = handle.create_dataset(
"fields",
shape=(time_steps, channels, height, width),
dtype=np.float32,
chunks=(chunk_time_steps, 1, height, width),
fillvalue=np.float32(fill_value),
compression="lzf",
)
fields.attrs["variables"] = np.asarray(variables, dtype=h5py.string_dtype())
fields.attrs["time_step"] = time_step_hours
handle.create_dataset(
"global_means", data=np.zeros((1, channels, 1, 1), dtype=np.float32)
)
handle.create_dataset(
"global_stds", data=np.ones((1, channels, 1, 1), dtype=np.float32)
)
if materialize_pattern:
latitude = np.linspace(1.0, -1.0, height, dtype=np.float32)[:, None]
longitude = np.linspace(
0.0, 2.0 * np.pi, width, endpoint=False, dtype=np.float32
)
base = latitude + np.sin(longitude)[None, :]
# Two frames are enough to exercise non-zero input and target reads.
for time_index in range(min(time_steps, 2)):
for channel_index in range(channels):
fields[time_index, channel_index] = (
base + channel_index / channels + time_index * 0.01
)
def main() -> None:
parser = argparse.ArgumentParser(description="Generate ERA5-compatible FCNv2 data")
parser.add_argument("--config", default=str(DEFAULT_CONFIG))
parser.add_argument("--no-pattern", action="store_true")
args = parser.parse_args()
config = load_config(args.config)
data = config["data"]
fake = config["fake_data"]
output_dir = resolve_path(config, data["dataset_dir"])
years = sorted(set(data["train_years"] + data["val_years"] + data["test_years"]))
height, width = data["grid_shape"]
for year in years:
path = output_dir / "data" / f"{year}.h5"
generate_year(
path,
data["variables"],
time_steps=fake["time_steps_per_year"],
height=height,
width=width,
time_step_hours=data["time_step_hours"],
chunk_time_steps=fake["chunk_time_steps"],
fill_value=fake["fill_value"],
materialize_pattern=fake["materialize_pattern"] and not args.no_pattern,
)
logical_gib = (
fake["time_steps_per_year"] * len(data["variables"]) * height * width * 4
) / 1024**3
print(f"{path}: logical={logical_gib:.2f} GiB, actual={path.stat().st_size / 1024**2:.2f} MiB")
if __name__ == "__main__":
main()
|