Index: gpaw/lrtddft/omega_matrix.py =================================================================== --- gpaw/lrtddft/omega_matrix.py (revisión: 12488) +++ gpaw/lrtddft/omega_matrix.py (copia de trabajo) @@ -172,13 +172,16 @@ else: nt_sg = paw.density.nt_sg # check if D_sp have been changed before - D_asp = self.paw.density.D_asp - for a, D_sp in D_asp.items(): + D_asp = {} + for a, D_sp in self.paw.density.D_asp.items(): if len(D_sp) != kss.npspins: if len(D_sp) == 1: + x = np.array([0.5 * D_sp[0], 0.5 * D_sp[0]]) D_asp[a] = np.array([0.5 * D_sp[0], 0.5 * D_sp[0]]) else: D_asp[a] = np.array([D_sp[0] + D_sp[1]]) + else: + D_asp[a] = D_sp.copy() # restrict the density if needed if fg: @@ -252,13 +255,13 @@ P_ii = np.outer(Pi_i, Pj_i) # we need the symmetric form, hence we can pack P_p = pack(P_ii) - D_sp = self.paw.density.D_asp[a].copy() + D_sp = D_asp[a].copy() D_sp[kss[ij].pspin] -= ns * P_p setup = wfs.setups[a] I_sp = np.zeros_like(D_sp) self.xc.calculate_paw_correction(setup, D_sp, I_sp) I_sp *= -1.0 - D_sp = self.paw.density.D_asp[a].copy() + D_sp = D_asp[a].copy() D_sp[kss[ij].pspin] += ns * P_p self.xc.calculate_paw_correction(setup, D_sp, I_sp) I_sp /= 2.0 * ns Index: gpaw/lrtddft/excited_state.py =================================================================== --- gpaw/lrtddft/excited_state.py (revisión: 12488) +++ gpaw/lrtddft/excited_state.py (copia de trabajo) @@ -346,9 +346,6 @@ self.gsdensity = calc.density self.gd = self.gsdensity.gd self.nbands = calc.wfs.bd.nbands - self.D_asp = {} - for a, D_sp in self.gsdensity.D_asp.items(): - self.D_asp[a] = 1. * D_sp # obtain weights ex = lrtddft[index] @@ -374,6 +371,13 @@ spos_ac = calc.get_atoms().get_scaled_positions() % 1.0 self.set_positions(spos_ac, calc.wfs.rank_a) + D_asp = {} + for a, D_sp in self.gsdensity.D_asp.items(): + repeats = self.ns // self.gsdensity.ns + # XXX does this work always? + D_asp[a] = (1. * D_sp).repeat(repeats, axis=0) + self.D_asp = D_asp + def update(self, wfs): self.timer.start('Density') self.timer.start('Pseudo density') Index: gpaw/arraydict.py =================================================================== --- gpaw/arraydict.py (revisión: 0) +++ gpaw/arraydict.py (revisión: 0) @@ -0,0 +1,82 @@ +import numpy as np + +class ArrayDict(dict): + def __init__(self, partition, shapes_a, dtype=float, d=None): + dict.__init__(self) + self.partition = partition + self.shapes_a = shapes_a # global + self.dtype = dtype + if d is None: + for a in partition.my_indices: + self[a] = np.empty(self.shapes_a[a], dtype=dtype) + else: + self.update(d) + self.check_consistency() + + def copy(self): + return ArrayDict(self.partition, self.shapes_a, self.dtype, self) + + def deepcopy(self): + copy = Arraydict(self.partition, self.shapes_a, self.dtype) + for a in self: + copy[a] = self[a].copy() + return copy + + def update(self, d): + dict.update(self, d) + self.check_consistency() + + def __getitem__(self, a): + value = dict.__getitem__(self, a) + assert value.shape == self.shapes_a[a] + assert value.dtype == self.dtype + return value + + def __setitem__(self, a, value): + assert value.shape == self.shapes_a[a], \ + 'defined shape %s vs new %s' % (self.shapes_a[a], value.shape) + assert value.dtype == self.dtype + dict.__setitem__(self, a, value) + + #def set(self, d): + # for a in d: + # assert a in self + # self[a][:] = d[a] + # self.check_consistency() + + def redistribute(self, partition): + def get_empty(a): + return np.empty(self.shapes_a[a]) + + self.partition.redistribute(partition, self, get_empty) + self.partition = partition # Better with immutable partition? + self.check_consistency() + + def check_consistency(self): + k1 = set(self.keys()) + k2 = set(self.partition.my_indices) + assert k1 == k2, 'Inconsistent keys %s vs from %s' % (k1, k2) + for a, array in self.items(): + assert array.dtype == self.dtype + assert array.shape == self.shapes_a[a], \ + 'array shape %s vs specified shape %s' % (array.shape, + self.shapes_a[a]) + + #def to_parent_comm(self): + # new_partition = self.partition.to_parent_comm() + # arraydict = ArrayDict(new_partition, self.shapes_a, dtype=float) + # assert arraydict.partition.comm.size >= self.partition.comm.size + # + # print self, arraydict + # # XXXX ? unfinished... + + def __str__(self): + tokens = [] + for key in sorted(self.keys()): + shapestr = 'x'.join(map(str, self.shapes_a[key])) + tokens.append('%s:%s' % (key, shapestr)) + text = ', '.join(tokens) + return '%s@rank%d/%d {%s}' % (self.__class__.__name__, + self.partition.comm.rank, + self.partition.comm.size, + text) Index: gpaw/utilities/partition.py =================================================================== --- gpaw/utilities/partition.py (revisión: 12488) +++ gpaw/utilities/partition.py (copia de trabajo) @@ -1,6 +1,6 @@ import numpy as np +from gpaw.arraydict import ArrayDict - class AtomicMatrixDistributor: """Class to distribute atomic dictionaries like dH_asp and D_asp.""" def __init__(self, atom_partition, setups, kptband_comm, ns): @@ -284,3 +284,11 @@ general_redistribute(self.comm, self.rank_a, new_partition.rank_a, Redist()) + if isinstance(atomdict_ax, ArrayDict): + atomdict_ax.partition = new_partition # XXX + atomdict_ax.check_consistency() + + def arraydict(self, shapes, dtype=float): + if callable(shapes): + shapes = [shapes(a) for a in self.natoms] + return ArrayDict(self, shapes, dtype) Index: gpaw/analyse/hirshfeld.py =================================================================== --- gpaw/analyse/hirshfeld.py (revisión: 12488) +++ gpaw/analyse/hirshfeld.py (copia de trabajo) @@ -8,6 +8,7 @@ from gpaw.setup import Setups from gpaw.xc import XC from gpaw.utilities.tools import coordinates +from gpaw.utilities.partition import AtomPartition from gpaw.mpi import rank @@ -27,6 +28,7 @@ """HirshfeldDensity builds a hack density object to calculate all electron density of atoms. This methods overrides the parallel distribution of atomic density matrices in density.py""" + self.atom_partition = AtomPartition(self.gd.comm, rank_a) self.nct.set_positions(spos_ac) self.ghat.set_positions(spos_ac) self.mixer.reset() @@ -75,15 +77,16 @@ setups = Setups(Z_a, par.setups, par.basis, par.lmax, XC(par.xc), self.calculator.wfs.world) - self.D_asp = D_asp # initialize self.initialize(setups, self.calculator.timer, np.zeros((len(atoms), 3)), False) self.set_mixer(None) + # FIXME nparray causes partitionong.py test to fail self.set_positions(spos_ac, np.array(rank_a)) + self.D_asp = D_asp basis_functions = BasisFunctions(self.gd, [setup.phit_j for setup in self.setups], Index: gpaw/density.py =================================================================== --- gpaw/density.py (revisión: 12488) +++ gpaw/density.py (copia de trabajo) @@ -18,9 +18,9 @@ from gpaw.utilities.timing import nulltimer from gpaw.io import read_atomic_matrices from gpaw.mpi import SerialCommunicator +from gpaw.arraydict import ArrayDict - -class Density: +class Density(object): """Density object. Attributes: @@ -95,19 +95,20 @@ # If both old and new atomic ranks are present, start a blank dict if # it previously didn't exist but it will needed for the new atoms. - assert rank_a is not None if (self.rank_a is not None and self.D_asp is None and (rank_a == self.gd.comm.rank).any()): - self.D_asp = {} + self.D_asp = self.setups.empty_asp(self.ns, self.atom_partition) if (self.rank_a is not None and self.D_asp is not None and not isinstance(self.gd.comm, SerialCommunicator)): self.timer.start('Redistribute') - def get_empty(a): - ni = self.setups[a].ni - return np.empty((self.ns, ni * (ni + 1) // 2)) - self.atom_partition.redistribute(atom_partition, self.D_asp, - get_empty) + #def get_empty(a): + # ni = self.setups[a].ni + # return np.empty((self.ns, ni * (ni + 1) // 2)) + #self.atom_partition.redistribute(atom_partition, self.D_asp, + # get_empty) + self.D_asp.redistribute(atom_partition) + self.D_asp.check_consistency() self.timer.stop('Redistribute') self.rank_a = rank_a @@ -122,6 +123,24 @@ wfs.calculate_density_contribution(self.nt_sG) self.nt_sG[:self.nspins] += self.nct_G + @property + def D_asp(self): + if self._D_asp is not None: + assert isinstance(self._D_asp, ArrayDict), type(self._D_asp) + self._D_asp.check_consistency() + return self._D_asp + + @D_asp.setter + def D_asp(self, value): + if isinstance(value, dict): + tmp = self.setups.empty_asp(self.ns, self.atom_partition) + tmp.update(value) + value = tmp + assert isinstance(value, ArrayDict) or value is None, type(value) + if value is not None: + value.check_consistency() + self._D_asp = value + def update(self, wfs): self.timer.start('Density') self.timer.start('Pseudo density') @@ -201,7 +220,7 @@ # but is not particularly efficient (not that this is a time # consuming step) - self.D_asp = {} + self.D_asp = self.setups.empty_asp(self.ns, self.atom_partition) f_asi = {} for a in basis_functions.atom_indices: c = self.charge / len(self.setups) # distribute on all atoms @@ -236,11 +255,8 @@ self.timer.start("Density initialize from wavefunctions") self.nt_sG = self.gd.empty(self.ns) self.calculate_pseudo_density(wfs) - self.D_asp = {} + self.D_asp = self.setups.empty_asp(self.ns, self.atom_partition) my_atom_indices = np.argwhere(wfs.rank_a == self.gd.comm.rank).ravel() - for a in my_atom_indices: - ni = self.setups[a].ni - self.D_asp[a] = np.empty((self.nspins, ni * (ni + 1) // 2)) wfs.calculate_atomic_density_matrices(self.D_asp) self.calculate_normalized_charges_and_mix() self.timer.stop("Density initialize from wavefunctions") @@ -249,6 +265,7 @@ """Set D_asp and nt_sG directly.""" self.nt_sG = nt_sG self.D_asp = D_asp + D_asp.check_consistency() #self.calculate_normalized_charges_and_mix() # No calculate multipole moments? Tests will fail because of # improperly initialized mixer @@ -485,12 +502,16 @@ nt_sG[s]) # Read atomic density matrices - D_asp = {} natoms = len(self.setups) + atom_partition = AtomPartition(self.gd.comm, np.zeros(natoms, int)) + D_asp = self.setups.empty_asp(self.ns, atom_partition) + self.atom_partition = atom_partition # XXXXXX + self.rank_a = np.zeros(natoms, int) all_D_sp = reader.get('AtomicDensityMatrices', broadcast=True) if self.gd.comm.rank == 0: - D_asp = read_atomic_matrices(all_D_sp, self.setups) + D_asp.update(read_atomic_matrices(all_D_sp, self.setups)) + D_asp.check_consistency() self.initialize_directly_from_arrays(nt_sG, D_asp) Index: gpaw/setup.py =================================================================== --- gpaw/setup.py (revisión: 12488) +++ gpaw/setup.py (copia de trabajo) @@ -1381,7 +1381,12 @@ for setup in self.setups.values(): setup.calculate_rotations(R_slmm) + def empty_asp(self, ns, atom_partition): + Dshapes_a = [(ns, setup.ni * (setup.ni + 1) // 2) + for setup in self] + return atom_partition.arraydict(Dshapes_a) + def types2atomtypes(symbols, types, default): """Map a types identifier to a list with a type id for each atom. Index: gpaw/hamiltonian.py =================================================================== --- gpaw/hamiltonian.py (revisión: 12488) +++ gpaw/hamiltonian.py (copia de trabajo) @@ -14,9 +14,10 @@ from gpaw.utilities import pack2, unpack, unpack2 from gpaw.io import read_atomic_matrices from gpaw.utilities.partition import AtomPartition, AtomicMatrixDistributor +from gpaw.arraydict import ArrayDict -class Hamiltonian: +class Hamiltonian(object): """Hamiltonian object. Attributes: @@ -91,7 +92,23 @@ self.ref_vt_sG = None self.ref_dH_asp = None + @property + def dH_asp(self): + assert isinstance(self._dH_asp, ArrayDict) or self._dH_asp is None, type(self._dH_asp) + self._dH_asp.check_consistency() + return self._dH_asp + @dH_asp.setter + def dH_asp(self, value): + if isinstance(value, dict): + tmp = self.setups.empty_asp(self.ns, self.atom_partition) + tmp.update(value) + value = tmp + assert isinstance(value, ArrayDict) or value is None, type(value) + if value is not None: + value.check_consistency() + self._dH_asp = value + def summary(self, fd): fd.write('XC and Coulomb potentials evaluated on a %d*%d*%d grid\n' % tuple(self.finegd.N_c)) @@ -505,10 +522,12 @@ self.gd.distribute(reader.get('PseudoPotential', s), self.vt_sG[s]) + natoms = len(self.setups) + self.rank_a = np.zeros(natoms, int) + self.atom_partition = AtomPartition(self.gd.comm, self.rank_a) + # Read non-local part of hamiltonian self.dH_asp = {} - natoms = len(self.setups) - self.rank_a = np.zeros(natoms, int) if version > 0.3: all_H_sp = reader.get('NonLocalPartOfHamiltonian', broadcast=True)