Coverage for ase / constraints / fix_internals.py: 92.53%
281 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-03-30 08:22 +0000
1from __future__ import annotations
3from typing import Sequence
4from warnings import warn
6import numpy as np
8from ase.constraints.constraint import FixConstraint, slice2enlist
9from ase.geometry import (
10 conditional_find_mic,
11 get_angles,
12 get_angles_derivatives,
13 get_dihedrals,
14 get_dihedrals_derivatives,
15 get_distances_derivatives,
16)
19# TODO: Better interface might be to use dictionaries in place of very
20# nested lists/tuples
21class FixInternals(FixConstraint):
22 """Constraint object for fixing multiple internal coordinates.
24 Allows fixing bonds, angles, dihedrals as well as linear combinations
25 of bonds (bondcombos).
27 Please provide angular units in degrees using `angles_deg` and
28 `dihedrals_deg`.
29 Fixing planar angles is not supported at the moment.
30 """
32 def __init__(
33 self,
34 bonds=None,
35 angles=None,
36 dihedrals=None,
37 angles_deg=None,
38 dihedrals_deg=None,
39 bondcombos=None,
40 mic=False,
41 epsilon=1.0e-7,
42 ):
43 """
44 A constrained internal coordinate is defined as a nested list:
45 '[value, [atom indices]]'. The constraint is initialized with a list of
46 constrained internal coordinates, i.e. '[[value, [atom indices]], ...]'.
47 If 'value' is None, the current value of the coordinate is constrained.
49 Parameters
50 ----------
51 bonds: nested python list, optional
52 List with targetvalue and atom indices defining the fixed bonds,
53 i.e. [[targetvalue, [index0, index1]], ...]
55 angles_deg: nested python list, optional
56 List with targetvalue and atom indices defining the fixedangles,
57 i.e. [[targetvalue, [index0, index1, index3]], ...]
59 dihedrals_deg: nested python list, optional
60 List with targetvalue and atom indices defining the fixed dihedrals,
61 i.e. [[targetvalue, [index0, index1, index3]], ...]
63 bondcombos: nested python list, optional
64 List with targetvalue, atom indices and linear coefficient defining
65 the fixed linear combination of bonds,
66 i.e. [[targetvalue, [[index0, index1, coefficient_for_bond],
67 [index1, index2, coefficient_for_bond]]], ...]
69 mic: bool, optional, default: False
70 Minimum image convention.
72 epsilon: float, optional, default: 1e-7
73 Convergence criterion.
74 """
75 warn_msg = 'Please specify {} in degrees using the {} argument.'
76 if angles:
77 warn(warn_msg.format('angles', 'angle_deg'), FutureWarning)
78 angles = np.asarray(angles)
79 angles[:, 0] = angles[:, 0] / np.pi * 180
80 angles = angles.tolist()
81 else:
82 angles = angles_deg
83 if dihedrals:
84 warn(warn_msg.format('dihedrals', 'dihedrals_deg'), FutureWarning)
85 dihedrals = np.asarray(dihedrals)
86 dihedrals[:, 0] = dihedrals[:, 0] / np.pi * 180
87 dihedrals = dihedrals.tolist()
88 else:
89 dihedrals = dihedrals_deg
91 self.bonds = bonds or []
92 self.angles = angles or []
93 self.dihedrals = dihedrals or []
94 self.bondcombos = bondcombos or []
95 self.mic = mic
96 self.epsilon = epsilon
98 self.n = (
99 len(self.bonds)
100 + len(self.angles)
101 + len(self.dihedrals)
102 + len(self.bondcombos)
103 )
105 # Initialize these at run-time:
106 self.constraints = []
107 self.initialized = False
109 def get_removed_dof(self, atoms):
110 return self.n
112 def initialize(self, atoms):
113 if self.initialized:
114 return
115 masses = np.repeat(atoms.get_masses(), 3)
116 cell = None
117 pbc = None
118 if self.mic:
119 cell = atoms.cell
120 pbc = atoms.pbc
121 self.constraints = []
122 for data, ConstrClass in [
123 (self.bonds, self.FixBondLengthAlt),
124 (self.angles, self.FixAngle),
125 (self.dihedrals, self.FixDihedral),
126 (self.bondcombos, self.FixBondCombo),
127 ]:
128 for datum in data:
129 targetvalue = datum[0]
130 if targetvalue is None: # set to current value
131 targetvalue = ConstrClass.get_value(
132 atoms, datum[1], self.mic
133 )
134 constr = ConstrClass(targetvalue, datum[1], masses, cell, pbc)
135 self.constraints.append(constr)
136 self.initialized = True
138 @staticmethod
139 def get_bondcombo(atoms, indices, mic=False):
140 """Convenience function to return the value of the bondcombo coordinate
141 (linear combination of bond lengths) for the given Atoms object 'atoms'.
142 Example: Get the current value of the linear combination of two bond
143 lengths defined as `bondcombo = [[0, 1, 1.0], [2, 3, -1.0]]`."""
144 c = sum(df[2] * atoms.get_distance(*df[:2], mic=mic) for df in indices)
145 return c
147 def get_subconstraint(self, atoms, definition):
148 """Get pointer to a specific subconstraint.
149 Identification by its definition via indices (and coefficients)."""
150 self.initialize(atoms)
151 for subconstr in self.constraints:
152 if isinstance(definition[0], Sequence): # Combo constraint
153 defin = [
154 d + [c] for d, c in zip(subconstr.indices, subconstr.coefs)
155 ]
156 if defin == definition:
157 return subconstr
158 else: # identify primitive constraints by their indices
159 if subconstr.indices == [definition]:
160 return subconstr
161 raise ValueError('Given `definition` not found on Atoms object.')
163 def shuffle_definitions(self, shuffle_dic, internal_type):
164 dfns = [] # definitions
165 for dfn in internal_type: # e.g. for bond in self.bonds
166 append = True
167 new_dfn = [dfn[0], list(dfn[1])]
168 for old in dfn[1]:
169 if old in shuffle_dic:
170 new_dfn[1][dfn[1].index(old)] = shuffle_dic[old]
171 else:
172 append = False
173 break
174 if append:
175 dfns.append(new_dfn)
176 return dfns
178 def shuffle_combos(self, shuffle_dic, internal_type):
179 dfns = [] # definitions
180 for dfn in internal_type: # i.e. for bondcombo in self.bondcombos
181 append = True
182 all_indices = [idx[0:-1] for idx in dfn[1]]
183 new_dfn = [dfn[0], list(dfn[1])]
184 for i, indices in enumerate(all_indices):
185 for old in indices:
186 if old in shuffle_dic:
187 new_dfn[1][i][indices.index(old)] = shuffle_dic[old]
188 else:
189 append = False
190 break
191 if not append:
192 break
193 if append:
194 dfns.append(new_dfn)
195 return dfns
197 def index_shuffle(self, atoms, ind):
198 # See docstring of superclass
199 self.initialize(atoms)
200 shuffle_dic = dict(slice2enlist(ind, len(atoms)))
201 shuffle_dic = {old: new for new, old in shuffle_dic.items()}
202 self.bonds = self.shuffle_definitions(shuffle_dic, self.bonds)
203 self.angles = self.shuffle_definitions(shuffle_dic, self.angles)
204 self.dihedrals = self.shuffle_definitions(shuffle_dic, self.dihedrals)
205 self.bondcombos = self.shuffle_combos(shuffle_dic, self.bondcombos)
206 self.initialized = False
207 self.initialize(atoms)
208 if len(self.constraints) == 0:
209 raise IndexError('Constraint not part of slice')
211 def get_indices(self):
212 cons = []
213 for dfn in self.bonds + self.dihedrals + self.angles:
214 cons.extend(dfn[1])
215 for dfn in self.bondcombos:
216 for partial_dfn in dfn[1]:
217 cons.extend(partial_dfn[0:-1]) # last index is the coefficient
218 return list(set(cons))
220 def todict(self):
221 return {
222 'name': 'FixInternals',
223 'kwargs': {
224 'bonds': self.bonds,
225 'angles_deg': self.angles,
226 'dihedrals_deg': self.dihedrals,
227 'bondcombos': self.bondcombos,
228 'mic': self.mic,
229 'epsilon': self.epsilon,
230 },
231 }
233 def adjust_positions(self, atoms, newpos):
234 self.initialize(atoms)
235 for constraint in self.constraints:
236 constraint.setup_jacobian(atoms.positions)
237 for _ in range(50):
238 maxerr = 0.0
239 for constraint in self.constraints:
240 constraint.adjust_positions(atoms.positions, newpos)
241 maxerr = max(abs(constraint.sigma), maxerr)
242 if maxerr < self.epsilon:
243 return
244 msg = 'FixInternals.adjust_positions did not converge.'
245 if any(
246 constr.targetvalue > 175.0 or constr.targetvalue < 5.0
247 for constr in self.constraints
248 if isinstance(constr, self.FixAngle)
249 ):
250 msg += (
251 ' This may be caused by an almost planar angle.'
252 ' Support for planar angles would require the'
253 ' implementation of ghost, i.e. dummy, atoms.'
254 ' See issue #868.'
255 )
256 raise ValueError(msg)
258 def adjust_forces(self, atoms, forces):
259 """Project out translations and rotations and all other constraints"""
260 self.initialize(atoms)
261 positions = atoms.positions
262 N = len(forces)
263 list2_constraints = list(np.zeros((6, N, 3)))
264 tx, ty, tz, rx, ry, rz = list2_constraints
266 list_constraints = [r.ravel() for r in list2_constraints]
268 tx[:, 0] = 1.0
269 ty[:, 1] = 1.0
270 tz[:, 2] = 1.0
271 ff = forces.ravel()
273 # Calculate the center of mass
274 center = positions.sum(axis=0) / N
276 rx[:, 1] = -(positions[:, 2] - center[2])
277 rx[:, 2] = positions[:, 1] - center[1]
278 ry[:, 0] = positions[:, 2] - center[2]
279 ry[:, 2] = -(positions[:, 0] - center[0])
280 rz[:, 0] = -(positions[:, 1] - center[1])
281 rz[:, 1] = positions[:, 0] - center[0]
283 # Normalizing transl., rotat. constraints
284 for r in list2_constraints:
285 r /= np.linalg.norm(r.ravel())
287 # Add all angle, etc. constraint vectors
288 for constraint in self.constraints:
289 constraint.setup_jacobian(positions)
290 constraint.adjust_forces(positions, forces)
291 list_constraints.insert(0, constraint.jacobian)
292 # QR DECOMPOSITION - GRAM SCHMIDT
294 list_constraints = [r.ravel() for r in list_constraints]
295 aa = np.column_stack(list_constraints)
296 (aa, _bb) = np.linalg.qr(aa)
297 # Projection
298 hh = []
299 for i, constraint in enumerate(self.constraints):
300 hh.append(aa[:, i] * np.vstack(aa[:, i]))
302 txx = aa[:, self.n] * np.vstack(aa[:, self.n])
303 tyy = aa[:, self.n + 1] * np.vstack(aa[:, self.n + 1])
304 tzz = aa[:, self.n + 2] * np.vstack(aa[:, self.n + 2])
305 rxx = aa[:, self.n + 3] * np.vstack(aa[:, self.n + 3])
306 ryy = aa[:, self.n + 4] * np.vstack(aa[:, self.n + 4])
307 rzz = aa[:, self.n + 5] * np.vstack(aa[:, self.n + 5])
308 T = txx + tyy + tzz + rxx + ryy + rzz
309 for vec in hh:
310 T += vec
311 ff = np.dot(T, np.vstack(ff))
312 forces[:, :] -= np.dot(T, np.vstack(ff)).reshape(-1, 3)
314 def __repr__(self):
315 constraints = [repr(constr) for constr in self.constraints]
316 return f'FixInternals(_copy_init={constraints}, epsilon={self.epsilon})'
318 # Classes for internal use in FixInternals
319 class FixInternalsBase:
320 """Base class for subclasses of FixInternals."""
322 def __init__(self, targetvalue, indices, masses, cell, pbc):
323 self.targetvalue = targetvalue # constant target value
324 self.indices = [defin[0:-1] for defin in indices] # indices, defs
325 self.coefs = np.asarray([defin[-1] for defin in indices])
326 self.masses = masses
327 self.jacobian = [] # geometric Jacobian matrix, Wilson B-matrix
328 self.sigma = 1.0 # difference between current and target value
329 self.projected_force = None # helps optimizers scan along constr.
330 self.cell = cell
331 self.pbc = pbc
333 def finalize_jacobian(self, pos, n_internals, n, derivs):
334 """Populate jacobian with derivatives for `n_internals` defined
335 internals. n = 2 (bonds), 3 (angles), 4 (dihedrals)."""
336 jacobian = np.zeros((n_internals, *pos.shape))
337 for i, idx in enumerate(self.indices):
338 for j in range(n):
339 jacobian[i, idx[j]] = derivs[i, j]
340 jacobian = jacobian.reshape((n_internals, 3 * len(pos)))
341 return self.coefs @ jacobian
343 def finalize_positions(self, newpos):
344 jacobian = self.jacobian / self.masses
345 lamda = -self.sigma / (jacobian @ self.get_jacobian(newpos))
346 dnewpos = lamda * jacobian
347 newpos += dnewpos.reshape(newpos.shape)
349 def adjust_forces(self, positions, forces):
350 self.projected_forces = (
351 self.jacobian @ forces.ravel()
352 ) * self.jacobian
353 self.jacobian /= np.linalg.norm(self.jacobian)
355 class FixBondCombo(FixInternalsBase):
356 """Constraint subobject for fixing linear combination of bond lengths
357 within FixInternals.
359 sum_i( coef_i * bond_length_i ) = constant
360 """
362 def get_jacobian(self, pos):
363 bondvectors = [pos[k] - pos[h] for h, k in self.indices]
364 derivs = get_distances_derivatives(
365 bondvectors, cell=self.cell, pbc=self.pbc
366 )
367 return self.finalize_jacobian(pos, len(bondvectors), 2, derivs)
369 def setup_jacobian(self, pos):
370 self.jacobian = self.get_jacobian(pos)
372 def adjust_positions(self, oldpos, newpos):
373 bondvectors = [newpos[k] - newpos[h] for h, k in self.indices]
374 (_,), (dists,) = conditional_find_mic(
375 [bondvectors], cell=self.cell, pbc=self.pbc
376 )
377 value = self.coefs @ dists
378 self.sigma = value - self.targetvalue
379 self.finalize_positions(newpos)
381 @staticmethod
382 def get_value(atoms, indices, mic):
383 return FixInternals.get_bondcombo(atoms, indices, mic)
385 def __repr__(self):
386 return (
387 f'FixBondCombo({self.targetvalue}, {self.indices}, '
388 '{self.coefs})'
389 )
391 class FixBondLengthAlt(FixBondCombo):
392 """Constraint subobject for fixing bond length within FixInternals.
393 Fix distance between atoms with indices a1, a2."""
395 def __init__(self, targetvalue, indices, masses, cell, pbc):
396 if targetvalue <= 0.0:
397 raise ZeroDivisionError('Invalid targetvalue for fixed bond')
398 indices = [list(indices) + [1.0]] # bond definition with coef 1.
399 super().__init__(targetvalue, indices, masses, cell=cell, pbc=pbc)
401 @staticmethod
402 def get_value(atoms, indices, mic):
403 return atoms.get_distance(*indices, mic=mic)
405 def __repr__(self):
406 return f'FixBondLengthAlt({self.targetvalue}, {self.indices})'
408 class FixAngle(FixInternalsBase):
409 """Constraint subobject for fixing an angle within FixInternals.
411 Convergence is potentially problematic for angles very close to
412 0 or 180 degrees as there is a singularity in the Cartesian derivative.
413 Fixing planar angles is therefore not supported at the moment.
414 """
416 def __init__(self, targetvalue, indices, masses, cell, pbc):
417 """Fix atom movement to construct a constant angle."""
418 if targetvalue <= 0.0 or targetvalue >= 180.0:
419 raise ZeroDivisionError('Invalid targetvalue for fixed angle')
420 indices = [list(indices) + [1.0]] # angle definition with coef 1.
421 super().__init__(targetvalue, indices, masses, cell=cell, pbc=pbc)
423 def gather_vectors(self, pos):
424 v0 = [pos[h] - pos[k] for h, k, l in self.indices]
425 v1 = [pos[l] - pos[k] for h, k, l in self.indices]
426 return v0, v1
428 def get_jacobian(self, pos):
429 v0, v1 = self.gather_vectors(pos)
430 derivs = get_angles_derivatives(
431 v0, v1, cell=self.cell, pbc=self.pbc
432 )
433 return self.finalize_jacobian(pos, len(v0), 3, derivs)
435 def setup_jacobian(self, pos):
436 self.jacobian = self.get_jacobian(pos)
438 def adjust_positions(self, oldpos, newpos):
439 v0, v1 = self.gather_vectors(newpos)
440 value = get_angles(v0, v1, cell=self.cell, pbc=self.pbc)
441 self.sigma = value - self.targetvalue
442 self.finalize_positions(newpos)
444 @staticmethod
445 def get_value(atoms, indices, mic):
446 return atoms.get_angle(*indices, mic=mic)
448 def __repr__(self):
449 return f'FixAngle({self.targetvalue}, {self.indices})'
451 class FixDihedral(FixInternalsBase):
452 """Constraint subobject for fixing a dihedral angle within FixInternals.
454 A dihedral becomes undefined when at least one of the inner two angles
455 becomes planar. Make sure to avoid this situation.
456 """
458 def __init__(self, targetvalue, indices, masses, cell, pbc):
459 indices = [list(indices) + [1.0]] # dihedral def. with coef 1.
460 super().__init__(targetvalue, indices, masses, cell=cell, pbc=pbc)
462 def gather_vectors(self, pos):
463 v0 = [pos[k] - pos[h] for h, k, l, m in self.indices]
464 v1 = [pos[l] - pos[k] for h, k, l, m in self.indices]
465 v2 = [pos[m] - pos[l] for h, k, l, m in self.indices]
466 return v0, v1, v2
468 def get_jacobian(self, pos):
469 v0, v1, v2 = self.gather_vectors(pos)
470 derivs = get_dihedrals_derivatives(
471 v0, v1, v2, cell=self.cell, pbc=self.pbc
472 )
473 return self.finalize_jacobian(pos, len(v0), 4, derivs)
475 def setup_jacobian(self, pos):
476 self.jacobian = self.get_jacobian(pos)
478 def adjust_positions(self, oldpos, newpos):
479 v0, v1, v2 = self.gather_vectors(newpos)
480 value = get_dihedrals(v0, v1, v2, cell=self.cell, pbc=self.pbc)
481 # apply minimum dihedral difference 'convention': (diff <= 180)
482 self.sigma = (value - self.targetvalue + 180) % 360 - 180
483 self.finalize_positions(newpos)
485 @staticmethod
486 def get_value(atoms, indices, mic):
487 return atoms.get_dihedral(*indices, mic=mic)
489 def __repr__(self):
490 return f'FixDihedral({self.targetvalue}, {self.indices})'