Chimie Théorique » QMX

from math import exp, sqrt

import numpy as np

from ase.atoms import Atoms

class STM:

    def __init__(self, atoms, symmetries=None):

        if isinstance(atoms, Atoms):

            calc = atoms.get_calculator()

        else:

            calc = atoms

            atoms = calc.get_atoms()

        self.nbands = calc.get_number_of_bands()

        self.weights = calc.get_k_point_weights()

        self.nkpts = len(self.weights)

        self.nspins = calc.get_number_of_spins()

        self.eigs = np.array([[calc.get_eigenvalues(k, s)

                               for k in range(self.nkpts)]

                              for s in range(self.nspins)])

        self.eigs -= calc.get_fermi_level()

        self.calc = calc

        self.cell = atoms.get_cell()

        assert not self.cell[2, :2].any() and not self.cell[:2, 2].any()

        self.ldos = None

        self.symmetries = symmetries or []

    def calculate_ldos(self, width=None):

        if self.ldos is not None and width == self.width:

            return

        if width is None:

            width = 0.1

        ldos = None

        for s in range(self.nspins):

            for k in range(self.nkpts):

                for n in range(self.nbands):

                    psi = self.calc.get_pseudo_wave_function(n, k, s)

                    if ldos is None:

                        ldos = np.zeros_like(psi)

                    f = (exp(-(self.eigs[s, k, n] / width)**2) *

                         self.weights[k])

                    ldos += f * (psi * np.conj(psi)).real

        if 0 in self.symmetries:

            # (x,y) -> (-x,y)

            ldos[1:] += ldos[:0:-1].copy()

            ldos[1:] *= 0.5

        if 1 in self.symmetries:

            # (x,y) -> (x,-y)

            ldos[:, 1:] += ldos[:, :0:-1].copy()

            ldos[:, 1:] *= 0.5

        if 2 in self.symmetries:

            # (x,y) -> (y,x)

            ldos += ldos.transpose((1, 0, 2)).copy()

            ldos *= 0.5

        self.ldos = ldos

        self.width = width

    #def save_ldos(self, filename='ldos.pckl'):

    def get_averaged_current(self, z, width=None):

        self.calculate_ldos(width)

        nz = self.ldos.shape[2]

        # Find grid point:

        n = z / self.cell[2, 2] * nz

        dn = n - np.floor(n)

        n = int(n) % nz

        print n,dn

        # Average and do linear interpolation:

        return ((1 - dn) * self.ldos[:, :, n].mean() +

                dn *       self.ldos[:, :, (n + 1) % nz].mean())

    def scan(self, current, z=None, width=None):

        self.calculate_ldos(width)

        L = self.cell[2, 2]

        if z is None:

            z = L / 2

        nz = self.ldos.shape[2]

        n = int(round(z / L * nz)) % nz

        h = L / nz

        ldos = self.ldos.reshape((-1, nz))

        heights = np.empty(ldos.shape[0])

        for i, a in enumerate(ldos):

            heights[i], z, n = find_height(a, current, z, n, nz, h)

        heights.shape = self.ldos.shape[:2]

        return heights

    def linescan(self, current, p1, p2, npoints=None, z=None, width=None):

        self.calculate_ldos(width)

        L = self.cell[2, 2]

        if z is None:

            z = L / 2

        nz = self.ldos.shape[2]

        n = int(round(z / L * nz)) % nz

        h = L / nz

        ldos = self.ldos.reshape((-1, nz))

        p1 = np.asarray(p1)

        p2 = np.asarray(p2)

        d = p2 - p1

        s = sqrt(np.dot(d, d))

        if npints == None:

            npoints = int(3 * s / h + 2)

        cell = self.cell[:2, :2]

        shape = np.array(self.ldos.shape[:2], float)

        M = cell.I

        heights = np.empty(npoints)

        for i in range(npoints):

            p = p1 + i * d / (npoints - 1)

            q = np.dot(M, p) * shape

            qi = q.astype(int)

            n0, n1 = qi

            f = q - qi

            g = 1 - f

            a = (g[0] * g[0] * ldos[n0,     n1    ] +

                 f[0] * g[0] * ldos[n0 + 1, n1    ] +

                 g[0] * f[0] * ldos[n0,     n1 + 1] +

                 f[0] * f[0] * ldos[n0 + 1, n1 + 1])

            heights[i], z, n = find_height(a, current, z, n, nz, h)

        return np.linspace(0, s, npoints), heights

    def cube(self, filename, atoms=None):

        pass

def find_height(array, current, z, n, nz, h):

    c1 = array[n]

    sign = cmp(c1, current)

    m = 0

    while m < nz:

        n = (n + sign) % nz

        z += sign * h

        c2 = array[n]

        if cmp(c2, current) != sign:

            break

        c1 = c2

        m += 1

    if m == nz:

        print z, n, nz, h, current, array

        raise RuntimeError('Tip crash!')

    return z - sign * h * (current - c2) / (c1 - c2), z, n