File size: 5,877 Bytes
41dd7e0 | 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 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 | """Pair manifests: which trajectories of ``trajectories_fine.npz`` a training run reads, and at which snapshots.
A manifest directory holds ``slice_<name>.json`` and ``master_order.npz`` (the same rows; the loaders re-hash it).
Rows are ``[group_row, member_col, gid]`` into ``vecs_train`` (segment pairs) or ``[block, group_row, member_col, gid]``
with block 0 = ``vecs_train`` (endpoint pairs). ``bifurcation.train`` writes the manifest of a run before training.
"""
from __future__ import annotations
import hashlib
import json
import os
from datetime import datetime, timezone
from pathlib import Path
import numpy as np
from .fine import ENDPOINT_SCHEMA as _fine_endpoint_schema
from .fine import FINE_FILE, _npz_member_memmap, _sha256_file
SEGMENT_SCHEMA = "segment_pairs_v1"
ENDPOINT_SCHEMA = _fine_endpoint_schema
def block_rows_segment(gids_train: np.ndarray, n_ic: int, members: int) -> np.ndarray:
"""Training conditions 0..n_ic-1 x realizations 0..members-1; rows [group_row, member_col, gid]."""
rows = np.empty((n_ic * members, 3), dtype=np.int64)
i = 0
for gi in range(n_ic):
for mj in range(members):
rows[i] = (gi, mj, int(gids_train[gi]))
i += 1
return rows
def block_rows_endpoint(gids_train: np.ndarray, n_ic: int, members: int) -> np.ndarray:
"""The same enumeration for endpoint pairs; rows [block, group_row, member_col, gid] with block 0 = vecs_train."""
rows = np.empty((n_ic * members, 4), dtype=np.int64)
i = 0
for gi in range(n_ic):
for mj in range(members):
rows[i] = (0, gi, mj, int(gids_train[gi]))
i += 1
return rows
def segment_geometry(T: float, S: int, t_values: np.ndarray) -> dict:
"""Uniform segmentation of [0, T] on the stored snapshot grid: n_T = index of T (an exact snapshot), S must
divide n_T; boundary_idx = [0, stride, ..., n_T]; the segment times are read from the stored grid."""
t_values = np.asarray(t_values, dtype=np.float64)
hits = np.nonzero(np.isclose(t_values, T, atol=1e-9))[0]
if len(hits) != 1:
raise ValueError(f"T={T} is not a stored snapshot time (grid {t_values[:3]}...)")
n_T = int(hits[0])
if n_T % S != 0:
raise ValueError(f"S={S} does not divide n_T={n_T} (T={T}); admissible S: "
f"{[s for s in range(1, n_T + 1) if n_T % s == 0]}")
stride = n_T // S
boundary_idx = [q * stride for q in range(S + 1)]
return {"T": float(T), "S": int(S), "n_T": n_T, "stride": stride, "delta": float(T) / S,
"boundary_idx": boundary_idx, "t_values": [float(t_values[b]) for b in boundary_idx]}
def _now() -> str:
return datetime.now(timezone.utc).isoformat(timespec="seconds")
def _write_master(slice_dir: Path, rows: np.ndarray, note: str) -> str:
slice_dir.mkdir(parents=True, exist_ok=True)
path = slice_dir / "master_order.npz"
if rows.shape[1] == 3:
payload = dict(group_row=rows[:, 0].astype(np.int32), member_col=rows[:, 1].astype(np.int32),
gid=rows[:, 2].astype(np.int32))
else:
payload = dict(block=rows[:, 0].astype(np.int32), group_row=rows[:, 1].astype(np.int32),
member_col=rows[:, 2].astype(np.int32), gid=rows[:, 3].astype(np.int32))
payload["algorithm"] = np.bytes_(note.encode())
np.savez(path, **payload)
return _sha256_file(str(path))
def dataset_provenance(packaged_dir: str | os.PathLike) -> dict:
packaged_dir = Path(packaged_dir)
return {"dataset_path": str(packaged_dir / FINE_FILE),
"split_sha256": _sha256_file(str(packaged_dir / "split.json")),
"normalization_sha256": _sha256_file(str(packaged_dir / "normalization.npz"))}
def write_segment_manifest(slice_dir: Path, rows3: np.ndarray, geometry: dict, name: str,
packaged_dir: str | os.PathLike, extra: dict | None = None) -> Path:
S = int(geometry["S"])
master_sha = _write_master(slice_dir, rows3, f"complete block ({name}): canonical enumeration")
man = {"schema": SEGMENT_SCHEMA, "arm": name, "cell": name, **geometry,
"N_traj": int(rows3.shape[0]), "B": int(rows3.shape[0]) * S,
"trajectories": rows3.tolist(), "master_order_sha256": master_sha,
**dataset_provenance(packaged_dir), "built_at": _now(), **(extra or {})}
out = slice_dir / f"slice_{name}.json"
out.write_text(json.dumps(man))
return out
def write_endpoint_manifest(slice_dir: Path, rows4: np.ndarray, idx_T: int, t_values: np.ndarray, name: str,
packaged_dir: str | os.PathLike, extra: dict | None = None) -> Path:
master_sha = _write_master(slice_dir, rows4, f"complete block ({name}): canonical enumeration, block 0 = vecs_train")
man = {"schema": ENDPOINT_SCHEMA, "arm": name, "S": 1,
"T": float(t_values[idx_T]), "n_T": int(idx_T),
"N_traj": int(rows4.shape[0]), "B": int(rows4.shape[0]),
"boundary_idx": [0, int(idx_T)], "t_values": [0.0, float(t_values[idx_T])],
"blocks": {"0": "vecs_train"},
"trajectories": rows4.tolist(), "master_order_sha256": master_sha,
**dataset_provenance(packaged_dir), "built_at": _now(), **(extra or {})}
out = slice_dir / f"slice_{name}.json"
out.write_text(json.dumps(man))
return out
def fine_grid(packaged_dir: str | os.PathLike) -> dict:
"""Key arrays of the fine file (memmapped; cheap)."""
fine = str(Path(packaged_dir) / FINE_FILE)
out = {"t_values": np.asarray(_npz_member_memmap(fine, "t_values")),
"gids_train": np.asarray(_npz_member_memmap(fine, "gids_train")),
"gids_eval": np.asarray(_npz_member_memmap(fine, "gids_eval"))}
for k in ("vecs_train", "vecs_eval"):
out[k + "_shape"] = tuple(_npz_member_memmap(fine, k).shape)
return out
|