OneScience-Group/oxides

1#!/usr/bin/env python32"""Validate the standardized OneScience oxide ASE databases."""3 4from __future__ import annotations5 6import argparse7import hashlib8import json9import math10from pathlib import Path11 12import numpy as np13from ase.db import connect14 15 16EXPECTED_COUNTS = {"train": 238, "val": 28, "test": 29}17REQUIRED_METADATA = {"oxide", "polymorph", "xc"}18 19 20def sha256_file(path: Path) -> str:21    digest = hashlib.sha256()22    with path.open("rb") as handle:23        for chunk in iter(lambda: handle.read(1024 * 1024), b""):24            digest.update(chunk)25    return digest.hexdigest()26 27 28def require_file(path: Path) -> None:29    if not path.is_file():30        raise FileNotFoundError(f"missing file: {path}")31 32 33def validate_split(path: Path, expected_count: int) -> dict[str, int]:34    require_file(path)35    database = connect(path)36    if database.count() != expected_count:37        raise ValueError(f"unexpected row count in {path}: {database.count()} != {expected_count}")38 39    oxides: set[str] = set()40    groups: set[tuple[str, str]] = set()41    for row in database.select():42        missing = REQUIRED_METADATA - set(row.key_value_pairs)43        if missing:44            raise ValueError(f"missing metadata in {path} row {row.id}: {sorted(missing)}")45        if row.xc != "PBE":46            raise ValueError(f"unexpected xc in {path} row {row.id}: {row.xc!r}")47        atoms = row.toatoms()48        if len(atoms) == 0 or not atoms.pbc.all() or abs(atoms.get_volume()) <= 0:49            raise ValueError(f"invalid periodic structure in {path} row {row.id}")50        forces = np.asarray(row.forces, dtype=float)51        stress = np.asarray(row.stress, dtype=float)52        if forces.shape != (len(atoms), 3):53            raise ValueError(f"invalid forces shape in {path} row {row.id}: {forces.shape}")54        if stress.shape != (6,):55            raise ValueError(f"invalid stress shape in {path} row {row.id}: {stress.shape}")56        values = np.concatenate(([float(row.energy)], forces.reshape(-1), stress))57        if not np.isfinite(values).all():58            raise ValueError(f"non-finite target in {path} row {row.id}")59        oxides.add(str(row.oxide))60        groups.add((str(row.oxide), str(row.polymorph)))61    return {"structures": database.count(), "oxide_count": len(oxides), "group_count": len(groups)}62 63 64def validate_checksums(package_root: Path, manifest: Path) -> int:65    require_file(manifest)66    count = 067    for line_number, raw in enumerate(manifest.read_text(encoding="utf-8").splitlines(), 1):68        if not raw.strip():69            continue70        digest, relative = raw.split(None, 1)71        target = package_root / relative72        require_file(target)73        if sha256_file(target) != digest:74            raise ValueError(f"checksum mismatch on line {line_number}: {relative}")75        count += 176    return count77 78 79def main() -> int:80    parser = argparse.ArgumentParser(description=__doc__)81    parser.add_argument("--dataset-root", default="data/OXIDES")82    parser.add_argument("--checksum-manifest", default="metadata/sha256_manifest.txt")83    parser.add_argument("--skip-checksum", action="store_true")84    args = parser.parse_args()85 86    root = Path(args.dataset_root)87    summary = {88        split: validate_split(root / "prepared" / f"{split}.db", count)89        for split, count in EXPECTED_COUNTS.items()90    }91    manifest = root / "manifest.json"92    require_file(manifest)93    metadata = json.loads(manifest.read_text(encoding="utf-8"))94    if metadata.get("counts") != EXPECTED_COUNTS:95        raise ValueError(f"manifest counts do not match expected counts: {metadata.get('counts')}")96    checksums = 0 if args.skip_checksum else validate_checksums(Path.cwd(), Path(args.checksum_manifest))97    print("Oxides dataset validation passed")98    print(json.dumps(summary, sort_keys=True))99    print(f"checksum entries verified: {checksums}")100    return 0101 102 103if __name__ == "__main__":104    raise SystemExit(main())105