Coverage for /builds/ase/ase/ase/io/opls.py: 53.80%

1# fmt: off 

2 

3import time 

4 

5import numpy as np 

6 

7from ase.atom import Atom 

8from ase.atoms import Atoms 

9from ase.calculators.lammpsrun import Prism 

10from ase.data import atomic_masses, chemical_symbols 

11from ase.io import read 

12from ase.neighborlist import NeighborList 

13 

14 

15def twochar(name): 

16 if len(name) > 1: 

17 return name[:2] 

18 else: 

19 return name + ' ' 

20 

21 

22class BondData: 

23 def __init__(self, name_value_hash): 

24 self.nvh = name_value_hash 

25 

26 def name_value(self, aname, bname): 

27 name1 = twochar(aname) + '-' + twochar(bname) 

28 name2 = twochar(bname) + '-' + twochar(aname) 

29 if name1 in self.nvh: 

30 return name1, self.nvh[name1] 

31 if name2 in self.nvh: 

32 return name2, self.nvh[name2] 

33 return None, None 

34 

35 def value(self, aname, bname): 

36 return self.name_value(aname, bname)[1] 

37 

38 

39class CutoffList(BondData): 

40 def max(self): 

41 return max(self.nvh.values()) 

42 

43 

44class AnglesData: 

45 def __init__(self, name_value_hash): 

46 self.nvh = name_value_hash 

47 

48 def name_value(self, aname, bname, cname): 

49 for name in [ 

50 (twochar(aname) + '-' + twochar(bname) + '-' + twochar(cname)), 

51 (twochar(cname) + '-' + twochar(bname) + '-' + twochar(aname))]: 

52 if name in self.nvh: 

53 return name, self.nvh[name] 

54 return None, None 

55 

56 

57class DihedralsData: 

58 def __init__(self, name_value_hash): 

59 self.nvh = name_value_hash 

60 

61 def name_value(self, aname, bname, cname, dname): 

62 for name in [ 

63 (twochar(aname) + '-' + twochar(bname) + '-' + 

64 twochar(cname) + '-' + twochar(dname)), 

65 (twochar(dname) + '-' + twochar(cname) + '-' + 

66 twochar(bname) + '-' + twochar(aname))]: 

67 if name in self.nvh: 

68 return name, self.nvh[name] 

69 return None, None 

70 

71 

72class OPLSff: 

73 def __init__(self, fileobj=None, warnings=0): 

74 self.warnings = warnings 

75 self.data = {} 

76 if fileobj is not None: 

77 self.read(fileobj) 

78 

79 def read(self, fileobj, comments='#'): 

80 

81 def read_block(name, symlen, nvalues): 

82 """Read a data block. 

83 

84 name: name of the block to store in self.data 

85 symlen: length of the symbol 

86 nvalues: number of values expected 

87 """ 

88 

89 if name not in self.data: 

90 self.data[name] = {} 

91 data = self.data[name] 

92 

93 def add_line(): 

94 line = fileobj.readline().strip() 

95 if not len(line): # end of the block 

96 return False 

97 line = line.split('#')[0] # get rid of comments 

98 if len(line) > symlen: 

99 symbol = line[:symlen] 

100 words = line[symlen:].split() 

101 if len(words) >= nvalues: 

102 if nvalues == 1: 

103 data[symbol] = float(words[0]) 

104 else: 

105 data[symbol] = [float(word) 

106 for word in words[:nvalues]] 

107 return True 

108 

109 while add_line(): 

110 pass 

111 

112 read_block('one', 2, 3) 

113 read_block('bonds', 5, 2) 

114 read_block('angles', 8, 2) 

115 read_block('dihedrals', 11, 4) 

116 read_block('cutoffs', 5, 1) 

117 

118 self.bonds = BondData(self.data['bonds']) 

119 self.angles = AnglesData(self.data['angles']) 

120 self.dihedrals = DihedralsData(self.data['dihedrals']) 

121 self.cutoffs = CutoffList(self.data['cutoffs']) 

122 

123 def write_lammps(self, atoms, prefix='lammps'): 

124 """Write input for a LAMMPS calculation.""" 

125 self.prefix = prefix 

126 

127 if hasattr(atoms, 'connectivities'): 

128 connectivities = atoms.connectivities 

129 else: 

130 btypes, blist = self.get_bonds(atoms) 

131 atypes, alist = self.get_angles() 

132 dtypes, dlist = self.get_dihedrals(alist, atypes) 

133 connectivities = { 

134 'bonds': blist, 

135 'bond types': btypes, 

136 'angles': alist, 

137 'angle types': atypes, 

138 'dihedrals': dlist, 

139 'dihedral types': dtypes} 

140 

141 self.write_lammps_definitions(atoms, btypes, atypes, dtypes) 

142 self.write_lammps_in() 

143 

144 self.write_lammps_atoms(atoms, connectivities) 

145 

146 def write_lammps_in(self): 

147 with open(self.prefix + '_in', 'w') as fileobj: 

148 self._write_lammps_in(fileobj) 

149 

150 def _write_lammps_in(self, fileobj): 

151 fileobj.write("""# LAMMPS relaxation (written by ASE) 

152 

153units metal 

154atom_style full 

155boundary p p p 

156#boundary p p f 

157 

158""") 

159 fileobj.write('read_data ' + self.prefix + '_atoms\n') 

160 fileobj.write('include ' + self.prefix + '_opls\n') 

161 fileobj.write(""" 

162kspace_style pppm 1e-5 

163#kspace_modify slab 3.0 

164 

165neighbor 1.0 bin 

166neigh_modify delay 0 every 1 check yes 

167 

168thermo 1000 

169thermo_style custom step temp press cpu pxx pyy pzz pxy pxz pyz ke pe etotal vol lx ly lz atoms 

170 

171dump 1 all xyz 1000 dump_relax.xyz 

172dump_modify 1 sort id 

173 

174restart 100000 test_relax 

175 

176min_style fire 

177minimize 1.0e-14 1.0e-5 100000 100000 

178""") # noqa: E501 

179 

180 def write_lammps_atoms(self, atoms, connectivities): 

181 """Write atoms input for LAMMPS""" 

182 with open(self.prefix + '_atoms', 'w') as fileobj: 

183 self._write_lammps_atoms(fileobj, atoms, connectivities) 

184 

185 def _write_lammps_atoms(self, fileobj, atoms, connectivities): 

186 # header 

187 fileobj.write(fileobj.name + ' (by ' + str(self.__class__) + ')\n\n') 

188 fileobj.write(str(len(atoms)) + ' atoms\n') 

189 fileobj.write(str(len(atoms.types)) + ' atom types\n') 

190 blist = connectivities['bonds'] 

191 if len(blist): 

192 btypes = connectivities['bond types'] 

193 fileobj.write(str(len(blist)) + ' bonds\n') 

194 fileobj.write(str(len(btypes)) + ' bond types\n') 

195 alist = connectivities['angles'] 

196 if len(alist): 

197 atypes = connectivities['angle types'] 

198 fileobj.write(str(len(alist)) + ' angles\n') 

199 fileobj.write(str(len(atypes)) + ' angle types\n') 

200 dlist = connectivities['dihedrals'] 

201 if len(dlist): 

202 dtypes = connectivities['dihedral types'] 

203 fileobj.write(str(len(dlist)) + ' dihedrals\n') 

204 fileobj.write(str(len(dtypes)) + ' dihedral types\n') 

205 

206 # cell 

207 p = Prism(atoms.get_cell()) 

208 xhi, yhi, zhi, xy, xz, yz = p.get_lammps_prism() 

209 fileobj.write(f'\n0.0 {xhi} xlo xhi\n') 

210 fileobj.write(f'0.0 {yhi} ylo yhi\n') 

211 fileobj.write(f'0.0 {zhi} zlo zhi\n') 

212 

213 if p.is_skewed(): 

214 fileobj.write(f"{xy} {xz} {yz} xy xz yz\n") 

215 

216 # atoms 

217 fileobj.write('\nAtoms\n\n') 

218 tag = atoms.get_tags() 

219 if atoms.has('molid'): 

220 molid = atoms.get_array('molid') 

221 else: 

222 molid = [1] * len(atoms) 

223 for i, r in enumerate( 

224 p.vector_to_lammps(atoms.get_positions())): 

225 atype = atoms.types[tag[i]] 

226 if len(atype) < 2: 

227 atype = atype + ' ' 

228 q = self.data['one'][atype][2] 

229 fileobj.write('%6d %3d %3d %s %s %s %s' % ((i + 1, molid[i], 

230 tag[i] + 1, 

231 q) + tuple(r))) 

232 fileobj.write(' # ' + atoms.types[tag[i]] + '\n') 

233 

234 # velocities 

235 velocities = atoms.get_velocities() 

236 if velocities is not None: 

237 velocities = p.vector_to_lammps(atoms.get_velocities()) 

238 fileobj.write('\nVelocities\n\n') 

239 for i, v in enumerate(velocities): 

240 fileobj.write('%6d %g %g %g\n' % 

241 (i + 1, v[0], v[1], v[2])) 

242 

243 # masses 

244 fileobj.write('\nMasses\n\n') 

245 for i, typ in enumerate(atoms.types): 

246 cs = atoms.split_symbol(typ)[0] 

247 fileobj.write('%6d %g # %s -> %s\n' % 

248 (i + 1, 

249 atomic_masses[chemical_symbols.index(cs)], 

250 typ, cs)) 

251 

252 # bonds 

253 if blist: 

254 fileobj.write('\nBonds\n\n') 

255 for ib, bvals in enumerate(blist): 

256 fileobj.write('%8d %6d %6d %6d ' % 

257 (ib + 1, bvals[0] + 1, bvals[1] + 1, 

258 bvals[2] + 1)) 

259 if bvals[0] in btypes: 

260 fileobj.write('# ' + btypes[bvals[0]]) 

261 fileobj.write('\n') 

262 

263 # angles 

264 if alist: 

265 fileobj.write('\nAngles\n\n') 

266 for ia, avals in enumerate(alist): 

267 fileobj.write('%8d %6d %6d %6d %6d ' % 

268 (ia + 1, avals[0] + 1, 

269 avals[1] + 1, avals[2] + 1, avals[3] + 1)) 

270 if avals[0] in atypes: 

271 fileobj.write('# ' + atypes[avals[0]]) 

272 fileobj.write('\n') 

273 

274 # dihedrals 

275 if dlist: 

276 fileobj.write('\nDihedrals\n\n') 

277 for i, dvals in enumerate(dlist): 

278 fileobj.write('%8d %6d %6d %6d %6d %6d ' % 

279 (i + 1, dvals[0] + 1, 

280 dvals[1] + 1, dvals[2] + 1, 

281 dvals[3] + 1, dvals[4] + 1)) 

282 if dvals[0] in dtypes: 

283 fileobj.write('# ' + dtypes[dvals[0]]) 

284 fileobj.write('\n') 

285 

286 def update_neighbor_list(self, atoms): 

287 cut = 0.5 * max(self.data['cutoffs'].values()) 

288 self.nl = NeighborList([cut] * len(atoms), skin=0, 

289 bothways=True, self_interaction=False) 

290 self.nl.update(atoms) 

291 self.atoms = atoms 

292 

293 def get_bonds(self, atoms): 

294 """Find bonds and return them and their types""" 

295 cutoffs = CutoffList(self.data['cutoffs']) 

296 self.update_neighbor_list(atoms) 

297 

298 types = atoms.get_types() 

299 tags = atoms.get_tags() 

300 cell = atoms.get_cell() 

301 bond_list = [] 

302 bond_types = [] 

303 for i, atom in enumerate(atoms): 

304 iname = types[tags[i]] 

305 indices, offsets = self.nl.get_neighbors(i) 

306 for j, offset in zip(indices, offsets): 

307 if j <= i: 

308 continue # do not double count 

309 jname = types[tags[j]] 

310 cut = cutoffs.value(iname, jname) 

311 if cut is None: 

312 if self.warnings > 1: 

313 print(f'Warning: cutoff {iname}-{jname} not found') 

314 continue # don't have it 

315 dist = np.linalg.norm(atom.position - atoms[j].position - 

316 np.dot(offset, cell)) 

317 if dist > cut: 

318 continue # too far away 

319 name, _val = self.bonds.name_value(iname, jname) 

320 if name is None: 

321 if self.warnings: 

322 print(f'Warning: potential {iname}-{jname} not found') 

323 continue # don't have it 

324 if name not in bond_types: 

325 bond_types.append(name) 

326 bond_list.append([bond_types.index(name), i, j]) 

327 return bond_types, bond_list 

328 

329 def get_angles(self, atoms=None): 

330 cutoffs = CutoffList(self.data['cutoffs']) 

331 if atoms is not None: 

332 self.update_neighbor_list(atoms) 

333 else: 

334 atoms = self.atoms 

335 

336 types = atoms.get_types() 

337 tags = atoms.get_tags() 

338 cell = atoms.get_cell() 

339 ang_list = [] 

340 ang_types = [] 

341 

342 # center atom *-i-* 

343 for i, atom in enumerate(atoms): 

344 iname = types[tags[i]] 

345 indicesi, offsetsi = self.nl.get_neighbors(i) 

346 

347 # search for first neighbor j-i-* 

348 for j, offsetj in zip(indicesi, offsetsi): 

349 jname = types[tags[j]] 

350 cut = cutoffs.value(iname, jname) 

351 if cut is None: 

352 continue # don't have it 

353 dist = np.linalg.norm(atom.position - atoms[j].position - 

354 np.dot(offsetj, cell)) 

355 if dist > cut: 

356 continue # too far away 

357 

358 # search for second neighbor j-i-k 

359 for k, offsetk in zip(indicesi, offsetsi): 

360 if k <= j: 

361 continue # avoid double count 

362 kname = types[tags[k]] 

363 cut = cutoffs.value(iname, kname) 

364 if cut is None: 

365 continue # don't have it 

366 dist = np.linalg.norm(atom.position - 

367 np.dot(offsetk, cell) - 

368 atoms[k].position) 

369 if dist > cut: 

370 continue # too far away 

371 name, _val = self.angles.name_value(jname, iname, 

372 kname) 

373 if name is None: 

374 if self.warnings > 1: 

375 print( 

376 f'Warning: angles {jname}-{iname}-{kname} not ' 

377 'found' 

378 ) 

379 continue # don't have it 

380 if name not in ang_types: 

381 ang_types.append(name) 

382 ang_list.append([ang_types.index(name), j, i, k]) 

383 

384 return ang_types, ang_list 

385 

386 def get_dihedrals(self, ang_types, ang_list): 

387 'Dihedrals derived from angles.' 

388 

389 cutoffs = CutoffList(self.data['cutoffs']) 

390 

391 atoms = self.atoms 

392 types = atoms.get_types() 

393 tags = atoms.get_tags() 

394 cell = atoms.get_cell() 

395 

396 dih_list = [] 

397 dih_types = [] 

398 

399 def append(name, i, j, k, L): 

400 if name not in dih_types: 

401 dih_types.append(name) 

402 index = dih_types.index(name) 

403 if (([index, i, j, k, L] not in dih_list) and 

404 ([index, L, k, j, i] not in dih_list)): 

405 dih_list.append([index, i, j, k, L]) 

406 

407 for angle in ang_types: 

408 L, i, j, k = angle 

409 iname = types[tags[i]] 

410 jname = types[tags[j]] 

411 kname = types[tags[k]] 

412 

413 # search for l-i-j-k 

414 indicesi, offsetsi = self.nl.get_neighbors(i) 

415 for L, offsetl in zip(indicesi, offsetsi): 

416 if L == j: 

417 continue # avoid double count 

418 lname = types[tags[L]] 

419 cut = cutoffs.value(iname, lname) 

420 if cut is None: 

421 continue # don't have it 

422 dist = np.linalg.norm(atoms[i].position - atoms[L].position - 

423 np.dot(offsetl, cell)) 

424 if dist > cut: 

425 continue # too far away 

426 name, _val = self.dihedrals.name_value(lname, iname, 

427 jname, kname) 

428 if name is None: 

429 continue # don't have it 

430 append(name, L, i, j, k) 

431 

432 # search for i-j-k-l 

433 indicesk, offsetsk = self.nl.get_neighbors(k) 

434 for L, offsetl in zip(indicesk, offsetsk): 

435 if L == j: 

436 continue # avoid double count 

437 lname = types[tags[L]] 

438 cut = cutoffs.value(kname, lname) 

439 if cut is None: 

440 continue # don't have it 

441 dist = np.linalg.norm(atoms[k].position - atoms[L].position - 

442 np.dot(offsetl, cell)) 

443 if dist > cut: 

444 continue # too far away 

445 name, _val = self.dihedrals.name_value(iname, jname, 

446 kname, lname) 

447 if name is None: 

448 continue # don't have it 

449 append(name, i, j, k, L) 

450 

451 return dih_types, dih_list 

452 

453 def write_lammps_definitions(self, atoms, btypes, atypes, dtypes): 

454 """Write force field definitions for LAMMPS.""" 

455 with open(self.prefix + '_opls', 'w') as fd: 

456 self._write_lammps_definitions(fd, atoms, btypes, atypes, dtypes) 

457 

458 def _write_lammps_definitions(self, fileobj, atoms, btypes, atypes, 

459 dtypes): 

460 fileobj.write('# OPLS potential\n') 

461 fileobj.write('# write_lammps' + 

462 str(time.asctime(time.localtime(time.time())))) 

463 

464 # bonds 

465 if len(btypes): 

466 fileobj.write('\n# bonds\n') 

467 fileobj.write('bond_style harmonic\n') 

468 for ib, btype in enumerate(btypes): 

469 fileobj.write('bond_coeff %6d' % (ib + 1)) 

470 for value in self.bonds.nvh[btype]: 

471 fileobj.write(' ' + str(value)) 

472 fileobj.write(' # ' + btype + '\n') 

473 

474 # angles 

475 if len(atypes): 

476 fileobj.write('\n# angles\n') 

477 fileobj.write('angle_style harmonic\n') 

478 for ia, atype in enumerate(atypes): 

479 fileobj.write('angle_coeff %6d' % (ia + 1)) 

480 for value in self.angles.nvh[atype]: 

481 fileobj.write(' ' + str(value)) 

482 fileobj.write(' # ' + atype + '\n') 

483 

484 # dihedrals 

485 if len(dtypes): 

486 fileobj.write('\n# dihedrals\n') 

487 fileobj.write('dihedral_style opls\n') 

488 for i, dtype in enumerate(dtypes): 

489 fileobj.write('dihedral_coeff %6d' % (i + 1)) 

490 for value in self.dihedrals.nvh[dtype]: 

491 fileobj.write(' ' + str(value)) 

492 fileobj.write(' # ' + dtype + '\n') 

493 

494 # Lennard Jones settings 

495 fileobj.write('\n# L-J parameters\n') 

496 fileobj.write('pair_style lj/cut/coul/long 10.0 7.4' + 

497 ' # consider changing these parameters\n') 

498 fileobj.write('special_bonds lj/coul 0.0 0.0 0.5\n') 

499 data = self.data['one'] 

500 for ia, atype in enumerate(atoms.types): 

501 if len(atype) < 2: 

502 atype = atype + ' ' 

503 fileobj.write('pair_coeff ' + str(ia + 1) + ' ' + str(ia + 1)) 

504 for value in data[atype][:2]: 

505 fileobj.write(' ' + str(value)) 

506 fileobj.write(' # ' + atype + '\n') 

507 fileobj.write('pair_modify shift yes mix geometric\n') 

508 

509 # Charges 

510 fileobj.write('\n# charges\n') 

511 for ia, atype in enumerate(atoms.types): 

512 if len(atype) < 2: 

513 atype = atype + ' ' 

514 fileobj.write('set type ' + str(ia + 1)) 

515 fileobj.write(' charge ' + str(data[atype][2])) 

516 fileobj.write(' # ' + atype + '\n') 

517 

518 

519class OPLSStructure(Atoms): 

520 default_map = { 

521 'BR': 'Br', 

522 'Be': 'Be', 

523 'C0': 'Ca', 

524 'Li': 'Li', 

525 'Mg': 'Mg', 

526 'Al': 'Al', 

527 'Ar': 'Ar'} 

528 

529 def __init__(self, filename=None, *args, **kwargs): 

530 Atoms.__init__(self, *args, **kwargs) 

531 if filename: 

532 self.read_extended_xyz(filename) 

533 else: 

534 self.types = [] 

535 for atom in self: 

536 if atom.symbol not in self.types: 

537 self.types.append(atom.symbol) 

538 atom.tag = self.types.index(atom.symbol) 

539 

540 def append(self, atom): 

541 """Append atom to end.""" 

542 self.extend(Atoms([atom])) 

543 

544 def read_extended_xyz(self, fileobj, map={}): 

545 """Read extended xyz file with labeled atoms.""" 

546 atoms = read(fileobj) 

547 self.set_cell(atoms.get_cell()) 

548 self.set_pbc(atoms.get_pbc()) 

549 

550 types = [] 

551 types_map = {} 

552 for atom, type in zip(atoms, atoms.get_array('type')): 

553 if type not in types: 

554 types_map[type] = len(types) 

555 types.append(type) 

556 atom.tag = types_map[type] 

557 self.append(atom) 

558 self.types = types 

559 

560 # copy extra array info 

561 for name, array in atoms.arrays.items(): 

562 if name not in self.arrays: 

563 self.new_array(name, array) 

564 

565 def split_symbol(self, string, translate=default_map): 

566 

567 if string in translate: 

568 return translate[string], string 

569 if len(string) < 2: 

570 return string, None 

571 return string[0], string[1] 

572 

573 def get_types(self): 

574 return self.types 

575 

576 def colored(self, elements): 

577 res = Atoms() 

578 res.set_cell(self.get_cell()) 

579 for atom in self: 

580 elem = self.types[atom.tag] 

581 if elem in elements: 

582 elem = elements[elem] 

583 res.append(Atom(elem, atom.position)) 

584 return res 

585 

586 def update_from_lammps_dump(self, fileobj, check=True): 

587 atoms = read(fileobj, format='lammps-dump') 

588 

589 if len(atoms) != len(self): 

590 raise RuntimeError('Structure in ' + str(fileobj) + 

591 ' has wrong length: %d != %d' % 

592 (len(atoms), len(self))) 

593 

594 if check: 

595 for a, b in zip(self, atoms): 

596 # check that the atom types match 

597 if a.tag + 1 != b.number: 

598 raise RuntimeError('Atoms index %d are of different ' 

599 'type (%d != %d)' 

600 % (a.index, a.tag + 1, b.number)) 

601 

602 self.set_cell(atoms.get_cell()) 

603 self.set_positions(atoms.get_positions()) 

604 if atoms.get_velocities() is not None: 

605 self.set_velocities(atoms.get_velocities()) 

606 # XXX what about energy and forces ??? 

607 

608 def read_connectivities(self, fileobj, update_types=False): 

609 """Read positions, connectivities, etc. 

610 

611 update_types: update atom types from the masses 

612 """ 

613 lines = fileobj.readlines() 

614 lines.pop(0) 

615 

616 def next_entry(): 

617 line = lines.pop(0).strip() 

618 if len(line) > 0: 

619 lines.insert(0, line) 

620 

621 def next_key(): 

622 while len(lines): 

623 line = lines.pop(0).strip() 

624 if len(line) > 0: 

625 lines.pop(0) 

626 return line 

627 return None 

628 

629 next_entry() 

630 header = {} 

631 while True: 

632 line = lines.pop(0).strip() 

633 if len(line): 

634 w = line.split() 

635 if len(w) == 2: 

636 header[w[1]] = int(w[0]) 

637 else: 

638 header[w[1] + ' ' + w[2]] = int(w[0]) 

639 else: 

640 break 

641 

642 while not lines.pop(0).startswith('Atoms'): 

643 pass 

644 lines.pop(0) 

645 

646 natoms = len(self) 

647 positions = np.empty((natoms, 3)) 

648 for i in range(natoms): 

649 w = lines.pop(0).split() 

650 assert int(w[0]) == (i + 1) 

651 positions[i] = np.array([float(w[4 + c]) for c in range(3)]) 

652 # print(w, positions[i]) 

653 

654 key = next_key() 

655 

656 velocities = None 

657 if key == 'Velocities': 

658 velocities = np.empty((natoms, 3)) 

659 for i in range(natoms): 

660 w = lines.pop(0).split() 

661 assert int(w[0]) == (i + 1) 

662 velocities[i] = np.array([float(w[1 + c]) for c in range(3)]) 

663 key = next_key() 

664 

665 if key == 'Masses': 

666 ntypes = len(self.types) 

667 masses = np.empty(ntypes) 

668 for i in range(ntypes): 

669 w = lines.pop(0).split() 

670 assert int(w[0]) == (i + 1) 

671 masses[i] = float(w[1]) 

672 

673 if update_types: 

674 # get the elements from the masses 

675 # this ensures that we have the right elements 

676 # even when reading from a lammps dump file 

677 def newtype(element, types): 

678 if len(element) > 1: 

679 # can not extend, we are restricted to 

680 # two characters 

681 return element 

682 count = 0 

683 for type in types: 

684 if type[0] == element: 

685 count += 1 

686 label = '0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ' 

687 return (element + label[count]) 

688 

689 symbolmap = {} 

690 typemap = {} 

691 types = [] 

692 ams = atomic_masses[:] 

693 ams[np.isnan(ams)] = 0 

694 for i, mass in enumerate(masses): 

695 m2 = (ams - mass)**2 

696 symbolmap[self.types[i]] = chemical_symbols[m2.argmin()] 

697 typemap[self.types[i]] = newtype( 

698 chemical_symbols[m2.argmin()], types) 

699 types.append(typemap[self.types[i]]) 

700 for atom in self: 

701 atom.symbol = symbolmap[atom.symbol]