1"""This module provides I/O functions for the MAGRES file format, introduced 

2by CASTEP as an output format to store structural data and ab-initio 

3calculated NMR parameters. 

4Authors: Simone Sturniolo (ase implementation), Tim Green (original magres 

5 parser code) 

6""" 

7 

8import re 

9import numpy as np 

10from collections import OrderedDict 

11 

12import ase.units 

13from ase.atoms import Atoms 

14from ase.spacegroup import Spacegroup 

15from ase.spacegroup.spacegroup import SpacegroupNotFoundError 

16from ase.calculators.singlepoint import SinglePointDFTCalculator 

17 

18_mprops = { 

19 'ms': ('sigma', 1), 

20 'sus': ('S', 0), 

21 'efg': ('V', 1), 

22 'isc': ('K', 2)} 

23# (matrix name, number of atoms in interaction) for various magres quantities 

24 

25 

26def read_magres(fd, include_unrecognised=False): 

27 """ 

28 Reader function for magres files. 

29 """ 

30 

31 blocks_re = re.compile(r'[\[<](?P<block_name>.*?)[>\]](.*?)[<\[]/' + 

32 r'(?P=block_name)[\]>]', re.M | re.S) 

33 

34 """ 

35 Here are defined the various functions required to parse 

36 different blocks. 

37 """ 

38 

39 def tensor33(x): 

40 return np.squeeze(np.reshape(x, (3, 3))).tolist() 

41 

42 def tensor31(x): 

43 return np.squeeze(np.reshape(x, (3, 1))).tolist() 

44 

45 def get_version(file_contents): 

46 """ 

47 Look for and parse the magres file format version line 

48 """ 

49 

50 lines = file_contents.split('\n') 

51 match = re.match(r'\#\$magres-abinitio-v([0-9]+).([0-9]+)', lines[0]) 

52 

53 if match: 

54 version = match.groups() 

55 version = tuple(vnum for vnum in version) 

56 else: 

57 version = None 

58 

59 return version 

60 

61 def parse_blocks(file_contents): 

62 """ 

63 Parse series of XML-like deliminated blocks into a list of 

64 (block_name, contents) tuples 

65 """ 

66 

67 blocks = blocks_re.findall(file_contents) 

68 

69 return blocks 

70 

71 def parse_block(block): 

72 """ 

73 Parse block contents into a series of (tag, data) records 

74 """ 

75 

76 def clean_line(line): 

77 # Remove comments and whitespace at start and ends of line 

78 line = re.sub('#(.*?)\n', '', line) 

79 line = line.strip() 

80 

81 return line 

82 

83 name, data = block 

84 

85 lines = [clean_line(line) for line in data.split('\n')] 

86 

87 records = [] 

88 

89 for line in lines: 

90 xs = line.split() 

91 

92 if len(xs) > 0: 

93 tag = xs[0] 

94 data = xs[1:] 

95 

96 records.append((tag, data)) 

97 

98 return (name, records) 

99 

100 def check_units(d): 

101 """ 

102 Verify that given units for a particular tag are correct. 

103 """ 

104 

105 allowed_units = {'lattice': 'Angstrom', 

106 'atom': 'Angstrom', 

107 'ms': 'ppm', 

108 'efg': 'au', 

109 'efg_local': 'au', 

110 'efg_nonlocal': 'au', 

111 'isc': '10^19.T^2.J^-1', 

112 'isc_fc': '10^19.T^2.J^-1', 

113 'isc_orbital_p': '10^19.T^2.J^-1', 

114 'isc_orbital_d': '10^19.T^2.J^-1', 

115 'isc_spin': '10^19.T^2.J^-1', 

116 'isc': '10^19.T^2.J^-1', 

117 'sus': '10^-6.cm^3.mol^-1', 

118 'calc_cutoffenergy': 'Hartree', } 

119 

120 if d[0] in d and d[1] == allowed_units[d[0]]: 

121 pass 

122 else: 

123 raise RuntimeError('Unrecognized units: %s %s' % (d[0], d[1])) 

124 

125 return d 

126 

127 def parse_magres_block(block): 

128 """ 

129 Parse magres block into data dictionary given list of record 

130 tuples. 

131 """ 

132 

133 name, records = block 

134 

135 # 3x3 tensor 

136 def ntensor33(name): 

137 return lambda d: {name: tensor33([float(x) for x in data])} 

138 

139 # Atom label, atom index and 3x3 tensor 

140 def sitensor33(name): 

141 return lambda d: {'atom': {'label': data[0], 

142 'index': int(data[1])}, 

143 name: tensor33([float(x) for x in data[2:]])} 

144 

145 # 2x(Atom label, atom index) and 3x3 tensor 

146 def sisitensor33(name): 

147 return lambda d: {'atom1': {'label': data[0], 

148 'index': int(data[1])}, 

149 'atom2': {'label': data[2], 

150 'index': int(data[3])}, 

151 name: tensor33([float(x) for x in data[4:]])} 

152 

153 tags = {'ms': sitensor33('sigma'), 

154 'sus': ntensor33('S'), 

155 'efg': sitensor33('V'), 

156 'efg_local': sitensor33('V'), 

157 'efg_nonlocal': sitensor33('V'), 

158 'isc': sisitensor33('K'), 

159 'isc_fc': sisitensor33('K'), 

160 'isc_spin': sisitensor33('K'), 

161 'isc_orbital_p': sisitensor33('K'), 

162 'isc_orbital_d': sisitensor33('K'), 

163 'units': check_units} 

164 

165 data_dict = {} 

166 

167 for record in records: 

168 tag, data = record 

169 

170 if tag not in data_dict: 

171 data_dict[tag] = [] 

172 

173 data_dict[tag].append(tags[tag](data)) 

174 

175 return data_dict 

176 

177 def parse_atoms_block(block): 

178 """ 

179 Parse atoms block into data dictionary given list of record tuples. 

180 """ 

181 

182 name, records = block 

183 

184 # Lattice record: a1, a2 a3, b1, b2, b3, c1, c2 c3 

185 def lattice(d): 

186 return tensor33([float(x) for x in data]) 

187 

188 # Atom record: label, index, x, y, z 

189 def atom(d): 

190 return {'species': data[0], 

191 'label': data[1], 

192 'index': int(data[2]), 

193 'position': tensor31([float(x) for x in data[3:]])} 

194 

195 def symmetry(d): 

196 return ' '.join(data) 

197 

198 tags = {'lattice': lattice, 

199 'atom': atom, 

200 'units': check_units, 

201 'symmetry': symmetry} 

202 

203 data_dict = {} 

204 

205 for record in records: 

206 tag, data = record 

207 if tag not in data_dict: 

208 data_dict[tag] = [] 

209 data_dict[tag].append(tags[tag](data)) 

210 

211 return data_dict 

212 

213 def parse_generic_block(block): 

214 """ 

215 Parse any other block into data dictionary given list of record 

216 tuples. 

217 """ 

218 

219 name, records = block 

220 

221 data_dict = {} 

222 

223 for record in records: 

224 tag, data = record 

225 

226 if tag not in data_dict: 

227 data_dict[tag] = [] 

228 

229 data_dict[tag].append(data) 

230 

231 return data_dict 

232 

233 """ 

234 Actual parser code. 

235 """ 

236 

237 block_parsers = {'magres': parse_magres_block, 

238 'atoms': parse_atoms_block, 

239 'calculation': parse_generic_block, } 

240 

241 file_contents = fd.read() 

242 

243 # This works as a validity check 

244 version = get_version(file_contents) 

245 if version is None: 

246 # This isn't even a .magres file! 

247 raise RuntimeError('File is not in standard Magres format') 

248 blocks = parse_blocks(file_contents) 

249 

250 data_dict = {} 

251 

252 for block_data in blocks: 

253 block = parse_block(block_data) 

254 

255 if block[0] in block_parsers: 

256 block_dict = block_parsers[block[0]](block) 

257 data_dict[block[0]] = block_dict 

258 else: 

259 # Throw in the text content of blocks we don't recognise 

260 if include_unrecognised: 

261 data_dict[block[0]] = block_data[1] 

262 

263 # Now the loaded data must be turned into an ASE Atoms object 

264 

265 # First check if the file is even viable 

266 if 'atoms' not in data_dict: 

267 raise RuntimeError('Magres file does not contain structure data') 

268 

269 # Allowed units handling. This is redundant for now but 

270 # could turn out useful in the future 

271 

272 magres_units = {'Angstrom': ase.units.Ang} 

273 

274 # Lattice parameters? 

275 if 'lattice' in data_dict['atoms']: 

276 try: 

277 u = dict(data_dict['atoms']['units'])['lattice'] 

278 except KeyError: 

279 raise RuntimeError('No units detected in file for lattice') 

280 u = magres_units[u] 

281 cell = np.array(data_dict['atoms']['lattice'][0]) * u 

282 pbc = True 

283 else: 

284 cell = None 

285 pbc = False 

286 

287 # Now the atoms 

288 symbols = [] 

289 positions = [] 

290 indices = [] 

291 labels = [] 

292 

293 if 'atom' in data_dict['atoms']: 

294 try: 

295 u = dict(data_dict['atoms']['units'])['atom'] 

296 except KeyError: 

297 raise RuntimeError('No units detected in file for atom positions') 

298 u = magres_units[u] 

299 # Now we have to account for the possibility of there being CASTEP 

300 # 'custom' species amongst the symbols 

301 custom_species = None 

302 for a in data_dict['atoms']['atom']: 

303 spec_custom = a['species'].split(':', 1) 

304 if len(spec_custom) > 1 and custom_species is None: 

305 # Add it to the custom info! 

306 custom_species = list(symbols) 

307 symbols.append(spec_custom[0]) 

308 positions.append(a['position']) 

309 indices.append(a['index']) 

310 labels.append(a['label']) 

311 if custom_species is not None: 

312 custom_species.append(a['species']) 

313 

314 atoms = Atoms(cell=cell, 

315 pbc=pbc, 

316 symbols=symbols, 

317 positions=positions) 

318 

319 # Add custom species if present 

320 if custom_species is not None: 

321 atoms.new_array('castep_custom_species', np.array(custom_species)) 

322 

323 # Add the spacegroup, if present and recognizable 

324 if 'symmetry' in data_dict['atoms']: 

325 try: 

326 spg = Spacegroup(data_dict['atoms']['symmetry'][0]) 

327 except SpacegroupNotFoundError: 

328 # Not found 

329 spg = Spacegroup(1) # Most generic one 

330 atoms.info['spacegroup'] = spg 

331 

332 # Set up the rest of the properties as arrays 

333 atoms.new_array('indices', np.array(indices)) 

334 atoms.new_array('labels', np.array(labels)) 

335 

336 # Now for the magres specific stuff 

337 li_list = list(zip(labels, indices)) 

338 

339 def create_magres_array(name, order, block): 

340 

341 if order == 1: 

342 u_arr = [None] * len(li_list) 

343 elif order == 2: 

344 u_arr = [[None] * (i + 1) for i in range(len(li_list))] 

345 else: 

346 raise ValueError( 

347 'Invalid order value passed to create_magres_array') 

348 

349 for s in block: 

350 # Find the atom index/indices 

351 if order == 1: 

352 # First find out which atom this is 

353 at = (s['atom']['label'], s['atom']['index']) 

354 try: 

355 ai = li_list.index(at) 

356 except ValueError: 

357 raise RuntimeError('Invalid data in magres block') 

358 # Then add the relevant quantity 

359 u_arr[ai] = s[mn] 

360 else: 

361 at1 = (s['atom1']['label'], s['atom1']['index']) 

362 at2 = (s['atom2']['label'], s['atom2']['index']) 

363 ai1 = li_list.index(at1) 

364 ai2 = li_list.index(at2) 

365 # Sort them 

366 ai1, ai2 = sorted((ai1, ai2), reverse=True) 

367 u_arr[ai1][ai2] = s[mn] 

368 

369 if order == 1: 

370 return np.array(u_arr) 

371 else: 

372 return np.array(u_arr, dtype=object) 

373 

374 if 'magres' in data_dict: 

375 if 'units' in data_dict['magres']: 

376 atoms.info['magres_units'] = dict(data_dict['magres']['units']) 

377 for u in atoms.info['magres_units']: 

378 # This bit to keep track of tags 

379 u0 = u.split('_')[0] 

380 

381 if u0 not in _mprops: 

382 raise RuntimeError('Invalid data in magres block') 

383 

384 mn, order = _mprops[u0] 

385 

386 if order > 0: 

387 u_arr = create_magres_array(mn, order, 

388 data_dict['magres'][u]) 

389 atoms.new_array(u, u_arr) 

390 else: 

391 # atoms.info['magres_data'] = atoms.info.get('magres_data', 

392 # {}) 

393 # # We only take element 0 because for this sort of data 

394 # # there should be only that 

395 # atoms.info['magres_data'][u] = \ 

396 # data_dict['magres'][u][0][mn] 

397 if atoms.calc is None: 

398 calc = SinglePointDFTCalculator(atoms) 

399 atoms.calc = calc 

400 atoms.calc.results[u] = data_dict['magres'][u][0][mn] 

401 

402 if 'calculation' in data_dict: 

403 atoms.info['magresblock_calculation'] = data_dict['calculation'] 

404 

405 if include_unrecognised: 

406 for b in data_dict: 

407 if b not in block_parsers: 

408 atoms.info['magresblock_' + b] = data_dict[b] 

409 

410 return atoms 

411 

412 

413def tensor_string(tensor): 

414 return ' '.join(' '.join(str(x) for x in xs) for xs in tensor) 

415 

416 

417def write_magres(fd, image): 

418 """ 

419 A writing function for magres files. Two steps: first data are arranged 

420 into structures, then dumped to the actual file 

421 """ 

422 

423 image_data = {} 

424 image_data['atoms'] = {'units': []} 

425 # Contains units, lattice and each individual atom 

426 if np.all(image.get_pbc()): 

427 # Has lattice! 

428 image_data['atoms']['units'].append(['lattice', 'Angstrom']) 

429 image_data['atoms']['lattice'] = [image.get_cell()] 

430 

431 # Now for the atoms 

432 if image.has('labels'): 

433 labels = image.get_array('labels') 

434 else: 

435 labels = image.get_chemical_symbols() 

436 

437 if image.has('indices'): 

438 indices = image.get_array('indices') 

439 else: 

440 indices = [labels[:i + 1].count(labels[i]) for i in range(len(labels))] 

441 

442 # Iterate over atoms 

443 symbols = (image.get_array('castep_custom_species') 

444 if image.has('castep_custom_species') 

445 else image.get_chemical_symbols()) 

446 

447 atom_info = list(zip(symbols, 

448 image.get_positions())) 

449 if len(atom_info) > 0: 

450 image_data['atoms']['units'].append(['atom', 'Angstrom']) 

451 image_data['atoms']['atom'] = [] 

452 

453 for i, a in enumerate(atom_info): 

454 image_data['atoms']['atom'].append({ 

455 'index': indices[i], 

456 'position': a[1], 

457 'species': a[0], 

458 'label': labels[i]}) 

459 

460 # Spacegroup, if present 

461 if 'spacegroup' in image.info: 

462 image_data['atoms']['symmetry'] = [image.info['spacegroup'] 

463 .symbol.replace(' ', '')] 

464 

465 # Now go on to do the same for magres information 

466 if 'magres_units' in image.info: 

467 

468 image_data['magres'] = {'units': []} 

469 

470 for u in image.info['magres_units']: 

471 # Get the type 

472 p = u.split('_')[0] 

473 if p in _mprops: 

474 image_data['magres']['units'].append( 

475 [u, image.info['magres_units'][u]]) 

476 image_data['magres'][u] = [] 

477 mn, order = _mprops[p] 

478 

479 if order == 0: 

480 # The case of susceptibility 

481 tens = { 

482 mn: image.calc.results[u] 

483 } 

484 image_data['magres'][u] = tens 

485 else: 

486 arr = image.get_array(u) 

487 li_tab = zip(labels, indices) 

488 for i, (lab, ind) in enumerate(li_tab): 

489 if order == 2: 

490 for j, (lab2, ind2) in enumerate(li_tab[:i + 1]): 

491 if arr[i][j] is not None: 

492 tens = {mn: arr[i][j], 

493 'atom1': {'label': lab, 

494 'index': ind}, 

495 'atom2': {'label': lab2, 

496 'index': ind2}} 

497 image_data['magres'][u].append(tens) 

498 elif order == 1: 

499 if arr[i] is not None: 

500 tens = {mn: arr[i], 

501 'atom': {'label': lab, 

502 'index': ind}} 

503 image_data['magres'][u].append(tens) 

504 

505 # Calculation block, if present 

506 if 'magresblock_calculation' in image.info: 

507 image_data['calculation'] = image.info['magresblock_calculation'] 

508 

509 def write_units(data, out): 

510 if 'units' in data: 

511 for tag, units in data['units']: 

512 out.append(' units %s %s' % (tag, units)) 

513 

514 def write_magres_block(data): 

515 """ 

516 Write out a <magres> block from its dictionary representation 

517 """ 

518 

519 out = [] 

520 

521 def nout(tag, tensor_name): 

522 if tag in data: 

523 out.append(' '.join([' ', tag, 

524 tensor_string(data[tag][tensor_name])])) 

525 

526 def siout(tag, tensor_name): 

527 if tag in data: 

528 for atom_si in data[tag]: 

529 out.append((' %s %s %d ' 

530 '%s') % (tag, 

531 atom_si['atom']['label'], 

532 atom_si['atom']['index'], 

533 tensor_string(atom_si[tensor_name]))) 

534 

535 write_units(data, out) 

536 

537 nout('sus', 'S') 

538 siout('ms', 'sigma') 

539 

540 siout('efg_local', 'V') 

541 siout('efg_nonlocal', 'V') 

542 siout('efg', 'V') 

543 

544 def sisiout(tag, tensor_name): 

545 if tag in data: 

546 for isc in data[tag]: 

547 out.append((' %s %s %d %s %d ' 

548 '%s') % (tag, 

549 isc['atom1']['label'], 

550 isc['atom1']['index'], 

551 isc['atom2']['label'], 

552 isc['atom2']['index'], 

553 tensor_string(isc[tensor_name]))) 

554 

555 sisiout('isc_fc', 'K') 

556 sisiout('isc_orbital_p', 'K') 

557 sisiout('isc_orbital_d', 'K') 

558 sisiout('isc_spin', 'K') 

559 sisiout('isc', 'K') 

560 

561 return '\n'.join(out) 

562 

563 def write_atoms_block(data): 

564 out = [] 

565 

566 write_units(data, out) 

567 

568 if 'lattice' in data: 

569 for lat in data['lattice']: 

570 out.append(" lattice %s" % tensor_string(lat)) 

571 

572 if 'symmetry' in data: 

573 for sym in data['symmetry']: 

574 out.append(' symmetry %s' % sym) 

575 

576 if 'atom' in data: 

577 for a in data['atom']: 

578 out.append((' atom %s %s %s ' 

579 '%s') % (a['species'], 

580 a['label'], 

581 a['index'], 

582 ' '.join(str(x) for x in a['position']))) 

583 

584 return '\n'.join(out) 

585 

586 def write_generic_block(data): 

587 out = [] 

588 

589 for tag, data in data.items(): 

590 for value in data: 

591 out.append('%s %s' % (tag, ' '.join(str(x) for x in value))) 

592 

593 return '\n'.join(out) 

594 

595 # Using this to preserve order 

596 block_writers = OrderedDict([('calculation', write_generic_block), 

597 ('atoms', write_atoms_block), 

598 ('magres', write_magres_block)]) 

599 

600 # First, write the header 

601 fd.write('#$magres-abinitio-v1.0\n') 

602 fd.write('# Generated by the Atomic Simulation Environment library\n') 

603 

604 for b in block_writers: 

605 if b in image_data: 

606 fd.write('[{0}]\n'.format(b)) 

607 fd.write(block_writers[b](image_data[b])) 

608 fd.write('\n[/{0}]\n'.format(b)) 

609 

610 # Now on to check for any non-standard blocks... 

611 for i in image.info: 

612 if '_' in i: 

613 ismag, b = i.split('_', 1) 

614 if ismag == 'magresblock' and b not in block_writers: 

615 fd.write('[{0}]\n'.format(b)) 

616 fd.write(image.info[i]) 

617 fd.write('[/{0}]\n'.format(b))