Source code for lcmd_db.adapters.ase

"""Convert Dataset[Molecule] entries to ASE Atoms objects."""

from __future__ import annotations

from typing import TYPE_CHECKING

from ..types import Molecule

if TYPE_CHECKING:
    import ase

    from ..dataset._base import Dataset




def convert(dataset: Dataset[Molecule[object]]) -> list[ase.Atoms]:
    """Convert each entry in a molecule dataset to an ``ase.Atoms`` object.

    Requires ``ase`` to be installed. Reads from structure files (XYZ).
    """
    try:
        import ase.io  # noqa: F401
    except ImportError as exc:
        raise ImportError(
            "The 'ase' package is required for this conversion. "
            "Install it with: uv add ase"
        ) from exc

    return [_entry_to_atoms(entry) for entry in dataset]



def _entry_to_atoms(entry: Molecule[object]) -> ase.Atoms:
    import ase.io

    if not (entry.structure_path and entry.structure_path.exists()):
        raise ValueError(
            f"Molecule {entry.id}: no structure file available. "
            "Include structures when loading the dataset."
        )

    atoms = ase.io.read(str(entry.structure_path))
    atoms.info.update(
        {k: v for k, v in entry.properties.items() if k != "id"}  # type: ignore[union-attr]
    )
    return atoms