| """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 |
|
|
