# -*- coding: utf-8 -*-
import matplotlib
matplotlib.rc('text', usetex=True)
import matplotlib.pyplot as plt
from ase import Atoms
from gpaw import GPAW
import numpy as np
from piirtoa import plot2d
a = 4.0
h = 0.2
nbands = 1
H = Atoms('H', [(a / 2, a / 2, a / 2)],
pbc=0,
cell=(a, a, a))
calc = GPAW(nbands=nbands, h=h, charge=0)
H.set_calculator(calc)
print H.get_potential_energy() + calc.get_reference_energy()
density = calc.get_pseudo_density()
print density.shape
nx = density.shape[0]
ny = density.shape[1]
density_xy = np.sum(density,axis=2)*h*h*h
plot2d(density_xy, h, h, nx, ny)
nx = 4
ny = 4
nz = 1
dx = 0.1
dy = 0.1
dz = 0.1
V_hartree = tee_gridi(nx,ny,nz)
#print "V_hartree", V_hartree
elektroni_tiheys = tee_gridi(nx,ny,nz)
elektroni_tiheys[1,1,0] = -1
elektroni_tiheys[2,2,0] = -1
#elektroni_tiheys[3,0,1] = -3
aseta_gridin_alkuarvo(0, V_hartree)
#print "elektroni_tiheys", elektroni_tiheys
#print "V_hartree", V_hartree
#sys.exit()
for step in range(10):
tee_gauss_seidel_yksi_askel(V_hartree, elektroni_tiheys, dx, dy, dz)
#print V_hartree
print "kineettinen energia", E_T(elektroni_tiheys)
print "vaihtokorrelaatio energia", E_vaihtokorrelaatio(elektroni_tiheys)
print "elektroni_elektroni energia", \
E_elektroni_elektroni(elektroni_tiheys,V_hartree)
print "elektroni_ydin energia", \
E_elektroni_ydin(elektroni_tiheys,ytimet,dx,dy,dz)
plot2d(V_hartree, dx, dy, nx, ny)
nx = 5
ny = 5
nz = 1
dx = 0.2
dy = 0.2
dz = 0.2
V_hartree_init = tee_gridi(nx,ny,nz)
elektroni_tiheys_init = tee_gridi(nx,ny,nz)
ydin_tiheys_init = tee_gridi(nx,ny,nz)
tasa_tiheys = 1/(nx*ny*nz*dx*dy*dz)
V_hartree = aseta_gridin_alkuarvo(0, V_hartree_init)
elektroni_tiheys = aseta_gridin_alkuarvo(-0.1, elektroni_tiheys_init)
ydin_tiheys = aseta_gridin_alkuarvo(0, ydin_tiheys_init)
ydin_tiheys[1,1,0] = 1
ydin_tiheys[3,3,0] = 1
elektroni_tiheys_flat = elektroni_tiheys.ravel().tolist()
xopt = fmin(func=E_tot, x0=elektroni_tiheys_flat, \
args = (V_hartree, ydin_tiheys, dx, dy, dz), \
V_hartree_xy = np.sum(V_hartree, axis=2)
print V_hartree_xy.shape
print V_hartree_xy
xopt.resize(V_hartree_xy.shape)
print xopt.shape
print xopt
#plot2d(V_hartree_xy, dx, dy, nx, ny)
plot2d(xopt, dx, dy, nx, ny)
#print "elektroneja oli ", np.sum(elektroni_tiheys_flat,axis=2)*dx*dy*dz
