Chimie Théorique » QMX

from ase import Atom, Atoms

from ase.data import covalent_radii, atomic_numbers

import numpy as np

class Embed(Atoms):

    #--- constructor of the Embed class ---

    def __init__(self, system=None, cluster=None):

        super(Embed, self).__init__()

        # define the atom map

        self.atom_map_sys_cl = []

        self.linkatoms = []

        # cluster dimensions

        self.xyz_cl_min = None

        self.xyz_cl_max = None

        # set the search radius for link atoms

        self.d = 10

        # define the systems for calculations

        self.set_system(system)

        self.set_cluster(cluster)

        #

        self.set_cell([10, 10, 10])

        return

    #--- set the cluster ---

    def set_cluster(self, atoms):

        import copy

        # set the min/max cluster dimensions

        self.xyz_cl_min = atoms[0].get_position()

        self.xyz_cl_max = atoms[0].get_position()

        for atom in atoms:

            # assign the label "Cluster (10)" in atom.TAG

            atom.set_tag(10)

            xyz=atom.get_position()

            for i in xrange(3):

                # set the min/max cluster dimensions

                if xyz[i] < self.xyz_cl_min[i]:

                    self.xyz_cl_min[i] = xyz[i]

                if xyz[i] > self.xyz_cl_max[i]:

                    self.xyz_cl_max[i] = xyz[i]

        # add self.d around min/max cluster dimensions

        self.xyz_cl_min -= [self.d, self.d, self.d]

        self.xyz_cl_max += [self.d, self.d, self.d]

        # set the cluster for low and high level calculation

        self.atoms_cluster = atoms

        return

    #--- set the system ---

    def set_system(self, atoms):

        self.atoms_system = atoms

        # assign the label "Cluster (10)" in atom.TAG

        for atom in atoms:

            atom.set_tag(0)

        # update search radius for link atoms

        dx = 0

        for atom in atoms:

            r = covalent_radii[atom.get_atomic_number()]

            if (r > dx):

                dx = r

        self.d = dx * 2.1

        return

    #--- return cluster ---

    def get_cluster(self):

        return self.atoms_cluster

    def get_system(self):

        return self.atoms_system

    #--- Embedding ---

    def embed(self):

        # is the cluster and the host system definied ?

        if self.atoms_cluster is None or self.atoms_system is None:

            return

        self.find_cluster()

        self.set_linkatoms()

        print "link atoms found: ", len(self.linkatoms)

    def find_cluster(self):

        # set tolerance

        d = 0.001

        #atoms

        xyzs_cl=[]

        for atom_cl in self.atoms_cluster:

            xyzs_cl.append(atom_cl.get_position())

        xyzs_sys=[]

        for atom_sys in self.atoms_system:

            xyzs_sys.append(atom_sys.get_position())

        self.atom_map_sys_cl=np.zeros(len(self.atoms_system), int)

        # loop over cluster atoms atom_sys

        for iat_sys in xrange(len(self.atoms_system)):

            # get the coordinates of the system atom atom_sys

            xyz_sys = xyzs_sys[iat_sys]

            # bSysOnly: no identical atom has been found

            bSysOnly = True

            # loop over system atoms atom_cl

            for iat_cl in xrange(len(self.atoms_cluster)):

                # difference vector between both atoms

                xyz_diff = np.abs(xyzs_sys[iat_sys]-xyzs_cl[iat_cl])

                # identical atoms

                if xyz_diff[0] < d and xyz_diff[1] < d and xyz_diff[2] < d:

                    # set tag (CLUSTER+HOST: 10) to atom_sys

                    self.atoms_system[iat_sys].set_tag(10)

                    # map the atom in the atom list

                    self.atom_map_sys_cl[iat_sys]=iat_cl

                    # atom has been identified

                    bSysOnly = False

                    break

            if bSysOnly:

                self.atom_map_sys_cl[iat_sys] = -1

    def set_linkatoms(self, tol=15., linkAtom=None, debug=False):

        # local copies of xyz coordinates to avoid massive copying of xyz objects

        xyzs_cl=[]

        for atom_cl in self.atoms_cluster:

            xyzs_cl.append(atom_cl.get_position())

        xyzs_sys=[]

        for atom_sys in self.atoms_system:

            xyzs_sys.append(atom_sys.get_position())

        # FIXME: mapping does not work when cluster atoms are outside of the box

        # set the standard link atom

        if linkAtom is None:

            linkAtom ='H'

        # number of atoms in the cluster and the system

        nat_cl=len(self.atoms_cluster)

        nat_sys=len(self.atoms_system)

        # system has pbc?

        pbc = self.get_pbc()

        # set the bond table

        bonds = []

        # set the 27 cell_vec, starting with the (0,0,0) vector for the unit cell

        cells_L = [(0.0,  0.0,  0.0)]

        # get the cell vectors

        cell = self.get_cell()

        # TODO: check 0D, 1D, 2D, 3D

        if False:

         for ix in xrange(-1, 2):

            for iy in xrange(-1, 2):

                for iz in xrange(-1, 2):

                    if ix == 0 and iy == 0 and iz == 0:

                        continue

                    cells_L.append(np.dot([ix, iy, iz], cell))

        # save the radius of system atoms

        rs_sys = []

        for atom in self.atoms_system:

            rs_sys.append(covalent_radii[atom.get_atomic_number()])

        # sum over cluster atoms (iat_cl)

        for iat_cl in xrange(nat_cl):

            # get the cluster atom

            atom_cl=self.atoms_cluster[iat_cl]

            # ignore link atoms

            if atom_cl.get_tag() == 50:

                continue

            # xyz coordinates and covalent radius of the cluster atom iat_cl

            xyz_cl = xyzs_cl[iat_cl]

            r_cl = covalent_radii[atom_cl.get_atomic_number()]

            # sum over system atoms (iat_sys)

            for iat_sys in xrange(nat_sys):

                # avoid cluster atoms

                if self.atoms_system[iat_sys].get_tag()==10:

                    continue

                # sum over all cell_L

                for cell_L in cells_L:

                    # xyz coordinates and covalent radius of the system atom iat_sys

                    xyz_sys = xyzs_sys[iat_sys]+cell_L

                    # go only in distance self.d around the cluster

                    lcont = True

                    for i in xrange(3):

                        if (xyz_sys[i] < self.xyz_cl_min[i] or

                            xyz_sys[i] > self.xyz_cl_max[i]):

                            lcont = False

                            break

                    if not lcont:

                        continue

                    # xyz coordinates and covalent radius of the system atom iat_sys

                    r_sys = rs_sys[iat_sys]

                    # diff vector

                    xyz_diff = xyz_sys - xyz_cl

                    # distance between the atoms

                    r = np.sqrt(np.dot(xyz_diff, xyz_diff))

                    # ratio of the distance to the sum of covalent radius

                    f = r / (r_cl + r_sys)

                    if debug:

                        print "Covalent radii = ",r_cl, r_sys

                        print "Distance ", f

                        print "tol = ",(1+tol/100.),(1-tol/100.),(1-2*tol/100.)

                    if f <= (1+tol/100.) and f >= (1-2*tol/100.):

                        s = cell_L, iat_cl, iat_sys, r_cl

                        bonds.append(s)

                        break

                    if f <= (1-2*tol/100.):

                        raise RuntimeError("QMX: The cluster atom", iat_cl, " and the system atom", iat_sys, "came too close")

        r_h = covalent_radii[atomic_numbers[linkAtom]]

        for bond in bonds:

            cell_L, iat_cl, iat_sys, r_cl = bond

            # assign the tags for the border atoms

            atom_sys.set_tag(1)

            atom_cl.set_tag(11)

            #difference vector for the link atom, scaling

            xyz_diff = xyzs_sys[iat_sys] + cell_L - xyzs_cl[iat_cl]

            r = (r_cl + r_h)

            xyz_diff *=  r / np.sqrt(np.dot(xyz_diff, xyz_diff))

            # determine position of the link atom

            xyz_diff += xyzs_cl[iat_cl]

            # create link atom

            atom = Atom(symbol=linkAtom, position=xyz_diff, tag=50)

            # add atom to cluster

            self.atoms_cluster.append(atom)

            # add atom to the linkatoms

            s = cell_L, iat_cl, iat_sys, r, len(self.atoms_cluster)-1

            self.linkatoms.append(s)

        return

    def set_positions(self, positions_new):

        # number of atoms

        nat_sys=len(self.atoms_system)

        # go over all pairs of atoms

        for iat_sys in xrange(nat_sys):

            xyz = positions_new[iat_sys]

            iat_cl = self.atom_map_sys_cl[iat_sys]

            self.atoms_system[iat_sys].set_position(xyz)

            if iat_cl > -1:

                self.atoms_cluster[iat_cl].set_position(xyz)

        for cell_L, iat_cl, iat_sys, r, iat in self.linkatoms:

            # determine position of the link atom

            xyz_cl = self.atoms_cluster[iat_cl].get_position()

            xyz = self.atoms_system[iat_sys].get_position() - xyz_cl + cell_L

            xyz *= r / np.sqrt(np.dot(xyz, xyz))

            xyz += xyz_cl

            # update xyz coordinates of the cluster

            self.atoms_cluster[iat].set_position(xyz)

    def __getitem__(self, i):

        return self.atoms_system.__getitem__(i)

    def get_positions(self):

        return self.atoms_system.get_positions()

    def __add__(self, other):

        return self.atoms_system.__add__(other)

    def __delitem__(self, i):

        return self.atoms_system.__delitem__(i)

    def __len__(self):

        return self.atoms_system.__len__()