Source code for ufp.analysis.inputs

"""Input-data analysis helpers for ASE-backed UFP workflows."""

from __future__ import annotations

from typing import Sequence

import ase
import ase.neighborlist
import numpy as np




def pair_distance_histogram(
    frames: Sequence[ase.Atoms],
    *,
    cutoff: float = 6.0,
    max_frames: int = 1000,
    seed: int = 7,
    bins: int = 300,
) -> tuple[np.ndarray, np.ndarray]:
    """Collect a pair-distance histogram from ASE structures."""
    rng = np.random.default_rng(seed)
    sample_size = min(max_frames, len(frames))
    frame_indices = rng.choice(len(frames), size=sample_size, replace=False)
    distances: list[np.ndarray] = []
    for frame_index in frame_indices:
        atoms = frames[int(frame_index)]
        i, j, d = ase.neighborlist.neighbor_list("ijd", atoms, cutoff=cutoff)
        mask = i < j
        if np.any(mask):
            distances.append(np.asarray(d[mask], dtype=float))
    if not distances:
        raise ValueError("no pair distances were collected")
    all_distances = np.concatenate(distances)
    histogram, bin_edges = np.histogram(
        all_distances,
        bins=np.linspace(0.0, cutoff, bins),
        density=True,
    )
    return histogram, bin_edges



__all__ = ["pair_distance_histogram"]