from ase import Atoms, PickleTrajectory
from ase.optimize.lbfgs import LBFGS, LBFGSLineSearch
from gpaw import GPAW, Mixer, MixerSum, ConvergenceError, PoissonSolver
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)
from gpaw.hs_operators import MatrixOperator
MatrixOperator.nblocks = 2

# 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


bulk = read('qn.traj')

calc = GPAW(gpts=(120, 120, 120),
            maxiter=100,
            mode = 'fd',
            poissonsolver = ps,
            nbands = 600,
            spinpol = False,
            xc = 'LDA',
            width = 0.01,
            mixer = Mixer(0.10, 5, 100.0), 
            # mixer = MixerSum(0.1, 5, 100.0), # spinpol = True
            eigensolver = es,
            txt = 'Au_bulk.out')

bulk.set_calculator(calc)
# spinpol = True
# magmom = 2.0
# mms = [magmom for i in range(len(bulk))]
# bulk.set_initial_magnetic_memonts(mms)

world.barrier() # make sure things are synchronized

mode = 'a'
traj = PickleTrajectory('qn.traj', mode, bulk)
relax = LBFGS(bulk, logfile = 'qn.log', trajectory = traj, restart= 'qn.pckl')
relax.replay_trajectory('qn.traj')
relax.run(fmax=0.05, steps=10)