"""ANI-1 dataset builder (pyanitools HDF5 from figshare).
The ANI-1 data set holds ~20 million off-equilibrium conformations and DFT
(wB97X/6-31G(d)) total energies for 57,462 small organic molecules (H, C, N, O)
built from GDB-11, generated by Normal Mode Sampling. It is distributed as eight
HDF5 files ``ani_gdb_s0X.h5`` (``X`` = number of heavy atoms, 1-8) bundled in a
single ~4.8 GB ``ANI-1_release.tar.gz`` archive.
Reference
---------
Smith, Isayev & Roitberg, "ANI-1, A data set of 20 million calculated
off-equilibrium conformations for organic molecules", *Sci. Data* **4**, 170193
(2017). Data: https://doi.org/10.6084/m9.figshare.c.3846712
Format / reader spec: https://github.com/isayev/ANI1_dataset
Notes
-----
* Upstream energies are in **Hartree** and positions in **angstrom**. By default
energies are converted to eV (the convention used elsewhere in xnn);
``units="hartree"`` keeps the raw values. The molecule contains no forces.
* The whole 4.8 GB archive is downloaded once (cached and MD5-verified), even to
read a single heavy-atom subset. Use ``heavy_atoms`` to select subsets and
``max_molecules`` / ``max_conformations`` to cap the amount materialised --
the full set is ~20 M conformations and will not fit in memory at once.
* ANI-1 ships no official split. ``split`` in ``{"train", "val", "test"}``
applies the paper's per-molecule 80/10/10 partition with a fixed seed, so
splits are disjoint and reproducible; ``split=None`` returns ``{"all": ...}``.
"""
from __future__ import annotations
from pathlib import Path
from typing import Optional, Union
import numpy as np
from tqdm.auto import tqdm
from ...config.coerce import SYMBOL_TO_Z
from ._download import download_file, extract_archive
from .base import DatasetBuilder, register_dataset
# figshare: ANI-1_release.tar.gz
_URL = "https://ndownloader.figshare.com/files/9057631"
_MD5 = "455e883a47ffd165f545c868ebaf95f9"
_ARCHIVE = "ANI-1_release.tar.gz"
# 1 Hartree in eV (CODATA 2018), matching ase.units.Hartree.
_HARTREE_TO_EV = 27.211386245988
def _iter_molecule_groups(group):
"""Yield every HDF5 group that holds a ``coordinates`` dataset.
pyanitools stores each molecule as a (possibly nested) group containing
``coordinates`` / ``energies`` / ``species`` datasets. This walks the tree
natively (no upstream reader needed) and yields the leaf molecule groups.
"""
import h5py
if "coordinates" in group and isinstance(group["coordinates"], h5py.Dataset):
yield group
return
for key in group:
item = group[key]
if isinstance(item, h5py.Group):
yield from _iter_molecule_groups(item)
def _decode_species(raw) -> np.ndarray:
"""Convert a pyanitools ``species`` dataset to an atomic-number array."""
out = []
for s in raw:
sym = s.decode() if isinstance(s, (bytes, bytearray)) else str(s)
out.append(SYMBOL_TO_Z[sym])
return np.asarray(out, dtype=np.int64)
[docs]
class ANI1Builder(DatasetBuilder):
"""Builder for the ANI-1 data set (20M DFT energies, GDB-11 subsets).
See the module docstring for the dataset description and citation. The one
4.8 GB archive is downloaded and extracted once; the requested heavy-atom
subsets are then parsed from their HDF5 files into xnn structure dicts.
"""
name = "ani1"
description = ("ANI-1: ~20M off-equilibrium conformations & wB97X energies "
"for H/C/N/O organic molecules (GDB-11).")
[docs]
def load(self, *, split: Optional[str] = None, cache_dir: Path,
heavy_atoms: Union[int, list, tuple, None] = None,
units: str = "eV", max_molecules: Optional[int] = None,
max_conformations: Optional[int] = None, seed: int = 1234,
quiet: bool = False) -> Union[dict[str, list[dict]], list[dict]]:
"""Download and preprocess the ANI-1 data set.
Parameters
----------
split : str or None
``None`` returns ``{"all": ...}``; ``"train"`` / ``"val"`` /
``"test"`` returns the paper's per-molecule 80/10/10 partition
(fixed ``seed``, disjoint splits).
cache_dir : pathlib.Path
Base cache directory; files live under ``cache_dir/"ani1"``.
heavy_atoms : int or sequence of int, optional
Which heavy-atom subset(s) ``ani_gdb_s0X.h5`` to load (``X`` in
1-8). Defaults to all eight. Fewer subsets = far less data.
units : str, optional
``"eV"`` (default) converts energies to eV; ``"hartree"`` keeps the
raw upstream values.
max_molecules : int, optional
Cap the number of molecules read per subset (useful for demos).
max_conformations : int, optional
Cap the number of conformations kept per molecule.
seed : int, optional
Seed for the reproducible 80/10/10 split. Defaults to ``1234``.
quiet : bool, optional
Suppress progress output. Defaults to ``False``.
Returns
-------
dict of {str: list of dict} or list of dict
Structure dicts with keys ``pos`` ``(N, 3)``, ``atomic_numbers``
``(N,)``, ``energy`` (scalar) and ``smiles`` (str). ANI-1 is
molecular and force-free, so no ``cell`` / ``pbc`` / ``forces``.
Raises
------
ValueError
If ``units`` or ``split`` is unrecognized, or ``heavy_atoms`` is out
of the 1-8 range.
ImportError
If ``h5py`` is not installed.
"""
if units not in ("eV", "hartree", "Hartree"):
raise ValueError(f"units must be 'eV' or 'hartree', got {units!r}")
subsets = self._resolve_subsets(heavy_atoms)
scale = 1.0 if units.lower() == "hartree" else _HARTREE_TO_EV
root = Path(cache_dir) / self.name
h5_dir = self._ensure_extracted(root, quiet)
structures: list[dict] = []
for x in subsets:
h5_path = h5_dir / f"ani_gdb_s{x:02d}.h5"
if not h5_path.exists():
raise FileNotFoundError(
f"expected {h5_path.name} in the extracted archive; "
f"found: {[p.name for p in h5_dir.glob('*.h5')]}")
structures.extend(self._read_h5(h5_path, scale, max_molecules,
max_conformations, quiet))
if split is None:
return {"all": structures}
if split not in ("train", "val", "test"):
raise ValueError(
f"unknown split {split!r}; use 'train', 'val', 'test', or None")
return self._partition(structures, seed)[split]
@staticmethod
def _resolve_subsets(heavy_atoms) -> list[int]:
"""Validate and normalize the ``heavy_atoms`` selection to a list."""
if heavy_atoms is None:
return list(range(1, 9))
vals = [heavy_atoms] if isinstance(heavy_atoms, int) else list(heavy_atoms)
for x in vals:
if not (1 <= int(x) <= 8):
raise ValueError(f"heavy_atoms entries must be 1-8, got {x!r}")
return [int(x) for x in vals]
def _ensure_extracted(self, root: Path, quiet: bool) -> Path:
"""Download (once) and extract the archive; return the h5 directory."""
raw = root / "raw"
# Locate the directory holding the ani_gdb_s0X.h5 files, if already there.
for cand in (raw / "ANI-1_release", raw):
if cand.is_dir() and any(cand.glob("ani_gdb_s0*.h5")):
return cand
archive = download_file(_URL, raw / _ARCHIVE, _MD5, quiet=quiet)
extract_archive(archive, raw)
for cand in (raw / "ANI-1_release", raw):
if cand.is_dir() and any(cand.glob("ani_gdb_s0*.h5")):
return cand
raise FileNotFoundError(
f"could not find ani_gdb_s0*.h5 after extracting {archive}")
@staticmethod
def _read_h5(h5_path: Path, scale: float, max_molecules: Optional[int],
max_conformations: Optional[int], quiet: bool) -> list[dict]:
"""Parse one ``ani_gdb_s0X.h5`` into structure dicts."""
import h5py
out: list[dict] = []
with h5py.File(h5_path, "r") as f:
groups = list(_iter_molecule_groups(f))
if max_molecules is not None:
groups = groups[:max_molecules]
for g in tqdm(groups, desc=f"ani1:{h5_path.stem}", unit=" mol",
disable=quiet, leave=False):
z = _decode_species(g["species"][()])
coords = np.asarray(g["coordinates"][()], dtype=np.float64)
energies = np.asarray(g["energies"][()], dtype=np.float64)
smiles = g["smiles"][()] if "smiles" in g else b""
if isinstance(smiles, np.ndarray):
smiles = b"".join(
s if isinstance(s, bytes) else str(s).encode()
for s in smiles.ravel())
smiles = smiles.decode() if isinstance(smiles, bytes) else str(smiles)
n = len(energies)
if max_conformations is not None:
n = min(n, max_conformations)
for i in range(n):
out.append({
"pos": coords[i],
"atomic_numbers": z,
"energy": float(energies[i]) * scale,
"smiles": smiles,
})
return out
@staticmethod
def _partition(structures: list[dict],
seed: int) -> dict[str, list[dict]]:
"""Deterministic 80/10/10 train/val/test partition (paper ratios)."""
idx = np.arange(len(structures))
np.random.default_rng(seed).shuffle(idx)
n_train = int(0.8 * len(idx))
n_val = int(0.1 * len(idx))
parts = {"train": idx[:n_train],
"val": idx[n_train:n_train + n_val],
"test": idx[n_train + n_val:]}
return {k: [structures[i] for i in v] for k, v in parts.items()}
register_dataset(ANI1Builder())
