kups.application.utils.particles ¶

def default_exclusion(n: int) -> Index[ExclusionId]:
    """Build a default per-particle exclusion index (each atom excludes itself).

    Args:
        n: Number of particles.

    Returns:
        Index mapping each particle to a unique ExclusionId.
    """
    return Index.integer(jnp.arange(n), n=n, label=ExclusionId)

def particles_from_arrays(
    *,
    positions: ArrayLike,
    cell_vectors: ArrayLike,
    periodicity: (
        Sequence[bool | np.bool_ | Array] | np.ndarray[Any, np.dtype[np.bool_]] | Array
    ),
    masses: ArrayLike,
    atomic_numbers: ArrayLike,
    labels: Sequence[str],
    charges: ArrayLike | None = None,
) -> tuple[
    Table[ParticleId, Particles], Cell[AnyPeriodicity], Callable[[Array], Array]
]:
    """Build particle data and cell from source-neutral arrays.

    Numeric inputs follow the active JAX dtype configuration; exact dtype
    preservation is not guaranteed.

    Args:
        positions: Cartesian positions in Å, shape ``(n, 3)``, expressed in the
            same Cartesian coordinate frame as ``cell_vectors``.
        cell_vectors: Three cell vectors in Å stored row-wise, shape ``(3, 3)``.
            The caller must provide a valid cell or bounding frame.
        periodicity: Three boolean values selecting periodic axes.
        masses: Particle masses in amu, shape ``(n,)``.
        atomic_numbers: Integer atomic numbers, shape ``(n,)``.
        labels: Non-string sequence of string labels, one per particle.
        charges: Partial charges in elementary-charge units, shape ``(n,)``.
            Omitted charges default to floating zeros.

    Returns:
        Tuple of ``(particles, cell, uc_transform)`` where ``uc_transform``
        rotates Cartesian coordinates into the lower-triangular cell frame.

    Raises:
        ValueError: If an input shape or particle/label count is invalid.
        TypeError: If conversion fails or an input has an invalid type or dtype.
    """
    positions_array = _as_jax_array("positions", positions)
    cell_vectors_array = _as_jax_array("cell_vectors", cell_vectors)
    masses_array = _as_jax_array("masses", masses)
    atomic_numbers_array = _as_jax_array("atomic_numbers", atomic_numbers)
    charges_array = None if charges is None else _as_jax_array("charges", charges)

    _require_real_dtype("positions", positions_array)
    _require_real_dtype("cell_vectors", cell_vectors_array)
    _require_real_dtype("masses", masses_array)
    _require_integer_dtype("atomic_numbers", atomic_numbers_array)
    if charges_array is not None:
        _require_real_dtype("charges", charges_array)

    if positions_array.ndim != 2 or positions_array.shape[1] != 3:
        raise ValueError(
            f"positions must have shape (n, 3); got {positions_array.shape}."
        )
    if cell_vectors_array.shape != (3, 3):
        raise ValueError(
            f"cell_vectors must have shape (3, 3); got {cell_vectors_array.shape}."
        )

    n_particles = positions_array.shape[0]
    _require_particle_shape("masses", masses_array, n_particles)
    _require_particle_shape("atomic_numbers", atomic_numbers_array, n_particles)
    if charges_array is not None:
        _require_particle_shape("charges", charges_array, n_particles)

    if isinstance(labels, (str, bytes)):
        raise TypeError("labels must be a non-string sequence of strings.")
    try:
        label_values = list(labels)
    except TypeError as error:
        raise TypeError("labels must be a sequence of strings.") from error
    if len(label_values) != n_particles:
        raise ValueError(
            f"labels must contain {n_particles} values; got {len(label_values)}."
        )
    if not all(isinstance(label, str) for label in label_values):
        raise TypeError("labels must contain only strings.")

    periodicity_tuple = _normalize_periodicity(periodicity)
    positions_array = _promote_integer_to_float(positions_array)
    cell_vectors_array = _promote_integer_to_float(cell_vectors_array)
    masses_array = _promote_integer_to_float(masses_array)
    if charges_array is None:
        charges_array = jnp.zeros(n_particles, dtype=jnp.result_type(float))
    else:
        charges_array = _promote_integer_to_float(charges_array)

    return _build_particles_and_cell(
        positions=positions_array,
        cell_vectors=cell_vectors_array,
        periodicity=periodicity_tuple,
        masses=masses_array,
        atomic_numbers=atomic_numbers_array,
        charges=charges_array,
        labels=list(map(Label, label_values)),
    )

def particles_from_ase(
    atoms: ase.Atoms | str | Path,
) -> tuple[
    Table[ParticleId, Particles], Cell[AnyPeriodicity], Callable[[Array], Array]
]:
    """Build particle data and cell from an ASE Atoms object or file path.

    Results are cached when ``atoms`` is a file path.

    Args:
        atoms: ASE Atoms object, or a file path (str/Path) readable by
            ``ase.io.read``.

    Returns:
        Tuple of (particles, cell, uc_transform) where uc_transform
        rotates Cartesian positions into the lower-triangular frame.
    """
    if isinstance(atoms, (str, Path)):
        return _particles_from_path(atoms)
    return _particles_from_atoms(atoms)