File size: 996 Bytes
1f88cea | 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 | from pathlib import Path
from boltzgen.data.data import Record, Structure
def load_record(record_id: str, record_dir: Path) -> Record:
"""Load the given record.
Parameters
----------
record_id : str
The record id to load.
record_dir : Path
The path to the record directory.
Returns
-------
Record
The loaded record.
"""
return Record.load(record_dir / f"{record_id}.json")
def load_structure(record: Record, struct_dir: Path) -> Structure:
"""Load the given input data.
Parameters
----------
record : str
The record to load.
target_dir : Path
The path to the data directory.
Returns
-------
Input
The loaded input.
"""
if (struct_dir / f"{record.id}.npz").exists():
structure_path = struct_dir / f"{record.id}.npz"
else:
structure_path = struct_dir / f"{record.id}" / f"{record.id}_model_0.npz"
return Structure.load(structure_path)
|