christian

#!/usr/bin/env pythonimport os, sysimport numpy as npfrom ase import Atom, Atomsfrom ase.units import Bohrfrom gpaw import GPAW, restartfrom gpaw.mpi import worldfrom gpaw.poisson import PoissonSolverfrom gpaw.mixer import Mixerdef output(txt, comm=world):    if comm.rank == 0:        sys.stdout.write(txt)        sys.stdout.flush()    comm.barrier()def prepare(comm=world, parallel=None):    a = 6.0 # Ang    d = 3.22 # Ang    atoms = Atoms([Atom('Na', (a, d/2, a))],                  pbc=(False, True, False),                  cell=(2*a, d, 2*a))    # Determine number of kpoints per axis    pbc_c = atoms.get_pbc()    cell_c = np.sum(atoms.get_cell()**2, axis=1)**0.5 / Bohr    cellprod = 50 / Bohr    kpts_c = 1 + (np.ceil(cellprod / cell_c).astype(int) - 1) * pbc_c    # Determine number of grid points    h = 0.2 / Bohr    div = 16    gpts_c = div * np.round(cell_c / (h*div))    xc = 'LDA'    if parallel is None:        parallel = {}    calc = GPAW(xc=xc,                gpts=tuple(gpts_c),                kpts=tuple(kpts_c),                spinpol=False,                nbands=4, #2,                basis='dzp',                poissonsolver=PoissonSolver(3, 'GS'),                mixer=Mixer(0.05, 5, 50.0),                convergence={'eigenstates': 1e-7, 'density': 1e-3},                communicator=comm,                parallel=parallel,                )    atoms.set_calculator(calc)    return atoms, calcif __name__ in ['__main__', '__builtin__']:    """    if world.rank == 0 and not os.path.isfile('Na.dzp.basis'):        from gpaw.atom.basis import BasisMaker        basis = BasisMaker('Na', 'dzp', xc=xc).generate(2, 1)        basis.write_xml()        del basis    world.barrier()    from gpaw import setup_paths    setup_paths.insert(0, '.')    ...    del setup_paths[0]    """    filename = 'parallel_io_%02d.hdf5' % world.size    atoms, calc = prepare()    e0 = atoms.get_potential_energy()    calc.write(filename, mode='all')    del atoms, calc    runs = {1: [(1,1)],            2: [(1,1), (1,2), (2,1)],            4: [(1,1), (1,2), (2,1), (1,4), (2,2), (4,1)],            8: [(1,1), (1,2), (2,1), (1,4), (2,2), (4,1), (1,8), (2,4), (4,2), (8,1)],            }    # 16 bands (~25 electrons) and 8 spin/kpts (2,2,2)        commsizes = 2**np.arange(np.floor(np.log2(world.size))+1, dtype=int)    output('\n\nRestarting for commsizes: %s\n' % commsizes)    for commsize in commsizes:        for parsize, parsize_bands in runs[commsize]:            #assert commsize % (parsize * parsize_bands) == 0            comm = world.new_communicator(np.arange(commsize, dtype=int))            if comm is not None:                output('world: %d, ' % comm.size, comm)                atoms, calc = restart(filename,                                      parallel={'domain': parsize,                                                'band': parsize_bands},                                      communicator=comm,                                      txt=None)                calc.initialize_positions()                calc.scf.reset()                output('gd: %d, bd: %d, kd: %d, ' % (calc.wfs.gd.comm.size, calc.wfs.bd.comm.size, calc.wfs.kpt_comm.size), comm)                e1 = atoms.get_potential_energy()                niter = calc.get_number_of_iterations()                output('niter: %d, e0: %f, e1: %f, diff: %f\n' % (niter, e0, e1, abs(e1-e0)), comm)            world.barrier()