Skip to content

kups.application.relaxation.data

Data structures and ASE initialisation for structure relaxation.

RelaxParameters

Bases: BaseModel

Optimiser configuration for relaxation.

Source code in src/kups/application/relaxation/data.py 
84
85
86
87
88
89
90
class RelaxParameters(BaseModel):
    """Optimiser configuration for relaxation."""

    optimizer: TransformationConfig
    """List of Optax transform specifications passed to `make_optimizer`."""
    optimize_cell: bool
    """Whether to also relax lattice vectors."""

optimize_cell instance-attribute

Whether to also relax lattice vectors.

optimizer instance-attribute

List of Optax transform specifications passed to make_optimizer.

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

Energy gradient w.r.t. positions, shape (n_atoms, 3).
Source code in src/kups/application/relaxation/data.py 
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
@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: Energy gradient w.r.t. positions, shape ``(n_atoms, 3)``.
    """

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

    def __post_init__(self):
        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 
71
72
73
74
75
76
77
78
79
80
81
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/Å)."""

force_tolerance instance-attribute

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

max_steps instance-attribute

Maximum number of optimisation steps.

out_file instance-attribute

Path to the HDF5 output file.

seed instance-attribute

Random seed. None for time-based.

RelaxSystems

System-level data for structure relaxation.

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

    unitcell: UnitCell
    """Unit cell geometry, batched with shape (1,)."""
    unitcell_gradients: UnitCell
    potential_energy: Array
    """Potential energy per system, shape (1,)."""

    @property
    def stress_tensor(self) -> Array:
        """Cauchy stress tensor, shape ``(..., 3, 3)``."""
        return (
            -self.unitcell_gradients.lattice_vectors
            / self.unitcell.volume[..., None, None]
        )

potential_energy instance-attribute

Potential energy per system, shape (1,).

stress_tensor property

Cauchy stress tensor, shape (..., 3, 3).

unitcell instance-attribute

Unit cell geometry, batched with 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 
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
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, unitcell, _ = 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),
        ),
    )
    unitcell = unitcell[None]
    systems = Table.arange(
        RelaxSystems(
            unitcell=unitcell,
            unitcell_gradients=tree_zeros_like(unitcell),
            potential_energy=jnp.array([0.0]),
        ),
        label=SystemId,
    )
    return particles, systems