Coverage for /builds/debichem-team/python-ase/ase/build/supercells.py: 80.17%
116 statements
« prev ^ index » next coverage.py v7.5.3, created at 2025-03-06 04:00 +0000
1"""Helper functions for creating supercells."""
3import warnings
5import numpy as np
7from ase import Atoms
10class SupercellError(Exception):
11 """Use if construction of supercell fails"""
14def get_deviation_from_optimal_cell_shape(cell, target_shape="sc", norm=None):
15 r"""Calculate the deviation from the target cell shape.
17 Calculates the deviation of the given cell metric from the ideal
18 cell metric defining a certain shape. Specifically, the function
19 evaluates the expression `\Delta = || Q \mathbf{h} -
20 \mathbf{h}_{target}||_2`, where `\mathbf{h}` is the input
21 metric (*cell*) and `Q` is a normalization factor (*norm*)
22 while the target metric `\mathbf{h}_{target}` (via
23 *target_shape*) represent simple cubic ('sc') or face-centered
24 cubic ('fcc') cell shapes.
26 Replaced with code from the `doped` defect simulation package
27 (https://doped.readthedocs.io) to be rotationally invariant,
28 boosting performance.
30 Parameters
31 ----------
32 cell : (..., 3, 3) array_like
33 Metric given as a 3x3 matrix of the input structure.
34 Multiple cells can also be given as a higher-dimensional array.
35 target_shape : {'sc', 'fcc'}
36 Desired supercell shape. Can be 'sc' for simple cubic or
37 'fcc' for face-centered cubic.
38 norm : float
39 Specify the normalization factor. This is useful to avoid
40 recomputing the normalization factor when computing the
41 deviation for a series of P matrices.
43 Returns
44 -------
45 float or ndarray
46 Cell metric(s) (0 is perfect score)
48 .. deprecated:: 3.24.0
49 `norm` is unused in ASE 3.24.0 and removed in ASE 3.25.0.
51 """
52 if norm is not None:
53 warnings.warn(
54 '`norm` is unused in ASE 3.24.0 and removed in ASE 3.25.0',
55 FutureWarning,
56 )
58 cell = np.asarray(cell)
59 cell_lengths = np.sqrt(np.add.reduce(cell**2, axis=-1))
60 eff_cubic_length = np.cbrt(np.abs(np.linalg.det(cell))) # 'a_0'
62 if target_shape == 'sc':
63 target_length = eff_cubic_length
65 elif target_shape == 'fcc':
66 # FCC is characterised by 60 degree angles & lattice vectors = 2**(1/6)
67 # times the eff cubic length:
68 target_length = eff_cubic_length * 2 ** (1 / 6)
70 else:
71 raise ValueError(target_shape)
73 inv_target_length = 1.0 / target_length
75 # rms difference to eff cubic/FCC length:
76 diffs = cell_lengths * inv_target_length[..., None] - 1.0
77 return np.sqrt(np.add.reduce(diffs**2, axis=-1))
80def find_optimal_cell_shape(
81 cell,
82 target_size,
83 target_shape,
84 lower_limit=-2,
85 upper_limit=2,
86 verbose=False,
87):
88 """Obtain the optimal transformation matrix for a supercell of target size
89 and shape.
91 Returns the transformation matrix that produces a supercell
92 corresponding to *target_size* unit cells with metric *cell* that
93 most closely approximates the shape defined by *target_shape*.
95 Updated with code from the `doped` defect simulation package
96 (https://doped.readthedocs.io) to be rotationally invariant and
97 allow transformation matrices with negative determinants, boosting
98 performance.
100 Parameters:
102 cell: 2D array of floats
103 Metric given as a (3x3 matrix) of the input structure.
104 target_size: integer
105 Size of desired supercell in number of unit cells.
106 target_shape: str
107 Desired supercell shape. Can be 'sc' for simple cubic or
108 'fcc' for face-centered cubic.
109 lower_limit: int
110 Lower limit of search range.
111 upper_limit: int
112 Upper limit of search range.
113 verbose: bool
114 Set to True to obtain additional information regarding
115 construction of transformation matrix.
117 Returns:
118 2D array of integers: Transformation matrix that produces the
119 optimal supercell.
120 """
121 cell = np.asarray(cell)
123 # Set up target metric
124 if target_shape == 'sc':
125 target_metric = np.eye(3)
126 elif target_shape == 'fcc':
127 target_metric = 0.5 * np.array([[0, 1, 1], [1, 0, 1], [1, 1, 0]],
128 dtype=float)
129 else:
130 raise ValueError(target_shape)
132 if verbose:
133 print("target metric (h_target):")
134 print(target_metric)
136 # Normalize cell metric to reduce computation time during looping
137 norm = (target_size * abs(np.linalg.det(cell)) /
138 np.linalg.det(target_metric)) ** (-1.0 / 3)
139 norm_cell = norm * cell
140 if verbose:
141 print(f"normalization factor (Q): {norm:g}")
143 # Approximate initial P matrix
144 ideal_P = np.dot(target_metric, np.linalg.inv(norm_cell))
145 if verbose:
146 print("idealized transformation matrix:")
147 print(ideal_P)
148 starting_P = np.array(np.around(ideal_P, 0), dtype=int)
149 if verbose:
150 print("closest integer transformation matrix (P_0):")
151 print(starting_P)
153 # Build a big matrix of all admissible integer matrix operations.
154 # (If this takes too much memory we could do blocking but there are
155 # too many for looping one by one.)
156 dimensions = [(upper_limit + 1) - lower_limit] * 9
157 operations = np.moveaxis(np.indices(dimensions), 0, -1).reshape(-1, 3, 3)
158 operations += lower_limit # Each element runs from lower to upper limits.
159 operations += starting_P
160 determinants = np.linalg.det(operations)
162 # screen supercells with the target size
163 good_indices = np.where(abs(determinants - target_size) < 1e-12)[0]
164 if not good_indices.size:
165 print("Failed to find a transformation matrix.")
166 return None
167 operations = operations[good_indices]
169 # evaluate derivations of the screened supercells
170 scores = get_deviation_from_optimal_cell_shape(
171 operations @ cell,
172 target_shape,
173 )
174 imin = np.argmin(scores)
175 best_score = scores[imin]
177 # screen candidates with the same best score
178 operations = operations[np.abs(scores - best_score) < 1e-6]
180 # select the one whose cell orientation is the closest to the target
181 # https://gitlab.com/ase/ase/-/merge_requests/3522
182 imin = np.argmin(np.add.reduce((operations - ideal_P)**2, axis=(-2, -1)))
183 optimal_P = operations[imin]
185 if np.linalg.det(optimal_P) <= 0:
186 optimal_P *= -1 # flip signs if negative determinant
188 # Finalize.
189 if verbose:
190 print(f"smallest score (|Q P h_p - h_target|_2): {best_score:f}")
191 print("optimal transformation matrix (P_opt):")
192 print(optimal_P)
193 print("supercell metric:")
194 print(np.round(np.dot(optimal_P, cell), 4))
195 det = np.linalg.det(optimal_P)
196 print(f"determinant of optimal transformation matrix: {det:g}")
197 return optimal_P
200def make_supercell(prim, P, *, wrap=True, order="cell-major", tol=1e-5):
201 r"""Generate a supercell by applying a general transformation (*P*) to
202 the input configuration (*prim*).
204 The transformation is described by a 3x3 integer matrix
205 `\mathbf{P}`. Specifically, the new cell metric
206 `\mathbf{h}` is given in terms of the metric of the input
207 configuration `\mathbf{h}_p` by `\mathbf{P h}_p =
208 \mathbf{h}`.
210 Parameters:
212 prim: ASE Atoms object
213 Input configuration.
214 P: 3x3 integer matrix
215 Transformation matrix `\mathbf{P}`.
216 wrap: bool
217 wrap in the end
218 order: str (default: "cell-major")
219 how to order the atoms in the supercell
221 "cell-major":
222 [atom1_shift1, atom2_shift1, ..., atom1_shift2, atom2_shift2, ...]
223 i.e. run first over all the atoms in cell1 and then move to cell2.
225 "atom-major":
226 [atom1_shift1, atom1_shift2, ..., atom2_shift1, atom2_shift2, ...]
227 i.e. run first over atom1 in all the cells and then move to atom2.
228 This may be the order preferred by most VASP users.
230 tol: float
231 tolerance for wrapping
232 """
234 supercell_matrix = P
235 supercell = clean_matrix(supercell_matrix @ prim.cell)
237 # cartesian lattice points
238 lattice_points_frac = lattice_points_in_supercell(supercell_matrix)
239 lattice_points = np.dot(lattice_points_frac, supercell)
240 N = len(lattice_points)
242 if order == "cell-major":
243 shifted = prim.positions[None, :, :] + lattice_points[:, None, :]
244 elif order == "atom-major":
245 shifted = prim.positions[:, None, :] + lattice_points[None, :, :]
246 else:
247 raise ValueError(f"invalid order: {order}")
248 shifted_reshaped = shifted.reshape(-1, 3)
250 superatoms = Atoms(positions=shifted_reshaped,
251 cell=supercell,
252 pbc=prim.pbc)
254 # Copy over any other possible arrays, inspired by atoms.__imul__
255 for name, arr in prim.arrays.items():
256 if name == "positions":
257 # This was added during construction of the super cell
258 continue
259 shape = (N * arr.shape[0], *arr.shape[1:])
260 if order == "cell-major":
261 new_arr = np.repeat(arr[None, :], N, axis=0).reshape(shape)
262 elif order == "atom-major":
263 new_arr = np.repeat(arr[:, None], N, axis=1).reshape(shape)
264 superatoms.set_array(name, new_arr)
266 # check number of atoms is correct
267 n_target = abs(int(np.round(np.linalg.det(supercell_matrix) * len(prim))))
268 if n_target != len(superatoms):
269 msg = "Number of atoms in supercell: {}, expected: {}".format(
270 n_target, len(superatoms))
271 raise SupercellError(msg)
273 if wrap:
274 superatoms.wrap(eps=tol)
276 return superatoms
279def lattice_points_in_supercell(supercell_matrix):
280 """Find all lattice points contained in a supercell.
282 Adapted from pymatgen, which is available under MIT license:
283 The MIT License (MIT) Copyright (c) 2011-2012 MIT & The Regents of the
284 University of California, through Lawrence Berkeley National Laboratory
285 """
287 diagonals = np.array([
288 [0, 0, 0],
289 [0, 0, 1],
290 [0, 1, 0],
291 [0, 1, 1],
292 [1, 0, 0],
293 [1, 0, 1],
294 [1, 1, 0],
295 [1, 1, 1],
296 ])
297 d_points = np.dot(diagonals, supercell_matrix)
299 mins = np.min(d_points, axis=0)
300 maxes = np.max(d_points, axis=0) + 1
302 ar = np.arange(mins[0], maxes[0])[:, None] * np.array([1, 0, 0])[None, :]
303 br = np.arange(mins[1], maxes[1])[:, None] * np.array([0, 1, 0])[None, :]
304 cr = np.arange(mins[2], maxes[2])[:, None] * np.array([0, 0, 1])[None, :]
306 all_points = ar[:, None, None] + br[None, :, None] + cr[None, None, :]
307 all_points = all_points.reshape((-1, 3))
309 frac_points = np.dot(all_points, np.linalg.inv(supercell_matrix))
311 tvects = frac_points[np.all(frac_points < 1 - 1e-10, axis=1)
312 & np.all(frac_points >= -1e-10, axis=1)]
313 assert len(tvects) == round(abs(np.linalg.det(supercell_matrix)))
314 return tvects
317def clean_matrix(matrix, eps=1e-12):
318 """ clean from small values"""
319 matrix = np.array(matrix)
320 for ij in np.ndindex(matrix.shape):
321 if abs(matrix[ij]) < eps:
322 matrix[ij] = 0
323 return matrix