Source code for ase.io.mustem

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