from math import *
import numpy as np
from supercell import Supercell
from netCDF4 import Dataset
import time
from ewald_cd import *
from utility import *
[docs]class Charge_dipole(Supercell):
    """
    Generate charge-dipole interaction matrix.
    The parameters are similar to that of 'charge-charge'.
 .. note::
    To avoid the directly ovelap between charges and dipoles, charges had
    been shift by (0.5,0.5,0.5) (relative coordinate) with respect to the origin
    of each unit cell, where the dipole stays. Therefore, the indexing in the
    generated 'netcdf' matrix is different from the others.
    """
    def __init__(self, n1, n2, nz, lattice):
        Supercell.__init__(self, n1, n2, nz, lattice)
        self.cdij = np.zeros((2 * n1 + 1, 2 * n2 + 1, 2 * nz + 1, 3))
        self.cd_matrix_calculated = False
    def write_cd_matrix(self, fn):
        if self.cd_matrix_calculated == False:
            print("Calculate the matrix first ...")
            self.generate_cd_matrix()
            self.cd_matrix_calculated = True
        dipm = Dataset(fn, "w", format="NETCDF4")
        direction = dipm.createDimension("direction", 3)
        directions = dipm.createVariable("direction", np.int32, ("direction"))
        iax = dipm.createDimension("iax", None)
        iaxs = dipm.createVariable("iax", np.int32, ("iax"))
        iay = dipm.createDimension("iay", None)
        iays = dipm.createVariable("iay", np.int32, ("iay"))
        iaz = dipm.createDimension("iaz", None)
        iazs = dipm.createVariable("iaz", np.int32, ("iaz"))
        # The actual 2-d varable.
        matrix = dipm.createVariable('matrix', np.float64, ('iax', 'iay', 'iaz', 'direction'))
        dipm.description = 'Charge dipole matrix: interaction matrix'
        dipm.history = 'Created at ' + time.ctime(time.time())
        matrix[:, :, :, :] = self.cdij
        dipm.close()
    def generate_cd_matrix(self):
        pi = 4.0 * atan(1.0)
        pi2 = pi * 2.0
        NN = 10
        tol = 1.0e-12
        eta = sqrt(-log(tol)) * 2
        gcut = 2.0 * eta ** 2
        gcut2 = gcut ** 2
        eta4 = 1.0 / (4 * eta ** 2)
        am = np.zeros((3))
        for i in range(3):
            for k in range(3):
                am[i] += self.a[k, i] ** 2
            am[i] = sqrt(am[i])
        mg1 = int(gcut * am[0] / pi2) + 1
        mg2 = int(gcut * am[1] / pi2) + 1
        mg3 = int(gcut * am[2] / pi2) + 1
        print('Gcut: ', gcut, ' mg1, mg2, mg3: ', mg1, mg2, mg3)
        # Set parameters to be used in the PYBIND11 C++ computation.
        cd_set_parameters(self.b, mg1, mg2, mg3, gcut2, eta4)
        flag_calculated = np.zeros((2 * self.n1 + 1, 2 * self.n2 + 1, 2 * self.n3 + 1), dtype=int)
        for ia in range(self.nsites):
            for ja in range(self.nsites):
                print('site: ', ia, ja)
                iax = self.ixa[ia]
                iay = self.iya[ia]
                iaz = self.iza[ia]
                jax = self.ixa[ja]
                jay = self.iya[ja]
                jaz = self.iza[ja]
                ix = jax - iax
                iy = jay - iay
                iz = jaz - iaz
                ixi = ix + self.n1
                iyi = iy + self.n2
                izi = iz + self.n3
                if (flag_calculated[ixi, iyi, izi] == 1):
                    continue
                pos = np.zeros(3)
                pos = (ix + 0.5) * self.lattice[0, :] \
                      
+ (iy + 0.5) * self.lattice[1, :] \
                      
+ (iz + 0.5) * self.lattice[2, :]
                rx = pos[0]
                ry = pos[1]
                rz = pos[2]
                c = 4.0 * pi / self.celvol
                dum = np.zeros(3)
                # print('Summing over k space')
                cd_sum_over_k(dum, rx, ry, rz)
                # The 2.0 comes from how the sum on k was done.
                self.cdij[ixi, iyi, izi, :] = dum[:] * c * 2.0
                # The following is an attempt to see what happens if
                # the k=0 term is added.
                # self.cdij[ixi,iyi,izi,:] = (dum[:]+pos)*c*2.0
                flag_calculated[ixi, iyi, izi] = 1
        self.cd_matrix_calculated = True