atomco.py

# -*- coding:utf-8 -*-
"""
===================================================================
Provide coordinate file class which do operations on these files.
===================================================================
Written by PytLab <shaozhengjiang@gmail.com>, November 2014
Updated by PytLab <shaozhengjiang@gmail.com>, May 2017

==============================================================

"""
import logging
import re
from collections import namedtuple
from math import acos, degrees
from itertools import combinations

import numpy as np

from . import VasPy
from .errors import CarfileValueError
from .functions import *
from .elements import chem_elements


class AtomCo(VasPy):
    "Base class to be inherited by atomco classes."
    def __init__(self, filename):
        VasPy.__init__(self, filename)

    def verify(self):
        if len(self.data) != self.natom:
            raise CarfileValueError('Atom numbers mismatch!')

    @property
    def atomco_dict(self):
        """
        Return the current atom type and coordinates mapping, 
        make sure the data in dict can be updated in time.
        """
        return self.get_atomco_dict(self.data)

    @property
    def tf_dict(self):
        """
        Return the current atom type and T/F mapping, make sure the data
        can be updated in time when returned.
        """
        return self.get_tf_dict(self.tf)

    def get_atomco_dict(self, data):
        """
        根据已获取的data和atoms, atoms_num, 获取atomco_dict
        """
        # [1, 1, 1, 16] -> [0, 1, 2, 3, 19]
        idx_list = [sum(self.atom_numbers[:i])
                    for i in range(1, len(self.atom_types)+1)]
        idx_list = [0] + idx_list

        data_list = data.tolist()
        atomco_dict = {}
        for atom_type, idx, next_idx in zip(self.atom_types,
                                            idx_list[:-1],
                                            idx_list[1:]):
            atomco_dict.setdefault(atom_type, data_list[idx: next_idx])

        return atomco_dict

    def get_tf_dict(self, tf):
        """
        根据已获取的tf和atoms, atoms_num, 获取tf_dict
        """
        # [1, 1, 1, 16] -> [0, 1, 2, 3, 19]
        idx_list = [sum(self.atom_numbers[:i])
                    for i in range(1, len(self.atom_types)+1)]
        idx_list = [0] + idx_list

        tf_list = tf.tolist()
        tf_dict = {}
        for atom_type, idx, next_idx in zip(self.atom_types,
                                            idx_list[:-1],
                                            idx_list[1:]):
            tf_dict.setdefault(atom_type, tf_list[idx: next_idx])

        return tf_dict

    def get_xyz_content(self, step=None, bases=None):

        """
        Get xyz file content.
        获取最新.xyz文件内容字符串

        Parameters:
        -----------
        step: The step number, int, optional, 1 by default.
        bases: If the bases is provided, default data is regarded as direct
               coordinates and would be converted to Cartesian coordinates using
               bases.
        """
        natom = "{:12d}\n".format(self.natom)
        try:
            step = self.step if step is None else step
        except AttributeError:
            step = 1
        step = "STEP ={:9d}\n".format(step)
        data = atomdict2str(self.atomco_dict, self.atom_types)
        data = ''
        for atom in self.atom_types:
            if bases is not None:
                coords = self.dir2cart(bases, np.array(self.atomco_dict[atom]))
                coords = coords.tolist()
            else:
                coords = self.atomco_dict[atom]

            template = '{:<3s}{:>16.7f}{:>16.7f}{:>16.7}\n'
            for i in range(len(coords)):
                data += template.format(atom, *coords[i])
        content = natom + step + data

        return content

    def get_poscar_content(self, **kwargs):
        """
        Get POSCAR content.
        根据对象数据获取poscar文件内容字符串

        Parameters:
        -----------
        bases_const: The constant for basis vectors, optional, 1.0 by default.

        bases: The basis vectors for the lattice, option, 3x3 np.array.
               [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]] by default.

        tf: The T/F info for all atoms. Nx3 np.array, n is the length of 0th axis of data.
        """
        content = 'Created by VASPy\n'

        # bases constant.
        try:
            bases_const = self.bases_const
        except AttributeError:
            bases_const = kwargs.get("bases_const", 1.0)
        bases_const = " {:.9f}\n".format(bases_const)

        # bases
        try:
            bases = self.bases
        except AttributeError:
            bases = kwargs.get("bases", np.array([[1.0, 0.0, 0.0],
                                                  [0.0, 1.0, 0.0],
                                                  [0.0, 0.0, 1.0]]))
        bases_list = bases.tolist()
        bases = ''
        for basis in bases_list:
            bases += "{:14.8f}{:14.8f}{:14.8f}\n".format(*basis)

        # atom info
        types, numbers = self.atom_types, self.atom_numbers
        atom_types = ("{:>5s}"*len(types) + "\n").format(*types)
        atom_numbers = ("{:>5d}"*len(numbers) + "\n").format(*numbers)

        # Direct or Cartesian
        coord_type = kwargs.get('coord_type', 'direct')

        # string
        if coord_type == 'direct':
            info = "Selective Dynamics\nDirect\n"
        else:
            info = "Cartensian\n"

        # data and tf
        try:
            tf = self.tf
        except AttributeError:
            # Initialize tf with 'T's.
            default_tf = np.full(self.data.shape, 'T', dtype=str)
            tf = kwargs.get("tf", default_tf)
        data_tf = ''
        if coord_type == 'direct':
            data = self.data.tolist()
        else:
            data = self.dir2cart(self.bases, self.data).tolist()

        for data, tf in zip(data, tf.tolist()):
            data_tf += ("{:18.12f}"*3 + "{:>5s}"*3 + "\n").format(*(data+tf))

        # merge all strings
        content += (bases_const + bases + atom_types + atom_numbers +
                    info + data_tf)

        return content

    def get_cif_content(self):
        """
        Get the cif file content.
        """
        content = 'data_VESTA_phase_1\n\n'

        # Phase name
        phase_name = ('xyz {}'*len(self.atom_types)).format(*self.atom_types)
        content += "{:<40s}'{}'\n".format('_pd_phase_name', phase_name)

        # Basis vectors lengths.
        length_a, length_b, length_c = [np.linalg.norm(basis) for basis in self.bases]
        content += '{:<40s}{:<.5f}\n'.format('_cell_length_a', length_a)
        content += '{:<40s}{:<.5f}\n'.format('_cell_length_b', length_b)
        content += '{:<40s}{:<.5f}\n'.format('_cell_length_c', length_c)

        # Angles between basis vectors.
        angle = lambda X, Y: degrees(acos(np.dot(X, Y)/(np.linalg.norm(X)*np.linalg.norm(Y))))
        alpha, beta, gamma = [angle(X, Y) for X, Y in combinations(self.bases, 2)]
        content += '{:<40s}{:<.2f}\n'.format('_cell_angle_alpha', alpha)
        content += '{:<40s}{:<.2f}\n'.format('_cell_angle_beta', beta)
        content += '{:<40s}{:<.2f}\n'.format('_cell_angle_gamma', gamma)

        # Other info.
        content += "{:<40s}'P 1'\n".format('_symmetry_space_group_name_H-M')
        content += '{:<40s}1\n\n'.format('_symmetry_Int_Tables_number')
        content += "loop_\n_symmetry_equiv_pos_as_xyz\n   'x, y, z'\n\n"

        # Atom info.
        content += ('loop_\n' +
                    '   _atom_site_label\n' +
                    '   _atom_site_occupancy\n' +
                    '   _atom_site_fract_x\n' +
                    '   _atom_site_fract_y\n' +
                    '   _atom_site_fract_z\n' +
                    '   _atom_site_adp_type\n' +
                    '   _atom_site_B_iso_or_equiv\n' +
                    '   _atom_site_type_symbol\n')

        # Atom coordinates.
        line_template = '   {:<9s}{:<7.1}{:<13.5f}{:<13.5f}{:<13.5f}{:<6s}{:<7.3f}{:s}\n'
        atom_count = 0
        for atom_type, coordinates in self.atomco_dict.items():
            for x, y, z in coordinates:
                atom_count += 1
                name = '{}{}'.format(atom_type, atom_count)
                content += line_template.format(name, 1.0, x, y, z, 'Biso', 1.0, atom_type)

        return content

    def get_lmp_content(self):
        """
        Get lammps data file content
        """
        content = '# Created by VASPy\n\n'

        # Info
        content += '{} atoms\n{} atom types\n\n'.format(len(self.data),
                                                        len(self.atom_types))

        # Basis info
        x, y, z = self.bases
        xhi, yhi, zhi = [np.linalg.norm(i) for i in self.bases]
        content += '0 {:.9f} xlo xhi\n0 {:.9f} ylo yhi\n0 {:.9f} zlo zhi\n'.format(xhi, yhi, zhi)
        xy = np.dot(x, y)/np.linalg.norm(x)
        xz = np.dot(x, z)/np.linalg.norm(x)
        yz = np.dot(y, z)/np.linalg.norm(y)
        content += '{}   {}   {} xy xz yz\n\n'.format(xy, xz, yz)

        # Masses
        content += 'Masses\n\n'
        for i, element in enumerate(self.atom_types):
            if element not in chem_elements:
                raise ValueError('element {} not in elements.py'.format(element))
            mass = chem_elements[element]['mass']
            content += '{} {:.4f}\n'.format(i+1, mass)

        # Coordinate
        content += '\nAtoms\n\n'
        cart_coords = self.dir2cart(self.bases, self.data).tolist()

        for i, (component, coord) in enumerate(zip(self.atom_components, cart_coords)):
            template = '{:>4d}{:>2d}{:>4.1f}{:>9.5f}{:>11.5f}{:>11.5f}\n'
            idx = i+1
            type_idx = self.atom_types.index(component) + 1
            x, y, z = coord
            content += template.format(idx, type_idx, 0.0, x, y, z)

        return content

    def get_volume(self):
        """
        Get volume of slab(Angstrom^3)
        获取晶格体积
        """
        if hasattr(self, 'bases_const') and hasattr(self, 'bases'):
            bases = self.bases_const*self.bases
            volume = np.linalg.det(bases)
            self.volume = volume
        else:
            raise AttributeError("Object has no bases and bases_const")

        return volume

    @staticmethod
    def dir2cart(bases, data):
        """
        Static method to convert direct coordinates to Cartisan coordinates.

        Parameters:
        -----------
        bases: The 3x3 array for basis vectors, 3x3 numpy.array.

        data: The direct coordinate data, Nx3 numpy.array.
        """
        A = np.matrix(bases).T
        x = np.matrix(data).T
        b = A*x

        b = np.array(b.T)

        if b.shape[0] == 1:
            b = b.reshape(3, )

        return b

    @staticmethod
    def cart2dir(bases, data):
        """
        Static method to convert Cartisian coordinates to direct coordinates.

        Parameters:
        -----------
        bases: The 3x3 array for basis vectors, 3x3 numpy.array.

        data: The Cartisan coordinate data, Nx3 numpy.array or a single 3D vector.
        """
        b = np.matrix(data).T
        A = np.matrix(bases).T
        x = A.I*b

        x = np.array(x.T)
        if x.shape[0] == 1:
            x = x.reshape(3, )

        return x


class XyzFile(AtomCo):
    """
    Create a .xyz file class.

    Example:

    >>> a = XyzFile(filename='ts.xyz')

    Class attributes descriptions
    =======================================================================
      Attribute      Description
      ============  =======================================================
      filename       string, name of the file the direct coordiante data
                     stored in
      natom          int, the number of total atom number
      step           int, STEP number in OUT.ANI file
      atom_types     list of string, atom types
      atom_numbers   list of int, atom number of atoms
      atomco_dict    dict, {atom name: coordinates}
      data           np.array, coordinates of atoms, dtype=float64
      ============  =======================================================
    """
    def __init__(self, **kwargs):
        filename = kwargs.pop("filename", None)
        content = kwargs.pop("content", None)
        content_list = kwargs.pop("content_list", None)

        if content_list is not None:
            content_list = content_list
        elif filename is not None:
            super(self.__class__, self).__init__(filename)
            with open(self.filename, 'r') as f:
                content_list = f.readlines()
        elif content is not None:
            content = content.strip()
            content_list = content.split("\n")

        self.load(content_list)
        self.verify()

    def load(self, content_list):
        """ Load all data in xyz file.
        """
        # Total number of all atoms.
        natom = int(content_list[0].strip())

        # The iteration step for this xyz file.
        step = int(str2list(content_list[1])[-1])

        # Get atom coordinate and number info
        data_list = [str2list(line) for line in content_list[2:]]
        data_array = np.array(data_list)      # dtype=np.string
        atoms_list = list(data_array[:, 0])   # 1st column
        data = np.float64(data_array[:, 1:])  # rest columns

        # Atom number for each atom
        atom_types = []
        for atom in atoms_list:
            if atom not in atom_types:
                atom_types.append(atom)
        atom_numbers = [atoms_list.count(atom) for atom in atom_types]

        # Set attributes.
        self.natom = natom
        self.step = step
        self.atom_types = atom_types
        self.atom_numbers = atom_numbers
        self.data = data

    def coordinate_transform(self, bases=None):
        "Use Ax=b to do coordinate transform cartesian to direct"
        if bases is None:
            bases = np.array([[1.0, 0.0, 0.0],
                              [0.0, 1.0, 0.0],
                              [0.0, 0.0, 1.0]])
        return self.cart2dir(bases, self.data)

    def get_content(self):
        "获取最新文件内容字符串"
        content = self.get_xyz_content()
        return content

    def tofile(self, filename='atomco.xyz'):
        "XyzFile object to .xyz file."
        content = self.get_content()

        with open(filename, 'w') as f:
            f.write(content)

        return


class PosCar(AtomCo):
    def __init__(self, filename='POSCAR'):
        """
        Class to generate POSCAR or CONTCAR-like objects.

        Example:

        >>> a = PosCar(filename='POSCAR')

        Class attributes descriptions
        =======================================================================
          Attribute      Description
          ============  =======================================================
          filename       string, name of the file the direct coordiante data
                         stored in
          bases_const    float, lattice bases constant
          bases          np.array, bases of POSCAR
          natom          int, the number of total atom number
          atom_types     list of strings, atom types
          atom_numbers   list of int, same shape with atoms
                         atom number of atoms in atoms
          tf             list of list, T&F info of atoms
          data           np.array, coordinates of atoms, dtype=float64
          ============  =======================================================
        """
        AtomCo.__init__(self, filename)

        # Load all data in file
        self.load()
        self.verify()

    def load(self):
        """ Load all information in POSCAR.
        """
        with open(self.filename, 'r') as f:
            content_list = f.readlines()

        # get scale factor
        bases_const = float(content_list[1])

        # bases
        bases = [str2list(basis) for basis in content_list[2:5]]

        # Atom info
        atom_types = str2list(content_list[5])
        # Atom number (str).
        atom_numbers = str2list(content_list[6])
        if content_list[7][0] in 'Ss':
            data_begin = 9
        else:
            data_begin = 8

        # get total number before load data
        atom_numbers = [int(i) for i in atom_numbers]
        natom = sum(atom_numbers)

        # data
        data, tf = [], []  # data and T or F info
        tf_dict = {}       # {tf: atom number}
        for line_str in content_list[data_begin: data_begin+natom]:
            line_list = str2list(line_str)
            data.append(line_list[:3])
            if len(line_list) > 3:
                tf_list = line_list[3:]
                tf.append(tf_list)
                # gather tf info to tf_dict
                tf_str = ','.join(tf_list)
                if tf_str not in tf_dict:
                    tf_dict[tf_str] = 1
                else:
                    tf_dict[tf_str] += 1
            else:
                tf.append(['T', 'T', 'T'])
                # gather tf info to tf_dict
                if 'T,T,T' not in tf_dict:
                    tf_dict['T,T,T'] = 1
                else:
                    tf_dict['T,T,T'] += 1

        # Data type convertion
        bases = np.float64(np.array(bases))  # to float
        data = np.float64(np.array(data))
        tf = np.array(tf)

        # set class attrs
        self.bases_const = bases_const
        self.bases = bases
        self.atom_types = atom_types
        self.atom_numbers = atom_numbers
        self.natom = natom
        self.data = data
        self.tf = tf
        self.totline = data_begin + natom  # total number of line

    def constrain_atom(self, atom, to='F', axis='all'):
        "修改某一类型原子的FT信息"
        # [1, 1, 1, 16] -> [0, 1, 2, 3, 19]
        idx_list = [sum(self.atom_numbers[:i])
                    for i in range(1, len(self.atom_types)+1)]
        idx_list = [0] + idx_list

        if to not in ['T', 'F']:
            raise CarfileValueError('Variable to must be T or F.')

        for atom_type, idx, next_idx in zip(self.atom_types,
                                            idx_list[:-1],
                                            idx_list[1:]):
            if atom_type == atom:
                if axis in ['x', 'X']:
                    self.tf[idx:next_idx, 0] = to
                elif axis in ['y', 'Y']:
                    self.tf[idx:next_idx, 1] = to
                elif axis in ['z', 'Z']:
                    self.tf[idx:next_idx, 2] = to
                else:
                    self.tf[idx:next_idx, :] = to
                break

        return self.tf

    def get_content(self):
        "根据对象数据获取文件内容字符串"
        content = self.get_poscar_content()
        return content

    def add_atom(self, atom_type, coordinate, fix=['T', 'T', 'T']):
        """
        Add a new atom to coordinate file.

        Parameters:
        -----------
        atom_type: element type of the atom, str.
        coordinate: position of the added atom, list of float.
        fix: flags for fixed atom in three directions, list of str.

        Example:
        --------
        >>> poscar.add_atom('C', [0.5, 0.5, 0.3])
        """
        atomco_dict = self.atomco_dict
        tf_dict = self.tf_dict

        self.natom += 1
        self.totline += 1

        if atom_type in self.atom_types:
            atomco_dict[atom_type].append(coordinate)
            tf_dict[atom_type].append(fix)
            idx = self.atom_types.index(atom_type)
            self.atom_numbers[idx] += 1
        else:
            self.atom_types.append(atom_type)
            atomco_dict[atom_type] = [coordinate]
            tf_dict[atom_type] = [fix]
            self.atom_numbers.append(1)

        # New data and fix info.
        data, tf = [], []
        for atom_type in self.atom_types:
            data += atomco_dict[atom_type]
            tf += tf_dict[atom_type]

        self.data = np.float64(np.array(data))
        self.tf = np.array(tf)

    def tofile(self, filename='POSCAR_c'):
        "生成文件"
        "PosCar object to POSCAR or CONTCAR."
        content = self.get_content()

        with open(filename, 'w') as f:
            f.write(content)

        return


class ContCar(PosCar):
    def __init__(self, filename='CONTCAR'):
        '''
        Class to generate POSCAR or CONTCAR-like objects.
        Totally same as PosCar class.

        Example:

        >>> a = ContCar(filename='POSCAR')
        '''
        PosCar.__init__(self, filename=filename)

    def tofile(self, filename='CONTCAR_c'):
        PosCar.tofile(self, filename=filename)


class XdatCar(AtomCo):
    def __init__(self, filename='XDATCAR'):
        """
        Class to generate XDATCAR objects.

        Example:

        >>> a = XdatCar()

        Class attributes descriptions
        =======================================================================
          Attribute      Description
          ============  =======================================================
          filename       string, name of the file the direct coordiante data
                         stored in
          bases_const    float, lattice bases constant
          bases          np.array, bases of POSCAR
          natom          int, the number of total atom number
          atom_types     list of strings, atom types
          tf             list of list, T&F info of atoms
          info_nline     int, line numbers of lattice info
          ============  =======================================================
        """
        AtomCo.__init__(self, filename)
        self.info_nline = 7  # line numbers of lattice info
        self.load()

    def load(self):
        with open(self.filename, 'r') as f:
            # read lattice info
            self.system = f.readline().strip()
            self.bases_const = float(f.readline().strip())

            # lattice basis
            self.bases = []
            for i in range(3):
                basis = line2list(f.readline())
                self.bases.append(basis)

            # atom info
            self.atom_types = str2list(f.readline())
            atoms_num = str2list(f.readline())
            self.atom_numbers = [int(i) for i in atoms_num]
            self.natom = sum(self.atom_numbers)

    def __iter__(self):
        """ Make the XdatCar object iterable.
        """
        # Define namedtuple for the item in iteration.
        XdatCarItem = namedtuple("XdatCarItem", ["step", "coordinates"])

        with open(self.filename, 'r') as f:
            # pass info lines
            for i in range(self.info_nline):
                f.readline()
            prompt = f.readline().strip()
            while '=' in prompt:
                step = int(prompt.split('=')[-1])
                data = []
                for i in range(self.natom):
                    data_line = f.readline()
                    data.append(line2list(data_line))
                prompt = f.readline().strip()

                yield XdatCarItem._make([step, np.array(data)])


class CifFile(AtomCo):
    def __init__(self, filename):
        """
        Class for *.cif files.

        Example:

        >>> a = CifFile(filename='ts.cif')

        Class attributes descriptions
        =======================================================================
          Attribute         Description
          ===============  ====================================================
          filename         string, name of the file the direct coordiante data
                           stored in
          natom            int, the number of total atom number
          atom_types       list of strings, atom types
          atom_numbers     list of int, atom number of atoms in atoms
          atom_names       list of string,
                           Value of attribute 'Name' in Atom3d tag.
          data             np.array, coordinates of atoms, dtype=float64

          cell_length_a    float, length of cell vector a
          cell_length_a    float, length of cell vector b
          cell_length_c    float, length of cell vector c
          cell_angle_alpha float, angle of cell alpha
          cell_angle_beta  float, angle of cell beta
          cell_angle_gamma float, angle of cell gamma
          ===============  ====================================================
        """
        super(CifFile, self).__init__(filename)
        self.__logger = logging.getLogger("vaspy.CifCar")
        self.load()

    def load(self):
        """
        Load data and attributes from *.cif file.
        """
        # Regular expression for attributes matching.
        regex = re.compile(r'^_(\w+)(?:\s+)(.+)$')

        with open(self.filename, 'r') as f:
            lines = f.readlines()

        # Split lines by 'loop_' indices.
        loop_indices = [i for i, line in enumerate(lines)
                        if line.startswith('loop_')]
        # [19, 23] -> [(0, 19), (20, 23), (24, line_length)]
        start_indices = [0] + [i + 1 for i in loop_indices]
        end_indices = loop_indices + [len(lines)]
        lines_groups = [lines[start: end] for start, end in
                        zip(start_indices, end_indices)]

        # Get attributes.
        float_candidates = ['cell_length_a', 'cell_length_b', 'cell_length_c',
                            'cell_angle_alpha', 'cell_angle_beta', 'cell_angle_gamma']
        for line in lines_groups[0]:
            line = line.strip()
            if line.startswith('_'):
                m = regex.match(line)
                if m:
                    attr, value = m.groups()
                    if attr in float_candidates:
                        value = float(value)
                    setattr(self, attr, value)
                    self.__logger.debug("{} = {}".format(attr, value))

        # Get coordinates data.
        titles = []
        data = []
        atom_names = []
        atom_types = []

        for line in lines_groups[-1]:
            line = line.strip()
            if line.startswith('_'):
                titles.append(line[1:])
            elif line:
                atom_name, _, x, y, z, _, _, atom_type = line2list(line, dtype=str)
                atom_names.append(atom_name)
                atom_types.append(atom_type)
                data.append([float(i) for i in (x, y, z)])

        # Set attributes.
        self.data = np.array(data)
        self.atom_names = atom_names
        self.atom_types = list(set(atom_types))
        self.atom_numbers = [atom_types.count(atom) for atom in self.atom_types]
        self.titles = titles
        self.natom = len(atom_names)