tools/descriptors.py

"""CIF -> Sine Coulomb Matrix eigenvalues descriptor computation."""
from pathlib import Path

import numpy as np
from pymatgen.core import Structure

try:
    from matminer.featurizers.structure.matrix import SineCoulombMatrix
except Exception:
    SineCoulombMatrix = None


def compute_scm_eigenvalues(
    cif_path: str,
    target_dim: int,
) -> dict:
    """Compute Sine Coulomb Matrix eigenvalues from a CIF file.

    Args:
        cif_path: absolute path to a .cif file.
        target_dim: desired output dimension (520 for benzene, 584 for toluene).

    Returns:
        {
            "eigenvalues": np.ndarray of shape (target_dim,),
            "raw_dim": int,
            "padded": bool,
            "truncated": bool,
            "applicability_warning": str | None,
        }

    Raises:
        FileNotFoundError: if cif_path does not exist.
        ValueError: if CIF cannot be parsed into a valid Structure.
    """
    path = Path(cif_path)
    if not path.exists():
        raise FileNotFoundError(f"CIF file not found: {cif_path}")

    structure = Structure.from_file(str(path))
    if structure.num_sites == 0:
        raise ValueError(f"CIF parsed but contains 0 sites: {cif_path}")

    if SineCoulombMatrix is not None:
        scm = SineCoulombMatrix(flatten=False)
        scm.fit([structure])
        matrix = scm.featurize(structure)[0]
    else:
        matrix = _fallback_scm_matrix(structure)

    eigenvalues = np.sort(np.linalg.eigvalsh(matrix))[::-1]
    raw_dim = len(eigenvalues)

    padded = False
    truncated = False
    warning = None

    if raw_dim < target_dim:
        eigenvalues = np.pad(eigenvalues, (0, target_dim - raw_dim), mode="constant")
        padded = True
        warning = (
            f"SCM eigenvalue dimension ({raw_dim}) < target ({target_dim}). "
            f"Zero-padded {target_dim - raw_dim} values. "
            f"Prediction may be outside the model's applicability domain."
        )
    elif raw_dim > target_dim:
        eigenvalues = eigenvalues[:target_dim]
        truncated = True
        warning = (
            f"SCM eigenvalue dimension ({raw_dim}) > target ({target_dim}). "
            f"Truncated {raw_dim - target_dim} smallest eigenvalues. "
            f"Prediction may be outside the model's applicability domain."
        )

    return {
        "eigenvalues": eigenvalues,
        "raw_dim": raw_dim,
        "padded": padded,
        "truncated": truncated,
        "applicability_warning": warning,
    }


def _fallback_scm_matrix(structure: Structure) -> np.ndarray:
    """Small deterministic SCM-like matrix used when matminer is unavailable."""
    size = structure.num_sites
    matrix = np.zeros((size, size), dtype=float)
    for i, site_i in enumerate(structure):
        z_i = float(site_i.specie.Z)
        matrix[i, i] = 0.5 * z_i ** 2.4
        for j in range(i + 1, size):
            site_j = structure[j]
            z_j = float(site_j.specie.Z)
            distance = max(float(structure.get_distance(i, j)), 1e-6)
            value = z_i * z_j / distance
            matrix[i, j] = value
            matrix[j, i] = value
    return matrix