Skip to content

kups.application.relaxation.data

Data structures and ASE initialisation for structure relaxation.

RelaxParticles

Bases: Particles

Particle data for structure relaxation.

Extends Particles with energy gradients and derived properties (forces, inclusion/exclusion indices) needed by relaxation propagators.

Attributes:

Name Type Description
position_gradients Array

Optimizer position-DOF gradient ∂E/∂u_pos (the relaxation filter's output), shape (n_atoms, 3): ∂E/∂q under cell_filter (reference-cartesian) or ∂E/∂r under positions_only. The force source and ASE-fmax convergence quantity.
Source code in src/kups/application/relaxation/data.py 
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
@dataclass
class RelaxParticles(Particles):
    """Particle data for structure relaxation.

    Extends ``Particles`` with energy gradients and derived properties
    (forces, inclusion/exclusion indices) needed by relaxation propagators.

    Attributes:
        position_gradients: Optimizer position-DOF gradient ``∂E/∂u_pos`` (the
            relaxation filter's output), shape ``(n_atoms, 3)``: ``∂E/∂q`` under
            ``cell_filter`` (reference-cartesian) or ``∂E/∂r`` under
            ``positions_only``. The force source and ASE-fmax convergence quantity.
    """

    position_gradients: Array
    exclusion: Index[ExclusionId] = field(default=None, kw_only=True)  # type: ignore

    def __post_init__(self) -> None:
        if self.exclusion is None:
            object.__setattr__(self, "exclusion", default_exclusion(len(self.charges)))

    @property
    def forces(self) -> Array:
        """Atomic forces, the negative position gradient."""
        return -self.position_gradients

forces property

Atomic forces, the negative position gradient.

RelaxRunConfig

Bases: BaseModel

Configuration for a relaxation run.

Source code in src/kups/application/relaxation/data.py 
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
class RelaxRunConfig(BaseModel):
    """Configuration for a relaxation run."""

    out_file: str | Path
    """Path to the HDF5 output file."""
    max_steps: int
    """Maximum number of optimisation steps."""
    seed: int | None
    """Random seed. None for time-based."""
    force_tolerance: float
    """Convergence threshold for max atomic force (eV/Å)."""
    optimizer: TransformationConfig
    """List of Optax transform specifications passed to `make_optimizer`."""
    optimize_cell: bool
    """Whether to also relax lattice vectors."""

force_tolerance instance-attribute

Convergence threshold for max atomic force (eV/Å).

max_steps instance-attribute

Maximum number of optimisation steps.

optimize_cell instance-attribute

Whether to also relax lattice vectors.

optimizer instance-attribute

List of Optax transform specifications passed to make_optimizer.

out_file instance-attribute

Path to the HDF5 output file.

seed instance-attribute

Random seed. None for time-based.

RelaxState

Force-field-agnostic relaxation state.

The potential is built with its parameters at construction time (via the adapters' parameters=), so no force-field field lives on the state.

Source code in src/kups/application/relaxation/data.py 
73
74
75
76
77
78
79
80
81
82
83
84
85
@dataclass
class RelaxState:
    """Force-field-agnostic relaxation state.

    The potential is built with its parameters at construction time (via the
    adapters' ``parameters=``), so no force-field field lives on the state.
    """

    particles: Table[ParticleId, RelaxParticles]
    systems: Table[SystemId, RelaxSystems]
    neighborlist_params: UniversalNeighborlistParameters
    opt_state: optax.OptState
    step: Array

RelaxSystems

System-level data for structure relaxation.

Source code in src/kups/application/relaxation/data.py 
58
59
60
61
62
63
64
65
66
67
68
69
70
@dataclass
class RelaxSystems:
    """System-level data for structure relaxation."""

    cell: Cell[AnyPeriodicity]
    """Cell geometry, batched with shape (1,)."""
    cell_gradients: Cell[AnyPeriodicity]
    """Optimizer cell-DOF gradient ``∂E/∂u_cell`` (the relaxation filter's output),
    stored on :attr:`cell`'s frame (the lower-triangular log-deformation entries
    under ``cell_filter``). The ASE-fmax convergence quantity for the cell; the
    atoms-ride-the-cell coupling is already folded in by the filter pullback."""
    potential_energy: Array
    """Potential energy per system, shape (1,)."""

cell instance-attribute

Cell geometry, batched with shape (1,).

cell_gradients instance-attribute

Optimizer cell-DOF gradient ∂E/∂u_cell (the relaxation filter's output), stored on :attr:cell's frame (the lower-triangular log-deformation entries under cell_filter). The ASE-fmax convergence quantity for the cell; the atoms-ride-the-cell coupling is already folded in by the filter pullback.

potential_energy instance-attribute

Potential energy per system, shape (1,).

relax_state_from_ase(atoms)

Build relaxation particle and system data from an ASE Atoms object or file.

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
tuple[Table[ParticleId, RelaxParticles], Table[SystemId, RelaxSystems]]

Tuple of (particles, systems) ready for relaxation propagators.
Source code in src/kups/application/relaxation/data.py 
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
def relax_state_from_ase(
    atoms: ase.Atoms | str | Path,
) -> tuple[Table[ParticleId, RelaxParticles], Table[SystemId, RelaxSystems]]:
    """Build relaxation particle and system data from an ASE Atoms object or file.

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

    Returns:
        Tuple of ``(particles, systems)`` ready for relaxation propagators.
    """
    p, cell, _ = particles_from_ase(atoms)
    particles = p.set_data(
        RelaxParticles(
            positions=p.data.positions,
            masses=p.data.masses,
            atomic_numbers=p.data.atomic_numbers,
            charges=p.data.charges,
            labels=p.data.labels,
            system=p.data.system,
            position_gradients=jnp.zeros_like(p.data.positions),
        ),
    )
    # cell_factor = per-system atom count (ASE's exp_cell_factor) balances the
    # extensive cell-virial gradient against the per-atom forces in the joint
    # optimiser. bincount over the system index gives one count per system.
    n_systems = p.data.system.num_labels
    cell_factor = jnp.bincount(p.data.system.indices, length=n_systems).astype(
        p.data.positions.dtype
    )
    cell = bind(cell[None], lambda x: x.frame).apply(
        lambda f: DeformedFrame.from_frame(
            f, cell_factor=cell_factor, deformation=MatrixLogFrame
        )
    )
    systems = Table.arange(
        RelaxSystems(
            cell=cell,
            cell_gradients=tree_zeros_like(cell),
            potential_energy=jnp.zeros(n_systems),
        ),
        label=SystemId,
    )
    return particles, systems