Skip to content

kups.potential.classical.harmonic

Harmonic potentials for bonded interactions.

This module provides harmonic bond and angle potentials commonly used in molecular mechanics force fields. These terms maintain molecular geometry and are typically applied to explicitly defined bonds and angles.

Bond potential: \(U(r) = k(r - r_0)^2\) Angle potential: \(U(\theta) = k(\theta - \theta_0)^2\)

HarmonicAngleParameters

Harmonic angle potential parameters.

Attributes:

Name Type Description
labels tuple[Label, ...]

Species labels, shape (n_species,)
theta0 Array

Equilibrium angles [degrees], shape (n_species, n_species, n_species)
k Array

Force constants [energy/degree²], shape (n_species, n_species, n_species)
Source code in src/kups/potential/classical/harmonic.py 
104
105
106
107
108
109
110
111
112
113
114
115
116
@dataclass
class HarmonicAngleParameters:
    """Harmonic angle potential parameters.

    Attributes:
        labels: Species labels, shape `(n_species,)`
        theta0: Equilibrium angles [degrees], shape `(n_species, n_species, n_species)`
        k: Force constants [energy/degree²], shape `(n_species, n_species, n_species)`
    """

    labels: tuple[Label, ...] = field(static=True)  # (n_species,)
    theta0: Array  # (n_species, n_species, n_species)
    k: Array  # (n_species, n_species, n_species)

HarmonicBondParameters

Harmonic bond potential parameters.

Attributes:

Name Type Description
labels tuple[Label, ...]

Species labels, shape (n_species,)
x0 Array

Equilibrium bond lengths [Ã…], shape (n_species, n_species)
k Array

Force constants [energy/Ų], shape (n_species, n_species)
Source code in src/kups/potential/classical/harmonic.py 
57
58
59
60
61
62
63
64
65
66
67
68
69
@dataclass
class HarmonicBondParameters:
    """Harmonic bond potential parameters.

    Attributes:
        labels: Species labels, shape `(n_species,)`
        x0: Equilibrium bond lengths [Ã…], shape `(n_species, n_species)`
        k: Force constants [energy/Ų], shape `(n_species, n_species)`
    """

    labels: tuple[Label, ...] = field(static=True)  # (n_species,)
    x0: Array  # (n_species, n_species)
    k: Array  # (n_species, n_species)

IsBondedParticles

Bases: HasPositionsAndLabels, IsRadiusGraphPoints, Protocol

Particle data with positions, labels, and system index.

Source code in src/kups/potential/classical/harmonic.py 
50
51
52
53
54
@runtime_checkable
class IsBondedParticles(HasPositionsAndLabels, IsRadiusGraphPoints, Protocol):
    """Particle data with positions, labels, and system index."""

    ...

harmonic_angle_energy(inp)

Compute harmonic angle energy for all angles.

Calculates energy as k(θ - θ₀)² for each angle triplet and sums over all systems. Angles are computed in degrees.

Parameters:

Name Type Description Default
inp HarmonicAngleInput

Graph potential input with harmonic angle parameters

 required

Returns:

Type Description
WithPatch[Table[SystemId, Energy], IdPatch[Any]]

Total angle energy per system
Source code in src/kups/potential/classical/harmonic.py 
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
150
151
152
153
154
155
156
157
def harmonic_angle_energy(
    inp: HarmonicAngleInput,
) -> WithPatch[Table[SystemId, Energy], IdPatch[Any]]:
    """Compute harmonic angle energy for all angles.

    Calculates energy as k(θ - θ₀)² for each angle triplet and sums over all systems.
    Angles are computed in degrees.

    Args:
        inp: Graph potential input with harmonic angle parameters

    Returns:
        Total angle energy per system
    """
    graph = inp.graph
    assert graph.edges.indices.indices.shape[1] == 3, (
        "Harmonic angle potential only supports triplet interactions (order=3)."
    )
    edg_species = graph.particles[graph.edges.indices].labels.indices_in(
        inp.parameters.labels
    )
    theta0 = inp.parameters.theta0[
        edg_species[:, 0], edg_species[:, 1], edg_species[:, 2]
    ]
    k = inp.parameters.k[edg_species[:, 0], edg_species[:, 1], edg_species[:, 2]]
    v1, v2 = graph.edge_shifts[:, 0], graph.edge_shifts[:, 1]
    angle = jnp.arccos(
        jnp.einsum("ij,ij->i", v1, v2)
        / (jnp.linalg.norm(v1, axis=-1) * jnp.linalg.norm(v2, axis=-1))
    )
    angle = jnp.rad2deg(angle)
    edge_energy = (angle - theta0) ** 2 * k
    total_energies = graph.edge_batch_mask.sum_over(edge_energy)
    return WithPatch(total_energies, IdPatch[Any]())

harmonic_bond_energy(inp)

Compute harmonic bond energy for all bonds.

Calculates energy as k(r - r₀)² for each bond and sums over all systems.

Parameters:

Name Type Description Default
inp HarmonicBondInput

Graph potential input with harmonic bond parameters

 required

Returns:

Type Description
WithPatch[Table[SystemId, Energy], IdPatch[Any]]

Total bond energy per system
Source code in src/kups/potential/classical/harmonic.py 
 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
def harmonic_bond_energy(
    inp: HarmonicBondInput,
) -> WithPatch[Table[SystemId, Energy], IdPatch[Any]]:
    """Compute harmonic bond energy for all bonds.

    Calculates energy as k(r - r₀)² for each bond and sums over all systems.

    Args:
        inp: Graph potential input with harmonic bond parameters

    Returns:
        Total bond energy per system
    """
    graph = inp.graph
    assert graph.edges.indices.indices.shape[1] == 2, (
        "Harmonic bond potential only supports pairwise interactions (order=2)."
    )
    edg_species = graph.particles[graph.edges.indices].labels.indices_in(
        inp.parameters.labels
    )
    x0 = inp.parameters.x0[edg_species[:, 0], edg_species[:, 1]]
    k = inp.parameters.k[edg_species[:, 0], edg_species[:, 1]]
    edge_energy = (jnp.linalg.norm(graph.edge_shifts[:, 0], axis=-1) - x0) ** 2 * k
    total_energies = graph.edge_batch_mask.sum_over(edge_energy)
    return WithPatch(total_energies, IdPatch[Any]())

make_harmonic_angle_potential(particles_view, edge_indices_view, systems_view, parameter_view, probe, gradient_lens, hessian_lens, hessian_idx_view, patch_idx_view=None, out_cache_lens=None)

Create harmonic angle potential for explicitly defined angles.

Applies harmonic restraints to specified atom triplets (angles). Angles must be explicitly provided via the input_view edge set as triplets (i-j-k).

Parameters:

Name Type Description Default
particles_view View[State, Table[ParticleId, IsBondedParticles]]

Extracts particle data (positions, species) with system index

 required
edge_indices_view View[State, Index[ParticleId]]

Extracts angle connectivity (triplets)

 required
systems_view View[State, Table[SystemId, HasCell[AnyPeriodicity]]]

Extracts indexed system data (cell)

 required
parameter_view View[State, HarmonicAngleParameters]

Extracts HarmonicAngleParameters

 required
probe Probe[State, P, IsGraphProbe[IsBondedParticles, Literal[3]]] | None

Graph probe for incremental particle and neighbor-list updates

 required
gradient_lens Lens[HarmonicAngleInput, Gradients]

Specifies gradients to compute

 required
hessian_lens Lens[Gradients, Hessians]

Specifies Hessians to compute

 required
hessian_idx_view View[State, Hessians]

Hessian index structure

 required
patch_idx_view View[State, PotentialOut[Gradients, Hessians]] | None

Cached output index structure

 None
out_cache_lens Lens[State, PotentialOut[Gradients, Hessians]] | None

Cache location lens

 None

Returns:

Type Description
Potential[State, Gradients, Hessians, P]

Harmonic angle Potential
Source code in src/kups/potential/classical/harmonic.py 
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
def make_harmonic_angle_potential[
    State,
    P: Patch[Any],
    Gradients,
    Hessians,
](
    particles_view: View[State, Table[ParticleId, IsBondedParticles]],
    edge_indices_view: View[State, Index[ParticleId]],
    systems_view: View[State, Table[SystemId, HasCell[AnyPeriodicity]]],
    parameter_view: View[State, HarmonicAngleParameters],
    probe: Probe[State, P, IsGraphProbe[IsBondedParticles, Literal[3]]] | None,
    gradient_lens: Lens[HarmonicAngleInput, Gradients],
    hessian_lens: Lens[Gradients, Hessians],
    hessian_idx_view: View[State, Hessians],
    patch_idx_view: View[State, PotentialOut[Gradients, Hessians]] | None = None,
    out_cache_lens: Lens[State, PotentialOut[Gradients, Hessians]] | None = None,
) -> Potential[State, Gradients, Hessians, P]:
    """Create harmonic angle potential for explicitly defined angles.

    Applies harmonic restraints to specified atom triplets (angles). Angles must be
    explicitly provided via the input_view edge set as triplets (i-j-k).

    Args:
        particles_view: Extracts particle data (positions, species) with system index
        edge_indices_view: Extracts angle connectivity (triplets)
        systems_view: Extracts indexed system data (cell)
        parameter_view: Extracts [HarmonicAngleParameters][kups.potential.classical.harmonic.HarmonicAngleParameters]
        probe: Graph probe for incremental particle and neighbor-list updates
        gradient_lens: Specifies gradients to compute
        hessian_lens: Specifies Hessians to compute
        hessian_idx_view: Hessian index structure
        patch_idx_view: Cached output index structure
        out_cache_lens: Cache location lens

    Returns:
        Harmonic angle [Potential][kups.core.potential.Potential]
    """
    graph_fn = GraphConstructor(
        particles=particles_view,
        systems=systems_view,
        neighborlist=lambda state: FixedEdgesNeighborList[Literal[3]](
            edge_indices_view(state)
        ),
        probe=probe,
    )
    composer = LocalGraphSumComposer(
        graph_constructor=graph_fn,
        parameter_view=parameter_view,
    )
    potential = PotentialFromEnergy(
        composer=composer,
        energy_fn=harmonic_angle_energy,
        gradient_lens=gradient_lens,
        hessian_lens=hessian_lens,
        hessian_idx_view=hessian_idx_view,
        cache_lens=out_cache_lens,
        patch_idx_view=patch_idx_view,
    )
    return potential

make_harmonic_bond_potential(particles_view, edge_indices_view, systems_view, parameter_view, probe, gradient_lens, hessian_lens, hessian_idx_view, patch_idx_view=None, out_cache_lens=None)

Create harmonic bond potential for explicitly defined bonds.

Applies harmonic restraints to specified atom pairs (bonds). Bonds must be explicitly provided via the input_view edge set.

Parameters:

Name Type Description Default
particles_view View[State, Table[ParticleId, IsBondedParticles]]

Extracts particle data (positions, species) with system index

 required
edge_indices_view View[State, Index[ParticleId]]

Extracts bond connectivity

 required
systems_view View[State, Table[SystemId, HasCell[AnyPeriodicity]]]

Extracts indexed system data (cell)

 required
parameter_view View[State, HarmonicBondParameters]

Extracts HarmonicBondParameters

 required
probe Probe[State, P, IsGraphProbe[IsBondedParticles, Literal[2]]] | None

Graph probe for incremental particle and neighbor-list updates

 required
gradient_lens Lens[HarmonicBondInput, Gradients]

Specifies gradients to compute

 required
hessian_lens Lens[Gradients, Hessians]

Specifies Hessians to compute

 required
hessian_idx_view View[State, Hessians]

Hessian index structure

 required
patch_idx_view View[State, PotentialOut[Gradients, Hessians]] | None

Cached output index structure

 None
out_cache_lens Lens[State, PotentialOut[Gradients, Hessians]] | None

Cache location lens

 None

Returns:

Type Description
Potential[State, Gradients, Hessians, P]

Harmonic bond Potential
Source code in src/kups/potential/classical/harmonic.py 
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
def make_harmonic_bond_potential[
    State,
    P: Patch[Any],
    Gradients,
    Hessians,
](
    particles_view: View[State, Table[ParticleId, IsBondedParticles]],
    edge_indices_view: View[State, Index[ParticleId]],
    systems_view: View[State, Table[SystemId, HasCell[AnyPeriodicity]]],
    parameter_view: View[State, HarmonicBondParameters],
    probe: Probe[State, P, IsGraphProbe[IsBondedParticles, Literal[2]]] | None,
    gradient_lens: Lens[HarmonicBondInput, Gradients],
    hessian_lens: Lens[Gradients, Hessians],
    hessian_idx_view: View[State, Hessians],
    patch_idx_view: View[State, PotentialOut[Gradients, Hessians]] | None = None,
    out_cache_lens: Lens[State, PotentialOut[Gradients, Hessians]] | None = None,
) -> Potential[State, Gradients, Hessians, P]:
    """Create harmonic bond potential for explicitly defined bonds.

    Applies harmonic restraints to specified atom pairs (bonds). Bonds must be
    explicitly provided via the input_view edge set.

    Args:
        particles_view: Extracts particle data (positions, species) with system index
        edge_indices_view: Extracts bond connectivity
        systems_view: Extracts indexed system data (cell)
        parameter_view: Extracts [HarmonicBondParameters][kups.potential.classical.harmonic.HarmonicBondParameters]
        probe: Graph probe for incremental particle and neighbor-list updates
        gradient_lens: Specifies gradients to compute
        hessian_lens: Specifies Hessians to compute
        hessian_idx_view: Hessian index structure
        patch_idx_view: Cached output index structure
        out_cache_lens: Cache location lens

    Returns:
        Harmonic bond [Potential][kups.core.potential.Potential]
    """
    graph_fn = GraphConstructor(
        particles=particles_view,
        systems=systems_view,
        neighborlist=lambda state: FixedEdgesNeighborList[Literal[2]](
            edge_indices_view(state)
        ),
        probe=probe,
    )
    composer = LocalGraphSumComposer(
        graph_constructor=graph_fn,
        parameter_view=parameter_view,
    )
    potential = PotentialFromEnergy(
        composer=composer,
        energy_fn=harmonic_bond_energy,
        gradient_lens=gradient_lens,
        hessian_lens=hessian_lens,
        hessian_idx_view=hessian_idx_view,
        cache_lens=out_cache_lens,
        patch_idx_view=patch_idx_view,
    )
    return potential