Skip to content

kups.application.utils.particles

Shared particle data structures and ASE loading utilities.

Particles

Particle state shared across simulation types.

Attributes:

Name Type Description
positions Array

Cartesian coordinates in the lower-triangular frame, shape (n_atoms, 3).
masses Array

Atomic masses (amu), shape (n_atoms,).
atomic_numbers Array

Atomic numbers, shape (n_atoms,).
charges Array

Partial charges, shape (n_atoms,).
labels Index[Label]

Per-atom string labels.
system Index[SystemId]

Index mapping each particle to a system.
Source code in src/kups/application/utils/particles.py 
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
@dataclass
class Particles:
    """Particle state shared across simulation types.

    Attributes:
        positions: Cartesian coordinates in the lower-triangular frame, shape (n_atoms, 3).
        masses: Atomic masses (amu), shape (n_atoms,).
        atomic_numbers: Atomic numbers, shape (n_atoms,).
        charges: Partial charges, shape (n_atoms,).
        labels: Per-atom string labels.
        system: Index mapping each particle to a system.
    """

    positions: Array
    masses: Array
    atomic_numbers: Array
    charges: Array
    labels: Index[Label]
    system: Index[SystemId]

    @property
    def inclusion(self) -> Index[InclusionId]:
        """System index re-labeled as InclusionId."""
        return Index(tuple(map(InclusionId, self.system.keys)), self.system.indices)

inclusion property

System index re-labeled as InclusionId.

default_exclusion(n)

Build a default per-particle exclusion index (each atom excludes itself).

Parameters:

Name Type Description Default
n int

Number of particles.

 required

Returns:

Type Description
Index[ExclusionId]

Index mapping each particle to a unique ExclusionId.
Source code in src/kups/application/utils/particles.py 
48
49
50
51
52
53
54
55
56
57
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)

particles_from_ase(atoms)

Build particle data and unit cell from an ASE Atoms object or file path.

Parameters:

Name Type Description Default
atoms Atoms | str | Path

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

 required

Returns:

Type Description
Table[ParticleId, Particles]

Tuple of (particles, unitcell, uc_transform) where uc_transform
UnitCell

rotates Cartesian positions into the lower-triangular frame.
Source code in src/kups/application/utils/particles.py 
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
def particles_from_ase(
    atoms: ase.Atoms | str | Path,
) -> tuple[Table[ParticleId, Particles], UnitCell, Callable[[Array], Array]]:
    """Build particle data and unit cell from an ASE Atoms object or file path.

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

    Returns:
        Tuple of (particles, unitcell, uc_transform) where uc_transform
        rotates Cartesian positions into the lower-triangular frame.
    """
    if isinstance(atoms, (str, Path)):
        atoms = next(ase.io.iread(atoms, index=-1, store_tags=True))
    L, uc_transform = to_lower_triangular(jnp.asarray(atoms.cell.array))
    unitcell = TriclinicUnitCell.from_matrix(L)
    # Rotate Cartesian positions into the lower-triangular frame.
    positions = uc_transform(jnp.asarray(atoms.positions))
    masses = jnp.asarray(atoms.get_masses())
    atomic_numbers = jnp.asarray(atoms.get_atomic_numbers())
    n_atoms = len(masses)
    charges = jnp.asarray(
        atoms.info.get(
            "_atom_type_partial_charge",
            atoms.info.get("_atom_site_charge", jnp.zeros((len(positions),))),
        )
    )
    labels = list(
        map(Label, atoms.info.get("_atom_site_label", atoms.get_chemical_symbols()))
    )
    particles = Table.arange(
        Particles(
            positions=positions,
            masses=masses,
            atomic_numbers=atomic_numbers,
            charges=charges,
            labels=Index.new(labels),
            system=Index.integer(jnp.zeros(n_atoms, dtype=int), label=SystemId),
        ),
        label=ParticleId,
    )
    return particles, unitcell, uc_transform