"""Module to read and write atoms in xtl file format for the muSTEM software.
See http://tcmp.ph.unimelb.edu.au/mustem/muSTEM.html for a few examples of
this format and the documentation of muSTEM.
See https://github.com/HamishGBrown/MuSTEM for the source code of muSTEM.
"""
import numpy as np
from ase.atoms import symbols2numbers
from ase.utils import basestring
[docs]def read_mustem(filename):
"""Import muSTEM input file.
Reads cell, atom positions, etc. from muSTEM xtl file
"""
from ase import Atoms
from ase.geometry import cellpar_to_cell
if isinstance(filename, basestring):
f = open(filename)
else: # Assume it's a file-like object
f = filename
# Read comment:
f.readline()
# Parse unit cell parameter
cellpar = [float(i) for i in f.readline().strip().split()[:3]]
cell = cellpar_to_cell(cellpar)
# beam energy
f.readline()
# Number of different type of atoms
element_number = int(f.readline().strip())
symbols = []
positions = []
for i in range(element_number):
# Read the element
symbol = str(f.readline().strip())
atoms_number = int(f.readline().split()[0])
# read all the position for each element
for j in range(atoms_number):
line = f.readline()
positions.append([float(i) for i in line.strip().split()])
symbols.append(symbol)
f.close()
atoms = Atoms(cell=cell, scaled_positions=positions)
atoms.set_chemical_symbols(symbols)
return atoms
class XtlmuSTEMWriter:
"""Write muSTEM input file.
Parameters:
atoms: Atoms object
keV: float
Energy of the electron beam required for the image simulation.
DW: float or dictionary of float with atom type as key
Debye-Waller factor of each atoms.
occupancy: float or dictionary of float with atom type as key (optional)
Occupancy of each atoms. Default value is `1.0`.
comment: str (optional)
Comments to be writen in the first line of the file. If not
provided, write the total number of atoms and the chemical formula.
fit_cell_to_atoms: bool (optional)
If `True`, fit the cell to the atoms positions. If negative coordinates
are present in the cell, the atoms are translated, so that all
positions are positive. If `False` (default), the atoms positions and
the cell are unchanged.
"""
def __init__(self, atoms, keV, DW, comment=None, occupancy=1.0,
fit_cell_to_atoms=False):
self.atoms = atoms.copy()
from collections import OrderedDict
self.atom_types = list(OrderedDict((element, None)
for element in self.atoms.get_chemical_symbols()))
self.keV = keV
self.DW = DW
self._check_key_dictionary(self.DW, 'DW')
self.comment = comment
if np.isscalar(occupancy):
self.occupancy = dict(zip(self.atom_types,
[occupancy] * len(self.atom_types)))
else:
self.occupancy = occupancy
self._check_key_dictionary(self.occupancy, 'occupancy')
self.numbers = symbols2numbers(self.atom_types)
if fit_cell_to_atoms:
self.atoms.translate(-self.atoms.positions.min(axis=0))
self.atoms.set_cell(self.atoms.positions.max(axis=0))
def _check_key_dictionary(self, d, dict_name):
# Check if we have enough key
for key in self.atom_types:
if not key in d:
raise ValueError('Missing the {0} key in the `{1}` dictionary.'
''.format(key, dict_name))
def _get_position_array_single_atom_type(self, number):
self._check_cell_is_orthorhombic()
# Get the scaled (reduced) position for a single atom type
return self.atoms.get_scaled_positions()[np.where(
self.atoms.numbers == number)]
def _check_cell_is_orthorhombic(self):
cell = self.atoms.get_cell()
# Diagonal element must different from 0
# Non-diagonal element must be zero
if ((cell != 0) != np.eye(3)).any():
raise ValueError('To export to this format, the cell need to be '
'orthorhombic.')
def _get_file_header(self):
# 1st line: comment line
if self.comment is None:
s = "{0} atoms with chemical formula: {1}\n".format(
len(self.atoms),
self.atoms.get_chemical_formula())
else:
s = self.comment
# 2nd line: lattice parameter
s += "{} {} {} {} {} {}\n".format(
*self.atoms.get_cell_lengths_and_angles().tolist())
# 3td line: acceleration voltage
s += "{}\n".format(self.keV)
# 4th line: number of different atom
s += "{}\n".format(len(self.atom_types))
return s
def _get_element_header(self, atom_type, number, atom_type_number,
occupancy, DW):
return "{0}\n{1} {2} {3} {4}\n".format(atom_type, number,
atom_type_number, occupancy, DW)
def _get_file_end(self):
return "Orientation\n 1 0 0\n 0 1 0\n 0 0 1\n"
def write_to_file(self, f):
if isinstance(f, basestring):
f = open(f, 'w')
f.write(self._get_file_header())
for atom_type, number, occupancy in zip(self.atom_types,
self.numbers,
self.occupancy):
positions = self._get_position_array_single_atom_type(number)
atom_type_number = positions.shape[0]
f.write(self._get_element_header(atom_type, atom_type_number,
number, self.occupancy[atom_type],
self.DW[atom_type]))
for pos in positions:
f.write('{0} {1} {2}\n'.format(pos[0], pos[1], pos[2]))
f.write(self._get_file_end())
[docs]def write_mustem(filename, *args, **kwargs):
writer = XtlmuSTEMWriter(*args, **kwargs)
writer.write_to_file(filename)