from ase import Atoms, PickleTrajectory, read
# from ase import Atoms
from gpaw import GPAW, Mixer, MixerSum, ConvergenceError, PoissonSolver
from gpaw.utilities.timing import TAUTimer
from gpaw.poisson import PoissonSolver
from gpaw.eigensolvers.rmm_diis import RMM_DIIS
# ps = PoissonSolver(nn='M', relax='GS', eps=1e-9)
es = RMM_DIIS(keep_htpsit=False)
# es = RMM_DIIS(keep_htpsit=False, blocksize=10)
from gpaw.hs_operators import MatrixOperator
MatrixOperator.nblocks = 1

# Makes Python traceback appear only on the root processor
import sys
from gpaw.mpi import world
from ase.utils import devnull
if world.rank != 0:
    sys.stderr = devnull

from gpaw import parsize_bands
cores = world.size
if parsize_bands is None:
    parsize_bands = 1
prefix = 'Au_bulk3x3x3_%d_%d' % (cores, parsize_bands)

class MyGPAW(GPAW):
    timer_class = TAUTimer
    
a = 4.08
bulk = Atoms('Au4', positions=((0, 0 ,0),
                               (0.5, 0.5, 0),
                               (0.5, 0, 0.5),
                               (0, 0.5, 0.5)), pbc=True)

bulk.set_cell((a,a,a), scale_atoms=True)
bulk = bulk.repeat((3, 3, 3))
calc = GPAW(h = 0.18,
            maxiter = 5,
            mode = 'fd',
#            poissonsolver = ps,
            nbands = 650,
            xc = 'LDA',
            width = 0.01,
            mixer = Mixer(0.10, 3),
            eigensolver = es,
            txt = prefix + '.out')

bulk.set_calculator(calc)

try:
    bulk.get_potential_energy()
except ConvergenceError:
    pass