Coverage for /builds/ase/ase/ase/build/supercells.py: 82.58%
155 statements
« prev ^ index » next coverage.py v7.5.3, created at 2025-08-02 00:12 +0000
1# fmt: off
3"""Helper functions for creating supercells."""
5import numpy as np
7from ase import Atoms
8from ase.utils import deprecated
11class SupercellError(Exception):
12 """Use if construction of supercell fails"""
15@deprecated('use `eval_length_deviation` instead.')
16def get_deviation_from_optimal_cell_shape(*args, **kwargs):
17 return eval_length_deviation(*args, **kwargs)
20def eval_shape_deviation(cell, target_shape="sc"):
21 r"""
22 Calculates the deviation of the given cell from the target cell metric.
24 Parameters
25 ----------
26 cell : (..., 3, 3) array_like
27 Metric given as a 3x3 matrix of the input structure.
28 Multiple cells can also be given as a higher-dimensional array.
29 target_shape : {'sc', 'fcc'}
30 Desired supercell shape.
32 Returns
33 -------
34 float or ndarray
35 Cell metric(s) (0 is perfect score)
37 """
39 cell = np.asarray(cell)
41 eff_cubic_length = np.cbrt(np.abs(np.linalg.det(cell))) # 'a_0'
43 if target_shape == 'sc':
44 target_len = eff_cubic_length
45 target_cos = 0.0 # cos(+-pi/2) = 0.0
46 target_metric = np.eye(3)
47 elif target_shape == 'fcc':
48 # FCC is characterised by 60 degree angles & lattice vectors = 2**(1/6)
49 # times the eff cubic length:
50 target_len = eff_cubic_length * 2 ** (1 / 6)
51 target_cos = 0.5 # cos(+-pi/3) = 0.5
52 target_metric = np.eye(3) + target_cos * (np.ones((3, 3)) - np.eye(3))
53 else:
54 raise ValueError(target_shape)
56 # calculate cell @ cell.T for (... , 3, 3)
57 # with cell -> C_mij
58 # and metric -> M_mkl
59 # M_mkl = (sum_j C_mkj * C_mlj) / leff**2
60 metric = cell @ np.swapaxes(cell, -2, -1)
61 normed = metric / target_len[..., None, None] ** 2
63 # offdiagonal ~ cos angle -> score = np.abs(cos angle - cos target_angle)
64 scores = np.add.reduce((normed - target_metric) ** 2, axis=(-2, -1))
66 return scores
69def eval_length_deviation(cell, target_shape="sc"):
70 r"""Calculate the deviation from the target cell shape.
72 Calculates the deviation of the given cell metric from the ideal
73 cell metric defining a certain shape. Specifically, the function
74 evaluates the expression `\Delta = || Q \mathbf{h} -
75 \mathbf{h}_{target}||_2`, where `\mathbf{h}` is the input
76 metric (*cell*) and `Q` is a normalization factor (*norm*)
77 while the target metric `\mathbf{h}_{target}` (via
78 *target_shape*) represent simple cubic ('sc') or face-centered
79 cubic ('fcc') cell shapes.
81 Replaced with code from the `doped` defect simulation package
82 (https://doped.readthedocs.io) to be rotationally invariant,
83 boosting performance.
85 Parameters
86 ----------
87 cell : (..., 3, 3) array_like
88 Metric given as a 3x3 matrix of the input structure.
89 Multiple cells can also be given as a higher-dimensional array.
90 target_shape : {'sc', 'fcc'}
91 Desired supercell shape. Can be 'sc' for simple cubic or
92 'fcc' for face-centered cubic.
94 Returns
95 -------
96 float or ndarray
97 Cell metric(s) (0 is perfect score)
99 .. deprecated:: 3.24.0
100 `norm` is unused in ASE 3.24.0 and removed in ASE 3.25.0.
102 """
104 cell = np.asarray(cell)
105 cell_lengths = np.sqrt(np.add.reduce(cell**2, axis=-1))
107 eff_cubic_length = np.cbrt(np.abs(np.linalg.det(cell))) # 'a_0'
109 if target_shape == 'sc':
110 target_len = eff_cubic_length
112 elif target_shape == 'fcc':
113 # FCC is characterised by 60 degree angles & lattice vectors = 2**(1/6)
114 # times the eff cubic length:
115 target_len = eff_cubic_length * 2 ** (1 / 6)
117 else:
118 raise ValueError(target_shape)
120 inv_target_len = 1.0 / target_len
122 # rms difference to eff cubic/FCC length:
123 diffs = cell_lengths * inv_target_len[..., None] - 1.0
124 scores = np.sqrt(np.add.reduce(diffs**2, axis=-1))
126 return scores
129def _guess_initial_transformation(cell, target_shape,
130 target_size, verbose=False):
132 # Set up target metric
133 if target_shape == 'sc':
134 target_metric = np.eye(3)
135 elif target_shape == 'fcc':
136 target_metric = 0.5 * np.array([[0, 1, 1], [1, 0, 1], [1, 1, 0]],
137 dtype=float)
138 else:
139 raise ValueError(target_shape)
141 if verbose:
142 print("target metric (h_target):")
143 print(target_metric)
145 # Normalize cell metric to reduce computation time during looping
146 norm = (target_size * abs(np.linalg.det(cell)) /
147 np.linalg.det(target_metric)) ** (-1.0 / 3)
148 norm_cell = norm * cell
149 if verbose:
150 print(f"normalization factor (Q): {norm:g}")
152 # Approximate initial P matrix
153 ideal_P = np.dot(target_metric, np.linalg.inv(norm_cell))
154 if verbose:
155 print("idealized transformation matrix:")
156 print(ideal_P)
157 starting_P = np.array(np.around(ideal_P, 0), dtype=int)
158 if verbose:
159 print("closest integer transformation matrix (P_0):")
160 print(starting_P)
162 return ideal_P, starting_P
165def _build_matrix_operations(starting_P, lower_limit, upper_limit):
166 mat_dim = starting_P.shape[0]
168 if not mat_dim == starting_P.shape[1]:
169 raise ValueError('Cell matrix should be quadratic.')
171 # Build a big matrix of all admissible integer matrix operations.
172 # (If this takes too much memory we could do blocking but there are
173 # too many for looping one by one.)
174 dimensions = [(upper_limit + 1) - lower_limit] * mat_dim**2
175 operations = np.moveaxis(np.indices(dimensions), 0, -1)
176 operations = operations.reshape(-1, mat_dim, mat_dim)
177 operations += lower_limit # Each element runs from lower to upper limits.
178 operations += starting_P
180 return operations
183def _screen_supercell_size(operations, target_size):
185 # screen supercells with the target size
186 determinants = np.round(np.linalg.det(operations), 0).astype(int)
187 good_indices = np.where(np.abs(determinants) == target_size)[0]
189 if not good_indices.size:
190 print("Failed to find a transformation matrix.")
191 return None
192 operations = operations[good_indices]
194 return operations
197def _optimal_transformation(operations, scores, ideal_P):
199 imin = np.argmin(scores)
200 best_score = scores[imin]
201 # screen candidates with the same best score
202 operations = operations[np.abs(scores - best_score) < 1e-6]
204 # select the one whose cell orientation is the closest to the target
205 # https://gitlab.com/ase/ase/-/merge_requests/3522
206 imin = np.argmin(np.add.reduce((operations - ideal_P)**2, axis=(-2, -1)))
208 optimal_P = operations[imin]
210 if np.linalg.det(optimal_P) <= 0:
211 optimal_P *= -1 # flip signs if negative determinant
213 return optimal_P, best_score
216all_score_funcs = {"length": eval_length_deviation,
217 "metric": eval_shape_deviation}
220def find_optimal_cell_shape(
221 cell,
222 target_size,
223 target_shape,
224 lower_limit=-2,
225 upper_limit=2,
226 verbose=False,
227 score_key='length'
228):
229 """Obtain the optimal transformation matrix for a supercell of target size
230 and shape.
232 Returns the transformation matrix that produces a supercell
233 corresponding to *target_size* unit cells with metric *cell* that
234 most closely approximates the shape defined by *target_shape*.
236 Updated with code from the `doped` defect simulation package
237 (https://doped.readthedocs.io) to be rotationally invariant and
238 allow transformation matrices with negative determinants, boosting
239 performance.
241 Parameters:
243 cell: 2D array of floats
244 Metric given as a (3x3 matrix) of the input structure.
245 target_size: integer
246 Size of desired supercell in number of unit cells.
247 target_shape: str
248 Desired supercell shape. Can be 'sc' for simple cubic or
249 'fcc' for face-centered cubic.
250 lower_limit: int
251 Lower limit of search range.
252 upper_limit: int
253 Upper limit of search range.
254 verbose: bool
255 Set to True to obtain additional information regarding
256 construction of transformation matrix.
257 score_key: str
258 key from all_score_funcs to select score function.
260 Returns:
261 2D array of integers: Transformation matrix that produces the
262 optimal supercell.
263 """
265 # transform to np.array
266 cell = np.asarray(cell)
268 # get starting transformation
269 # ideal_P ... transformation: target_cell = ideal_P @ cell
270 # starting_P ... integer rounded (ideal_P)
271 ideal_P, starting_P = _guess_initial_transformation(cell, target_shape,
272 target_size,
273 verbose=verbose)
275 # build all admissible matrix operations 'centered' at starting_P
276 operations = _build_matrix_operations(starting_P,
277 lower_limit, upper_limit)
279 # pre-screen operations based on target_size
280 operations = _screen_supercell_size(operations, target_size)
282 # evaluate derivations of the screened supercells
283 if score_key in all_score_funcs:
284 get_deviation_score = all_score_funcs[score_key]
285 else:
286 msg = (f'Score func key {score_key} not implemented.'
287 + f'Please select from {all_score_funcs}.')
288 raise SupercellError(msg)
290 scores = get_deviation_score(operations @ cell,
291 target_shape)
293 # obtain optimal transformation from scores
294 optimal_P, best_score = _optimal_transformation(operations, scores, ideal_P)
296 # Finalize.
297 if verbose:
298 print(f"smallest score (|Q P h_p - h_target|_2): {best_score:f}")
299 print("optimal transformation matrix (P_opt):")
300 print(optimal_P)
301 print("supercell metric:")
302 print(np.round(np.dot(optimal_P, cell), 4))
303 det = np.linalg.det(optimal_P)
304 print(f"determinant of optimal transformation matrix: {det:g}")
306 return optimal_P
309def make_supercell(prim, P, *, wrap=True, order="cell-major", tol=1e-5):
310 r"""Generate a supercell by applying a general transformation (*P*) to
311 the input configuration (*prim*).
313 The transformation is described by a 3x3 integer matrix
314 `\mathbf{P}`. Specifically, the new cell metric
315 `\mathbf{h}` is given in terms of the metric of the input
316 configuration `\mathbf{h}_p` by `\mathbf{P h}_p =
317 \mathbf{h}`.
319 Parameters:
321 prim: ASE Atoms object
322 Input configuration.
323 P: 3x3 integer matrix
324 Transformation matrix `\mathbf{P}`.
325 wrap: bool
326 wrap in the end
327 order: str (default: "cell-major")
328 how to order the atoms in the supercell
330 "cell-major":
331 [atom1_shift1, atom2_shift1, ..., atom1_shift2, atom2_shift2, ...]
332 i.e. run first over all the atoms in cell1 and then move to cell2.
334 "atom-major":
335 [atom1_shift1, atom1_shift2, ..., atom2_shift1, atom2_shift2, ...]
336 i.e. run first over atom1 in all the cells and then move to atom2.
337 This may be the order preferred by most VASP users.
339 tol: float
340 tolerance for wrapping
341 """
343 supercell_matrix = P
344 supercell = clean_matrix(supercell_matrix @ prim.cell)
346 # cartesian lattice points
347 lattice_points_frac = lattice_points_in_supercell(supercell_matrix)
348 lattice_points = np.dot(lattice_points_frac, supercell)
349 N = len(lattice_points)
351 if order == "cell-major":
352 shifted = prim.positions[None, :, :] + lattice_points[:, None, :]
353 elif order == "atom-major":
354 shifted = prim.positions[:, None, :] + lattice_points[None, :, :]
355 else:
356 raise ValueError(f"invalid order: {order}")
357 shifted_reshaped = shifted.reshape(-1, 3)
359 superatoms = Atoms(positions=shifted_reshaped,
360 cell=supercell,
361 pbc=prim.pbc)
363 # Copy over any other possible arrays, inspired by atoms.__imul__
364 for name, arr in prim.arrays.items():
365 if name == "positions":
366 # This was added during construction of the super cell
367 continue
368 shape = (N * arr.shape[0], *arr.shape[1:])
369 if order == "cell-major":
370 new_arr = np.repeat(arr[None, :], N, axis=0).reshape(shape)
371 elif order == "atom-major":
372 new_arr = np.repeat(arr[:, None], N, axis=1).reshape(shape)
373 superatoms.set_array(name, new_arr)
375 # check number of atoms is correct
376 n_target = abs(int(np.round(np.linalg.det(supercell_matrix) * len(prim))))
377 if n_target != len(superatoms):
378 msg = "Number of atoms in supercell: {}, expected: {}".format(
379 n_target, len(superatoms))
380 raise SupercellError(msg)
382 if wrap:
383 superatoms.wrap(eps=tol)
385 return superatoms
388def lattice_points_in_supercell(supercell_matrix):
389 """Find all lattice points contained in a supercell.
391 Adapted from pymatgen, which is available under MIT license:
392 The MIT License (MIT) Copyright (c) 2011-2012 MIT & The Regents of the
393 University of California, through Lawrence Berkeley National Laboratory
394 """
396 diagonals = np.array([
397 [0, 0, 0],
398 [0, 0, 1],
399 [0, 1, 0],
400 [0, 1, 1],
401 [1, 0, 0],
402 [1, 0, 1],
403 [1, 1, 0],
404 [1, 1, 1],
405 ])
406 d_points = np.dot(diagonals, supercell_matrix)
408 mins = np.min(d_points, axis=0)
409 maxes = np.max(d_points, axis=0) + 1
411 ar = np.arange(mins[0], maxes[0])[:, None] * np.array([1, 0, 0])[None, :]
412 br = np.arange(mins[1], maxes[1])[:, None] * np.array([0, 1, 0])[None, :]
413 cr = np.arange(mins[2], maxes[2])[:, None] * np.array([0, 0, 1])[None, :]
415 all_points = ar[:, None, None] + br[None, :, None] + cr[None, None, :]
416 all_points = all_points.reshape((-1, 3))
418 frac_points = np.dot(all_points, np.linalg.inv(supercell_matrix))
420 tvects = frac_points[np.all(frac_points < 1 - 1e-10, axis=1)
421 & np.all(frac_points >= -1e-10, axis=1)]
422 assert len(tvects) == round(abs(np.linalg.det(supercell_matrix)))
423 return tvects
426def clean_matrix(matrix, eps=1e-12):
427 """ clean from small values"""
428 matrix = np.array(matrix)
429 for ij in np.ndindex(matrix.shape):
430 if abs(matrix[ij]) < eps:
431 matrix[ij] = 0
432 return matrix