Coverage for /builds/debichem-team/python-ase/ase/ga/slab_operators.py: 68.48%
330 statements
« prev ^ index » next coverage.py v7.5.3, created at 2025-03-06 04:00 +0000
1"""Operators that work on slabs.
2Allowed compositions are respected.
3Identical indexing of the slabs are assumed for the cut-splice operator."""
4from collections import Counter
5from itertools import permutations
6from operator import itemgetter
8import numpy as np
10from ase.ga.element_mutations import get_periodic_table_distance
11from ase.ga.offspring_creator import OffspringCreator
12from ase.utils import atoms_to_spglib_cell
14try:
15 import spglib
16except ImportError:
17 spglib = None
20def permute2(atoms, rng=np.random):
21 i1 = rng.choice(range(len(atoms)))
22 sym1 = atoms[i1].symbol
23 i2 = rng.choice([a.index for a in atoms if a.symbol != sym1])
24 atoms[i1].symbol = atoms[i2].symbol
25 atoms[i2].symbol = sym1
28def replace_element(atoms, element_out, element_in):
29 syms = np.array(atoms.get_chemical_symbols())
30 syms[syms == element_out] = element_in
31 atoms.set_chemical_symbols(syms)
34def get_add_remove_lists(**kwargs):
35 to_add, to_rem = [], []
36 for s, amount in kwargs.items():
37 if amount > 0:
38 to_add.extend([s] * amount)
39 elif amount < 0:
40 to_rem.extend([s] * abs(amount))
41 return to_add, to_rem
44def get_minority_element(atoms):
45 counter = Counter(atoms.get_chemical_symbols())
46 return sorted(counter.items(), key=itemgetter(1), reverse=False)[0][0]
49def minority_element_segregate(atoms, layer_tag=1, rng=np.random):
50 """Move the minority alloy element to the layer specified by the layer_tag,
51 Atoms object should contain atoms with the corresponding tag."""
52 sym = get_minority_element(atoms)
53 layer_indices = {a.index for a in atoms if a.tag == layer_tag}
54 minority_indices = {a.index for a in atoms if a.symbol == sym}
55 change_indices = minority_indices - layer_indices
56 in_layer_not_sym = list(layer_indices - minority_indices)
57 rng.shuffle(in_layer_not_sym)
58 if len(change_indices) > 0:
59 for i, ai in zip(change_indices, in_layer_not_sym):
60 atoms[i].symbol = atoms[ai].symbol
61 atoms[ai].symbol = sym
64def same_layer_comp(atoms, rng=np.random):
65 unique_syms, comp = np.unique(sorted(atoms.get_chemical_symbols()),
66 return_counts=True)
67 layer = get_layer_comps(atoms)
68 sym_dict = {s: int(np.array(c) / len(layer))
69 for s, c in zip(unique_syms, comp)}
70 for la in layer:
71 correct_by = sym_dict.copy()
72 lcomp = dict(
73 zip(*np.unique([atoms[i].symbol for i in la], return_counts=True)))
74 for s, num in lcomp.items():
75 correct_by[s] -= num
76 to_add, to_rem = get_add_remove_lists(**correct_by)
77 for add, rem in zip(to_add, to_rem):
78 ai = rng.choice([i for i in la if atoms[i].symbol == rem])
79 atoms[ai].symbol = add
82def get_layer_comps(atoms, eps=1e-2):
83 lc = []
84 old_z = np.inf
85 for z, ind in sorted([(a.z, a.index) for a in atoms]):
86 if abs(old_z - z) < eps:
87 lc[-1].append(ind)
88 else:
89 lc.append([ind])
90 old_z = z
92 return lc
95def get_ordered_composition(syms, pools=None):
96 if pools is None:
97 pool_index = {sym: 0 for sym in set(syms)}
98 else:
99 pool_index = {}
100 for i, pool in enumerate(pools):
101 if isinstance(pool, str):
102 pool_index[pool] = i
103 else:
104 for sym in set(syms):
105 if sym in pool:
106 pool_index[sym] = i
107 syms = [(sym, pool_index[sym], c)
108 for sym, c in zip(*np.unique(syms, return_counts=True))]
109 unique_syms, pn, comp = zip(
110 *sorted(syms, key=lambda k: (k[1] - k[2], k[0])))
111 return (unique_syms, pn, comp)
114def dummy_func(*args):
115 return
118class SlabOperator(OffspringCreator):
119 def __init__(self, verbose=False, num_muts=1,
120 allowed_compositions=None,
121 distribution_correction_function=None,
122 element_pools=None,
123 rng=np.random):
124 OffspringCreator.__init__(self, verbose, num_muts=num_muts, rng=rng)
126 self.allowed_compositions = allowed_compositions
127 self.element_pools = element_pools
128 if distribution_correction_function is None:
129 self.dcf = dummy_func
130 else:
131 self.dcf = distribution_correction_function
132 # Number of different elements i.e. [2, 1] if len(element_pools) == 2
133 # then 2 different elements in pool 1 is allowed but only 1 from pool 2
135 def get_symbols_to_use(self, syms):
136 """Get the symbols to use for the offspring candidate. The returned
137 list of symbols will respect self.allowed_compositions"""
138 if self.allowed_compositions is None:
139 return syms
141 unique_syms, counts = np.unique(syms, return_counts=True)
142 comp, unique_syms = zip(*sorted(zip(counts, unique_syms),
143 reverse=True))
145 for cc in self.allowed_compositions:
146 comp += (0,) * (len(cc) - len(comp))
147 if comp == tuple(sorted(cc)):
148 return syms
150 comp_diff = self.get_closest_composition_diff(comp)
151 to_add, to_rem = get_add_remove_lists(
152 **dict(zip(unique_syms, comp_diff)))
153 for add, rem in zip(to_add, to_rem):
154 tbc = [i for i in range(len(syms)) if syms[i] == rem]
155 ai = self.rng.choice(tbc)
156 syms[ai] = add
157 return syms
159 def get_add_remove_elements(self, syms):
160 if self.element_pools is None or self.allowed_compositions is None:
161 return [], []
162 unique_syms, pool_number, comp = get_ordered_composition(
163 syms, self.element_pools)
164 stay_comp, stay_syms = [], []
165 add_rem = {}
166 per_pool = len(self.allowed_compositions[0]) / len(self.element_pools)
167 pool_count = np.zeros(len(self.element_pools), dtype=int)
168 for pn, num, sym in zip(pool_number, comp, unique_syms):
169 pool_count[pn] += 1
170 if pool_count[pn] <= per_pool:
171 stay_comp.append(num)
172 stay_syms.append(sym)
173 else:
174 add_rem[sym] = -num
175 # collect elements from individual pools
177 diff = self.get_closest_composition_diff(stay_comp)
178 add_rem.update({s: c for s, c in zip(stay_syms, diff)})
179 return get_add_remove_lists(**add_rem)
181 def get_closest_composition_diff(self, c):
182 comp = np.array(c)
183 mindiff = 1e10
184 allowed_list = list(self.allowed_compositions)
185 self.rng.shuffle(allowed_list)
186 for ac in allowed_list:
187 diff = self.get_composition_diff(comp, ac)
188 numdiff = sum(abs(i) for i in diff)
189 if numdiff < mindiff:
190 mindiff = numdiff
191 ccdiff = diff
192 return ccdiff
194 def get_composition_diff(self, c1, c2):
195 difflen = len(c1) - len(c2)
196 if difflen > 0:
197 c2 += (0,) * difflen
198 return np.array(c2) - c1
200 def get_possible_mutations(self, a):
201 unique_syms, comp = np.unique(sorted(a.get_chemical_symbols()),
202 return_counts=True)
203 min_num = min(
204 i for i in np.ravel(list(self.allowed_compositions)) if i > 0
205 )
206 muts = set()
207 for i, n in enumerate(comp):
208 if n != 0:
209 muts.add((unique_syms[i], n))
210 if n % min_num >= 0:
211 for j in range(1, n // min_num):
212 muts.add((unique_syms[i], min_num * j))
213 return list(muts)
215 def get_all_element_mutations(self, a):
216 """Get all possible mutations for the supplied atoms object given
217 the element pools."""
218 muts = []
219 symset = set(a.get_chemical_symbols())
220 for sym in symset:
221 for pool in self.element_pools:
222 if sym in pool:
223 muts.extend([(sym, s) for s in pool if s not in symset])
224 return muts
226 def finalize_individual(self, indi):
227 atoms_string = ''.join(indi.get_chemical_symbols())
228 indi.info['key_value_pairs']['atoms_string'] = atoms_string
229 return OffspringCreator.finalize_individual(self, indi)
232class CutSpliceSlabCrossover(SlabOperator):
233 def __init__(self, allowed_compositions=None, element_pools=None,
234 verbose=False,
235 num_muts=1, tries=1000, min_ratio=0.25,
236 distribution_correction_function=None, rng=np.random):
237 SlabOperator.__init__(self, verbose, num_muts,
238 allowed_compositions,
239 distribution_correction_function,
240 element_pools=element_pools,
241 rng=rng)
243 self.tries = tries
244 self.min_ratio = min_ratio
245 self.descriptor = 'CutSpliceSlabCrossover'
247 def get_new_individual(self, parents):
248 f, m = parents
250 indi = self.initialize_individual(f, self.operate(f, m))
251 indi.info['data']['parents'] = [i.info['confid'] for i in parents]
253 parent_message = ': Parents {} {}'.format(f.info['confid'],
254 m.info['confid'])
255 return (self.finalize_individual(indi),
256 self.descriptor + parent_message)
258 def operate(self, f, m):
259 child = f.copy()
260 fp = f.positions
261 ma = np.max(fp.transpose(), axis=1)
262 mi = np.min(fp.transpose(), axis=1)
264 for _ in range(self.tries):
265 # Find center point of cut
266 rv = [self.rng.random() for _ in range(3)] # random vector
267 midpoint = (ma - mi) * rv + mi
269 # Determine cut plane
270 theta = self.rng.random() * 2 * np.pi # 0,2pi
271 phi = self.rng.random() * np.pi # 0,pi
272 e = np.array((np.sin(phi) * np.cos(theta),
273 np.sin(theta) * np.sin(phi),
274 np.cos(phi)))
276 # Cut structures
277 d2fp = np.dot(fp - midpoint, e)
278 fpart = d2fp > 0
279 ratio = float(np.count_nonzero(fpart)) / len(f)
280 if ratio < self.min_ratio or ratio > 1 - self.min_ratio:
281 continue
282 syms = np.where(fpart, f.get_chemical_symbols(),
283 m.get_chemical_symbols())
284 dists2plane = abs(d2fp)
286 # Correct the composition
287 # What if only one element pool is represented in the offspring
288 to_add, to_rem = self.get_add_remove_elements(syms)
290 # Change elements closest to the cut plane
291 for add, rem in zip(to_add, to_rem):
292 tbc = [(dists2plane[i], i)
293 for i in range(len(syms)) if syms[i] == rem]
294 ai = sorted(tbc)[0][1]
295 syms[ai] = add
297 child.set_chemical_symbols(syms)
298 break
300 self.dcf(child)
302 return child
305# Mutations: Random, MoveUp/Down/Left/Right, six or all elements
307class RandomCompositionMutation(SlabOperator):
308 """Change the current composition to another of the allowed compositions.
309 The allowed compositions should be input in the same order as the element
310 pools, for example:
311 element_pools = [['Au', 'Cu'], ['In', 'Bi']]
312 allowed_compositions = [(6, 2), (5, 3)]
313 means that there can be 5 or 6 Au and Cu, and 2 or 3 In and Bi.
314 """
316 def __init__(self, verbose=False, num_muts=1, element_pools=None,
317 allowed_compositions=None,
318 distribution_correction_function=None, rng=np.random):
319 SlabOperator.__init__(self, verbose, num_muts,
320 allowed_compositions,
321 distribution_correction_function,
322 element_pools=element_pools,
323 rng=rng)
325 self.descriptor = 'RandomCompositionMutation'
327 def get_new_individual(self, parents):
328 f = parents[0]
329 parent_message = ': Parent {}'.format(f.info['confid'])
331 if self.allowed_compositions is None:
332 if len(set(f.get_chemical_symbols())) == 1:
333 if self.element_pools is None:
334 # We cannot find another composition without knowledge of
335 # other allowed elements or compositions
336 return None, self.descriptor + parent_message
338 # Do the operation
339 indi = self.initialize_individual(f, self.operate(f))
340 indi.info['data']['parents'] = [i.info['confid'] for i in parents]
342 return (self.finalize_individual(indi),
343 self.descriptor + parent_message)
345 def operate(self, atoms):
346 allowed_comps = self.allowed_compositions
347 if allowed_comps is None:
348 n_elems = len(set(atoms.get_chemical_symbols()))
349 n_atoms = len(atoms)
350 allowed_comps = [c for c in permutations(range(1, n_atoms),
351 n_elems)
352 if sum(c) == n_atoms]
354 # Sorting the composition to have the same order as in element_pools
355 syms = atoms.get_chemical_symbols()
356 unique_syms, _, comp = get_ordered_composition(syms,
357 self.element_pools)
359 # Choose the composition to change to
360 for i, allowed in enumerate(allowed_comps):
361 if comp == tuple(allowed):
362 allowed_comps = np.delete(allowed_comps, i, axis=0)
363 break
364 chosen = self.rng.randint(len(allowed_comps))
365 comp_diff = self.get_composition_diff(comp, allowed_comps[chosen])
367 # Get difference from current composition
368 to_add, to_rem = get_add_remove_lists(
369 **dict(zip(unique_syms, comp_diff)))
371 # Correct current composition
372 syms = atoms.get_chemical_symbols()
373 for add, rem in zip(to_add, to_rem):
374 tbc = [i for i in range(len(syms)) if syms[i] == rem]
375 ai = self.rng.choice(tbc)
376 syms[ai] = add
378 atoms.set_chemical_symbols(syms)
379 self.dcf(atoms)
380 return atoms
383class RandomElementMutation(SlabOperator):
384 def __init__(self, element_pools, verbose=False, num_muts=1,
385 allowed_compositions=None,
386 distribution_correction_function=None, rng=np.random):
387 SlabOperator.__init__(self, verbose, num_muts,
388 allowed_compositions,
389 distribution_correction_function,
390 element_pools=element_pools,
391 rng=rng)
393 self.descriptor = 'RandomElementMutation'
395 def get_new_individual(self, parents):
396 f = parents[0]
398 # Do the operation
399 indi = self.initialize_individual(f, self.operate(f))
400 indi.info['data']['parents'] = [i.info['confid'] for i in parents]
402 parent_message = ': Parent {}'.format(f.info['confid'])
403 return (self.finalize_individual(indi),
404 self.descriptor + parent_message)
406 def operate(self, atoms):
407 poss_muts = self.get_all_element_mutations(atoms)
408 chosen = self.rng.randint(len(poss_muts))
409 replace_element(atoms, *poss_muts[chosen])
410 self.dcf(atoms)
411 return atoms
414class NeighborhoodElementMutation(SlabOperator):
415 def __init__(self, element_pools, verbose=False, num_muts=1,
416 allowed_compositions=None,
417 distribution_correction_function=None, rng=np.random):
418 SlabOperator.__init__(self, verbose, num_muts,
419 allowed_compositions,
420 distribution_correction_function,
421 element_pools=element_pools,
422 rng=rng)
424 self.descriptor = 'NeighborhoodElementMutation'
426 def get_new_individual(self, parents):
427 f = parents[0]
429 indi = self.initialize_individual(f, f)
430 indi.info['data']['parents'] = [i.info['confid'] for i in parents]
432 indi = self.operate(indi)
434 parent_message = ': Parent {}'.format(f.info['confid'])
435 return (self.finalize_individual(indi),
436 self.descriptor + parent_message)
438 def operate(self, atoms):
439 least_diff = 1e22
440 for mut in self.get_all_element_mutations(atoms):
441 dist = get_periodic_table_distance(*mut)
442 if dist < least_diff:
443 poss_muts = [mut]
444 least_diff = dist
445 elif dist == least_diff:
446 poss_muts.append(mut)
448 chosen = self.rng.randint(len(poss_muts))
449 replace_element(atoms, *poss_muts[chosen])
450 self.dcf(atoms)
451 return atoms
454class SymmetrySlabPermutation(SlabOperator):
455 """Permutes the atoms in the slab until it has a higher symmetry number."""
457 def __init__(self, verbose=False, num_muts=1, sym_goal=100, max_tries=50,
458 allowed_compositions=None,
459 distribution_correction_function=None, rng=np.random):
460 SlabOperator.__init__(self, verbose, num_muts,
461 allowed_compositions,
462 distribution_correction_function,
463 rng=rng)
464 if spglib is None:
465 print("SymmetrySlabPermutation needs spglib to function")
467 assert sym_goal >= 1
468 self.sym_goal = sym_goal
469 self.max_tries = max_tries
470 self.descriptor = 'SymmetrySlabPermutation'
472 def get_new_individual(self, parents):
473 f = parents[0]
474 # Permutation only makes sense if two different elements are present
475 if len(set(f.get_chemical_symbols())) == 1:
476 f = parents[1]
477 if len(set(f.get_chemical_symbols())) == 1:
478 return None, '{1} not possible in {0}'.format(f.info['confid'],
479 self.descriptor)
481 indi = self.initialize_individual(f, self.operate(f))
482 indi.info['data']['parents'] = [i.info['confid'] for i in parents]
484 parent_message = ': Parent {}'.format(f.info['confid'])
485 return (self.finalize_individual(indi),
486 self.descriptor + parent_message)
488 def operate(self, atoms):
489 from ase.spacegroup.symmetrize import spglib_get_symmetry_dataset
490 # Do the operation
491 sym_num = 1
492 sg = self.sym_goal
493 while sym_num < sg:
494 for _ in range(self.max_tries):
495 for _ in range(2):
496 permute2(atoms, rng=self.rng)
497 self.dcf(atoms)
498 sym_num = spglib_get_symmetry_dataset(
499 atoms_to_spglib_cell(atoms))['number']
500 if sym_num >= sg:
501 break
502 sg -= 1
503 return atoms
506class RandomSlabPermutation(SlabOperator):
507 def __init__(self, verbose=False, num_muts=1,
508 allowed_compositions=None,
509 distribution_correction_function=None, rng=np.random):
510 SlabOperator.__init__(self, verbose, num_muts,
511 allowed_compositions,
512 distribution_correction_function,
513 rng=rng)
515 self.descriptor = 'RandomSlabPermutation'
517 def get_new_individual(self, parents):
518 f = parents[0]
519 # Permutation only makes sense if two different elements are present
520 if len(set(f.get_chemical_symbols())) == 1:
521 f = parents[1]
522 if len(set(f.get_chemical_symbols())) == 1:
523 return None, '{1} not possible in {0}'.format(f.info['confid'],
524 self.descriptor)
526 indi = self.initialize_individual(f, f)
527 indi.info['data']['parents'] = [i.info['confid'] for i in parents]
529 indi = self.operate(indi)
531 parent_message = ': Parent {}'.format(f.info['confid'])
532 return (self.finalize_individual(indi),
533 self.descriptor + parent_message)
535 def operate(self, atoms):
536 # Do the operation
537 for _ in range(self.num_muts):
538 permute2(atoms, rng=self.rng)
539 self.dcf(atoms)
540 return atoms