In [1]:
####Implementing the one line poisson solver.
"""This notebook has a code written in Python2.
The code solves the poisson's equation numerically with in a unit cell
using the method described in Tomas Arias lecture.""";

In [2]:
"""Importing all the necessary funtions required for this code"""
# numpy for math calculations
import numpy as np  
# For pplotting or visualising 
import matplotlib.pylab as pylab 
%pylab notebook
pylab.rcParams['figure.figsize'] = (10, 6)
import matplotlib.pyplot as plt
# For performing linear algebra operations using matrices
from numpy import linalg as LA
from numpy.linalg import inv
import operator
# for Generating tuples list.
import itertools


Populating the interactive namespace from numpy and matplotlib
/usr/local/lib/python2.7/site-packages/IPython/kernel/__init__.py:13: ShimWarning: The `IPython.kernel` package has been deprecated. You should import from ipykernel or jupyter_client instead.
  "You should import from ipykernel or jupyter_client instead.", ShimWarning)

In [3]:
### Creating the setup.m file. All these inputs are setup.

"""Creating the Cell and Geometry

    Defining the variables 
    
    R=[a,b,c] => Defines the lattice vectors. (a,b,c) are any +ve integers.
    s=[a,b,c] => Defines the number of sample points along the
                    lattice vectors.  (a,b,c) are any integers.
    sig=[a,b,c,...] => Defining sigma for each gaussian funtion to be summed up.
    coeff=array of {+/-}coefficients. The factor is used for gaussian.  
         Gaussian => coeff[0].Gaussian(sig[0])+coeff[1].Gaussian(sig[1])+....so on.
"""
#Inputs for Geometry
s = [4,4,4]
R= np.diag([4,4,4])
#Input for Charge distribution 
sig=[0.75,0.5]
coeff=[-1,1]

In [4]:
def _M(s):
    """Defining the 'M' matrix.
    'M' has coordinates in matrix form for each sample point in real space.
    We use M to get list of real space sample  points 'r'
    r=M(Diag(S))^(-1)R^(Transpose)
    
     Args:
        input:
            s => number of sample points along the lattice vectors.
        output: 
            M matrix. 
    """
    M = tuple(itertools.product(*[range(0,s[0]),range(0,s[1]),range(0,s[2])]))
    M=[list(i) for i in M]
    return M

In [5]:
"""This does not work here. It works in terminal code."""
def test_M():
    """Tests that the generate M subroutine works.                                                                                                                          
    """
    assert _M([2,2,2])==[[0, 0, 0],[0, 0, 1], [0, 1, 0], [0, 1, 1], [1, 0, 0], [1, 0, 1], [1, 1, 0], [1, 1, 1]]
    assert _M([3,2,1])==[[0, 0, 0], [0, 1, 0], [1, 0, 0], [1, 1, 0], [2, 0, 0], [2, 1, 0]]

In [6]:
def _N(s):
    """Defining the 'N' matrix.
    'N' has coordinates in matrix form for each sample point in reciprocal
    space.
    We use N to get list of reciprocal lattice vectos 'G'
    G= 2*pi*N*R^(inverse)
    
    Args:
        input:
            s => number of sample points along the lattice vectors.
        output: 
            N matrix. 
    """
    M=_M(s)
    m1=np.transpose(M)[0]
    m2=np.transpose(M)[1]
    m3=np.transpose(M)[2]
    n1=[ m1[i]-s[0] if m1[i] > s[0]/2 else m1[i] for i in range(len(m1)) ]
    n2=[ m2[i]-s[1] if m2[i] > s[1]/2 else m2[i] for i in range(len(m2)) ]
    n3=[ m3[i]-s[2] if m3[i] > s[2]/2 else m3[i] for i in range(len(m3)) ]
    N=np.transpose([n1,n2,n3])
    return N

In [7]:
def test_N():
    """Tests that the generate N subroutine works.                                                                                                                          
    """
    assert _N([2,2,2])==array([[0, 0, 0],[0, 0, 1],[0, 1, 0],[0, 1, 1],[1, 0, 0],[1, 0, 1],[1, 1, 0],[1, 1, 1]])
    assert _N([3,2,1])==array([[ 0,  0,  0],[ 0,  1,  0],[ 1,  0,  0],[ 1,  1,  0],[-1,  0,  0],[-1,  1,  0]])

In [8]:
def test_N():
    """Tests that the generate N subroutine works.                                                                                                                          
    This is another way to test arrays. using np.allclose                                                                                                                   
    """
    assert np.allclose(_N([3,2,1]),np.array([[ 0,  0,  0],[ 0,  1,  0],[ 1,  0,  0],[ 1,  1,  0],[-1,  0,  0],[-1,  1,  0]])) == True
    assert np.allclose(_N([2,2,2]),np.array([[0, 0, 0],
                                             [0, 0, 1],
                                             [0, 1, 0],
                                             [0, 1, 1],
                                             [1, 0, 0],
                                             [1, 0, 1],
                                             [1, 1, 0],
                                             [1, 1, 1]]))

In [9]:
def _gen_r_G_G2(s,R,p=None):
    """Generates r, G and G2 matrices.
    
    Args: input:
           s => number of sample points along the lattice vectors. 
           R=[a,b,c] => Defines the lattice vectors. (a,b,c) are any +ve integers.
           p = 'r', 'G' or 'G2'(square of G matrix) 
           
          Output:
           p='r' => Outputs r matrix.
           p='G' => Outputs G matrix.
           p='2' => Outputs G2 matrix.
       
       Example: s = [15,15,15]
                R= np.diag([6,6,6])
                _gen_r_G_G2(s,R,"r")
    """
    if p is None:
        return "Please pass a third argument asking for 'r', 'G' or 'G2'."
    if p is "r":
        r=np.dot(np.dot(_M(s),inv(np.diag(s))),np.transpose(R))
        return r
    if p is "G":
        G=2*np.pi*np.dot(_N(s),inv(R))
        return G
    if p is "G2":
        G=2*np.pi*np.dot(_N(s),inv(R))
        G2=np.sum(np.transpose(G*G),axis=0) 
        return G2
    else:
        return "Error: Passed an argument other than r, G or G2."

In [10]:
"""REMEMBER I CHANGED THE NAME OF FUNCTION TO gen_r_G_G2 in real code."""

def test_gen_r_G_G2():
    """ Tests weather _gen_r_G_G2 subroutine works                                                                                                                          
    """
    assert np.allclose(_gen_r_G_G2([2,2,2],np.diag([2,2,2]),'r'),np.array([[ 0.,  0.,  0.],
                                                                           [ 0.,  0.,  1.],
                                                                           [ 0.,  1.,  0.],
                                                                           [ 0.,  1.,  1.],
                                                                           [ 1.,  0.,  0.],
                                                                           [ 1.,  0.,  1.],
                                                                           [ 1.,  1.,  0.],
                                                                           [ 1.,  1.,  1.]]))
    assert np.allclose(_gen_r_G_G2([2,2,2],np.diag([2,2,2]),'G'),np.array([[ 0.        ,  0.        ,  0.        ],
                                                                           [ 0.        ,  0.        ,  3.14159265],
                                                                           [ 0.        ,  3.14159265,  0.        ],
                                                                           [ 0.        ,  3.14159265,  3.14159265],
                                                                           [ 3.14159265,  0.        ,  0.        ],
                                                                           [ 3.14159265,  0.        ,  3.14159265],
                                                                           [ 3.14159265,  3.14159265,  0.        ],
                                                                           [ 3.14159265,  3.14159265,  3.14159265]]))
    assert np.allclose(_gen_r_G_G2([2,2,2],np.diag([2,2,2]),'G2'),np.array([  0.       ,   9.8696044,   9.8696044,  19.7392088,   9.8696044,
                                                                              19.7392088,  19.7392088,  29.6088132]))

In [11]:
def gaussian(sig,coeff,x):
    """ Defining Gaussian charge distribution at point x in 3D space.
    Args:
        inputs: 
            sig => sigmas for each Gaussian.
            coeff => coefficients for each gaussian funtion.
                Note: All the gaussians will be summed up.
            x (np.array()): List of input points 
                Note: x in 3D space is x=Sqrt(a^2+b^2+c^2) where
                {a,b,c} are the cartesian coordinates.
            
        Output: Charge distribution at points x. 
        
    """
    g=lambda x,sigma: (np.exp(-(x**2)/(2*sigma**2)))/(2*np.pi*sigma**2)**(3./2)
    n=lambda x: np.sum([coeff[i]*g(x,sig[i]) for i in range(len(sig))],axis=0)
    if len(sig)!=len(coeff):
        return "Error: Length of 'sig' not equal to length of 'coeff'."
    else:
        return n(x)

In [12]:
def test_gaussian():
    """ Tests the gaussian charge distribution routine in charge_dis.py                                                                                                     
    """
    assert np.allclose(gaussian([0.75,0.5],[1,-1],np.linspace(0,1,10)),np.array([-0.35744565, -0.34669986, -0.31613924, -0.27038447, -0.21590651,
       -0.15959883, -0.10743921, -0.06358218, -0.03003532, -0.00686962])) == True
    assert np.allclose(gaussian([0.75,0.5,0.3],[1,-1,0.2],np.linspace(0,1,10)),np.array([ 0.11287757,  0.09244664,  0.04133802, -0.01668859, -0.05894094,
       -0.07493042, -0.06762259, -0.04725793, -0.02420049, -0.00505139]))

In [13]:
def charge_dis(s,R,sig,coeff):
    """Generates the charge distribution at each sample point for the unit cell.
        We change the midpoint of the unitcell to origin here.
    
    Args:
        input:
            R=[a,b,c] => Defines the lattice vectors. (a,b,c) are any +ve integers.
            s=[a,b,c] => Defines the number of sample points along the
                        lattice vectors.  (a,b,c) are any integers.
    """
    mid = sum(R,axis=1)/2.
    dr=np.sqrt(sum((_gen_r_G_G2(s,R,"r")-mid)**2,axis=1))
    return gaussian(sig,coeff,dr)

In [14]:
def test_charge_dis():
    """ Tests the charge_dis subroutine in charge_dis.py                                                                                                                    
    """
    assert np.allclose(charge_dis([2,2,2],np.diag([2,2,2]),[0.75,0.5],[1,-1]),np.array([ 0.00919842,  0.01613367,  0.01613367, -0.00686962,  0.01613367,
      -0.00686962, -0.00686962, -0.35744565])) == True

In [15]:
def Ope(inp,R):
    """Outputs the overlap matrix. In our case here we are specific to plane wave basis.
    Args:
        input:
            inp (numpy array): The overlap matrix operator acts on this input.
            R=[a,b,c] => Defines the lattice vectors. (a,b,c) are any +ve integers.
    """
    return np.linalg.det(R)*np.dot(np.identity(len(inp)),inp)

In [16]:
def test_Ope():
    """Tests the Ope() subroutine in fourierBasis.py                                                                                                                        
    """

    assert np.allclose(Ope([2,3,4,5],np.diag([2,2,2])),np.array([ 16.,  24.,  32.,  40.]))
    assert np.allclose(Ope([1,2,3,4,5,5.5,6,3],np.diag([2,1,6])),np.array([ 12.,  24.,  36.,  48.,  60.,  66.,  72.,  36.]))

In [17]:
def Lope(inp,s,R):
    """Lope is the L operator as defined in Tomas arias Lecture.
     Note:Works only for plane wave basis.
    Args:
        input:
            inp (numpy array): The overlap matrix operator acts on this input.
            s=[a,b,c] => Defines the number of sample points along the
                        lattice vectors.  (a,b,c) are any integers.
            R=[a,b,c] => Defines the lattice vectors. (a,b,c) are any +ve integers.
    """
    if len(_gen_r_G_G2(s,R,"G2")) == len(inp):
        return -np.linalg.det(R)*np.dot(np.diag(_gen_r_G_G2(s,R,"G2")),inp)
    else:
        return "Error: Input vector size is not same as G2"

In [18]:
def test_Lope():
    """Tests the Lope() subroutine in fourierBasis.py                                                                                                                       
                                                                                                                                                                            
    Function: Lope(inp,s,R)                                                                                                                                                 
                                                                                                                                                                            
    input:                                                                                                                                                                  
            inp (numpy array): The overlap matrix operator acts on this input.                                                                                              
            s=[a,b,c] => Defines the number of sample points along the                                                                                                      
                        lattice vectors.  (a,b,c) are any integers.                                                                                                         
            R=[a,b,c] => Defines the lattice vectors. (a,b,c) are any +ve integers.                                                                                         
    """

    assert Lope([2,3,4,5],[2,2,2],np.diag([2,2,2])) == 'Error: Input vector size is not same as G2'
    assert np.allclose(Lope(_gen_r_G_G2([2,2,2],np.diag([2,2,2]),'G2'),[2,2,2],np.diag([2,2,2])),np.array([   -0.        ,  -779.27272827,  -779.27272827, -3117.09091309,
       -779.27272827, -3117.09091309, -3117.09091309, -7013.45455445])) == True

In [19]:
def LinvOpe(inp,s,R):
    """LinvOpe is the inverse of L operator as defined in Tomas arias Lecture.
     Note:Works only for plane wave basis.
    Args:
        input:
            inp (numpy array): The overlap matrix operator acts on this input.
            s=[a,b,c] => Defines the number of sample points along the
                        lattice vectors.  (a,b,c) are any integers.
            R=[a,b,c] => Defines the lattice vectors. (a,b,c) are any +ve integers.
    """
    if len(_gen_r_G_G2(s,R,"G2")) == len(inp):
        Lele=-np.linalg.det(R)*np.diag(_gen_r_G_G2(s,R,"G2"))
        Lele[0][0]=1
        Leleinv=inv(Lele)
        Leleinv[0][0]=0
        return np.dot(Leleinv,inp)
    else:
        return "Error: Input vector size is not same as G2"

In [20]:
def test_LinvOpe():
    """ Testing the LinvOpe operator in fourierBasis.py                                                                       
                                                                                                                              
    LinvOpe is the inverse of L operator as defined in Tomas arias Lecture.                                                   
     Note:Works only for plane wave basis.                                                                                    
    Args:                                                                                                                     
        input:                                                                                                                
            inp (numpy array): The overlap matrix operator acts on this input.                                                
            s=[a,b,c] => Defines the number of sample points along the                                                        
                        lattice vectors.  (a,b,c) are any integers.                                                           
            R=[a,b,c] => Defines the lattice vectors. (a,b,c) are any +ve integers.                                           
    """
    assert LinvOpe([2,3,4,5],[2,2,2],np.diag([2,2,2])) == 'Error: Input vector size is not same as G2'
    assert np.allclose(LinvOpe(gen_r_G_G2([2,2,2],np.diag([2,2,2]),'G2'),[2,2,2],np.diag([2,2,2])),np.array([ 0.   , -0.125, \
-0.125, -0.125, -0.125, -0.125, -0.125, -0.125]))

In [21]:
def cI(inp,s,R,p=None):
    """cI is the fourier series expansion coeff Matrix. It is the forward transform matrix for the problem.
        Note:Works only for plane wave basis.
        Args:
        input:
            inp (numpy array): The overlap matrix operator acts on this input.
            s=[a,b,c] => Defines the number of sample points along the
                        lattice vectors.  (a,b,c) are any integers.
            R=[a,b,c] => Defines the lattice vectors. (a,b,c) are any +ve integers.
            p (default): Fast fourier transform. #Written with Wiley's code help.
                p='s' => slow fourier transform.
    """
    if p is None:
        f=np.fft.fftn(inp.reshape(s,order="F")).reshape(np.prod(s),order="F")
    elif p is "s": 
        f=np.dot(np.exp(1j*np.dot(_gen_r_G_G2(s,R,"G"),np.transpose(_gen_r_G_G2(s,R,"r")))),inp)
    return f

In [22]:
def test_cI():
    """ Testing the cI subroutine in fourierBasis.py                                                                                                                                                        
                                                                                                                                                                                                            
    Function: cI                                                                                                                                                                                            
    cI is the fourier series expansion coeff Matrix. It is the forward transform matrix for the problem.                                                                                                    
        Note:Works only for plane wave basis.                                                                                                                                                               
        Args:                                                                                                                                                                                               
        input:                                                                                                                                                                                              
            inp (numpy array): The overlap matrix operator acts on this input.                                                                                                                              
            s=[a,b,c] => Defines the number of sample points along the                                                                                                                                      
                        lattice vectors.  (a,b,c) are any integers.                                                                                                                                         
            R=[a,b,c] => Defines the lattice vectors. (a,b,c) are any +ve integers.                                                                                                                         
            p (default): Fast fourier transform. #Written with Wiley's code help.                                                                                                                           
                p='s' => slow fourier transform.                                                                                                                                                            
    """

    assert np.allclose(cI(_gen_r_G_G2([2,2,2],np.diag([3,3,3]),'r').T[1],[2,2,2],np.diag([3,3,3])),np.array([ 6.+0.j,  0.+0.j, -6.+0.j,  0.+0.j,  0.+0.j,  0.+0.j,  0.+0.j,0.+0.j]))
    assert np.allclose(cI(_gen_r_G_G2([2,2,2],np.diag([3,3,3]),'r').T[1],[2,2,2],np.diag([3,3,3]),'s'),np.array([ 6. +0.00000000e+00j,  0. +3.67394040e-16j, -6. +7.34788079e-16j,
                                                                    0. -3.67394040e-16j,  0. +3.67394040e-16j,  0. -2.46519033e-32j,0. -3.67394040e-16j,  0. -2.46519033e-32j]))

In [23]:
def cJ(inp,s,R,p=None):
    """cJ is the fourier series expansion coeff Matrix. You can get the fast fourier transform too.
      Note:Works only for plane wave basis.
    Args:
        input:
            inp (numpy array): The overlap matrix operator acts on this input.
            s=[a,b,c] => Defines the number of sample points along the
                        lattice vectors.  (a,b,c) are any integers.
            R=[a,b,c] => Defines the lattice vectors. (a,b,c) are any +ve integers.
    """
    if p is None:
        f=np.fft.ifftn(inp.reshape(s,order="F")).reshape(np.prod(s),order="F")
        return f
    elif p is "s":
        f=np.exp(1j*np.dot(_gen_r_G_G2(s,R,"G"),np.transpose(_gen_r_G_G2(s,R,"r"))))
        return (1./np.prod(s))*np.dot(np.transpose(f.conjugate()),inp)

In [24]:
def test_cJ():
    """                                                                                                                                                                                                     
    This funtion tests the cJ funtion in fourierBasis.py                                                                                                                                                    
                                                                                                                                                                                                            
    Funtion: cJ                                                                                                                                                                                             
    cJ is the fourier series expansion coeff Matrix. You can get the fast fourier transform too.                                                                                                            
      Note:Works only for plane wave basis.                                                                                                                                                                 
    Args:                                                                                                                                                                                                   
        input:                                                                                                                                                                                              
            inp (numpy array): The overlap matrix operator acts on this input.                                                                                                                              
            s=[a,b,c] => Defines the number of sample points along the                                                                                                                                      
                        lattice vectors.  (a,b,c) are any integers.                                                                                                                                         
            R=[a,b,c] => Defines the lattice vectors. (a,b,c) are any +ve integers.                                                                                                                         
    """

    assert np.allclose(cJ(_gen_r_G_G2([2,2,2],np.diag([3,3,3]),'r').T[1],[2,2,2],np.diag([3,3,3])),np.array([ 0.75+0.j,  0.00+0.j, -0.75+0.j,  0.00+0.j,  0.00+0.j,  0.00+0.j,0.00+0.j,  0.00+0.j]))
    assert np.allclose(cJ(_gen_r_G_G2([2,2,2],np.diag([3,3,3]),'r').T[1],[2,2,2],np.diag([3,3,3]),'s'),np.array([ 0.75 +0.00000000e+00j,  0.00 -4.59242550e-17j,
                                                                                                                  -0.75 -9.18485099e-17j,  0.00 +4.59242550e-17j,
                                                                                                                  0.00 -4.59242550e-17j,  0.00 +3.08148791e-33j,
                                                                                                                  0.00 +4.59242550e-17j,  0.00 +6.16297582e-33j]))

In [25]:
def phi(inp,s,R):
    """Poisson solver in a single line of code."""
    return cI(LinvOpe(-4.0*np.pi*Ope(cJ(inp,s,R),R),s,R),s,R)

In [26]:
def test_phi():
    """                                                                                                                                                                                                     
     Tests the oneline poisson solver funtion.                                                                                                                                                              
                                                                                                                                                                                                            
    Poisson solver in a single line of code."""

    assert np.allclose(phi(charge_dis([2,2,2],np.diag([2,2,2]),[0.75,0.5],[1,-1]),[2,2,2],np.diag([2,2,2])),np.array([ 0.11648337+0.j,  0.07467416+0.j,  0.07467416+0.j, -0.01777448+0.j,
    0.07467416+0.j, -0.01777448+0.j, -0.01777448+0.j, -0.28718243+0.j]))

In [27]:
Uanalytic=((1/sig[0] + 1/sig[1])/2 -np.sqrt(2)/np.sqrt(sig[0]**2+sig[1]**2))/np.sqrt(np.pi)
Uanalytic


Out[27]:
0.055142527694733337

In [28]:
Unumerical=0.5*np.real(np.dot(np.conjugate(cJ(np.real(phi(charge_dis(s,R,sig,coeff),s,R)),s,R)),Ope(cJ(charge_dis(s,R,sig,coeff),s,R),R)))
Unumerical


Out[28]:
0.11031018403072157

In [29]:
#################code ends ;

In [30]:
# Just to check how long the code takes to run.
import time
start_time = time.time()
Unumerical=0.5*np.real(np.dot(np.conjugate(cJ(np.real(phi(charge_dis(s,R,sig,coeff),s,R)),s,R)),Ope(cJ(charge_dis(s,R,sig,coeff),s,R),R)))
print("--- %s seconds ---" % (time.time() - start_time))


--- 0.00365304946899 seconds ---

In [31]:
def cIm(inp,s,R,p=None):
    """Difference between cI and cI2: Input is a matrix instead of an array or coloumn vector. 
        
        cIm is the fourier series expansion coeff Matrix. It is the forward transform matrix for the problem.
        Note:Works only for plane wave basis. 
        Args:
        input:
            inp (numpy matrix array): The overlap matrix operator acts on this input.
            s=[a,b,c] => Defines the number of sample points along the
                        lattice vectors.  (a,b,c) are any integers.
            R=[a,b,c] => Defines the lattice vectors. (a,b,c) are any +ve integers.
            p (default): Fast fourier transform. #Written with Wiley's code help.
                p='s' => slow fourier transform.
                
        Example: cIm(_gen_r_G_G2([2,2,2],np.diag([3,3,3]),'r').T,[2,2,2],np.diag([3,3,3]))
            
    """
    f=[]
    inpp=[]
    inpp=np.array([inp.tolist()])
    f= [cI(inpp[i],s,R) for i in range(len(inpp))]
    return f

In [38]:
x1=charge_dis(s,R,sig,coeff).tolist()
x2=charge_dis(s,R,sig,coeff).tolist()
e=np.array(np.array([x2,x1]).tolist())
len(np.array(e.tolist()))


Out[38]:
2

In [ ]: