Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1import time 

2import numpy as np 

3 

4from ase.atom import Atom 

5from ase.atoms import Atoms 

6from ase.calculators.lammpsrun import Prism 

7from ase.neighborlist import NeighborList 

8from ase.data import atomic_masses, chemical_symbols 

9from ase.io import read 

10 

11 

12def twochar(name): 

13 if len(name) > 1: 

14 return name[:2] 

15 else: 

16 return name + ' ' 

17 

18 

19class BondData: 

20 def __init__(self, name_value_hash): 

21 self.nvh = name_value_hash 

22 

23 def name_value(self, aname, bname): 

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

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

26 if name1 in self.nvh: 

27 return name1, self.nvh[name1] 

28 if name2 in self.nvh: 

29 return name2, self.nvh[name2] 

30 return None, None 

31 

32 def value(self, aname, bname): 

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

34 

35 

36class CutoffList(BondData): 

37 def max(self): 

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

39 

40 

41class AnglesData: 

42 def __init__(self, name_value_hash): 

43 self.nvh = name_value_hash 

44 

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

46 for name in [ 

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

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

49 if name in self.nvh: 

50 return name, self.nvh[name] 

51 return None, None 

52 

53 

54class DihedralsData: 

55 def __init__(self, name_value_hash): 

56 self.nvh = name_value_hash 

57 

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

59 for name in [ 

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

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

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

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

64 if name in self.nvh: 

65 return name, self.nvh[name] 

66 return None, None 

67 

68 

69class OPLSff: 

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

71 self.warnings = warnings 

72 self.data = {} 

73 if fileobj is not None: 

74 self.read(fileobj) 

75 

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

77 

78 def read_block(name, symlen, nvalues): 

79 """Read a data block. 

80 

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

82 symlen: length of the symbol 

83 nvalues: number of values expected 

84 """ 

85 

86 if name not in self.data: 

87 self.data[name] = {} 

88 data = self.data[name] 

89 

90 def add_line(): 

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

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

93 return False 

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

95 if len(line) > symlen: 

96 symbol = line[:symlen] 

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

98 if len(words) >= nvalues: 

99 if nvalues == 1: 

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

101 else: 

102 data[symbol] = [float(word) 

103 for word in words[:nvalues]] 

104 return True 

105 

106 while add_line(): 

107 pass 

108 

109 read_block('one', 2, 3) 

110 read_block('bonds', 5, 2) 

111 read_block('angles', 8, 2) 

112 read_block('dihedrals', 11, 4) 

113 read_block('cutoffs', 5, 1) 

114 

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

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

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

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

119 

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

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

122 self.prefix = prefix 

123 

124 if hasattr(atoms, 'connectivities'): 

125 connectivities = atoms.connectivities 

126 else: 

127 btypes, blist = self.get_bonds(atoms) 

128 atypes, alist = self.get_angles() 

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

130 connectivities = { 

131 'bonds': blist, 

132 'bond types': btypes, 

133 'angles': alist, 

134 'angle types': atypes, 

135 'dihedrals': dlist, 

136 'dihedral types': dtypes} 

137 

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

139 self.write_lammps_in() 

140 

141 self.write_lammps_atoms(atoms, connectivities) 

142 

143 def write_lammps_in(self): 

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

145 self._write_lammps_in(fileobj) 

146 

147 def _write_lammps_in(self, fileobj): 

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

149 

150units metal 

151atom_style full 

152boundary p p p 

153#boundary p p f 

154 

155""") 

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

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

158 fileobj.write(""" 

159kspace_style pppm 1e-5 

160#kspace_modify slab 3.0 

161 

162neighbor 1.0 bin 

163neigh_modify delay 0 every 1 check yes 

164 

165thermo 1000 

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

167 

168dump 1 all xyz 1000 dump_relax.xyz 

169dump_modify 1 sort id 

170 

171restart 100000 test_relax 

172 

173min_style fire 

174minimize 1.0e-14 1.0e-5 100000 100000 

175""") # noqa: E501 

176 

177 def write_lammps_atoms(self, atoms, connectivities): 

178 """Write atoms input for LAMMPS""" 

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

180 self._write_lammps_atoms(fileobj, atoms, connectivities) 

181 

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

183 # header 

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

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

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

187 blist = connectivities['bonds'] 

188 if len(blist): 

189 btypes = connectivities['bond types'] 

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

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

192 alist = connectivities['angles'] 

193 if len(alist): 

194 atypes = connectivities['angle types'] 

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

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

197 dlist = connectivities['dihedrals'] 

198 if len(dlist): 

199 dtypes = connectivities['dihedral types'] 

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

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

202 

203 # cell 

204 p = Prism(atoms.get_cell()) 

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

206 fileobj.write('\n0.0 %s xlo xhi\n' % xhi) 

207 fileobj.write('0.0 %s ylo yhi\n' % yhi) 

208 fileobj.write('0.0 %s zlo zhi\n' % zhi) 

209 

210 if p.is_skewed(): 

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

212 

213 # atoms 

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

215 tag = atoms.get_tags() 

216 if atoms.has('molid'): 

217 molid = atoms.get_array('molid') 

218 else: 

219 molid = [1] * len(atoms) 

220 for i, r in enumerate( 

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

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

223 if len(atype) < 2: 

224 atype = atype + ' ' 

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

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

227 tag[i] + 1, 

228 q) + tuple(r))) 

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

230 

231 # velocities 

232 velocities = atoms.get_velocities() 

233 if velocities is not None: 

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

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

236 for i, v in enumerate(velocities): 

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

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

239 

240 # masses 

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

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

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

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

245 (i + 1, 

246 atomic_masses[chemical_symbols.index(cs)], 

247 typ, cs)) 

248 

249 # bonds 

250 if blist: 

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

252 for ib, bvals in enumerate(blist): 

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

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

255 bvals[2] + 1)) 

256 if bvals[0] in btypes: 

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

258 fileobj.write('\n') 

259 

260 # angles 

261 if alist: 

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

263 for ia, avals in enumerate(alist): 

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

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

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

267 if avals[0] in atypes: 

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

269 fileobj.write('\n') 

270 

271 # dihedrals 

272 if dlist: 

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

274 for i, dvals in enumerate(dlist): 

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

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

277 dvals[1] + 1, dvals[2] + 1, 

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

279 if dvals[0] in dtypes: 

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

281 fileobj.write('\n') 

282 

283 def update_neighbor_list(self, atoms): 

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

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

286 bothways=True, self_interaction=False) 

287 self.nl.update(atoms) 

288 self.atoms = atoms 

289 

290 def get_bonds(self, atoms): 

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

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

293 self.update_neighbor_list(atoms) 

294 

295 types = atoms.get_types() 

296 tags = atoms.get_tags() 

297 cell = atoms.get_cell() 

298 bond_list = [] 

299 bond_types = [] 

300 for i, atom in enumerate(atoms): 

301 iname = types[tags[i]] 

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

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

304 if j <= i: 

305 continue # do not double count 

306 jname = types[tags[j]] 

307 cut = cutoffs.value(iname, jname) 

308 if cut is None: 

309 if self.warnings > 1: 

310 print('Warning: cutoff %s-%s not found' 

311 % (iname, jname)) 

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('Warning: potential %s-%s not found' 

321 % (iname, jname)) 

322 continue # don't have it 

323 if name not in bond_types: 

324 bond_types.append(name) 

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

326 return bond_types, bond_list 

327 

328 def get_angles(self, atoms=None): 

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

330 if atoms is not None: 

331 self.update_neighbor_list(atoms) 

332 else: 

333 atoms = self.atoms 

334 

335 types = atoms.get_types() 

336 tags = atoms.get_tags() 

337 cell = atoms.get_cell() 

338 ang_list = [] 

339 ang_types = [] 

340 

341 # center atom *-i-* 

342 for i, atom in enumerate(atoms): 

343 iname = types[tags[i]] 

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

345 

346 # search for first neighbor j-i-* 

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

348 jname = types[tags[j]] 

349 cut = cutoffs.value(iname, jname) 

350 if cut is None: 

351 continue # don't have it 

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

353 np.dot(offsetj, cell)) 

354 if dist > cut: 

355 continue # too far away 

356 

357 # search for second neighbor j-i-k 

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

359 if k <= j: 

360 continue # avoid double count 

361 kname = types[tags[k]] 

362 cut = cutoffs.value(iname, kname) 

363 if cut is None: 

364 continue # don't have it 

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

366 np.dot(offsetk, cell) - 

367 atoms[k].position) 

368 if dist > cut: 

369 continue # too far away 

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

371 kname) 

372 if name is None: 

373 if self.warnings > 1: 

374 print('Warning: angles %s-%s-%s not found' 

375 % (jname, iname, kname)) 

376 continue # don't have it 

377 if name not in ang_types: 

378 ang_types.append(name) 

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

380 

381 return ang_types, ang_list 

382 

383 def get_dihedrals(self, ang_types, ang_list): 

384 'Dihedrals derived from angles.' 

385 

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

387 

388 atoms = self.atoms 

389 types = atoms.get_types() 

390 tags = atoms.get_tags() 

391 cell = atoms.get_cell() 

392 

393 dih_list = [] 

394 dih_types = [] 

395 

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

397 if name not in dih_types: 

398 dih_types.append(name) 

399 index = dih_types.index(name) 

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

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

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

403 

404 for angle in ang_types: 

405 L, i, j, k = angle 

406 iname = types[tags[i]] 

407 jname = types[tags[j]] 

408 kname = types[tags[k]] 

409 

410 # search for l-i-j-k 

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

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

413 if L == j: 

414 continue # avoid double count 

415 lname = types[tags[L]] 

416 cut = cutoffs.value(iname, lname) 

417 if cut is None: 

418 continue # don't have it 

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

420 np.dot(offsetl, cell)) 

421 if dist > cut: 

422 continue # too far away 

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

424 jname, kname) 

425 if name is None: 

426 continue # don't have it 

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

428 

429 # search for i-j-k-l 

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

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

432 if L == j: 

433 continue # avoid double count 

434 lname = types[tags[L]] 

435 cut = cutoffs.value(kname, lname) 

436 if cut is None: 

437 continue # don't have it 

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

439 np.dot(offsetl, cell)) 

440 if dist > cut: 

441 continue # too far away 

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

443 kname, lname) 

444 if name is None: 

445 continue # don't have it 

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

447 

448 return dih_types, dih_list 

449 

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

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

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

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

454 

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

456 dtypes): 

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

458 fileobj.write('# write_lammps' + 

459 str(time.asctime(time.localtime(time.time())))) 

460 

461 # bonds 

462 if len(btypes): 

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

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

465 for ib, btype in enumerate(btypes): 

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

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

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

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

470 

471 # angles 

472 if len(atypes): 

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

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

475 for ia, atype in enumerate(atypes): 

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

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

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

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

480 

481 # dihedrals 

482 if len(dtypes): 

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

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

485 for i, dtype in enumerate(dtypes): 

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

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

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

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

490 

491 # Lennard Jones settings 

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

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

494 ' # consider changing these parameters\n') 

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

496 data = self.data['one'] 

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

498 if len(atype) < 2: 

499 atype = atype + ' ' 

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

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

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

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

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

505 

506 # Charges 

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

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

509 if len(atype) < 2: 

510 atype = atype + ' ' 

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

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

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

514 

515 

516class OPLSStructure(Atoms): 

517 default_map = { 

518 'BR': 'Br', 

519 'Be': 'Be', 

520 'C0': 'Ca', 

521 'Li': 'Li', 

522 'Mg': 'Mg', 

523 'Al': 'Al', 

524 'Ar': 'Ar'} 

525 

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

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

528 if filename: 

529 self.read_extended_xyz(filename) 

530 else: 

531 self.types = [] 

532 for atom in self: 

533 if atom.symbol not in self.types: 

534 self.types.append(atom.symbol) 

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

536 

537 def append(self, atom): 

538 """Append atom to end.""" 

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

540 

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

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

543 atoms = read(fileobj) 

544 self.set_cell(atoms.get_cell()) 

545 self.set_pbc(atoms.get_pbc()) 

546 

547 types = [] 

548 types_map = {} 

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

550 if type not in types: 

551 types_map[type] = len(types) 

552 types.append(type) 

553 atom.tag = types_map[type] 

554 self.append(atom) 

555 self.types = types 

556 

557 # copy extra array info 

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

559 if name not in self.arrays: 

560 self.new_array(name, array) 

561 

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

563 

564 if string in translate: 

565 return translate[string], string 

566 if len(string) < 2: 

567 return string, None 

568 return string[0], string[1] 

569 

570 def get_types(self): 

571 return self.types 

572 

573 def colored(self, elements): 

574 res = Atoms() 

575 res.set_cell(self.get_cell()) 

576 for atom in self: 

577 elem = self.types[atom.tag] 

578 if elem in elements: 

579 elem = elements[elem] 

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

581 return res 

582 

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

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

585 

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

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

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

589 (len(atoms), len(self))) 

590 

591 if check: 

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

593 # check that the atom types match 

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

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

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

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

598 

599 self.set_cell(atoms.get_cell()) 

600 self.set_positions(atoms.get_positions()) 

601 if atoms.get_velocities() is not None: 

602 self.set_velocities(atoms.get_velocities()) 

603 # XXX what about energy and forces ??? 

604 

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

606 """Read positions, connectivities, etc. 

607 

608 update_types: update atom types from the masses 

609 """ 

610 lines = fileobj.readlines() 

611 lines.pop(0) 

612 

613 def next_entry(): 

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

615 if(len(line) > 0): 

616 lines.insert(0, line) 

617 

618 def next_key(): 

619 while(len(lines)): 

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

621 if(len(line) > 0): 

622 lines.pop(0) 

623 return line 

624 return None 

625 

626 next_entry() 

627 header = {} 

628 while True: 

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

630 if len(line): 

631 w = line.split() 

632 if len(w) == 2: 

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

634 else: 

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

636 else: 

637 break 

638 

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

640 pass 

641 lines.pop(0) 

642 

643 natoms = len(self) 

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

645 for i in range(natoms): 

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

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

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

649 # print(w, positions[i]) 

650 

651 key = next_key() 

652 

653 velocities = None 

654 if key == 'Velocities': 

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

656 for i in range(natoms): 

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

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

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

660 key = next_key() 

661 

662 if key == 'Masses': 

663 ntypes = len(self.types) 

664 masses = np.empty((ntypes)) 

665 for i in range(ntypes): 

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

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

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

669 

670 if update_types: 

671 # get the elements from the masses 

672 # this ensures that we have the right elements 

673 # even when reading from a lammps dump file 

674 def newtype(element, types): 

675 if len(element) > 1: 

676 # can not extend, we are restricted to 

677 # two characters 

678 return element 

679 count = 0 

680 for type in types: 

681 if type[0] == element: 

682 count += 1 

683 label = '0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ' 

684 return (element + label[count]) 

685 

686 symbolmap = {} 

687 typemap = {} 

688 types = [] 

689 ams = atomic_masses[:] 

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

691 for i, mass in enumerate(masses): 

692 m2 = (ams - mass)**2 

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

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

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

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

697 for atom in self: 

698 atom.symbol = symbolmap[atom.symbol] 

699 self.types = types 

700 

701 key = next_key() 

702 

703 def read_list(key_string, length, debug=False): 

704 if key != key_string: 

705 return [], key 

706 

707 lst = [] 

708 while(len(lines)): 

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

710 if len(w) > length: 

711 lst.append([(int(w[1 + c]) - 1) for c in range(length)]) 

712 else: 

713 return lst, next_key() 

714 return lst, None 

715 

716 bonds, key = read_list('Bonds', 3) 

717 angles, key = read_list('Angles', 4) 

718 dihedrals, key = read_list('Dihedrals', 5, True) 

719 

720 self.connectivities = { 

721 'bonds': bonds, 

722 'angles': angles, 

723 'dihedrals': dihedrals 

724 } 

725 

726 if 'bonds' in header: 

727 assert(len(bonds) == header['bonds']) 

728 self.connectivities['bond types'] = list( 

729 range(header['bond types'])) 

730 if 'angles' in header: 

731 assert(len(angles) == header['angles']) 

732 self.connectivities['angle types'] = list( 

733 range(header['angle types'])) 

734 if 'dihedrals' in header: 

735 assert(len(dihedrals) == header['dihedrals']) 

736 self.connectivities['dihedral types'] = list(range( 

737 header['dihedral types']))