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

1import time 

2 

3import numpy as np 

4 

5from ase.atom import Atom 

6from ase.atoms import Atoms 

7from ase.calculators.lammpsrun import Prism 

8from ase.data import atomic_masses, chemical_symbols 

9from ase.io import read 

10from ase.neighborlist import NeighborList 

11 

12 

13def twochar(name): 

14 if len(name) > 1: 

15 return name[:2] 

16 else: 

17 return name + ' ' 

18 

19 

20class BondData: 

21 def __init__(self, name_value_hash): 

22 self.nvh = name_value_hash 

23 

24 def name_value(self, aname, bname): 

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

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

27 if name1 in self.nvh: 

28 return name1, self.nvh[name1] 

29 if name2 in self.nvh: 

30 return name2, self.nvh[name2] 

31 return None, None 

32 

33 def value(self, aname, bname): 

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

35 

36 

37class CutoffList(BondData): 

38 def max(self): 

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

40 

41 

42class AnglesData: 

43 def __init__(self, name_value_hash): 

44 self.nvh = name_value_hash 

45 

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

47 for name in [ 

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

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

50 if name in self.nvh: 

51 return name, self.nvh[name] 

52 return None, None 

53 

54 

55class DihedralsData: 

56 def __init__(self, name_value_hash): 

57 self.nvh = name_value_hash 

58 

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

60 for name in [ 

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

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

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

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

65 if name in self.nvh: 

66 return name, self.nvh[name] 

67 return None, None 

68 

69 

70class OPLSff: 

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

72 self.warnings = warnings 

73 self.data = {} 

74 if fileobj is not None: 

75 self.read(fileobj) 

76 

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

78 

79 def read_block(name, symlen, nvalues): 

80 """Read a data block. 

81 

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

83 symlen: length of the symbol 

84 nvalues: number of values expected 

85 """ 

86 

87 if name not in self.data: 

88 self.data[name] = {} 

89 data = self.data[name] 

90 

91 def add_line(): 

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

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

94 return False 

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

96 if len(line) > symlen: 

97 symbol = line[:symlen] 

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

99 if len(words) >= nvalues: 

100 if nvalues == 1: 

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

102 else: 

103 data[symbol] = [float(word) 

104 for word in words[:nvalues]] 

105 return True 

106 

107 while add_line(): 

108 pass 

109 

110 read_block('one', 2, 3) 

111 read_block('bonds', 5, 2) 

112 read_block('angles', 8, 2) 

113 read_block('dihedrals', 11, 4) 

114 read_block('cutoffs', 5, 1) 

115 

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

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

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

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

120 

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

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

123 self.prefix = prefix 

124 

125 if hasattr(atoms, 'connectivities'): 

126 connectivities = atoms.connectivities 

127 else: 

128 btypes, blist = self.get_bonds(atoms) 

129 atypes, alist = self.get_angles() 

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

131 connectivities = { 

132 'bonds': blist, 

133 'bond types': btypes, 

134 'angles': alist, 

135 'angle types': atypes, 

136 'dihedrals': dlist, 

137 'dihedral types': dtypes} 

138 

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

140 self.write_lammps_in() 

141 

142 self.write_lammps_atoms(atoms, connectivities) 

143 

144 def write_lammps_in(self): 

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

146 self._write_lammps_in(fileobj) 

147 

148 def _write_lammps_in(self, fileobj): 

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

150 

151units metal 

152atom_style full 

153boundary p p p 

154#boundary p p f 

155 

156""") 

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

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

159 fileobj.write(""" 

160kspace_style pppm 1e-5 

161#kspace_modify slab 3.0 

162 

163neighbor 1.0 bin 

164neigh_modify delay 0 every 1 check yes 

165 

166thermo 1000 

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

168 

169dump 1 all xyz 1000 dump_relax.xyz 

170dump_modify 1 sort id 

171 

172restart 100000 test_relax 

173 

174min_style fire 

175minimize 1.0e-14 1.0e-5 100000 100000 

176""") # noqa: E501 

177 

178 def write_lammps_atoms(self, atoms, connectivities): 

179 """Write atoms input for LAMMPS""" 

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

181 self._write_lammps_atoms(fileobj, atoms, connectivities) 

182 

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

184 # header 

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

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

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

188 blist = connectivities['bonds'] 

189 if len(blist): 

190 btypes = connectivities['bond types'] 

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

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

193 alist = connectivities['angles'] 

194 if len(alist): 

195 atypes = connectivities['angle types'] 

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

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

198 dlist = connectivities['dihedrals'] 

199 if len(dlist): 

200 dtypes = connectivities['dihedral types'] 

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

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

203 

204 # cell 

205 p = Prism(atoms.get_cell()) 

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

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

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

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

210 

211 if p.is_skewed(): 

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

213 

214 # atoms 

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

216 tag = atoms.get_tags() 

217 if atoms.has('molid'): 

218 molid = atoms.get_array('molid') 

219 else: 

220 molid = [1] * len(atoms) 

221 for i, r in enumerate( 

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

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

224 if len(atype) < 2: 

225 atype = atype + ' ' 

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

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

228 tag[i] + 1, 

229 q) + tuple(r))) 

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

231 

232 # velocities 

233 velocities = atoms.get_velocities() 

234 if velocities is not None: 

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

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

237 for i, v in enumerate(velocities): 

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

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

240 

241 # masses 

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

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

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

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

246 (i + 1, 

247 atomic_masses[chemical_symbols.index(cs)], 

248 typ, cs)) 

249 

250 # bonds 

251 if blist: 

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

253 for ib, bvals in enumerate(blist): 

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

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

256 bvals[2] + 1)) 

257 if bvals[0] in btypes: 

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

259 fileobj.write('\n') 

260 

261 # angles 

262 if alist: 

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

264 for ia, avals in enumerate(alist): 

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

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

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

268 if avals[0] in atypes: 

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

270 fileobj.write('\n') 

271 

272 # dihedrals 

273 if dlist: 

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

275 for i, dvals in enumerate(dlist): 

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

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

278 dvals[1] + 1, dvals[2] + 1, 

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

280 if dvals[0] in dtypes: 

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

282 fileobj.write('\n') 

283 

284 def update_neighbor_list(self, atoms): 

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

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

287 bothways=True, self_interaction=False) 

288 self.nl.update(atoms) 

289 self.atoms = atoms 

290 

291 def get_bonds(self, atoms): 

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

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

294 self.update_neighbor_list(atoms) 

295 

296 types = atoms.get_types() 

297 tags = atoms.get_tags() 

298 cell = atoms.get_cell() 

299 bond_list = [] 

300 bond_types = [] 

301 for i, atom in enumerate(atoms): 

302 iname = types[tags[i]] 

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

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

305 if j <= i: 

306 continue # do not double count 

307 jname = types[tags[j]] 

308 cut = cutoffs.value(iname, jname) 

309 if cut is None: 

310 if self.warnings > 1: 

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

312 continue # don't have it 

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

314 np.dot(offset, cell)) 

315 if dist > cut: 

316 continue # too far away 

317 name, val = self.bonds.name_value(iname, jname) 

318 if name is None: 

319 if self.warnings: 

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

321 continue # don't have it 

322 if name not in bond_types: 

323 bond_types.append(name) 

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

325 return bond_types, bond_list 

326 

327 def get_angles(self, atoms=None): 

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

329 if atoms is not None: 

330 self.update_neighbor_list(atoms) 

331 else: 

332 atoms = self.atoms 

333 

334 types = atoms.get_types() 

335 tags = atoms.get_tags() 

336 cell = atoms.get_cell() 

337 ang_list = [] 

338 ang_types = [] 

339 

340 # center atom *-i-* 

341 for i, atom in enumerate(atoms): 

342 iname = types[tags[i]] 

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

344 

345 # search for first neighbor j-i-* 

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

347 jname = types[tags[j]] 

348 cut = cutoffs.value(iname, jname) 

349 if cut is None: 

350 continue # don't have it 

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

352 np.dot(offsetj, cell)) 

353 if dist > cut: 

354 continue # too far away 

355 

356 # search for second neighbor j-i-k 

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

358 if k <= j: 

359 continue # avoid double count 

360 kname = types[tags[k]] 

361 cut = cutoffs.value(iname, kname) 

362 if cut is None: 

363 continue # don't have it 

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

365 np.dot(offsetk, cell) - 

366 atoms[k].position) 

367 if dist > cut: 

368 continue # too far away 

369 name, val = self.angles.name_value(jname, iname, 

370 kname) 

371 if name is None: 

372 if self.warnings > 1: 

373 print( 

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

375 'found' 

376 ) 

377 continue # don't have it 

378 if name not in ang_types: 

379 ang_types.append(name) 

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

381 

382 return ang_types, ang_list 

383 

384 def get_dihedrals(self, ang_types, ang_list): 

385 'Dihedrals derived from angles.' 

386 

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

388 

389 atoms = self.atoms 

390 types = atoms.get_types() 

391 tags = atoms.get_tags() 

392 cell = atoms.get_cell() 

393 

394 dih_list = [] 

395 dih_types = [] 

396 

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

398 if name not in dih_types: 

399 dih_types.append(name) 

400 index = dih_types.index(name) 

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

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

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

404 

405 for angle in ang_types: 

406 L, i, j, k = angle 

407 iname = types[tags[i]] 

408 jname = types[tags[j]] 

409 kname = types[tags[k]] 

410 

411 # search for l-i-j-k 

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

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

414 if L == j: 

415 continue # avoid double count 

416 lname = types[tags[L]] 

417 cut = cutoffs.value(iname, lname) 

418 if cut is None: 

419 continue # don't have it 

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

421 np.dot(offsetl, cell)) 

422 if dist > cut: 

423 continue # too far away 

424 name, val = self.dihedrals.name_value(lname, iname, 

425 jname, kname) 

426 if name is None: 

427 continue # don't have it 

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

429 

430 # search for i-j-k-l 

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

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

433 if L == j: 

434 continue # avoid double count 

435 lname = types[tags[L]] 

436 cut = cutoffs.value(kname, lname) 

437 if cut is None: 

438 continue # don't have it 

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

440 np.dot(offsetl, cell)) 

441 if dist > cut: 

442 continue # too far away 

443 name, val = self.dihedrals.name_value(iname, jname, 

444 kname, lname) 

445 if name is None: 

446 continue # don't have it 

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

448 

449 return dih_types, dih_list 

450 

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

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

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

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

455 

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

457 dtypes): 

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

459 fileobj.write('# write_lammps' + 

460 str(time.asctime(time.localtime(time.time())))) 

461 

462 # bonds 

463 if len(btypes): 

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

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

466 for ib, btype in enumerate(btypes): 

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

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

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

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

471 

472 # angles 

473 if len(atypes): 

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

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

476 for ia, atype in enumerate(atypes): 

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

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

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

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

481 

482 # dihedrals 

483 if len(dtypes): 

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

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

486 for i, dtype in enumerate(dtypes): 

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

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

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

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

491 

492 # Lennard Jones settings 

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

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

495 ' # consider changing these parameters\n') 

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

497 data = self.data['one'] 

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

499 if len(atype) < 2: 

500 atype = atype + ' ' 

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

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

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

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

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

506 

507 # Charges 

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

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

510 if len(atype) < 2: 

511 atype = atype + ' ' 

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

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

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

515 

516 

517class OPLSStructure(Atoms): 

518 default_map = { 

519 'BR': 'Br', 

520 'Be': 'Be', 

521 'C0': 'Ca', 

522 'Li': 'Li', 

523 'Mg': 'Mg', 

524 'Al': 'Al', 

525 'Ar': 'Ar'} 

526 

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

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

529 if filename: 

530 self.read_extended_xyz(filename) 

531 else: 

532 self.types = [] 

533 for atom in self: 

534 if atom.symbol not in self.types: 

535 self.types.append(atom.symbol) 

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

537 

538 def append(self, atom): 

539 """Append atom to end.""" 

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

541 

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

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

544 atoms = read(fileobj) 

545 self.set_cell(atoms.get_cell()) 

546 self.set_pbc(atoms.get_pbc()) 

547 

548 types = [] 

549 types_map = {} 

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

551 if type not in types: 

552 types_map[type] = len(types) 

553 types.append(type) 

554 atom.tag = types_map[type] 

555 self.append(atom) 

556 self.types = types 

557 

558 # copy extra array info 

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

560 if name not in self.arrays: 

561 self.new_array(name, array) 

562 

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

564 

565 if string in translate: 

566 return translate[string], string 

567 if len(string) < 2: 

568 return string, None 

569 return string[0], string[1] 

570 

571 def get_types(self): 

572 return self.types 

573 

574 def colored(self, elements): 

575 res = Atoms() 

576 res.set_cell(self.get_cell()) 

577 for atom in self: 

578 elem = self.types[atom.tag] 

579 if elem in elements: 

580 elem = elements[elem] 

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

582 return res 

583 

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

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

586 

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

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

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

590 (len(atoms), len(self))) 

591 

592 if check: 

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

594 # check that the atom types match 

595 if not (a.tag + 1 == b.number): 

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

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

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

599 

600 self.set_cell(atoms.get_cell()) 

601 self.set_positions(atoms.get_positions()) 

602 if atoms.get_velocities() is not None: 

603 self.set_velocities(atoms.get_velocities()) 

604 # XXX what about energy and forces ??? 

605 

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

607 """Read positions, connectivities, etc. 

608 

609 update_types: update atom types from the masses 

610 """ 

611 lines = fileobj.readlines() 

612 lines.pop(0) 

613 

614 def next_entry(): 

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

616 if len(line) > 0: 

617 lines.insert(0, line) 

618 

619 def next_key(): 

620 while len(lines): 

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

622 if len(line) > 0: 

623 lines.pop(0) 

624 return line 

625 return None 

626 

627 next_entry() 

628 header = {} 

629 while True: 

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

631 if len(line): 

632 w = line.split() 

633 if len(w) == 2: 

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

635 else: 

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

637 else: 

638 break 

639 

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

641 pass 

642 lines.pop(0) 

643 

644 natoms = len(self) 

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

646 for i in range(natoms): 

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

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

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

650 # print(w, positions[i]) 

651 

652 key = next_key() 

653 

654 velocities = None 

655 if key == 'Velocities': 

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

657 for i in range(natoms): 

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

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

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

661 key = next_key() 

662 

663 if key == 'Masses': 

664 ntypes = len(self.types) 

665 masses = np.empty(ntypes) 

666 for i in range(ntypes): 

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

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

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

670 

671 if update_types: 

672 # get the elements from the masses 

673 # this ensures that we have the right elements 

674 # even when reading from a lammps dump file 

675 def newtype(element, types): 

676 if len(element) > 1: 

677 # can not extend, we are restricted to 

678 # two characters 

679 return element 

680 count = 0 

681 for type in types: 

682 if type[0] == element: 

683 count += 1 

684 label = '0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ' 

685 return (element + label[count]) 

686 

687 symbolmap = {} 

688 typemap = {} 

689 types = [] 

690 ams = atomic_masses[:] 

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

692 for i, mass in enumerate(masses): 

693 m2 = (ams - mass)**2 

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

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

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

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

698 for atom in self: 

699 atom.symbol = symbolmap[atom.symbol] 

700 self.types = types