1import time 

2import warnings 

3 

4from ase.units import Ang, fs 

5from ase.utils import reader, writer 

6 

7 

8v_unit = Ang / (1000.0 * fs) 

9 

10 

11@reader 

12def read_aims(fd, apply_constraints=True): 

13 """Import FHI-aims geometry type files. 

14 

15 Reads unitcell, atom positions and constraints from 

16 a geometry.in file. 

17 

18 If geometric constraint (symmetry parameters) are in the file 

19 include that information in atoms.info["symmetry_block"] 

20 """ 

21 

22 lines = fd.readlines() 

23 return parse_geometry_lines(lines, apply_constraints=apply_constraints) 

24 

25 

26def parse_geometry_lines(lines, apply_constraints=True): 

27 

28 from ase import Atoms 

29 from ase.constraints import ( 

30 FixAtoms, 

31 FixCartesian, 

32 FixScaledParametricRelations, 

33 FixCartesianParametricRelations, 

34 ) 

35 import numpy as np 

36 

37 atoms = Atoms() 

38 

39 positions = [] 

40 cell = [] 

41 symbols = [] 

42 velocities = [] 

43 magmoms = [] 

44 symmetry_block = [] 

45 charges = [] 

46 fix = [] 

47 fix_cart = [] 

48 xyz = np.array([0, 0, 0]) 

49 i = -1 

50 n_periodic = -1 

51 periodic = np.array([False, False, False]) 

52 cart_positions, scaled_positions = False, False 

53 for line in lines: 

54 inp = line.split() 

55 if inp == []: 

56 continue 

57 if inp[0] in ["atom", "atom_frac"]: 

58 

59 if inp[0] == "atom": 

60 cart_positions = True 

61 else: 

62 scaled_positions = True 

63 

64 if xyz.all(): 

65 fix.append(i) 

66 elif xyz.any(): 

67 fix_cart.append(FixCartesian(i, xyz)) 

68 floatvect = float(inp[1]), float(inp[2]), float(inp[3]) 

69 positions.append(floatvect) 

70 symbols.append(inp[4]) 

71 magmoms.append(0.0) 

72 charges.append(0.0) 

73 xyz = np.array([0, 0, 0]) 

74 i += 1 

75 

76 elif inp[0] == "lattice_vector": 

77 floatvect = float(inp[1]), float(inp[2]), float(inp[3]) 

78 cell.append(floatvect) 

79 n_periodic = n_periodic + 1 

80 periodic[n_periodic] = True 

81 

82 elif inp[0] == "initial_moment": 

83 magmoms[-1] = float(inp[1]) 

84 

85 elif inp[0] == "initial_charge": 

86 charges[-1] = float(inp[1]) 

87 

88 elif inp[0] == "constrain_relaxation": 

89 if inp[1] == ".true.": 

90 fix.append(i) 

91 elif inp[1] == "x": 

92 xyz[0] = 1 

93 elif inp[1] == "y": 

94 xyz[1] = 1 

95 elif inp[1] == "z": 

96 xyz[2] = 1 

97 

98 elif inp[0] == "velocity": 

99 floatvect = [v_unit * float(l) for l in inp[1:4]] 

100 velocities.append(floatvect) 

101 

102 elif inp[0] in [ 

103 "symmetry_n_params", 

104 "symmetry_params", 

105 "symmetry_lv", 

106 "symmetry_frac", 

107 ]: 

108 symmetry_block.append(" ".join(inp)) 

109 

110 if xyz.all(): 

111 fix.append(i) 

112 elif xyz.any(): 

113 fix_cart.append(FixCartesian(i, xyz)) 

114 

115 if cart_positions and scaled_positions: 

116 raise Exception( 

117 "Can't specify atom positions with mixture of " 

118 "Cartesian and fractional coordinates" 

119 ) 

120 elif scaled_positions and periodic.any(): 

121 atoms = Atoms( 

122 symbols, scaled_positions=positions, cell=cell, pbc=periodic 

123 ) 

124 else: 

125 atoms = Atoms(symbols, positions) 

126 

127 if len(velocities) > 0: 

128 if len(velocities) != len(positions): 

129 raise Exception( 

130 "Number of positions and velocities have to coincide." 

131 ) 

132 atoms.set_velocities(velocities) 

133 

134 fix_params = [] 

135 

136 if len(symmetry_block) > 5: 

137 params = symmetry_block[1].split()[1:] 

138 

139 lattice_expressions = [] 

140 lattice_params = [] 

141 

142 atomic_expressions = [] 

143 atomic_params = [] 

144 

145 n_lat_param = int(symmetry_block[0].split(" ")[2]) 

146 

147 lattice_params = params[:n_lat_param] 

148 atomic_params = params[n_lat_param:] 

149 

150 for ll, line in enumerate(symmetry_block[2:]): 

151 expression = " ".join(line.split(" ")[1:]) 

152 if ll < 3: 

153 lattice_expressions += expression.split(",") 

154 else: 

155 atomic_expressions += expression.split(",") 

156 

157 fix_params.append( 

158 FixCartesianParametricRelations.from_expressions( 

159 list(range(3)), 

160 lattice_params, 

161 lattice_expressions, 

162 use_cell=True, 

163 ) 

164 ) 

165 

166 fix_params.append( 

167 FixScaledParametricRelations.from_expressions( 

168 list(range(len(atoms))), atomic_params, atomic_expressions 

169 ) 

170 ) 

171 

172 if any(magmoms): 

173 atoms.set_initial_magnetic_moments(magmoms) 

174 if any(charges): 

175 atoms.set_initial_charges(charges) 

176 

177 if periodic.any(): 

178 atoms.set_cell(cell) 

179 atoms.set_pbc(periodic) 

180 if len(fix): 

181 atoms.set_constraint([FixAtoms(indices=fix)] + fix_cart + fix_params) 

182 else: 

183 atoms.set_constraint(fix_cart + fix_params) 

184 

185 if fix_params and apply_constraints: 

186 atoms.set_positions(atoms.get_positions()) 

187 return atoms 

188 

189 

190@writer 

191def write_aims( 

192 fd, 

193 atoms, 

194 scaled=False, 

195 geo_constrain=False, 

196 velocities=False, 

197 ghosts=None, 

198 info_str=None, 

199 wrap=False, 

200): 

201 """Method to write FHI-aims geometry files. 

202 

203 Writes the atoms positions and constraints (only FixAtoms is 

204 supported at the moment). 

205 

206 Args: 

207 fd: file object 

208 File to output structure to 

209 atoms: ase.atoms.Atoms 

210 structure to output to the file 

211 scaled: bool 

212 If True use fractional coordinates instead of Cartesian coordinates 

213 symmetry_block: list of str 

214 List of geometric constraints as defined in: 

215 https://arxiv.org/abs/1908.01610 

216 velocities: bool 

217 If True add the atomic velocity vectors to the file 

218 ghosts: list of Atoms 

219 A list of ghost atoms for the system 

220 info_str: str 

221 A string to be added to the header of the file 

222 wrap: bool 

223 Wrap atom positions to cell before writing 

224 """ 

225 

226 from ase.constraints import FixAtoms, FixCartesian 

227 

228 import numpy as np 

229 

230 if geo_constrain: 

231 if not scaled: 

232 warnings.warn( 

233 "Setting scaled to True because a symmetry_block is detected." 

234 ) 

235 scaled = True 

236 

237 fd.write("#=======================================================\n") 

238 if hasattr(fd, 'name'): 

239 fd.write("# FHI-aims file: " + fd.name + "\n") 

240 fd.write("# Created using the Atomic Simulation Environment (ASE)\n") 

241 fd.write("# " + time.asctime() + "\n") 

242 

243 # If writing additional information is requested via info_str: 

244 if info_str is not None: 

245 fd.write("\n# Additional information:\n") 

246 if isinstance(info_str, list): 

247 fd.write("\n".join(["# {}".format(s) for s in info_str])) 

248 else: 

249 fd.write("# {}".format(info_str)) 

250 fd.write("\n") 

251 

252 fd.write("#=======================================================\n") 

253 

254 i = 0 

255 if atoms.get_pbc().any(): 

256 for n, vector in enumerate(atoms.get_cell()): 

257 fd.write("lattice_vector ") 

258 for i in range(3): 

259 fd.write("%16.16f " % vector[i]) 

260 fd.write("\n") 

261 fix_cart = np.zeros([len(atoms), 3]) 

262 

263 # else aims crashes anyways 

264 # better be more explicit 

265 # write_magmoms = np.any([a.magmom for a in atoms]) 

266 

267 if atoms.constraints: 

268 for constr in atoms.constraints: 

269 if isinstance(constr, FixAtoms): 

270 fix_cart[constr.index] = [1, 1, 1] 

271 elif isinstance(constr, FixCartesian): 

272 fix_cart[constr.a] = -constr.mask + 1 

273 

274 if ghosts is None: 

275 ghosts = np.zeros(len(atoms)) 

276 else: 

277 assert len(ghosts) == len(atoms) 

278 

279 if geo_constrain: 

280 wrap = False 

281 scaled_positions = atoms.get_scaled_positions(wrap=wrap) 

282 

283 for i, atom in enumerate(atoms): 

284 if ghosts[i] == 1: 

285 atomstring = "empty " 

286 elif scaled: 

287 atomstring = "atom_frac " 

288 else: 

289 atomstring = "atom " 

290 fd.write(atomstring) 

291 if scaled: 

292 for pos in scaled_positions[i]: 

293 fd.write("%16.16f " % pos) 

294 else: 

295 for pos in atom.position: 

296 fd.write("%16.16f " % pos) 

297 fd.write(atom.symbol) 

298 fd.write("\n") 

299 # (1) all coords are constrained: 

300 if fix_cart[i].all(): 

301 fd.write(" constrain_relaxation .true.\n") 

302 # (2) some coords are constrained: 

303 elif fix_cart[i].any(): 

304 xyz = fix_cart[i] 

305 for n in range(3): 

306 if xyz[n]: 

307 fd.write(" constrain_relaxation %s\n" % "xyz"[n]) 

308 if atom.charge: 

309 fd.write(" initial_charge %16.6f\n" % atom.charge) 

310 if atom.magmom: 

311 fd.write(" initial_moment %16.6f\n" % atom.magmom) 

312 

313 # Write velocities if this is wanted 

314 if velocities and atoms.get_velocities() is not None: 

315 fd.write( 

316 " velocity {:.16f} {:.16f} {:.16f}\n".format( 

317 *atoms.get_velocities()[i] / v_unit 

318 ) 

319 ) 

320 

321 if geo_constrain: 

322 for line in get_sym_block(atoms): 

323 fd.write(line) 

324 

325 

326def get_sym_block(atoms): 

327 """Get symmetry block for Parametric constraints in atoms.constraints""" 

328 import numpy as np 

329 

330 from ase.constraints import ( 

331 FixScaledParametricRelations, 

332 FixCartesianParametricRelations, 

333 ) 

334 

335 # Initialize param/expressions lists 

336 atomic_sym_params = [] 

337 lv_sym_params = [] 

338 atomic_param_constr = np.zeros((len(atoms),), dtype="<U100") 

339 lv_param_constr = np.zeros((3,), dtype="<U100") 

340 

341 # Populate param/expressions list 

342 for constr in atoms.constraints: 

343 if isinstance(constr, FixScaledParametricRelations): 

344 atomic_sym_params += constr.params 

345 

346 if np.any(atomic_param_constr[constr.indices] != ""): 

347 warnings.warn( 

348 "multiple parametric constraints defined for the same " 

349 "atom, using the last one defined" 

350 ) 

351 

352 atomic_param_constr[constr.indices] = [ 

353 ", ".join(expression) for expression in constr.expressions 

354 ] 

355 elif isinstance(constr, FixCartesianParametricRelations): 

356 lv_sym_params += constr.params 

357 

358 if np.any(lv_param_constr[constr.indices] != ""): 

359 warnings.warn( 

360 "multiple parametric constraints defined for the same " 

361 "lattice vector, using the last one defined" 

362 ) 

363 

364 lv_param_constr[constr.indices] = [ 

365 ", ".join(expression) for expression in constr.expressions 

366 ] 

367 

368 if np.all(atomic_param_constr == "") and np.all(lv_param_constr == ""): 

369 return [] 

370 

371 # Check Constraint Parameters 

372 if len(atomic_sym_params) != len(np.unique(atomic_sym_params)): 

373 warnings.warn( 

374 "Some parameters were used across constraints, they will be " 

375 "combined in the aims calculations" 

376 ) 

377 atomic_sym_params = np.unique(atomic_sym_params) 

378 

379 if len(lv_sym_params) != len(np.unique(lv_sym_params)): 

380 warnings.warn( 

381 "Some parameters were used across constraints, they will be " 

382 "combined in the aims calculations" 

383 ) 

384 lv_sym_params = np.unique(lv_sym_params) 

385 

386 if np.any(atomic_param_constr == ""): 

387 raise IOError( 

388 "FHI-aims input files require all atoms have defined parametric " 

389 "constraints" 

390 ) 

391 

392 cell_inds = np.where(lv_param_constr == "")[0] 

393 for ind in cell_inds: 

394 lv_param_constr[ind] = "{:.16f}, {:.16f}, {:.16f}".format( 

395 *atoms.cell[ind] 

396 ) 

397 

398 n_atomic_params = len(atomic_sym_params) 

399 n_lv_params = len(lv_sym_params) 

400 n_total_params = n_atomic_params + n_lv_params 

401 

402 sym_block = [] 

403 if n_total_params > 0: 

404 sym_block.append("#" + "="*55 + "\n") 

405 sym_block.append("# Parametric constraints\n") 

406 sym_block.append("#" + "="*55 + "\n") 

407 sym_block.append( 

408 "symmetry_n_params {:d} {:d} {:d}\n".format( 

409 n_total_params, n_lv_params, n_atomic_params 

410 ) 

411 ) 

412 sym_block.append( 

413 "symmetry_params %s\n" 

414 % " ".join(lv_sym_params + atomic_sym_params) 

415 ) 

416 

417 for constr in lv_param_constr: 

418 sym_block.append("symmetry_lv {:s}\n".format(constr)) 

419 

420 for constr in atomic_param_constr: 

421 sym_block.append("symmetry_frac {:s}\n".format(constr)) 

422 return sym_block 

423 

424 

425def _parse_atoms(fd, n_atoms, molecular_dynamics=False): 

426 """parse structure information from aims output to Atoms object""" 

427 from ase import Atoms, Atom 

428 

429 next(fd) 

430 atoms = Atoms() 

431 for i in range(n_atoms): 

432 inp = next(fd).split() 

433 if "lattice_vector" in inp[0]: 

434 cell = [] 

435 for i in range(3): 

436 cell += [[float(inp[1]), float(inp[2]), float(inp[3])]] 

437 inp = next(fd).split() 

438 atoms.set_cell(cell) 

439 inp = next(fd).split() 

440 atoms.append(Atom(inp[4], (inp[1], inp[2], inp[3]))) 

441 if molecular_dynamics: 

442 inp = next(fd).split() 

443 

444 return atoms 

445 

446 

447@reader 

448def read_aims_output(fd, index=-1): 

449 """Import FHI-aims output files with all data available, i.e. 

450 relaxations, MD information, force information etc etc etc.""" 

451 from ase import Atoms, Atom 

452 from ase.calculators.singlepoint import SinglePointCalculator 

453 from ase.constraints import FixAtoms, FixCartesian 

454 

455 molecular_dynamics = False 

456 cell = [] 

457 images = [] 

458 fix = [] 

459 fix_cart = [] 

460 f = None 

461 pbc = False 

462 found_aims_calculator = False 

463 stress = None 

464 for line in fd: 

465 # if "List of parameters used to initialize the calculator:" in line: 

466 # next(fd) 

467 # calc = read_aims_calculator(fd) 

468 # calc.out = filename 

469 # found_aims_calculator = True 

470 if "| Number of atoms :" in line: 

471 inp = line.split() 

472 n_atoms = int(inp[5]) 

473 if "| Unit cell:" in line: 

474 if not pbc: 

475 pbc = True 

476 for i in range(3): 

477 inp = next(fd).split() 

478 cell.append([inp[1], inp[2], inp[3]]) 

479 if "Found relaxation constraint for atom" in line: 

480 xyz = [0, 0, 0] 

481 ind = int(line.split()[5][:-1]) - 1 

482 if "All coordinates fixed" in line: 

483 if ind not in fix: 

484 fix.append(ind) 

485 if "coordinate fixed" in line: 

486 coord = line.split()[6] 

487 if coord == "x": 

488 xyz[0] = 1 

489 elif coord == "y": 

490 xyz[1] = 1 

491 elif coord == "z": 

492 xyz[2] = 1 

493 keep = True 

494 for n, c in enumerate(fix_cart): 

495 if ind == c.a: 

496 keep = False 

497 if keep: 

498 fix_cart.append(FixCartesian(ind, xyz)) 

499 else: 

500 fix_cart[n].mask[xyz.index(1)] = 0 

501 

502 if "Atomic structure:" in line and not molecular_dynamics: 

503 next(fd) 

504 atoms = Atoms() 

505 for _ in range(n_atoms): 

506 inp = next(fd).split() 

507 atoms.append(Atom(inp[3], (inp[4], inp[5], inp[6]))) 

508 

509 if "Complete information for previous time-step:" in line: 

510 molecular_dynamics = True 

511 

512 if "Updated atomic structure:" in line and not molecular_dynamics: 

513 atoms = _parse_atoms(fd, n_atoms=n_atoms) 

514 elif "Atomic structure (and velocities)" in line: 

515 next(fd) 

516 atoms = Atoms() 

517 velocities = [] 

518 for i in range(n_atoms): 

519 inp = next(fd).split() 

520 atoms.append(Atom(inp[4], (inp[1], inp[2], inp[3]))) 

521 inp = next(fd).split() 

522 floatvect = [v_unit * float(l) for l in inp[1:4]] 

523 velocities.append(floatvect) 

524 atoms.set_velocities(velocities) 

525 if len(fix): 

526 atoms.set_constraint([FixAtoms(indices=fix)] + fix_cart) 

527 else: 

528 atoms.set_constraint(fix_cart) 

529 images.append(atoms) 

530 

531 # if we enter here, the SocketIO/PIMD Wrapper was used 

532 elif ( 

533 "Atomic structure that " 

534 "was used in the preceding time step of the wrapper" 

535 ) in line: 

536 # parse atoms and add calculator information, i.e., the energies 

537 # and forces that were already collected 

538 atoms = _parse_atoms(fd, n_atoms=n_atoms) 

539 results = images[-1].calc.results 

540 atoms.calc = SinglePointCalculator(atoms, **results) 

541 

542 # replace last image with updated atoms 

543 images[-1] = atoms 

544 

545 # make sure `atoms` does not point to `images[-1` later on 

546 atoms = atoms.copy() 

547 

548 # FlK: add analytical stress and replace stress=None 

549 if "Analytical stress tensor - Symmetrized" in line: 

550 # scroll to significant lines 

551 for _ in range(4): 

552 next(fd) 

553 stress = [] 

554 for _ in range(3): 

555 inp = next(fd) 

556 stress.append([float(i) for i in inp.split()[2:5]]) 

557 

558 if "Total atomic forces" in line: 

559 f = [] 

560 for i in range(n_atoms): 

561 inp = next(fd).split() 

562 # FlK: use inp[-3:] instead of inp[1:4] to make sure this works 

563 # when atom number is not preceded by a space. 

564 f.append([float(i) for i in inp[-3:]]) 

565 if not found_aims_calculator: 

566 e = images[-1].get_potential_energy() 

567 # FlK: Add the stress if it has been computed 

568 if stress is None: 

569 calc = SinglePointCalculator(atoms, energy=e, forces=f) 

570 else: 

571 calc = SinglePointCalculator( 

572 atoms, energy=e, forces=f, stress=stress 

573 ) 

574 images[-1].calc = calc 

575 e = None 

576 f = None 

577 

578 if "Total energy corrected" in line: 

579 e = float(line.split()[5]) 

580 if pbc: 

581 atoms.set_cell(cell) 

582 atoms.pbc = True 

583 if not found_aims_calculator: 

584 atoms.calc = SinglePointCalculator(atoms, energy=e) 

585 if not molecular_dynamics: 

586 if len(fix): 

587 atoms.set_constraint([FixAtoms(indices=fix)] + fix_cart) 

588 else: 

589 atoms.set_constraint(fix_cart) 

590 images.append(atoms) 

591 e = None 

592 # if found_aims_calculator: 

593 # calc.set_results(images[-1]) 

594 # images[-1].calc = calc 

595 

596 # FlK: add stress per atom 

597 if "Per atom stress (eV) used for heat flux calculation" in line: 

598 # scroll to boundary 

599 next(l for l in fd if "-------------" in l) 

600 

601 stresses = [] 

602 for l in [next(fd) for _ in range(n_atoms)]: 

603 # Read stresses 

604 xx, yy, zz, xy, xz, yz = [float(d) for d in l.split()[2:8]] 

605 stresses.append([[xx, xy, xz], [xy, yy, yz], [xz, yz, zz]]) 

606 

607 if not found_aims_calculator: 

608 e = images[-1].get_potential_energy() 

609 f = images[-1].get_forces() 

610 stress = images[-1].get_stress(voigt=False) 

611 

612 calc = SinglePointCalculator( 

613 atoms, energy=e, forces=f, stress=stress, stresses=stresses 

614 ) 

615 images[-1].calc = calc 

616 

617 fd.close() 

618 if molecular_dynamics: 

619 images = images[1:] 

620 

621 # return requested images, code borrowed from ase/io/trajectory.py 

622 if isinstance(index, int): 

623 return images[index] 

624 else: 

625 step = index.step or 1 

626 if step > 0: 

627 start = index.start or 0 

628 if start < 0: 

629 start += len(images) 

630 stop = index.stop or len(images) 

631 if stop < 0: 

632 stop += len(images) 

633 else: 

634 if index.start is None: 

635 start = len(images) - 1 

636 else: 

637 start = index.start 

638 if start < 0: 

639 start += len(images) 

640 if index.stop is None: 

641 stop = -1 

642 else: 

643 stop = index.stop 

644 if stop < 0: 

645 stop += len(images) 

646 return [images[i] for i in range(start, stop, step)]