Coverage for ase / constraints / fix_parametric_relations.py: 88.62%
167 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-03-30 08:22 +0000
1import numpy as np
3from ase.constraints.constraint import FixConstraint
4from ase.stress import full_3x3_to_voigt_6_stress, voigt_6_to_full_3x3_stress
5from ase.utils.parsemath import eval_expression
8class FixParametricRelations(FixConstraint):
9 def __init__(
10 self,
11 indices,
12 Jacobian,
13 const_shift,
14 params=None,
15 eps=1e-12,
16 use_cell=False,
17 ):
18 """Constrains the degrees of freedom to act in a reduced parameter
19 space defined by the Jacobian
21 These constraints are based off the work in:
22 https://arxiv.org/abs/1908.01610
24 The constraints linearly maps the full 3N degrees of freedom,
25 where N is number of active lattice vectors/atoms onto a
26 reduced subset of M free parameters, where M <= 3*N. The
27 Jacobian matrix and constant shift vector map the full set of
28 degrees of freedom onto the reduced parameter space.
30 Currently the constraint is set up to handle either atomic
31 positions or lattice vectors at one time, but not both. To do
32 both simply add a two constraints for each set. This is done
33 to keep the mathematics behind the operations separate.
35 It would be possible to extend these constraints to allow
36 non-linear transformations if functionality to update the
37 Jacobian at each position update was included. This would
38 require passing an update function evaluate it every time
39 adjust_positions is callled. This is currently NOT supported,
40 and there are no plans to implement it in the future.
42 Args:
43 indices (list of int): indices of the constrained atoms
44 (if not None or empty then cell_indices must be None or Empty)
45 Jacobian (np.ndarray(shape=(3*len(indices), len(params)))):
46 The Jacobian describing
47 the parameter space transformation
48 const_shift (np.ndarray(shape=(3*len(indices)))):
49 A vector describing the constant term
50 in the transformation not accounted for in the Jacobian
51 params (list of str):
52 parameters used in the parametric representation
53 if None a list is generated based on the shape of the Jacobian
54 eps (float): a small number to compare the similarity of
55 numbers and set the precision used
56 to generate the constraint expressions
57 use_cell (bool): if True then act on the cell object
59 """
60 self.indices = np.array(indices)
61 self.Jacobian = np.array(Jacobian)
62 self.const_shift = np.array(const_shift)
64 assert self.const_shift.shape[0] == 3 * len(self.indices)
65 assert self.Jacobian.shape[0] == 3 * len(self.indices)
67 self.eps = eps
68 self.use_cell = use_cell
70 if params is None:
71 params = []
72 if self.Jacobian.shape[1] > 0:
73 int_fmt_str = (
74 '{:0'
75 + str(int(np.ceil(np.log10(self.Jacobian.shape[1]))))
76 + 'd}'
77 )
78 for param_ind in range(self.Jacobian.shape[1]):
79 params.append('param_' + int_fmt_str.format(param_ind))
80 else:
81 assert len(params) == self.Jacobian.shape[-1]
83 self.params = params
85 self.Jacobian_inv = (
86 np.linalg.inv(self.Jacobian.T @ self.Jacobian) @ self.Jacobian.T
87 )
89 @classmethod
90 def from_expressions(
91 cls, indices, params, expressions, eps=1e-12, use_cell=False
92 ):
93 """Converts the expressions into a Jacobian Matrix/const_shift
94 vector and constructs a FixParametricRelations constraint
96 The expressions must be a list like object of size 3*N and
97 elements must be ordered as:
98 [n_0,i; n_0,j; n_0,k; n_1,i; n_1,j; .... ; n_N-1,i; n_N-1,j; n_N-1,k],
99 where i, j, and k are the first, second and third
100 component of the atomic position/lattice
101 vector. Currently only linear operations are allowed to be
102 included in the expressions so
103 only terms like:
104 - const * param_0
105 - sqrt[const] * param_1
106 - const * param_0 +/- const * param_1 +/- ... +/- const * param_M
107 where const is any real number and param_0, param_1, ..., param_M are
108 the parameters passed in
109 params, are allowed.
111 For example, fractional atomic position constraints for wurtzite are:
112 params = ["z1", "z2"]
113 expressions = [
114 "1.0/3.0", "2.0/3.0", "z1",
115 "2.0/3.0", "1.0/3.0", "0.5 + z1",
116 "1.0/3.0", "2.0/3.0", "z2",
117 "2.0/3.0", "1.0/3.0", "0.5 + z2",
118 ]
120 For diamond are:
121 params = []
122 expressions = [
123 "0.0", "0.0", "0.0",
124 "0.25", "0.25", "0.25",
125 ],
127 and for stannite are
128 params=["x4", "z4"]
129 expressions = [
130 "0.0", "0.0", "0.0",
131 "0.0", "0.5", "0.5",
132 "0.75", "0.25", "0.5",
133 "0.25", "0.75", "0.5",
134 "x4 + z4", "x4 + z4", "2*x4",
135 "x4 - z4", "x4 - z4", "-2*x4",
136 "0.0", "-1.0 * (x4 + z4)", "x4 - z4",
137 "0.0", "x4 - z4", "-1.0 * (x4 + z4)",
138 ]
140 Args:
141 indices (list of int): indices of the constrained atoms
142 (if not None or empty then cell_indices must be None or Empty)
143 params (list of str): parameters used in the
144 parametric representation
145 expressions (list of str): expressions used to convert from the
146 parametric to the real space representation
147 eps (float): a small number to compare the similarity of
148 numbers and set the precision used
149 to generate the constraint expressions
150 use_cell (bool): if True then act on the cell object
152 Returns
153 -------
154 cls(
155 indices,
156 Jacobian generated from expressions,
157 const_shift generated from expressions,
158 params,
159 eps-12,
160 use_cell,
161 )
162 """
163 Jacobian = np.zeros((3 * len(indices), len(params)))
164 const_shift = np.zeros(3 * len(indices))
166 for expr_ind, expression in enumerate(expressions):
167 expression = expression.strip()
169 # Convert subtraction to addition
170 expression = expression.replace('-', '+(-1.0)*')
171 if expression[0] == '+':
172 expression = expression[1:]
173 elif expression[:2] == '(+':
174 expression = '(' + expression[2:]
176 # Explicitly add leading zeros so when replacing param_1 with 0.0
177 # param_11 does not become 0.01
178 int_fmt_str = (
179 '{:0' + str(int(np.ceil(np.log10(len(params) + 1)))) + 'd}'
180 )
182 param_dct = {}
183 param_map = {}
185 # Construct a standardized param template for A/B filling
186 for param_ind, param in enumerate(params):
187 param_str = 'param_' + int_fmt_str.format(param_ind)
188 param_map[param] = param_str
189 param_dct[param_str] = 0.0
191 # Replace the parameters according to the map
192 # Sort by string length (long to short) to prevent cases like x11
193 # becoming f"{param_map["x1"]}1"
194 for param in sorted(params, key=lambda s: -1.0 * len(s)):
195 expression = expression.replace(param, param_map[param])
197 # Partial linearity check
198 for express_sec in expression.split('+'):
199 in_sec = [param in express_sec for param in param_dct]
200 n_params_in_sec = len(np.where(np.array(in_sec))[0])
201 if n_params_in_sec > 1:
202 raise ValueError(
203 'FixParametricRelations expressions must be linear.'
204 )
206 const_shift[expr_ind] = float(
207 eval_expression(expression, param_dct)
208 )
210 for param_ind in range(len(params)):
211 param_str = 'param_' + int_fmt_str.format(param_ind)
212 if param_str not in expression:
213 Jacobian[expr_ind, param_ind] = 0.0
214 continue
215 param_dct[param_str] = 1.0
216 test_1 = float(eval_expression(expression, param_dct))
217 test_1 -= const_shift[expr_ind]
218 Jacobian[expr_ind, param_ind] = test_1
220 param_dct[param_str] = 2.0
221 test_2 = float(eval_expression(expression, param_dct))
222 test_2 -= const_shift[expr_ind]
223 if abs(test_2 / test_1 - 2.0) > eps:
224 raise ValueError(
225 'FixParametricRelations expressions must be linear.'
226 )
227 param_dct[param_str] = 0.0
229 args = [
230 indices,
231 Jacobian,
232 const_shift,
233 params,
234 eps,
235 use_cell,
236 ]
237 if cls is FixScaledParametricRelations:
238 args = args[:-1]
239 return cls(*args)
241 @property
242 def expressions(self):
243 """Generate the expressions represented by the current self.Jacobian
244 and self.const_shift objects"""
245 expressions = []
246 per = int(round(-1 * np.log10(self.eps)))
247 fmt_str = '{:.' + str(per + 1) + 'g}'
248 for index, shift_val in enumerate(self.const_shift):
249 exp = ''
250 if (
251 np.all(np.abs(self.Jacobian[index]) < self.eps)
252 or np.abs(shift_val) > self.eps
253 ):
254 exp += fmt_str.format(shift_val)
256 param_exp = ''
257 for param_index, jacob_val in enumerate(self.Jacobian[index]):
258 abs_jacob_val = np.round(np.abs(jacob_val), per + 1)
259 if abs_jacob_val < self.eps:
260 continue
262 param = self.params[param_index]
263 if param_exp or exp:
264 if jacob_val > -1.0 * self.eps:
265 param_exp += ' + '
266 else:
267 param_exp += ' - '
268 elif (
269 (not exp)
270 and (not param_exp)
271 and (jacob_val < -1.0 * self.eps)
272 ):
273 param_exp += '-'
275 if np.abs(abs_jacob_val - 1.0) <= self.eps:
276 param_exp += f'{param:s}'
277 else:
278 param_exp += (fmt_str + '*{:s}').format(
279 abs_jacob_val, param
280 )
282 exp += param_exp
284 expressions.append(exp)
285 return np.array(expressions).reshape((-1, 3))
287 def todict(self):
288 """Create a dictionary representation of the constraint"""
289 return {
290 'name': type(self).__name__,
291 'kwargs': {
292 'indices': self.indices,
293 'params': self.params,
294 'Jacobian': self.Jacobian,
295 'const_shift': self.const_shift,
296 'eps': self.eps,
297 'use_cell': self.use_cell,
298 },
299 }
301 def __repr__(self):
302 """The str representation of the constraint"""
303 if len(self.indices) > 1:
304 indices_str = '[{:d}, ..., {:d}]'.format(
305 self.indices[0], self.indices[-1]
306 )
307 else:
308 indices_str = f'[{self.indices[0]:d}]'
310 if len(self.params) > 1:
311 params_str = '[{:s}, ..., {:s}]'.format(
312 self.params[0], self.params[-1]
313 )
314 elif len(self.params) == 1:
315 params_str = f'[{self.params[0]:s}]'
316 else:
317 params_str = '[]'
319 return '{:s}({:s}, {:s}, ..., {:e})'.format(
320 type(self).__name__, indices_str, params_str, self.eps
321 )
324class FixScaledParametricRelations(FixParametricRelations):
325 def __init__(
326 self,
327 indices,
328 Jacobian,
329 const_shift,
330 params=None,
331 eps=1e-12,
332 ):
333 """The fractional coordinate version of FixParametricRelations
335 All arguments are the same, but since this is for fractional
336 coordinates use_cell is false"""
337 super().__init__(
338 indices,
339 Jacobian,
340 const_shift,
341 params,
342 eps,
343 False,
344 )
346 def adjust_contravariant(self, cell, vecs, B):
347 """Adjust the values of a set of vectors that are contravariant
348 with the unit transformation"""
349 scaled = cell.scaled_positions(vecs).flatten()
350 scaled = self.Jacobian_inv @ (scaled - B)
351 scaled = ((self.Jacobian @ scaled) + B).reshape((-1, 3))
353 return cell.cartesian_positions(scaled)
355 def adjust_positions(self, atoms, positions):
356 """Adjust positions of the atoms to match the constraints"""
357 positions[self.indices] = self.adjust_contravariant(
358 atoms.cell,
359 positions[self.indices],
360 self.const_shift,
361 )
362 positions[self.indices] = self.adjust_B(
363 atoms.cell, positions[self.indices]
364 )
366 def adjust_B(self, cell, positions):
367 """Wraps the positions back to the unit cell and adjust B to
368 keep track of this change"""
369 fractional = cell.scaled_positions(positions)
370 wrapped_fractional = (fractional % 1.0) % 1.0
371 self.const_shift += np.round(wrapped_fractional - fractional).flatten()
372 return cell.cartesian_positions(wrapped_fractional)
374 def adjust_momenta(self, atoms, momenta):
375 """Adjust momenta of the atoms to match the constraints"""
376 momenta[self.indices] = self.adjust_contravariant(
377 atoms.cell,
378 momenta[self.indices],
379 np.zeros(self.const_shift.shape),
380 )
382 def adjust_forces(self, atoms, forces):
383 """Adjust forces of the atoms to match the constraints"""
384 # Forces are coavarient to the coordinate transformation, use the
385 # inverse transformations
386 cart2frac_jacob = np.zeros(2 * (3 * len(atoms),))
387 for i_atom in range(len(atoms)):
388 cart2frac_jacob[
389 3 * i_atom : 3 * (i_atom + 1), 3 * i_atom : 3 * (i_atom + 1)
390 ] = atoms.cell.T
392 jacobian = cart2frac_jacob @ self.Jacobian
393 jacobian_inv = np.linalg.inv(jacobian.T @ jacobian) @ jacobian.T
395 reduced_forces = jacobian.T @ forces.flatten()
396 forces[self.indices] = (jacobian_inv.T @ reduced_forces).reshape(-1, 3)
398 def todict(self):
399 """Create a dictionary representation of the constraint"""
400 dct = super().todict()
401 del dct['kwargs']['use_cell']
402 return dct
405class FixCartesianParametricRelations(FixParametricRelations):
406 def __init__(
407 self,
408 indices,
409 Jacobian,
410 const_shift,
411 params=None,
412 eps=1e-12,
413 use_cell=False,
414 ):
415 """The Cartesian coordinate version of FixParametricRelations"""
416 super().__init__(
417 indices,
418 Jacobian,
419 const_shift,
420 params,
421 eps,
422 use_cell,
423 )
425 def adjust_contravariant(self, vecs, B):
426 """Adjust the values of a set of vectors that are contravariant with
427 the unit transformation"""
428 vecs = self.Jacobian_inv @ (vecs.flatten() - B)
429 vecs = ((self.Jacobian @ vecs) + B).reshape((-1, 3))
430 return vecs
432 def adjust_positions(self, atoms, positions):
433 """Adjust positions of the atoms to match the constraints"""
434 if self.use_cell:
435 return
436 positions[self.indices] = self.adjust_contravariant(
437 positions[self.indices],
438 self.const_shift,
439 )
441 def adjust_momenta(self, atoms, momenta):
442 """Adjust momenta of the atoms to match the constraints"""
443 if self.use_cell:
444 return
445 momenta[self.indices] = self.adjust_contravariant(
446 momenta[self.indices],
447 np.zeros(self.const_shift.shape),
448 )
450 def adjust_forces(self, atoms, forces):
451 """Adjust forces of the atoms to match the constraints"""
452 if self.use_cell:
453 return
455 forces_reduced = self.Jacobian.T @ forces[self.indices].flatten()
456 forces[self.indices] = (self.Jacobian_inv.T @ forces_reduced).reshape(
457 -1, 3
458 )
460 def adjust_cell(self, atoms, cell):
461 """Adjust the cell of the atoms to match the constraints"""
462 if not self.use_cell:
463 return
464 cell[self.indices] = self.adjust_contravariant(
465 cell[self.indices],
466 np.zeros(self.const_shift.shape),
467 )
469 def adjust_stress(self, atoms, stress):
470 """Adjust the stress of the atoms to match the constraints"""
471 if not self.use_cell:
472 return
474 stress_3x3 = voigt_6_to_full_3x3_stress(stress)
475 stress_reduced = self.Jacobian.T @ stress_3x3[self.indices].flatten()
476 stress_3x3[self.indices] = (
477 self.Jacobian_inv.T @ stress_reduced
478 ).reshape(-1, 3)
480 stress[:] = full_3x3_to_voigt_6_stress(stress_3x3)