def main(): from gpaw.grid_descriptor import GridDescriptor from gpaw.mpi import world serial = world.new_communicator([world.rank]) # Generator which must run on all ranks gen = np.random.RandomState(0) # This one is just used by master gen_serial = np.random.RandomState(17) maxsize = 5 for i in range(1): N1_c = gen.randint(1, maxsize, 3) N2_c = gen.randint(1, maxsize, 3) gd1 = GridDescriptor(N1_c, N1_c) gd2 = GridDescriptor(N2_c, N2_c) serial_gd1 = gd1.new_descriptor(comm=serial) serial_gd2 = gd2.new_descriptor(comm=serial) a1_serial = serial_gd1.empty() a1_serial.flat[:] = gen_serial.rand(a1_serial.size) if world.rank == 0: print('r0: a1 serial', a1_serial.ravel()) a1 = gd1.empty() a1[:] = -1 grid2grid(world, serial_gd1, gd1, a1_serial, a1) print(world.rank, 'a1 distributed', a1.ravel()) world.barrier() a2 = gd2.zeros() a2[:] = -2 grid2grid(world, gd1, gd2, a1, a2) print(world.rank, 'a2 distributed', a2.ravel()) world.barrier() #grid2grid(world, gd2, gd2_serial gd1 = GridDescriptor(N1_c, N1_c * 0.2) #serialgd = gd2.new_descriptor( a1 = gd1.empty() a1.flat[:] = gen.rand(a1.size) #print a1 grid2grid(world, gd1, gd2, a1, a2)
class BaseInducedField(object): """Partially virtual base class for induced field calculations. Attributes: ----------- omega_w: ndarray Frequencies for Fourier transform in atomic units folding: string Folding type ('Gauss' or 'Lorentz' or None) width: float Width parameter for folding Frho_wg: ndarray (complex) Fourier transform of induced charge density Fphi_wg: ndarray (complex) Fourier transform of induced electric potential Fef_wvg: ndarray (complex) Fourier transform of induced electric field Ffe_wg: ndarray (float) Fourier transform of field enhancement Fbgef_v: ndarray (float) Fourier transform of background electric field """ def __init__(self, filename=None, paw=None, frequencies=None, folding='Gauss', width=0.08, readmode=''): """ Parameters: ----------- filename: string Filename of a previous ``InducedField`` to be loaded. Setting filename disables parameters ``frequencies``, ``folding`` and ``width``. paw: PAW object PAW object for InducedField frequencies: ndarray or list of floats Frequencies in eV for Fourier transforms. This parameter is neglected if ``filename`` is given. folding: string Folding type: 'Gauss' or 'Lorentz' or None: Gaussian or Lorentzian folding or no folding. This parameter is neglected if ``filename`` is given. width: float Width in eV for the Gaussian (sigma) or Lorentzian (eta) folding This parameter is neglected if ``filename`` is given. """ self.dtype = complex # These are allocated when calculated self.fieldgd = None self.Frho_wg = None self.Fphi_wg = None self.Fef_wvg = None self.Ffe_wg = None self.Fbgef_v = None # has variables self.has_paw = False self.has_field = False if not hasattr(self, 'readwritemode_str_to_list'): self.readwritemode_str_to_list = \ {'': ['Frho', 'atoms'], 'all': ['Frho', 'Fphi', 'Fef', 'Ffe', 'atoms'], 'field': ['Frho', 'Fphi', 'Fef', 'Ffe', 'atoms']} if filename is not None: self.initialize(paw, allocate=False) self.read(filename, mode=readmode) return self.folding = folding self.width = width * eV_to_aufrequency self.set_folding(folding, width * eV_to_aufrequency) self.omega_w = np.asarray(frequencies) * eV_to_aufrequency self.nw = len(self.omega_w) self.nv = 3 # dimensionality of the space self.initialize(paw, allocate=True) def initialize(self, paw, allocate=True): self.allocated = False self.has_paw = paw is not None if self.has_paw: # If no paw is given, then the variables # marked with "# !" are created in _read(). # Other variables are not accessible without # paw (TODO: could be implemented...). self.paw = paw self.world = paw.wfs.world # ! self.domain_comm = paw.wfs.gd.comm # ! self.band_comm = paw.wfs.bd.comm # ! self.kpt_comm = paw.wfs.kd.comm # ! self.rank_a = paw.wfs.atom_partition.rank_a self.nspins = paw.density.nspins # ! self.setups = paw.wfs.setups self.density = paw.density self.atoms = paw.atoms # ! self.na = len(self.atoms.get_atomic_numbers()) # ! self.gd = self.density.gd # ! self.stencil = self.density.stencil if allocate: self.allocate() def allocate(self): self.allocated = True def deallocate(self): self.fieldgd = None self.Frho_wg = None self.Fphi_wg = None self.Fef_wvg = None self.Ffe_wg = None self.Fbgef_v = None def set_folding(self, folding, width): """ width: float Width in atomic units """ if width is None: folding = None self.folding = folding if self.folding is None: self.width = None else: self.width = width def get_induced_density(self, from_density, gridrefinement): raise RuntimeError('Virtual member function called') def calculate_induced_field(self, from_density='comp', gridrefinement=2, extend_N_cd=None, deextend=False, poissonsolver=None, gradient_n=3): if self.has_field and \ from_density == self.field_from_density and \ self.Frho_wg is not None: Frho_wg = self.Frho_wg gd = self.fieldgd else: Frho_wg, gd = self.get_induced_density(from_density, gridrefinement) # Always extend a bit to get field without jumps if extend_N_cd is None: extend_N = max(8, 2**int(np.ceil(np.log(gradient_n) / np.log(2.)))) extend_N_cd = extend_N * np.ones(shape=(3, 2), dtype=np.int) deextend = True # Extend grid oldgd = gd egd, cell_cv, move_c = \ extended_grid_descriptor(gd, extend_N_cd=extend_N_cd) Frho_we = egd.zeros((self.nw, ), dtype=self.dtype) for w in range(self.nw): extend_array(gd, egd, Frho_wg[w], Frho_we[w]) Frho_wg = Frho_we gd = egd if not deextend: # TODO: this will make atoms unusable with original grid self.atoms.set_cell(cell_cv, scale_atoms=False) move_atoms(self.atoms, move_c) # Allocate arrays Fphi_wg = gd.zeros((self.nw, ), dtype=self.dtype) Fef_wvg = gd.zeros(( self.nw, self.nv, ), dtype=self.dtype) Ffe_wg = gd.zeros((self.nw, ), dtype=float) # Poissonsolver if poissonsolver is None: poissonsolver = PoissonSolver(name='fd', eps=1e-20) poissonsolver.set_grid_descriptor(gd) for w in range(self.nw): # TODO: better output of progress # parprint('%d' % w) calculate_field(gd, Frho_wg[w], self.Fbgef_v, Fphi_wg[w], Fef_wvg[w], Ffe_wg[w], poissonsolver=poissonsolver, nv=self.nv, gradient_n=gradient_n) # De-extend grid if deextend: Frho_wo = oldgd.zeros((self.nw, ), dtype=self.dtype) Fphi_wo = oldgd.zeros((self.nw, ), dtype=self.dtype) Fef_wvo = oldgd.zeros(( self.nw, self.nv, ), dtype=self.dtype) Ffe_wo = oldgd.zeros((self.nw, ), dtype=float) for w in range(self.nw): deextend_array(oldgd, gd, Frho_wo[w], Frho_wg[w]) deextend_array(oldgd, gd, Fphi_wo[w], Fphi_wg[w]) deextend_array(oldgd, gd, Ffe_wo[w], Ffe_wg[w]) for v in range(self.nv): deextend_array(oldgd, gd, Fef_wvo[w][v], Fef_wvg[w][v]) Frho_wg = Frho_wo Fphi_wg = Fphi_wo Fef_wvg = Fef_wvo Ffe_wg = Ffe_wo gd = oldgd # Store results self.has_field = True self.field_from_density = from_density self.fieldgd = gd self.Frho_wg = Frho_wg self.Fphi_wg = Fphi_wg self.Fef_wvg = Fef_wvg self.Ffe_wg = Ffe_wg def _parse_readwritemode(self, mode): if isinstance(mode, str): try: readwrites = self.readwritemode_str_to_list[mode] except KeyError: raise IOError('unknown readwrite mode string') elif isinstance(mode, list): readwrites = mode else: raise IOError('unknown readwrite mode type') if any(k in readwrites for k in ['Frho', 'Fphi', 'Fef', 'Ffe']): readwrites.append('field') return readwrites def read(self, filename, mode='', idiotproof=True): if idiotproof and not filename.endswith('.ind'): raise IOError('Filename must end with `.ind`.') reads = self._parse_readwritemode(mode) # Open reader from gpaw.io import Reader reader = Reader(filename) self._read(reader, reads) reader.close() self.world.barrier() def _read(self, reader, reads): r = reader # Test data type dtype = {'float': float, 'complex': complex}[r.dtype] if dtype != self.dtype: raise IOError('Data is an incompatible type.') # Read dimensions na = r.na self.nv = r.nv nspins = r.nspins ng = r.ng # Background electric field Fbgef_v = r.Fbgef_v if self.has_paw: # Test dimensions if na != self.na: raise IOError('natoms is incompatible with calculator') if nspins != self.nspins: raise IOError('nspins is incompatible with calculator') if (ng != self.gd.get_size_of_global_array()).any(): raise IOError('grid is incompatible with calculator') if (Fbgef_v != self.Fbgef_v).any(): raise IOError('kick is incompatible with calculator') else: # Construct objects / assign values without paw self.na = na self.nspins = nspins self.Fbgef_v = Fbgef_v from ase.io.trajectory import read_atoms self.atoms = read_atoms(r.atoms) self.world = mpi.world self.gd = GridDescriptor(ng + 1, self.atoms.get_cell() / Bohr, pbc_c=False, comm=self.world) self.domain_comm = self.gd.comm self.band_comm = mpi.SerialCommunicator() self.kpt_comm = mpi.SerialCommunicator() # Folding folding = r.folding width = r.width self.set_folding(folding, width) # Frequencies self.omega_w = r.omega_w self.nw = len(self.omega_w) # Read field if 'field' in reads: nfieldg = r.nfieldg self.has_field = True self.field_from_density = r.field_from_density self.fieldgd = self.gd.new_descriptor(N_c=nfieldg + 1) def readarray(name, shape, dtype): if name.split('_')[0] in reads: setattr(self, name, self.fieldgd.empty(shape, dtype=dtype)) self.fieldgd.distribute(r.get(name), getattr(self, name)) readarray('Frho_wg', (self.nw, ), self.dtype) readarray('Fphi_wg', (self.nw, ), self.dtype) readarray('Fef_wvg', (self.nw, self.nv), self.dtype) readarray('Ffe_wg', (self.nw, ), float) def write(self, filename, mode='', idiotproof=True): """ Parameters ---------- mode: string or list of strings """ if idiotproof and not filename.endswith('.ind'): raise IOError('Filename must end with `.ind`.') writes = self._parse_readwritemode(mode) if 'field' in writes and self.fieldgd is None: raise IOError('field variables cannot be written ' + 'before they are calculated') from gpaw.io import Writer writer = Writer(filename, self.world, tag='INDUCEDFIELD') # Actual write self._write(writer, writes) # Make sure slaves don't return before master is done writer.close() self.world.barrier() def _write(self, writer, writes): # Write parameters/dimensions writer.write(dtype={ float: 'float', complex: 'complex' }[self.dtype], folding=self.folding, width=self.width, na=self.na, nv=self.nv, nspins=self.nspins, ng=self.gd.get_size_of_global_array()) # Write field grid if 'field' in writes: writer.write(field_from_density=self.field_from_density, nfieldg=self.fieldgd.get_size_of_global_array()) # Write frequencies writer.write(omega_w=self.omega_w) # Write background electric field writer.write(Fbgef_v=self.Fbgef_v) from ase.io.trajectory import write_atoms write_atoms(writer.child('atoms'), self.atoms) if 'field' in writes: def writearray(name, shape, dtype): if name.split('_')[0] in writes: writer.add_array(name, shape, dtype) a_wxg = getattr(self, name) for w in range(self.nw): writer.fill(self.fieldgd.collect(a_wxg[w])) shape = (self.nw, ) + tuple( self.fieldgd.get_size_of_global_array()) writearray('Frho_wg', shape, self.dtype) writearray('Fphi_wg', shape, self.dtype) writearray('Ffe_wg', shape, float) shape3 = shape[:1] + (self.nv, ) + shape[1:] writearray('Fef_wvg', shape3, self.dtype)