1"""Operators that work on slabs. 

2Allowed compositions are respected. 

3Identical indexing of the slabs are assumed for the cut-splice operator.""" 

4from operator import itemgetter 

5from collections import Counter 

6from itertools import permutations 

7import numpy as np 

8 

9from ase.ga.offspring_creator import OffspringCreator 

10from ase.ga.element_mutations import get_periodic_table_distance 

11from ase.utils import atoms_to_spglib_cell 

12 

13try: 

14 import spglib 

15except ImportError: 

16 spglib = None 

17 

18 

19def permute2(atoms, rng=np.random): 

20 i1 = rng.choice(range(len(atoms))) 

21 sym1 = atoms[i1].symbol 

22 i2 = rng.choice([a.index for a in atoms if a.symbol != sym1]) 

23 atoms[i1].symbol = atoms[i2].symbol 

24 atoms[i2].symbol = sym1 

25 

26 

27def replace_element(atoms, element_out, element_in): 

28 syms = np.array(atoms.get_chemical_symbols()) 

29 syms[syms == element_out] = element_in 

30 atoms.set_chemical_symbols(syms) 

31 

32 

33def get_add_remove_lists(**kwargs): 

34 to_add, to_rem = [], [] 

35 for s, amount in kwargs.items(): 

36 if amount > 0: 

37 to_add.extend([s] * amount) 

38 elif amount < 0: 

39 to_rem.extend([s] * abs(amount)) 

40 return to_add, to_rem 

41 

42 

43def get_minority_element(atoms): 

44 counter = Counter(atoms.get_chemical_symbols()) 

45 return sorted(counter.items(), key=itemgetter(1), reverse=False)[0][0] 

46 

47 

48def minority_element_segregate(atoms, layer_tag=1, rng=np.random): 

49 """Move the minority alloy element to the layer specified by the layer_tag, 

50 Atoms object should contain atoms with the corresponding tag.""" 

51 sym = get_minority_element(atoms) 

52 layer_indices = set([a.index for a in atoms if a.tag == layer_tag]) 

53 minority_indices = set([a.index for a in atoms if a.symbol == sym]) 

54 change_indices = minority_indices - layer_indices 

55 in_layer_not_sym = list(layer_indices - minority_indices) 

56 rng.shuffle(in_layer_not_sym) 

57 if len(change_indices) > 0: 

58 for i, ai in zip(change_indices, in_layer_not_sym): 

59 atoms[i].symbol = atoms[ai].symbol 

60 atoms[ai].symbol = sym 

61 

62 

63def same_layer_comp(atoms, rng=np.random): 

64 unique_syms, comp = np.unique(sorted(atoms.get_chemical_symbols()), 

65 return_counts=True) 

66 l = get_layer_comps(atoms) 

67 sym_dict = dict((s, int(np.array(c) / len(l))) 

68 for s, c in zip(unique_syms, comp)) 

69 for la in l: 

70 correct_by = sym_dict.copy() 

71 lcomp = dict( 

72 zip(*np.unique([atoms[i].symbol for i in la], return_counts=True))) 

73 for s, num in lcomp.items(): 

74 correct_by[s] -= num 

75 to_add, to_rem = get_add_remove_lists(**correct_by) 

76 for add, rem in zip(to_add, to_rem): 

77 ai = rng.choice([i for i in la if atoms[i].symbol == rem]) 

78 atoms[ai].symbol = add 

79 

80 

81def get_layer_comps(atoms, eps=1e-2): 

82 lc = [] 

83 old_z = np.inf 

84 for z, ind in sorted([(a.z, a.index) for a in atoms]): 

85 if abs(old_z - z) < eps: 

86 lc[-1].append(ind) 

87 else: 

88 lc.append([ind]) 

89 old_z = z 

90 

91 return lc 

92 

93 

94def get_ordered_composition(syms, pools=None): 

95 if pools is None: 

96 pool_index = dict((sym, 0) for sym in set(syms)) 

97 else: 

98 pool_index = {} 

99 for i, pool in enumerate(pools): 

100 if isinstance(pool, str): 

101 pool_index[pool] = i 

102 else: 

103 for sym in set(syms): 

104 if sym in pool: 

105 pool_index[sym] = i 

106 syms = [(sym, pool_index[sym], c) 

107 for sym, c in zip(*np.unique(syms, return_counts=True))] 

108 unique_syms, pn, comp = zip( 

109 *sorted(syms, key=lambda k: (k[1] - k[2], k[0]))) 

110 return (unique_syms, pn, comp) 

111 

112 

113def dummy_func(*args): 

114 return 

115 

116 

117class SlabOperator(OffspringCreator): 

118 def __init__(self, verbose=False, num_muts=1, 

119 allowed_compositions=None, 

120 distribution_correction_function=None, 

121 element_pools=None, 

122 rng=np.random): 

123 OffspringCreator.__init__(self, verbose, num_muts=num_muts, rng=rng) 

124 

125 self.allowed_compositions = allowed_compositions 

126 self.element_pools = element_pools 

127 if distribution_correction_function is None: 

128 self.dcf = dummy_func 

129 else: 

130 self.dcf = distribution_correction_function 

131 # Number of different elements i.e. [2, 1] if len(element_pools) == 2 

132 # then 2 different elements in pool 1 is allowed but only 1 from pool 2 

133 

134 def get_symbols_to_use(self, syms): 

135 """Get the symbols to use for the offspring candidate. The returned 

136 list of symbols will respect self.allowed_compositions""" 

137 if self.allowed_compositions is None: 

138 return syms 

139 

140 unique_syms, counts = np.unique(syms, return_counts=True) 

141 comp, unique_syms = zip(*sorted(zip(counts, unique_syms), 

142 reverse=True)) 

143 

144 for cc in self.allowed_compositions: 

145 comp += (0,) * (len(cc) - len(comp)) 

146 if comp == tuple(sorted(cc)): 

147 return syms 

148 

149 comp_diff = self.get_closest_composition_diff(comp) 

150 to_add, to_rem = get_add_remove_lists( 

151 **dict(zip(unique_syms, comp_diff))) 

152 for add, rem in zip(to_add, to_rem): 

153 tbc = [i for i in range(len(syms)) if syms[i] == rem] 

154 ai = self.rng.choice(tbc) 

155 syms[ai] = add 

156 return syms 

157 

158 def get_add_remove_elements(self, syms): 

159 if self.element_pools is None or self.allowed_compositions is None: 

160 return [], [] 

161 unique_syms, pool_number, comp = get_ordered_composition( 

162 syms, self.element_pools) 

163 stay_comp, stay_syms = [], [] 

164 add_rem = {} 

165 per_pool = len(self.allowed_compositions[0]) / len(self.element_pools) 

166 pool_count = np.zeros(len(self.element_pools), dtype=int) 

167 for pn, num, sym in zip(pool_number, comp, unique_syms): 

168 pool_count[pn] += 1 

169 if pool_count[pn] <= per_pool: 

170 stay_comp.append(num) 

171 stay_syms.append(sym) 

172 else: 

173 add_rem[sym] = -num 

174 # collect elements from individual pools 

175 

176 diff = self.get_closest_composition_diff(stay_comp) 

177 add_rem.update(dict((s, c) for s, c in zip(stay_syms, diff))) 

178 return get_add_remove_lists(**add_rem) 

179 

180 def get_closest_composition_diff(self, c): 

181 comp = np.array(c) 

182 mindiff = 1e10 

183 allowed_list = list(self.allowed_compositions) 

184 self.rng.shuffle(allowed_list) 

185 for ac in allowed_list: 

186 diff = self.get_composition_diff(comp, ac) 

187 numdiff = sum([abs(i) for i in diff]) 

188 if numdiff < mindiff: 

189 mindiff = numdiff 

190 ccdiff = diff 

191 return ccdiff 

192 

193 def get_composition_diff(self, c1, c2): 

194 difflen = len(c1) - len(c2) 

195 if difflen > 0: 

196 c2 += (0,) * difflen 

197 return np.array(c2) - c1 

198 

199 def get_possible_mutations(self, a): 

200 unique_syms, comp = np.unique(sorted(a.get_chemical_symbols()), 

201 return_counts=True) 

202 min_num = min([i for i in np.ravel(list(self.allowed_compositions)) 

203 if i > 0]) 

204 muts = set() 

205 for i, n in enumerate(comp): 

206 if n != 0: 

207 muts.add((unique_syms[i], n)) 

208 if n % min_num >= 0: 

209 for j in range(1, n // min_num): 

210 muts.add((unique_syms[i], min_num * j)) 

211 return list(muts) 

212 

213 def get_all_element_mutations(self, a): 

214 """Get all possible mutations for the supplied atoms object given 

215 the element pools.""" 

216 muts = [] 

217 symset = set(a.get_chemical_symbols()) 

218 for sym in symset: 

219 for pool in self.element_pools: 

220 if sym in pool: 

221 muts.extend([(sym, s) for s in pool if s not in symset]) 

222 return muts 

223 

224 def finalize_individual(self, indi): 

225 atoms_string = ''.join(indi.get_chemical_symbols()) 

226 indi.info['key_value_pairs']['atoms_string'] = atoms_string 

227 return OffspringCreator.finalize_individual(self, indi) 

228 

229 

230class CutSpliceSlabCrossover(SlabOperator): 

231 def __init__(self, allowed_compositions=None, element_pools=None, 

232 verbose=False, 

233 num_muts=1, tries=1000, min_ratio=0.25, 

234 distribution_correction_function=None, rng=np.random): 

235 SlabOperator.__init__(self, verbose, num_muts, 

236 allowed_compositions, 

237 distribution_correction_function, 

238 element_pools=element_pools, 

239 rng=rng) 

240 

241 self.tries = tries 

242 self.min_ratio = min_ratio 

243 self.descriptor = 'CutSpliceSlabCrossover' 

244 

245 def get_new_individual(self, parents): 

246 f, m = parents 

247 

248 indi = self.initialize_individual(f, self.operate(f, m)) 

249 indi.info['data']['parents'] = [i.info['confid'] for i in parents] 

250 

251 parent_message = ': Parents {0} {1}'.format(f.info['confid'], 

252 m.info['confid']) 

253 return (self.finalize_individual(indi), 

254 self.descriptor + parent_message) 

255 

256 def operate(self, f, m): 

257 child = f.copy() 

258 fp = f.positions 

259 ma = np.max(fp.transpose(), axis=1) 

260 mi = np.min(fp.transpose(), axis=1) 

261 

262 for _ in range(self.tries): 

263 # Find center point of cut 

264 rv = [self.rng.rand() for _ in range(3)] # random vector 

265 midpoint = (ma - mi) * rv + mi 

266 

267 # Determine cut plane 

268 theta = self.rng.rand() * 2 * np.pi # 0,2pi 

269 phi = self.rng.rand() * np.pi # 0,pi 

270 e = np.array((np.sin(phi) * np.cos(theta), 

271 np.sin(theta) * np.sin(phi), 

272 np.cos(phi))) 

273 

274 # Cut structures 

275 d2fp = np.dot(fp - midpoint, e) 

276 fpart = d2fp > 0 

277 ratio = float(np.count_nonzero(fpart)) / len(f) 

278 if ratio < self.min_ratio or ratio > 1 - self.min_ratio: 

279 continue 

280 syms = np.where(fpart, f.get_chemical_symbols(), 

281 m.get_chemical_symbols()) 

282 dists2plane = abs(d2fp) 

283 

284 # Correct the composition 

285 # What if only one element pool is represented in the offspring 

286 to_add, to_rem = self.get_add_remove_elements(syms) 

287 

288 # Change elements closest to the cut plane 

289 for add, rem in zip(to_add, to_rem): 

290 tbc = [(dists2plane[i], i) 

291 for i in range(len(syms)) if syms[i] == rem] 

292 ai = sorted(tbc)[0][1] 

293 syms[ai] = add 

294 

295 child.set_chemical_symbols(syms) 

296 break 

297 

298 self.dcf(child) 

299 

300 return child 

301 

302 

303# Mutations: Random, MoveUp/Down/Left/Right, six or all elements 

304 

305class RandomCompositionMutation(SlabOperator): 

306 """Change the current composition to another of the allowed compositions. 

307 The allowed compositions should be input in the same order as the element pools, 

308 for example: 

309 element_pools = [['Au', 'Cu'], ['In', 'Bi']] 

310 allowed_compositions = [(6, 2), (5, 3)] 

311 means that there can be 5 or 6 Au and Cu, and 2 or 3 In and Bi. 

312 """ 

313 

314 def __init__(self, verbose=False, num_muts=1, element_pools=None, 

315 allowed_compositions=None, 

316 distribution_correction_function=None, rng=np.random): 

317 SlabOperator.__init__(self, verbose, num_muts, 

318 allowed_compositions, 

319 distribution_correction_function, 

320 element_pools=element_pools, 

321 rng=rng) 

322 

323 self.descriptor = 'RandomCompositionMutation' 

324 

325 def get_new_individual(self, parents): 

326 f = parents[0] 

327 parent_message = ': Parent {0}'.format(f.info['confid']) 

328 

329 if self.allowed_compositions is None: 

330 if len(set(f.get_chemical_symbols())) == 1: 

331 if self.element_pools is None: 

332 # We cannot find another composition without knowledge of 

333 # other allowed elements or compositions 

334 return None, self.descriptor + parent_message 

335 

336 # Do the operation 

337 indi = self.initialize_individual(f, self.operate(f)) 

338 indi.info['data']['parents'] = [i.info['confid'] for i in parents] 

339 

340 return (self.finalize_individual(indi), 

341 self.descriptor + parent_message) 

342 

343 def operate(self, atoms): 

344 allowed_comps = self.allowed_compositions 

345 if allowed_comps is None: 

346 n_elems = len(set(atoms.get_chemical_symbols())) 

347 n_atoms = len(atoms) 

348 allowed_comps = [c for c in permutations(range(1, n_atoms), 

349 n_elems) 

350 if sum(c) == n_atoms] 

351 

352 # Sorting the composition to have the same order as in element_pools 

353 syms = atoms.get_chemical_symbols() 

354 unique_syms, _, comp = get_ordered_composition(syms, 

355 self.element_pools) 

356 

357 # Choose the composition to change to 

358 for i, allowed in enumerate(allowed_comps): 

359 if comp == tuple(allowed): 

360 allowed_comps = np.delete(allowed_comps, i, axis=0) 

361 break 

362 chosen = self.rng.randint(len(allowed_comps)) 

363 comp_diff = self.get_composition_diff(comp, allowed_comps[chosen]) 

364 

365 # Get difference from current composition 

366 to_add, to_rem = get_add_remove_lists( 

367 **dict(zip(unique_syms, comp_diff))) 

368 

369 # Correct current composition 

370 syms = atoms.get_chemical_symbols() 

371 for add, rem in zip(to_add, to_rem): 

372 tbc = [i for i in range(len(syms)) if syms[i] == rem] 

373 ai = self.rng.choice(tbc) 

374 syms[ai] = add 

375 

376 atoms.set_chemical_symbols(syms) 

377 self.dcf(atoms) 

378 return atoms 

379 

380 

381class RandomElementMutation(SlabOperator): 

382 def __init__(self, element_pools, verbose=False, num_muts=1, 

383 allowed_compositions=None, 

384 distribution_correction_function=None, rng=np.random): 

385 SlabOperator.__init__(self, verbose, num_muts, 

386 allowed_compositions, 

387 distribution_correction_function, 

388 element_pools=element_pools, 

389 rng=rng) 

390 

391 self.descriptor = 'RandomElementMutation' 

392 

393 def get_new_individual(self, parents): 

394 f = parents[0] 

395 

396 # Do the operation 

397 indi = self.initialize_individual(f, self.operate(f)) 

398 indi.info['data']['parents'] = [i.info['confid'] for i in parents] 

399 

400 parent_message = ': Parent {0}'.format(f.info['confid']) 

401 return (self.finalize_individual(indi), 

402 self.descriptor + parent_message) 

403 

404 def operate(self, atoms): 

405 poss_muts = self.get_all_element_mutations(atoms) 

406 chosen = self.rng.randint(len(poss_muts)) 

407 replace_element(atoms, *poss_muts[chosen]) 

408 self.dcf(atoms) 

409 return atoms 

410 

411 

412class NeighborhoodElementMutation(SlabOperator): 

413 def __init__(self, element_pools, verbose=False, num_muts=1, 

414 allowed_compositions=None, 

415 distribution_correction_function=None, rng=np.random): 

416 SlabOperator.__init__(self, verbose, num_muts, 

417 allowed_compositions, 

418 distribution_correction_function, 

419 element_pools=element_pools, 

420 rng=rng) 

421 

422 self.descriptor = 'NeighborhoodElementMutation' 

423 

424 def get_new_individual(self, parents): 

425 f = parents[0] 

426 

427 indi = self.initialize_individual(f, f) 

428 indi.info['data']['parents'] = [i.info['confid'] for i in parents] 

429 

430 indi = self.operate(indi) 

431 

432 parent_message = ': Parent {0}'.format(f.info['confid']) 

433 return (self.finalize_individual(indi), 

434 self.descriptor + parent_message) 

435 

436 def operate(self, atoms): 

437 least_diff = 1e22 

438 for mut in self.get_all_element_mutations(atoms): 

439 dist = get_periodic_table_distance(*mut) 

440 if dist < least_diff: 

441 poss_muts = [mut] 

442 least_diff = dist 

443 elif dist == least_diff: 

444 poss_muts.append(mut) 

445 

446 chosen = self.rng.randint(len(poss_muts)) 

447 replace_element(atoms, *poss_muts[chosen]) 

448 self.dcf(atoms) 

449 return atoms 

450 

451 

452class SymmetrySlabPermutation(SlabOperator): 

453 """Permutes the atoms in the slab until it has a higher symmetry number.""" 

454 

455 def __init__(self, verbose=False, num_muts=1, sym_goal=100, max_tries=50, 

456 allowed_compositions=None, 

457 distribution_correction_function=None, rng=np.random): 

458 SlabOperator.__init__(self, verbose, num_muts, 

459 allowed_compositions, 

460 distribution_correction_function, 

461 rng=rng) 

462 if spglib is None: 

463 print("SymmetrySlabPermutation needs spglib to function") 

464 

465 assert sym_goal >= 1 

466 self.sym_goal = sym_goal 

467 self.max_tries = max_tries 

468 self.descriptor = 'SymmetrySlabPermutation' 

469 

470 def get_new_individual(self, parents): 

471 f = parents[0] 

472 # Permutation only makes sense if two different elements are present 

473 if len(set(f.get_chemical_symbols())) == 1: 

474 f = parents[1] 

475 if len(set(f.get_chemical_symbols())) == 1: 

476 return None, '{1} not possible in {0}'.format(f.info['confid'], 

477 self.descriptor) 

478 

479 indi = self.initialize_individual(f, self.operate(f)) 

480 indi.info['data']['parents'] = [i.info['confid'] for i in parents] 

481 

482 parent_message = ': Parent {0}'.format(f.info['confid']) 

483 return (self.finalize_individual(indi), 

484 self.descriptor + parent_message) 

485 

486 def operate(self, atoms): 

487 # Do the operation 

488 sym_num = 1 

489 sg = self.sym_goal 

490 while sym_num < sg: 

491 for _ in range(self.max_tries): 

492 for _ in range(2): 

493 permute2(atoms, rng=self.rng) 

494 self.dcf(atoms) 

495 sym_num = spglib.get_symmetry_dataset( 

496 atoms_to_spglib_cell(atoms))['number'] 

497 if sym_num >= sg: 

498 break 

499 sg -= 1 

500 return atoms 

501 

502 

503class RandomSlabPermutation(SlabOperator): 

504 def __init__(self, verbose=False, num_muts=1, 

505 allowed_compositions=None, 

506 distribution_correction_function=None, rng=np.random): 

507 SlabOperator.__init__(self, verbose, num_muts, 

508 allowed_compositions, 

509 distribution_correction_function, 

510 rng=rng) 

511 

512 self.descriptor = 'RandomSlabPermutation' 

513 

514 def get_new_individual(self, parents): 

515 f = parents[0] 

516 # Permutation only makes sense if two different elements are present 

517 if len(set(f.get_chemical_symbols())) == 1: 

518 f = parents[1] 

519 if len(set(f.get_chemical_symbols())) == 1: 

520 return None, '{1} not possible in {0}'.format(f.info['confid'], 

521 self.descriptor) 

522 

523 indi = self.initialize_individual(f, f) 

524 indi.info['data']['parents'] = [i.info['confid'] for i in parents] 

525 

526 indi = self.operate(indi) 

527 

528 parent_message = ': Parent {0}'.format(f.info['confid']) 

529 return (self.finalize_individual(indi), 

530 self.descriptor + parent_message) 

531 

532 def operate(self, atoms): 

533 # Do the operation 

534 for _ in range(self.num_muts): 

535 permute2(atoms, rng=self.rng) 

536 self.dcf(atoms) 

537 return atoms