Example #1
0
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)
Example #2
0
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)