Coverage for /builds/kinetik161/ase/ase/io/magres.py: 84.52%

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"""

8import re

9from collections import OrderedDict

11import numpy as np

13import ase.units

14from ase.atoms import Atoms

15from ase.calculators.singlepoint import SinglePointDFTCalculator

16from ase.spacegroup import Spacegroup

17from ase.spacegroup.spacegroup import SpacegroupNotFoundError

19_mprops = {

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

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

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

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

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

27def read_magres(fd, include_unrecognised=False):

28 """

29 Reader function for magres files.

30 """

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

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

35 """

36 Here are defined the various functions required to parse

37 different blocks.

38 """

40 def tensor33(x):

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

43 def tensor31(x):

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

46 def get_version(file_contents):

47 """

48 Look for and parse the magres file format version line

49 """

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

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

54 if match:

55 version = match.groups()

56 version = tuple(vnum for vnum in version)

57 else:

58 version = None

60 return version

62 def parse_blocks(file_contents):

63 """

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

65 (block_name, contents) tuples

66 """

68 blocks = blocks_re.findall(file_contents)

70 return blocks

72 def parse_block(block):

73 """

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

75 """

77 def clean_line(line):

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

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

80 line = line.strip()

82 return line

84 name, data = block

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

88 records = []

90 for line in lines:

91 xs = line.split()

93 if len(xs) > 0:

94 tag = xs[0]

95 data = xs[1:]

97 records.append((tag, data))

99 return (name, records)

100

101 def check_units(d):

102 """

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

104 """

105

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

107 'atom': 'Angstrom',

108 'ms': 'ppm',

109 'efg': 'au',

110 'efg_local': 'au',

111 'efg_nonlocal': 'au',

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

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

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

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

116 'isc_spin': '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(f'Unrecognized units: {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 # Solve incompatibility for atomic indices above 100

244 pattern_ms = r'(ms [a-zA-Z]{1,2})(\d+)'

245 file_contents = re.sub(pattern_ms, r'\g<1> \g<2>', file_contents)

246 pattern_efg = r'(efg [a-zA-Z]{1,2})(\d+)'

247 file_contents = re.sub(pattern_efg, r'\g<1> \g<2>', file_contents)

248

249 # This works as a validity check

250 version = get_version(file_contents)

251 if version is None:

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

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

254 blocks = parse_blocks(file_contents)

255

256 data_dict = {}

257

258 for block_data in blocks:

259 block = parse_block(block_data)

260

261 if block[0] in block_parsers:

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

263 data_dict[block[0]] = block_dict

264 else:

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

266 if include_unrecognised:

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

268

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

270

271 # First check if the file is even viable

272 if 'atoms' not in data_dict:

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

274

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

276 # could turn out useful in the future

277

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

279

280 # Lattice parameters?

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

282 try:

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

284 except KeyError:

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

286 u = magres_units[u]

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

288 pbc = True

289 else:

290 cell = None

291 pbc = False

292

293 # Now the atoms

294 symbols = []

295 positions = []

296 indices = []

297 labels = []

298

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

300 try:

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

302 except KeyError:

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

304 u = magres_units[u]

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

306 # 'custom' species amongst the symbols

307 custom_species = None

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

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

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

311 # Add it to the custom info!

312 custom_species = list(symbols)

313 symbols.append(spec_custom[0])

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

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

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

317 if custom_species is not None:

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

319

320 atoms = Atoms(cell=cell,

321 pbc=pbc,

322 symbols=symbols,

323 positions=positions)

324

325 # Add custom species if present

326 if custom_species is not None:

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

328

329 # Add the spacegroup, if present and recognizable

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

331 try:

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

333 except SpacegroupNotFoundError:

334 # Not found

335 spg = Spacegroup(1) # Most generic one

336 atoms.info['spacegroup'] = spg

337

338 # Set up the rest of the properties as arrays

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

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

341

342 # Now for the magres specific stuff

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

344

345 def create_magres_array(name, order, block):

346

347 if order == 1:

348 u_arr = [None] * len(li_list)

349 elif order == 2:

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

351 else:

352 raise ValueError(

353 'Invalid order value passed to create_magres_array')

354

355 for s in block:

356 # Find the atom index/indices

357 if order == 1:

358 # First find out which atom this is

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

360 try:

361 ai = li_list.index(at)

362 except ValueError:

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

364 # Then add the relevant quantity

365 u_arr[ai] = s[mn]

366 else:

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

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

369 ai1 = li_list.index(at1)

370 ai2 = li_list.index(at2)

371 # Sort them

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

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

374

375 if order == 1:

376 return np.array(u_arr)

377 else:

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

379

380 if 'magres' in data_dict:

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

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

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

384 # This bit to keep track of tags

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

386

387 if u0 not in _mprops:

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

389

390 mn, order = _mprops[u0]

391

392 if order > 0:

393 u_arr = create_magres_array(mn, order,

394 data_dict['magres'][u])

395 atoms.new_array(u, u_arr)

396 else:

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

398 # {})

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

400 # # there should be only that

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

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

403 if atoms.calc is None:

404 calc = SinglePointDFTCalculator(atoms)

405 atoms.calc = calc

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

407

408 if 'calculation' in data_dict:

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

410

411 if include_unrecognised:

412 for b in data_dict:

413 if b not in block_parsers:

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

415

416 return atoms

417

418

419def tensor_string(tensor):

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

421

422

423def write_magres(fd, image):

424 """

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

426 into structures, then dumped to the actual file

427 """

428

429 image_data = {}

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

431 # Contains units, lattice and each individual atom

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

433 # Has lattice!

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

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

436

437 # Now for the atoms

438 if image.has('labels'):

439 labels = image.get_array('labels')

440 else:

441 labels = image.get_chemical_symbols()

442

443 if image.has('indices'):

444 indices = image.get_array('indices')

445 else:

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

447

448 # Iterate over atoms

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

450 if image.has('castep_custom_species')

451 else image.get_chemical_symbols())

452

453 atom_info = list(zip(symbols,

454 image.get_positions()))

455 if len(atom_info) > 0:

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

457 image_data['atoms']['atom'] = []

458

459 for i, a in enumerate(atom_info):

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

461 'index': indices[i],

462 'position': a[1],

463 'species': a[0],

464 'label': labels[i]})

465

466 # Spacegroup, if present

467 if 'spacegroup' in image.info:

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

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

470

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

472 if 'magres_units' in image.info:

473

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

475

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

477 # Get the type

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

479 if p in _mprops:

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

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

482 image_data['magres'][u] = []

483 mn, order = _mprops[p]

484

485 if order == 0:

486 # The case of susceptibility

487 tens = {

488 mn: image.calc.results[u]

489 }

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

491 else:

492 arr = image.get_array(u)

493 li_tab = zip(labels, indices)

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

495 if order == 2:

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

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

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

499 'atom1': {'label': lab,

500 'index': ind},

501 'atom2': {'label': lab2,

502 'index': ind2}}

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

504 elif order == 1:

505 if arr[i] is not None:

506 tens = {mn: arr[i],

507 'atom': {'label': lab,

508 'index': ind}}

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

510

511 # Calculation block, if present

512 if 'magresblock_calculation' in image.info:

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

514

515 def write_units(data, out):

516 if 'units' in data:

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

518 out.append(f' units {tag} {units}')

519

520 def write_magres_block(data):

521 """

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

523 """

524

525 out = []

526

527 def nout(tag, tensor_name):

528 if tag in data:

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

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

531

532 def siout(tag, tensor_name):

533 if tag in data:

534 for atom_si in data[tag]:

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

536 '%s') % (tag,

537 atom_si['atom']['label'],

538 atom_si['atom']['index'],

539 tensor_string(atom_si[tensor_name])))

540

541 write_units(data, out)

542

543 nout('sus', 'S')

544 siout('ms', 'sigma')

545

546 siout('efg_local', 'V')

547 siout('efg_nonlocal', 'V')

548 siout('efg', 'V')

549

550 def sisiout(tag, tensor_name):

551 if tag in data:

552 for isc in data[tag]:

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

554 '%s') % (tag,

555 isc['atom1']['label'],

556 isc['atom1']['index'],

557 isc['atom2']['label'],

558 isc['atom2']['index'],

559 tensor_string(isc[tensor_name])))

560

561 sisiout('isc_fc', 'K')

562 sisiout('isc_orbital_p', 'K')

563 sisiout('isc_orbital_d', 'K')

564 sisiout('isc_spin', 'K')

565 sisiout('isc', 'K')

566

567 return '\n'.join(out)

568

569 def write_atoms_block(data):

570 out = []

571

572 write_units(data, out)

573

574 if 'lattice' in data:

575 for lat in data['lattice']:

576 out.append(f" lattice {tensor_string(lat)}")

577

578 if 'symmetry' in data:

579 for sym in data['symmetry']:

580 out.append(f' symmetry {sym}')

581

582 if 'atom' in data:

583 for a in data['atom']:

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

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

586 a['label'],

587 a['index'],

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

589

590 return '\n'.join(out)

591

592 def write_generic_block(data):

593 out = []

594

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

596 for value in data:

597 out.append('{} {}'.format(tag, ' '.join(str(x) for x in value)))

598

599 return '\n'.join(out)

600

601 # Using this to preserve order

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

603 ('atoms', write_atoms_block),

604 ('magres', write_magres_block)])

605

606 # First, write the header

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

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

609

610 for b in block_writers:

611 if b in image_data:

612 fd.write(f'[{b}]\n')

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

614 fd.write(f'\n[/{b}]\n')

615

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

617 for i in image.info:

618 if '_' in i:

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

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

621 fd.write(f'[{b}]\n')

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

623 fd.write(f'[/{b}]\n')