Source code for atom_hifi.corpus

"""Focus-corpus helpers for Atom-HiFi runner orchestration."""
from __future__ import annotations

import numpy as np




def filter_atoms(atoms, desc, elements):
    """Return (idx, desc_focused, atoms_focused) for the focus atoms.

    The ``is_focus`` per-atom array (extxyz column, int 0/1) and the
    ``elements`` (focus_elements) list are combined with AND: an atom is a focus
    atom iff its element is requested and its is_focus flag is 1. With
    focus_elements=null the element mask is all-True, so is_focus alone selects
    focus atoms.
    """
    if elements:
        sym = np.array(atoms.get_chemical_symbols())
        mask = np.isin(sym, elements)
    else:
        mask = np.ones(len(atoms), dtype=bool)
    if 'is_focus' in atoms.arrays:
        mask &= atoms.arrays['is_focus'].astype(bool)
    idx = np.where(mask)[0]
    return idx, desc[idx], atoms[idx]





def focus_arrays(lib, desc_list, focus_elements):
    """Return focus-only structures, descriptor array, structure map, desc list."""
    focus_desc_list, lib_focus, focus_sizes = [], [], []
    for atoms, desc in zip(lib, desc_list):
        _, foc, foc_atoms = filter_atoms(atoms, desc, focus_elements)
        focus_desc_list.append(foc)
        lib_focus.append(foc_atoms)
        focus_sizes.append(len(foc))
    return (
        lib_focus,
        np.vstack(focus_desc_list),
        np.repeat(np.arange(len(focus_desc_list)), focus_sizes),
        focus_desc_list,
    )





def validate_focus_species(lib, lib_focus, focus_elements, focus_explicit):
    """Fail fast if an explicit focus element has zero focus atoms."""
    focus_present = sorted({s for a in lib_focus for s in a.get_chemical_symbols()})
    if focus_explicit:
        missing = [el for el in focus_elements if el not in focus_present]
        if missing:
            has_isfocus = any('is_focus' in a.arrays for a in lib)
            hint = (
                ' -- the is_focus column excludes every atom of these elements. '
                'Regenerate it so the focus element(s) are flagged is_focus=1, '
                'or drop them from focus_elements.'
                if has_isfocus else
                ' -- no atoms of these elements are present in the library.'
            )
            raise ValueError(
                f'focus_elements {missing} have zero focus atoms after '
                f'filtering{hint}')
    return [el for el in focus_elements if el in focus_present]