Coverage for /builds/kinetik161/ase/ase/ga/slab_operators.py: 68.69%
329 statements
« prev ^ index » next coverage.py v7.2.7, created at 2023-12-10 11:04 +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([i for i in np.ravel(list(self.allowed_compositions))
204 if i > 0])
205 muts = set()
206 for i, n in enumerate(comp):
207 if n != 0:
208 muts.add((unique_syms[i], n))
209 if n % min_num >= 0:
210 for j in range(1, n // min_num):
211 muts.add((unique_syms[i], min_num * j))
212 return list(muts)
214 def get_all_element_mutations(self, a):
215 """Get all possible mutations for the supplied atoms object given
216 the element pools."""
217 muts = []
218 symset = set(a.get_chemical_symbols())
219 for sym in symset:
220 for pool in self.element_pools:
221 if sym in pool:
222 muts.extend([(sym, s) for s in pool if s not in symset])
223 return muts
225 def finalize_individual(self, indi):
226 atoms_string = ''.join(indi.get_chemical_symbols())
227 indi.info['key_value_pairs']['atoms_string'] = atoms_string
228 return OffspringCreator.finalize_individual(self, indi)
231class CutSpliceSlabCrossover(SlabOperator):
232 def __init__(self, allowed_compositions=None, element_pools=None,
233 verbose=False,
234 num_muts=1, tries=1000, min_ratio=0.25,
235 distribution_correction_function=None, rng=np.random):
236 SlabOperator.__init__(self, verbose, num_muts,
237 allowed_compositions,
238 distribution_correction_function,
239 element_pools=element_pools,
240 rng=rng)
242 self.tries = tries
243 self.min_ratio = min_ratio
244 self.descriptor = 'CutSpliceSlabCrossover'
246 def get_new_individual(self, parents):
247 f, m = parents
249 indi = self.initialize_individual(f, self.operate(f, m))
250 indi.info['data']['parents'] = [i.info['confid'] for i in parents]
252 parent_message = ': Parents {} {}'.format(f.info['confid'],
253 m.info['confid'])
254 return (self.finalize_individual(indi),
255 self.descriptor + parent_message)
257 def operate(self, f, m):
258 child = f.copy()
259 fp = f.positions
260 ma = np.max(fp.transpose(), axis=1)
261 mi = np.min(fp.transpose(), axis=1)
263 for _ in range(self.tries):
264 # Find center point of cut
265 rv = [self.rng.random() for _ in range(3)] # random vector
266 midpoint = (ma - mi) * rv + mi
268 # Determine cut plane
269 theta = self.rng.random() * 2 * np.pi # 0,2pi
270 phi = self.rng.random() * np.pi # 0,pi
271 e = np.array((np.sin(phi) * np.cos(theta),
272 np.sin(theta) * np.sin(phi),
273 np.cos(phi)))
275 # Cut structures
276 d2fp = np.dot(fp - midpoint, e)
277 fpart = d2fp > 0
278 ratio = float(np.count_nonzero(fpart)) / len(f)
279 if ratio < self.min_ratio or ratio > 1 - self.min_ratio:
280 continue
281 syms = np.where(fpart, f.get_chemical_symbols(),
282 m.get_chemical_symbols())
283 dists2plane = abs(d2fp)
285 # Correct the composition
286 # What if only one element pool is represented in the offspring
287 to_add, to_rem = self.get_add_remove_elements(syms)
289 # Change elements closest to the cut plane
290 for add, rem in zip(to_add, to_rem):
291 tbc = [(dists2plane[i], i)
292 for i in range(len(syms)) if syms[i] == rem]
293 ai = sorted(tbc)[0][1]
294 syms[ai] = add
296 child.set_chemical_symbols(syms)
297 break
299 self.dcf(child)
301 return child
304# Mutations: Random, MoveUp/Down/Left/Right, six or all elements
306class RandomCompositionMutation(SlabOperator):
307 """Change the current composition to another of the allowed compositions.
308 The allowed compositions should be input in the same order as the element
309 pools, for example:
310 element_pools = [['Au', 'Cu'], ['In', 'Bi']]
311 allowed_compositions = [(6, 2), (5, 3)]
312 means that there can be 5 or 6 Au and Cu, and 2 or 3 In and Bi.
313 """
315 def __init__(self, verbose=False, num_muts=1, element_pools=None,
316 allowed_compositions=None,
317 distribution_correction_function=None, rng=np.random):
318 SlabOperator.__init__(self, verbose, num_muts,
319 allowed_compositions,
320 distribution_correction_function,
321 element_pools=element_pools,
322 rng=rng)
324 self.descriptor = 'RandomCompositionMutation'
326 def get_new_individual(self, parents):
327 f = parents[0]
328 parent_message = ': Parent {}'.format(f.info['confid'])
330 if self.allowed_compositions is None:
331 if len(set(f.get_chemical_symbols())) == 1:
332 if self.element_pools is None:
333 # We cannot find another composition without knowledge of
334 # other allowed elements or compositions
335 return None, self.descriptor + parent_message
337 # Do the operation
338 indi = self.initialize_individual(f, self.operate(f))
339 indi.info['data']['parents'] = [i.info['confid'] for i in parents]
341 return (self.finalize_individual(indi),
342 self.descriptor + parent_message)
344 def operate(self, atoms):
345 allowed_comps = self.allowed_compositions
346 if allowed_comps is None:
347 n_elems = len(set(atoms.get_chemical_symbols()))
348 n_atoms = len(atoms)
349 allowed_comps = [c for c in permutations(range(1, n_atoms),
350 n_elems)
351 if sum(c) == n_atoms]
353 # Sorting the composition to have the same order as in element_pools
354 syms = atoms.get_chemical_symbols()
355 unique_syms, _, comp = get_ordered_composition(syms,
356 self.element_pools)
358 # Choose the composition to change to
359 for i, allowed in enumerate(allowed_comps):
360 if comp == tuple(allowed):
361 allowed_comps = np.delete(allowed_comps, i, axis=0)
362 break
363 chosen = self.rng.randint(len(allowed_comps))
364 comp_diff = self.get_composition_diff(comp, allowed_comps[chosen])
366 # Get difference from current composition
367 to_add, to_rem = get_add_remove_lists(
368 **dict(zip(unique_syms, comp_diff)))
370 # Correct current composition
371 syms = atoms.get_chemical_symbols()
372 for add, rem in zip(to_add, to_rem):
373 tbc = [i for i in range(len(syms)) if syms[i] == rem]
374 ai = self.rng.choice(tbc)
375 syms[ai] = add
377 atoms.set_chemical_symbols(syms)
378 self.dcf(atoms)
379 return atoms
382class RandomElementMutation(SlabOperator):
383 def __init__(self, element_pools, verbose=False, num_muts=1,
384 allowed_compositions=None,
385 distribution_correction_function=None, rng=np.random):
386 SlabOperator.__init__(self, verbose, num_muts,
387 allowed_compositions,
388 distribution_correction_function,
389 element_pools=element_pools,
390 rng=rng)
392 self.descriptor = 'RandomElementMutation'
394 def get_new_individual(self, parents):
395 f = parents[0]
397 # Do the operation
398 indi = self.initialize_individual(f, self.operate(f))
399 indi.info['data']['parents'] = [i.info['confid'] for i in parents]
401 parent_message = ': Parent {}'.format(f.info['confid'])
402 return (self.finalize_individual(indi),
403 self.descriptor + parent_message)
405 def operate(self, atoms):
406 poss_muts = self.get_all_element_mutations(atoms)
407 chosen = self.rng.randint(len(poss_muts))
408 replace_element(atoms, *poss_muts[chosen])
409 self.dcf(atoms)
410 return atoms
413class NeighborhoodElementMutation(SlabOperator):
414 def __init__(self, element_pools, verbose=False, num_muts=1,
415 allowed_compositions=None,
416 distribution_correction_function=None, rng=np.random):
417 SlabOperator.__init__(self, verbose, num_muts,
418 allowed_compositions,
419 distribution_correction_function,
420 element_pools=element_pools,
421 rng=rng)
423 self.descriptor = 'NeighborhoodElementMutation'
425 def get_new_individual(self, parents):
426 f = parents[0]
428 indi = self.initialize_individual(f, f)
429 indi.info['data']['parents'] = [i.info['confid'] for i in parents]
431 indi = self.operate(indi)
433 parent_message = ': Parent {}'.format(f.info['confid'])
434 return (self.finalize_individual(indi),
435 self.descriptor + parent_message)
437 def operate(self, atoms):
438 least_diff = 1e22
439 for mut in self.get_all_element_mutations(atoms):
440 dist = get_periodic_table_distance(*mut)
441 if dist < least_diff:
442 poss_muts = [mut]
443 least_diff = dist
444 elif dist == least_diff:
445 poss_muts.append(mut)
447 chosen = self.rng.randint(len(poss_muts))
448 replace_element(atoms, *poss_muts[chosen])
449 self.dcf(atoms)
450 return atoms
453class SymmetrySlabPermutation(SlabOperator):
454 """Permutes the atoms in the slab until it has a higher symmetry number."""
456 def __init__(self, verbose=False, num_muts=1, sym_goal=100, max_tries=50,
457 allowed_compositions=None,
458 distribution_correction_function=None, rng=np.random):
459 SlabOperator.__init__(self, verbose, num_muts,
460 allowed_compositions,
461 distribution_correction_function,
462 rng=rng)
463 if spglib is None:
464 print("SymmetrySlabPermutation needs spglib to function")
466 assert sym_goal >= 1
467 self.sym_goal = sym_goal
468 self.max_tries = max_tries
469 self.descriptor = 'SymmetrySlabPermutation'
471 def get_new_individual(self, parents):
472 f = parents[0]
473 # Permutation only makes sense if two different elements are present
474 if len(set(f.get_chemical_symbols())) == 1:
475 f = parents[1]
476 if len(set(f.get_chemical_symbols())) == 1:
477 return None, '{1} not possible in {0}'.format(f.info['confid'],
478 self.descriptor)
480 indi = self.initialize_individual(f, self.operate(f))
481 indi.info['data']['parents'] = [i.info['confid'] for i in parents]
483 parent_message = ': Parent {}'.format(f.info['confid'])
484 return (self.finalize_individual(indi),
485 self.descriptor + parent_message)
487 def operate(self, atoms):
488 # Do the operation
489 sym_num = 1
490 sg = self.sym_goal
491 while sym_num < sg:
492 for _ in range(self.max_tries):
493 for _ in range(2):
494 permute2(atoms, rng=self.rng)
495 self.dcf(atoms)
496 sym_num = spglib.get_symmetry_dataset(
497 atoms_to_spglib_cell(atoms))['number']
498 if sym_num >= sg:
499 break
500 sg -= 1
501 return atoms
504class RandomSlabPermutation(SlabOperator):
505 def __init__(self, verbose=False, num_muts=1,
506 allowed_compositions=None,
507 distribution_correction_function=None, rng=np.random):
508 SlabOperator.__init__(self, verbose, num_muts,
509 allowed_compositions,
510 distribution_correction_function,
511 rng=rng)
513 self.descriptor = 'RandomSlabPermutation'
515 def get_new_individual(self, parents):
516 f = parents[0]
517 # Permutation only makes sense if two different elements are present
518 if len(set(f.get_chemical_symbols())) == 1:
519 f = parents[1]
520 if len(set(f.get_chemical_symbols())) == 1:
521 return None, '{1} not possible in {0}'.format(f.info['confid'],
522 self.descriptor)
524 indi = self.initialize_individual(f, f)
525 indi.info['data']['parents'] = [i.info['confid'] for i in parents]
527 indi = self.operate(indi)
529 parent_message = ': Parent {}'.format(f.info['confid'])
530 return (self.finalize_individual(indi),
531 self.descriptor + parent_message)
533 def operate(self, atoms):
534 # Do the operation
535 for _ in range(self.num_muts):
536 permute2(atoms, rng=self.rng)
537 self.dcf(atoms)
538 return atoms