index | modules | gitlab

# 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

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

    def write_mustem(filename, *args, **kwargs):

        writer = XtlmuSTEMWriter(*args, **kwargs)
        writer.write_to_file(filename)

Documentation v: 3.19.0

Versions
  3.20.0b1 (development)
  3.19.0 (latest stable)