"""This module provides I/O functions for the MAGRES file format, introduced
by CASTEP as an output format to store structural data and ab-initio
calculated NMR parameters.
Authors: Simone Sturniolo (ase implementation), Tim Green (original magres
parser code)
"""
import re
import numpy as np
from collections import OrderedDict
import ase.units
from ase.atoms import Atoms
from ase.spacegroup import Spacegroup
from ase.spacegroup.spacegroup import SpacegroupNotFoundError
from ase.calculators.singlepoint import SinglePointDFTCalculator
_mprops = {
'ms': ('sigma', 1),
'sus': ('S', 0),
'efg': ('V', 1),
'isc': ('K', 2)}
# (matrix name, number of atoms in interaction) for various magres quantities
[docs]def read_magres(fd, include_unrecognised=False):
"""
Reader function for magres files.
"""
blocks_re = re.compile(r'[\[<](?P<block_name>.*?)[>\]](.*?)[<\[]/' +
r'(?P=block_name)[\]>]', re.M | re.S)
"""
Here are defined the various functions required to parse
different blocks.
"""
def tensor33(x):
return np.squeeze(np.reshape(x, (3, 3))).tolist()
def tensor31(x):
return np.squeeze(np.reshape(x, (3, 1))).tolist()
def get_version(file_contents):
"""
Look for and parse the magres file format version line
"""
lines = file_contents.split('\n')
match = re.match(r'\#\$magres-abinitio-v([0-9]+).([0-9]+)', lines[0])
if match:
version = match.groups()
version = tuple(vnum for vnum in version)
else:
version = None
return version
def parse_blocks(file_contents):
"""
Parse series of XML-like deliminated blocks into a list of
(block_name, contents) tuples
"""
blocks = blocks_re.findall(file_contents)
return blocks
def parse_block(block):
"""
Parse block contents into a series of (tag, data) records
"""
def clean_line(line):
# Remove comments and whitespace at start and ends of line
line = re.sub('#(.*?)\n', '', line)
line = line.strip()
return line
name, data = block
lines = [clean_line(line) for line in data.split('\n')]
records = []
for line in lines:
xs = line.split()
if len(xs) > 0:
tag = xs[0]
data = xs[1:]
records.append((tag, data))
return (name, records)
def check_units(d):
"""
Verify that given units for a particular tag are correct.
"""
allowed_units = {'lattice': 'Angstrom',
'atom': 'Angstrom',
'ms': 'ppm',
'efg': 'au',
'efg_local': 'au',
'efg_nonlocal': 'au',
'isc': '10^19.T^2.J^-1',
'isc_fc': '10^19.T^2.J^-1',
'isc_orbital_p': '10^19.T^2.J^-1',
'isc_orbital_d': '10^19.T^2.J^-1',
'isc_spin': '10^19.T^2.J^-1',
'isc': '10^19.T^2.J^-1',
'sus': '10^-6.cm^3.mol^-1',
'calc_cutoffenergy': 'Hartree', }
if d[0] in d and d[1] == allowed_units[d[0]]:
pass
else:
raise RuntimeError('Unrecognized units: %s %s' % (d[0], d[1]))
return d
def parse_magres_block(block):
"""
Parse magres block into data dictionary given list of record
tuples.
"""
name, records = block
# 3x3 tensor
def ntensor33(name):
return lambda d: {name: tensor33([float(x) for x in data])}
# Atom label, atom index and 3x3 tensor
def sitensor33(name):
return lambda d: {'atom': {'label': data[0],
'index': int(data[1])},
name: tensor33([float(x) for x in data[2:]])}
# 2x(Atom label, atom index) and 3x3 tensor
def sisitensor33(name):
return lambda d: {'atom1': {'label': data[0],
'index': int(data[1])},
'atom2': {'label': data[2],
'index': int(data[3])},
name: tensor33([float(x) for x in data[4:]])}
tags = {'ms': sitensor33('sigma'),
'sus': ntensor33('S'),
'efg': sitensor33('V'),
'efg_local': sitensor33('V'),
'efg_nonlocal': sitensor33('V'),
'isc': sisitensor33('K'),
'isc_fc': sisitensor33('K'),
'isc_spin': sisitensor33('K'),
'isc_orbital_p': sisitensor33('K'),
'isc_orbital_d': sisitensor33('K'),
'units': check_units}
data_dict = {}
for record in records:
tag, data = record
if tag not in data_dict:
data_dict[tag] = []
data_dict[tag].append(tags[tag](data))
return data_dict
def parse_atoms_block(block):
"""
Parse atoms block into data dictionary given list of record tuples.
"""
name, records = block
# Lattice record: a1, a2 a3, b1, b2, b3, c1, c2 c3
def lattice(d):
return tensor33([float(x) for x in data])
# Atom record: label, index, x, y, z
def atom(d):
return {'species': data[0],
'label': data[1],
'index': int(data[2]),
'position': tensor31([float(x) for x in data[3:]])}
def symmetry(d):
return ' '.join(data)
tags = {'lattice': lattice,
'atom': atom,
'units': check_units,
'symmetry': symmetry}
data_dict = {}
for record in records:
tag, data = record
if tag not in data_dict:
data_dict[tag] = []
data_dict[tag].append(tags[tag](data))
return data_dict
def parse_generic_block(block):
"""
Parse any other block into data dictionary given list of record
tuples.
"""
name, records = block
data_dict = {}
for record in records:
tag, data = record
if tag not in data_dict:
data_dict[tag] = []
data_dict[tag].append(data)
return data_dict
"""
Actual parser code.
"""
block_parsers = {'magres': parse_magres_block,
'atoms': parse_atoms_block,
'calculation': parse_generic_block, }
file_contents = fd.read()
# This works as a validity check
version = get_version(file_contents)
if version is None:
# This isn't even a .magres file!
raise RuntimeError('File is not in standard Magres format')
blocks = parse_blocks(file_contents)
data_dict = {}
for block_data in blocks:
block = parse_block(block_data)
if block[0] in block_parsers:
block_dict = block_parsers[block[0]](block)
data_dict[block[0]] = block_dict
else:
# Throw in the text content of blocks we don't recognise
if include_unrecognised:
data_dict[block[0]] = block_data[1]
# Now the loaded data must be turned into an ASE Atoms object
# First check if the file is even viable
if 'atoms' not in data_dict:
raise RuntimeError('Magres file does not contain structure data')
# Allowed units handling. This is redundant for now but
# could turn out useful in the future
magres_units = {'Angstrom': ase.units.Ang}
# Lattice parameters?
if 'lattice' in data_dict['atoms']:
try:
u = dict(data_dict['atoms']['units'])['lattice']
except KeyError:
raise RuntimeError('No units detected in file for lattice')
u = magres_units[u]
cell = np.array(data_dict['atoms']['lattice'][0]) * u
pbc = True
else:
cell = None
pbc = False
# Now the atoms
symbols = []
positions = []
indices = []
labels = []
if 'atom' in data_dict['atoms']:
try:
u = dict(data_dict['atoms']['units'])['atom']
except KeyError:
raise RuntimeError('No units detected in file for atom positions')
u = magres_units[u]
# Now we have to account for the possibility of there being CASTEP
# 'custom' species amongst the symbols
custom_species = None
for a in data_dict['atoms']['atom']:
spec_custom = a['species'].split(':', 1)
if len(spec_custom) > 1 and custom_species is None:
# Add it to the custom info!
custom_species = list(symbols)
symbols.append(spec_custom[0])
positions.append(a['position'])
indices.append(a['index'])
labels.append(a['label'])
if custom_species is not None:
custom_species.append(a['species'])
atoms = Atoms(cell=cell,
pbc=pbc,
symbols=symbols,
positions=positions)
# Add custom species if present
if custom_species is not None:
atoms.new_array('castep_custom_species', np.array(custom_species))
# Add the spacegroup, if present and recognizable
if 'symmetry' in data_dict['atoms']:
try:
spg = Spacegroup(data_dict['atoms']['symmetry'][0])
except SpacegroupNotFoundError:
# Not found
spg = Spacegroup(1) # Most generic one
atoms.info['spacegroup'] = spg
# Set up the rest of the properties as arrays
atoms.new_array('indices', np.array(indices))
atoms.new_array('labels', np.array(labels))
# Now for the magres specific stuff
li_list = list(zip(labels, indices))
def create_magres_array(name, order, block):
if order == 1:
u_arr = [None] * len(li_list)
elif order == 2:
u_arr = [[None] * (i + 1) for i in range(len(li_list))]
else:
raise ValueError(
'Invalid order value passed to create_magres_array')
for s in block:
# Find the atom index/indices
if order == 1:
# First find out which atom this is
at = (s['atom']['label'], s['atom']['index'])
try:
ai = li_list.index(at)
except ValueError:
raise RuntimeError('Invalid data in magres block')
# Then add the relevant quantity
u_arr[ai] = s[mn]
else:
at1 = (s['atom1']['label'], s['atom1']['index'])
at2 = (s['atom2']['label'], s['atom2']['index'])
ai1 = li_list.index(at1)
ai2 = li_list.index(at2)
# Sort them
ai1, ai2 = sorted((ai1, ai2), reverse=True)
u_arr[ai1][ai2] = s[mn]
if order == 1:
return np.array(u_arr)
else:
return np.array(u_arr, dtype=object)
if 'magres' in data_dict:
if 'units' in data_dict['magres']:
atoms.info['magres_units'] = dict(data_dict['magres']['units'])
for u in atoms.info['magres_units']:
# This bit to keep track of tags
u0 = u.split('_')[0]
if u0 not in _mprops:
raise RuntimeError('Invalid data in magres block')
mn, order = _mprops[u0]
if order > 0:
u_arr = create_magres_array(mn, order,
data_dict['magres'][u])
atoms.new_array(u, u_arr)
else:
# atoms.info['magres_data'] = atoms.info.get('magres_data',
# {})
# # We only take element 0 because for this sort of data
# # there should be only that
# atoms.info['magres_data'][u] = \
# data_dict['magres'][u][0][mn]
if atoms.calc is None:
calc = SinglePointDFTCalculator(atoms)
atoms.calc = calc
atoms.calc.results[u] = data_dict['magres'][u][0][mn]
if 'calculation' in data_dict:
atoms.info['magresblock_calculation'] = data_dict['calculation']
if include_unrecognised:
for b in data_dict:
if b not in block_parsers:
atoms.info['magresblock_' + b] = data_dict[b]
return atoms
def tensor_string(tensor):
return ' '.join(' '.join(str(x) for x in xs) for xs in tensor)
[docs]def write_magres(fd, image):
"""
A writing function for magres files. Two steps: first data are arranged
into structures, then dumped to the actual file
"""
image_data = {}
image_data['atoms'] = {'units': []}
# Contains units, lattice and each individual atom
if np.all(image.get_pbc()):
# Has lattice!
image_data['atoms']['units'].append(['lattice', 'Angstrom'])
image_data['atoms']['lattice'] = [image.get_cell()]
# Now for the atoms
if image.has('labels'):
labels = image.get_array('labels')
else:
labels = image.get_chemical_symbols()
if image.has('indices'):
indices = image.get_array('indices')
else:
indices = [labels[:i + 1].count(labels[i]) for i in range(len(labels))]
# Iterate over atoms
symbols = (image.get_array('castep_custom_species')
if image.has('castep_custom_species')
else image.get_chemical_symbols())
atom_info = list(zip(symbols,
image.get_positions()))
if len(atom_info) > 0:
image_data['atoms']['units'].append(['atom', 'Angstrom'])
image_data['atoms']['atom'] = []
for i, a in enumerate(atom_info):
image_data['atoms']['atom'].append({
'index': indices[i],
'position': a[1],
'species': a[0],
'label': labels[i]})
# Spacegroup, if present
if 'spacegroup' in image.info:
image_data['atoms']['symmetry'] = [image.info['spacegroup']
.symbol.replace(' ', '')]
# Now go on to do the same for magres information
if 'magres_units' in image.info:
image_data['magres'] = {'units': []}
for u in image.info['magres_units']:
# Get the type
p = u.split('_')[0]
if p in _mprops:
image_data['magres']['units'].append(
[u, image.info['magres_units'][u]])
image_data['magres'][u] = []
mn, order = _mprops[p]
if order == 0:
# The case of susceptibility
tens = {
mn: image.calc.results[u]
}
image_data['magres'][u] = tens
else:
arr = image.get_array(u)
li_tab = zip(labels, indices)
for i, (lab, ind) in enumerate(li_tab):
if order == 2:
for j, (lab2, ind2) in enumerate(li_tab[:i + 1]):
if arr[i][j] is not None:
tens = {mn: arr[i][j],
'atom1': {'label': lab,
'index': ind},
'atom2': {'label': lab2,
'index': ind2}}
image_data['magres'][u].append(tens)
elif order == 1:
if arr[i] is not None:
tens = {mn: arr[i],
'atom': {'label': lab,
'index': ind}}
image_data['magres'][u].append(tens)
# Calculation block, if present
if 'magresblock_calculation' in image.info:
image_data['calculation'] = image.info['magresblock_calculation']
def write_units(data, out):
if 'units' in data:
for tag, units in data['units']:
out.append(' units %s %s' % (tag, units))
def write_magres_block(data):
"""
Write out a <magres> block from its dictionary representation
"""
out = []
def nout(tag, tensor_name):
if tag in data:
out.append(' '.join([' ', tag,
tensor_string(data[tag][tensor_name])]))
def siout(tag, tensor_name):
if tag in data:
for atom_si in data[tag]:
out.append((' %s %s %d '
'%s') % (tag,
atom_si['atom']['label'],
atom_si['atom']['index'],
tensor_string(atom_si[tensor_name])))
write_units(data, out)
nout('sus', 'S')
siout('ms', 'sigma')
siout('efg_local', 'V')
siout('efg_nonlocal', 'V')
siout('efg', 'V')
def sisiout(tag, tensor_name):
if tag in data:
for isc in data[tag]:
out.append((' %s %s %d %s %d '
'%s') % (tag,
isc['atom1']['label'],
isc['atom1']['index'],
isc['atom2']['label'],
isc['atom2']['index'],
tensor_string(isc[tensor_name])))
sisiout('isc_fc', 'K')
sisiout('isc_orbital_p', 'K')
sisiout('isc_orbital_d', 'K')
sisiout('isc_spin', 'K')
sisiout('isc', 'K')
return '\n'.join(out)
def write_atoms_block(data):
out = []
write_units(data, out)
if 'lattice' in data:
for lat in data['lattice']:
out.append(" lattice %s" % tensor_string(lat))
if 'symmetry' in data:
for sym in data['symmetry']:
out.append(' symmetry %s' % sym)
if 'atom' in data:
for a in data['atom']:
out.append((' atom %s %s %s '
'%s') % (a['species'],
a['label'],
a['index'],
' '.join(str(x) for x in a['position'])))
return '\n'.join(out)
def write_generic_block(data):
out = []
for tag, data in data.items():
for value in data:
out.append('%s %s' % (tag, ' '.join(str(x) for x in value)))
return '\n'.join(out)
# Using this to preserve order
block_writers = OrderedDict([('calculation', write_generic_block),
('atoms', write_atoms_block),
('magres', write_magres_block)])
# First, write the header
fd.write('#$magres-abinitio-v1.0\n')
fd.write('# Generated by the Atomic Simulation Environment library\n')
for b in block_writers:
if b in image_data:
fd.write('[{0}]\n'.format(b))
fd.write(block_writers[b](image_data[b]))
fd.write('\n[/{0}]\n'.format(b))
# Now on to check for any non-standard blocks...
for i in image.info:
if '_' in i:
ismag, b = i.split('_', 1)
if ismag == 'magresblock' and b not in block_writers:
fd.write('[{0}]\n'.format(b))
fd.write(image.info[i])
fd.write('[/{0}]\n'.format(b))