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
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]:
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]:
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))
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]:
In [ ]: