Exemplo n.º 1
0
    def test_op4(self):
        S = SparseCSR((10, 100), dtype=np.int32)
        # Create initial stuff
        for i in range(10):
            j = range(i * 4, i * 4 + 3)
            S[0, j] = i

        s = 1 + S
        assert_equal(s.dtype, np.int32)
        s = 1. + S
        assert_equal(s.dtype, np.float64)
        s = 1.j + S
        assert_equal(s.dtype, np.complex128)

        s = 1 - S
        assert_equal(s.dtype, np.int32)
        s = 1. - S
        assert_equal(s.dtype, np.float64)
        s = 1.j - S
        assert_equal(s.dtype, np.complex128)

        s = 1 * S
        assert_equal(s.dtype, np.int32)
        s = 1. * S
        assert_equal(s.dtype, np.float64)
        s = 1.j * S
        assert_equal(s.dtype, np.complex128)

        s = 1**S
        assert_equal(s.dtype, np.int32)
        s = 1.**S
        assert_equal(s.dtype, np.float64)
        s = 1.j**S
        assert_equal(s.dtype, np.complex128)
Exemplo n.º 2
0
 def test_init3(self):
     csr = sc.sparse.csr_matrix((10, 10), dtype=np.int32)
     csr[0, 1] = 1
     csr[0, 2] = 2
     sp = SparseCSR(csr)
     assert_equal(sp.dtype, np.int32)
     assert_equal(sp.shape, (10, 10, 1))
     assert_equal(len(sp), 2)
     assert_equal(sp[0, 1], 1)
     assert_equal(sp[0, 2], 2)
     sp = SparseCSR(csr, dtype=np.float64)
     assert_equal(sp.shape, (10, 10, 1))
     assert_equal(sp.dtype, np.float64)
     assert_equal(len(sp), 2)
     assert_equal(sp[0, 1], 1)
     assert_equal(sp[0, 2], 2)
Exemplo n.º 3
0
 def test_init4(self):
     csr = sc.sparse.csr_matrix((10, 10), dtype=np.int32)
     csr[0, 1] = 1
     csr[0, 2] = 2
     print(csr.indices, csr.indptr)
     sp = SparseCSR((csr.data, csr.indices, csr.indptr))
     assert_equal(sp.dtype, np.int32)
     assert_equal(sp.shape, (10, 10, 1))
     assert_equal(len(sp), 2)
     assert_equal(sp[0, 1], 1)
     assert_equal(sp[0, 2], 2)
     sp = SparseCSR((csr.data, csr.indices, csr.indptr), dtype=np.float64)
     assert_equal(sp.shape, (10, 10, 1))
     assert_equal(sp.dtype, np.float64)
     assert_equal(len(sp), 2)
     assert_equal(sp[0, 1], 1)
     assert_equal(sp[0, 2], 2)
Exemplo n.º 4
0
 def test_init2(self):
     SparseCSR((10, 100))
     for d in [np.int32, np.float64, np.complex128]:
         s = SparseCSR((10, 100), dtype=d)
         assert_equal(s.shape, (10, 100, 1))
         assert_equal(s.dim, 1)
         assert_equal(s.dtype, d)
         for k in [1, 2]:
             s = SparseCSR((10, 100, k), dtype=d)
             assert_equal(s.shape, (10, 100, k))
             assert_equal(s.dim, k)
             s = SparseCSR((10, 100), dim=k, dtype=d)
             assert_equal(s.shape, (10, 100, k))
             assert_equal(s.dim, k)
             s = SparseCSR((10, 100, 3), dim=k, dtype=d)
             assert_equal(s.shape, (10, 100, 3))
             assert_equal(s.dim, 3)
Exemplo n.º 5
0
    def test_op3(self):
        S = SparseCSR((10, 100), dtype=np.int32)
        # Create initial stuff
        for i in range(10):
            j = range(i * 4, i * 4 + 3)
            S[0, j] = i

        for op in ['add', 'sub', 'mul', 'pow']:
            func = getattr(S, '__{}__'.format(op))
            s = func(1)
            assert_equal(s.dtype, np.int32)
            s = func(1.)
            assert_equal(s.dtype, np.float64)
            if op != 'pow':
                s = func(1.j)
                assert_equal(s.dtype, np.complex128)

        S = S.copy(dtype=np.float64)
        for op in ['add', 'sub', 'mul', 'pow']:
            func = getattr(S, '__{}__'.format(op))
            s = func(1)
            assert_equal(s.dtype, np.float64)
            s = func(1.)
            assert_equal(s.dtype, np.float64)
            if op != 'pow':
                s = func(1.j)
                assert_equal(s.dtype, np.complex128)

        S = S.copy(dtype=np.complex128)
        for op in ['add', 'sub', 'mul', 'pow']:
            func = getattr(S, '__{}__'.format(op))
            s = func(1)
            assert_equal(s.dtype, np.complex128)
            s = func(1.)
            assert_equal(s.dtype, np.complex128)
            if op != 'pow':
                s = func(1.j)
                assert_equal(s.dtype, np.complex128)
Exemplo n.º 6
0
    def test_op3(self):
        S = SparseCSR((10,100), dtype=np.int32)
        # Create initial stuff
        for i in range(10):
            j = range(i*4, i*4+3)
            S[0, j] = i

        for op in ['add', 'sub', 'mul', 'pow']:
            func = getattr(S, '__{}__'.format(op))
            s = func(1)
            assert_equal(s.dtype, np.int32)
            s = func(1.)
            assert_equal(s.dtype, np.float64)
            if op != 'pow':
                s = func(1.j)
                assert_equal(s.dtype, np.complex128)

        S = S.copy(dtype=np.float64)
        for op in ['add', 'sub', 'mul', 'pow']:
            func = getattr(S, '__{}__'.format(op))
            s = func(1)
            assert_equal(s.dtype, np.float64)
            s = func(1.)
            assert_equal(s.dtype, np.float64)
            if op != 'pow':
                s = func(1.j)
                assert_equal(s.dtype, np.complex128)

        S = S.copy(dtype=np.complex128)
        for op in ['add', 'sub', 'mul', 'pow']:
            func = getattr(S, '__{}__'.format(op))
            s = func(1)
            assert_equal(s.dtype, np.complex128)
            s = func(1.)
            assert_equal(s.dtype, np.complex128)
            if op != 'pow':
                s = func(1.j)
                assert_equal(s.dtype, np.complex128)
    def density(self, grid, spinor=None, tol=1e-7, eta=False):
        r""" Expand the density matrix to the charge density on a grid

        This routine calculates the real-space density components on a specified grid.

        This is an *in-place* operation that *adds* to the current values in the grid.

        Note: To calculate :math:`\rho(\mathbf r)` in a unit-cell different from the
        originating geometry, simply pass a grid with a unit-cell different than the originating
        supercell.

        The real-space density is calculated as:

        .. math::
            \rho(\mathbf r) = \sum_{\nu\mu}\phi_\nu(\mathbf r)\phi_\mu(\mathbf r) D_{\nu\mu}

        While for non-collinear/spin-orbit calculations the density is determined from the
        spinor component (`spinor`) by

        .. math::
           \rho_{\boldsymbol\sigma}(\mathbf r) = \sum_{\nu\mu}\phi_\nu(\mathbf r)\phi_\mu(\mathbf r) \sum_\alpha [\boldsymbol\sigma \mathbf \rho_{\nu\mu}]_{\alpha\alpha}

        Here :math:`\boldsymbol\sigma` corresponds to a spinor operator to extract relevant quantities. By passing the identity matrix the total charge is added. By using the Pauli matrix :math:`\boldsymbol\sigma_x`
        only the :math:`x` component of the density is added to the grid (see `Spin.X`).

        Parameters
        ----------
        grid : Grid
           the grid on which to add the density (the density is in ``e/Ang^3``)
        spinor : (2,) or (2, 2), optional
           the spinor matrix to obtain the diagonal components of the density. For un-polarized density matrices
           this keyword has no influence. For spin-polarized it *has* to be either 1 integer or a vector of
           length 2 (defaults to total density).
           For non-collinear/spin-orbit density matrices it has to be a 2x2 matrix (defaults to total density).
        tol : float, optional
           DM tolerance for accepted values. For all density matrix elements with absolute values below
           the tolerance, they will be treated as strictly zeros.
        eta: bool, optional
           show a progressbar on stdout
        """
        try:
            # Once unique has the axis keyword, we know we can safely
            # use it in this routine
            # Otherwise we raise an ImportError
            unique([[0, 1], [2, 3]], axis=0)
        except:
            raise NotImplementedError(
                self.__class__.__name__ +
                '.density requires numpy >= 1.13, either update '
                'numpy or do not use this function!')

        geometry = self.geometry
        # Check that the atomic coordinates, really are all within the intrinsic supercell.
        # If not, it may mean that the DM does not conform to the primary unit-cell paradigm
        # of matrix elements. It complicates things.
        fxyz = geometry.fxyz
        f_min = fxyz.min()
        f_max = fxyz.max()
        if f_min < 0 or 1. < f_max:
            warn(
                self.__class__.__name__ +
                '.density has been passed a geometry where some coordinates are '
                'outside the primary unit-cell. This may potentially lead to problems! '
                'Double check the charge density!')
        del fxyz, f_min, f_max

        # Extract sub variables used throughout the loop
        shape = _a.asarrayi(grid.shape)
        dcell = grid.dcell

        # Sparse matrix data
        csr = self._csr

        # In the following we don't care about division
        # So 1) save error state, 2) turn off divide by 0, 3) calculate, 4) turn on old error state
        old_err = np.seterr(divide='ignore', invalid='ignore')

        # Placeholder for the resulting coefficients
        DM = None
        if self.spin.kind > Spin.POLARIZED:
            if spinor is None:
                # Default to the total density
                spinor = np.identity(2, dtype=np.complex128)
            else:
                spinor = _a.arrayz(spinor)
            if spinor.size != 4 or spinor.ndim != 2:
                raise ValueError(
                    self.__class__.__name__ +
                    '.density with NC/SO spin, requires a 2x2 matrix.')

            DM = _a.emptyz([self.nnz, 2, 2])
            idx = array_arange(csr.ptr[:-1], n=csr.ncol)
            if self.spin.kind == Spin.NONCOLINEAR:
                # non-collinear
                DM[:, 0, 0] = csr._D[idx, 0]
                DM[:, 1, 1] = csr._D[idx, 1]
                DM[:, 1,
                   0] = csr._D[idx,
                               2] - 1j * csr._D[idx, 3]  #TODO check sign here!
                DM[:, 0, 1] = np.conj(DM[:, 1, 0])
            else:
                # spin-orbit
                DM[:, 0, 0] = csr._D[idx, 0] + 1j * csr._D[idx, 4]
                DM[:, 1, 1] = csr._D[idx, 1] + 1j * csr._D[idx, 5]
                DM[:, 1,
                   0] = csr._D[idx,
                               2] - 1j * csr._D[idx, 3]  #TODO check sign here!
                DM[:, 0, 1] = csr._D[idx, 6] + 1j * csr._D[idx, 7]

            # Perform dot-product with spinor, and take out the diagonal real part
            DM = dot(DM, spinor.T)[:, [0, 1], [0, 1]].sum(1).real

        elif self.spin.kind == Spin.POLARIZED:
            if spinor is None:
                spinor = _a.onesd(2)

            elif isinstance(spinor, Integral):
                # extract the provided spin-polarization
                s = _a.zerosd(2)
                s[spinor] = 1.
                spinor = s
            else:
                spinor = _a.arrayd(spinor)

            if spinor.size != 2 or spinor.ndim != 1:
                raise ValueError(
                    self.__class__.__name__ +
                    '.density with polarized spin, requires spinor '
                    'argument as an integer, or a vector of length 2')

            idx = array_arange(csr.ptr[:-1], n=csr.ncol)
            DM = csr._D[idx, 0] * spinor[0] + csr._D[idx, 1] * spinor[1]

        else:
            idx = array_arange(csr.ptr[:-1], n=csr.ncol)
            DM = csr._D[idx, 0]

        # Create the DM csr matrix.
        csrDM = csr_matrix(
            (DM, csr.col[idx], np.insert(np.cumsum(csr.ncol), 0, 0)),
            shape=(self.shape[:2]),
            dtype=DM.dtype)

        # Clean-up
        del idx, DM

        # To heavily speed up the construction of the density we can recreate
        # the sparse csrDM matrix by summing the lower and upper triangular part.
        # This means we only traverse the sparse UPPER part of the DM matrix
        # I.e.:
        #    psi_i * DM_{ij} * psi_j + psi_j * DM_{ji} * psi_i
        # is equal to:
        #    psi_i * (DM_{ij} + DM_{ji}) * psi_j
        # Secondly, to ease the loops we extract the main diagonal (on-site terms)
        # and store this for separate usage
        csr_sum = [None] * geometry.n_s
        no = geometry.no
        primary_i_s = geometry.sc_index([0, 0, 0])
        for i_s in range(geometry.n_s):
            # Extract the csr matrix
            o_start, o_end = i_s * no, (i_s + 1) * no
            csr = csrDM[:, o_start:o_end]
            if i_s == primary_i_s:
                csr_sum[i_s] = triu(csr) + tril(csr, -1).transpose()
            else:
                csr_sum[i_s] = csr

        # Recreate the column-stacked csr matrix
        csrDM = ss_hstack(csr_sum, format='csr')
        del csr, csr_sum

        # Remove all zero elements (note we use the tolerance here!)
        csrDM.data = np.where(np.fabs(csrDM.data) > tol, csrDM.data, 0.)

        # Eliminate zeros and sort indices etc.
        csrDM.eliminate_zeros()
        csrDM.sort_indices()
        csrDM.prune()

        # 1. Ensure the grid has a geometry associated with it
        sc = grid.sc.copy()
        if grid.geometry is None:
            # Create the actual geometry that encompass the grid
            ia, xyz, _ = geometry.within_inf(sc)
            if len(ia) > 0:
                grid.set_geometry(Geometry(xyz, geometry.atom[ia], sc=sc))

        # Instead of looping all atoms in the supercell we find the exact atoms
        # and their supercell indices.
        add_R = _a.zerosd(3) + geometry.maxR()
        # Calculate the required additional vectors required to increase the fictitious
        # supercell by add_R in each direction.
        # For extremely skewed lattices this will be way too much, hence we make
        # them square.
        o = sc.toCuboid(True)
        sc = SuperCell(o._v, origo=o.origo) + np.diag(2 * add_R)
        sc.origo -= add_R

        # Retrieve all atoms within the grid supercell
        # (and the neighbours that connect into the cell)
        IA, XYZ, ISC = geometry.within_inf(sc)

        # Retrieve progressbar
        eta = tqdm_eta(len(IA), self.__class__.__name__ + '.density', 'atom',
                       eta)

        cell = geometry.cell
        atom = geometry.atom
        axyz = geometry.axyz
        a2o = geometry.a2o

        def xyz2spherical(xyz, offset):
            """ Calculate the spherical coordinates from indices """
            rx = xyz[:, 0] - offset[0]
            ry = xyz[:, 1] - offset[1]
            rz = xyz[:, 2] - offset[2]

            # Calculate radius ** 2
            xyz_to_spherical_cos_phi(rx, ry, rz)
            return rx, ry, rz

        def xyz2sphericalR(xyz, offset, R):
            """ Calculate the spherical coordinates from indices """
            rx = xyz[:, 0] - offset[0]
            idx = indices_fabs_le(rx, R)
            ry = xyz[idx, 1] - offset[1]
            ix = indices_fabs_le(ry, R)
            ry = ry[ix]
            idx = idx[ix]
            rz = xyz[idx, 2] - offset[2]
            ix = indices_fabs_le(rz, R)
            ry = ry[ix]
            rz = rz[ix]
            idx = idx[ix]
            if len(idx) == 0:
                return [], [], [], []
            rx = rx[idx]

            # Calculate radius ** 2
            ix = indices_le(rx**2 + ry**2 + rz**2, R**2)
            idx = idx[ix]
            if len(idx) == 0:
                return [], [], [], []
            rx = rx[ix]
            ry = ry[ix]
            rz = rz[ix]
            xyz_to_spherical_cos_phi(rx, ry, rz)
            return idx, rx, ry, rz

        # Looping atoms in the sparse pattern is better since we can pre-calculate
        # the radial parts and then add them.
        # First create a SparseOrbital matrix, then convert to SparseAtom
        spO = SparseOrbital(geometry, dtype=np.int16)
        spO._csr = SparseCSR(csrDM)
        spA = spO.toSparseAtom(dtype=np.int16)
        del spO
        na = geometry.na
        # Remove the diagonal part of the sparse atom matrix
        off = na * primary_i_s
        for ia in range(na):
            del spA[ia, off + ia]

        # Get pointers and delete the atomic sparse pattern
        # The below complexity is because we are not finalizing spA
        csr = spA._csr
        a_ptr = np.insert(_a.cumsumi(csr.ncol), 0, 0)
        a_col = csr.col[array_arange(csr.ptr, n=csr.ncol)]
        del spA, csr

        # Get offset in supercell in orbitals
        off = geometry.no * primary_i_s
        origo = grid.origo
        # TODO sum the non-origo atoms to the csrDM matrix
        #      this would further decrease the loops required.

        # Loop over all atoms in the grid-cell
        for ia, ia_xyz, isc in zip(IA, XYZ - origo.reshape(1, 3), ISC):
            # Get current atom
            ia_atom = atom[ia]
            IO = a2o(ia)
            IO_range = range(ia_atom.no)
            cell_offset = (cell * isc.reshape(3, 1)).sum(0) - origo

            # Extract maximum R
            R = ia_atom.maxR()
            if R <= 0.:
                warn("Atom '{}' does not have a wave-function, skipping atom.".
                     format(ia_atom))
                eta.update()
                continue

            # Retrieve indices of the grid for the atomic shape
            idx = grid.index(ia_atom.toSphere(ia_xyz))

            # Now we have the indices for the largest orbital on the atom

            # Subsequently we have to loop the orbitals and the
            # connecting orbitals
            # Then we find the indices that overlap with these indices
            # First reduce indices to inside the grid-cell
            idx[idx[:, 0] < 0, 0] = 0
            idx[shape[0] <= idx[:, 0], 0] = shape[0] - 1
            idx[idx[:, 1] < 0, 1] = 0
            idx[shape[1] <= idx[:, 1], 1] = shape[1] - 1
            idx[idx[:, 2] < 0, 2] = 0
            idx[shape[2] <= idx[:, 2], 2] = shape[2] - 1

            # Remove duplicates, requires numpy >= 1.13
            idx = unique(idx, axis=0)
            if len(idx) == 0:
                eta.update()
                continue

            # Get real-space coordinates for the current atom
            # as well as the radial parts
            grid_xyz = dot(idx, dcell)

            # Perform loop on connection atoms
            # Allocate the DM_pj arrays
            # This will have a size equal to number of elements times number of
            # orbitals on this atom
            # In this way we do not have to calculate the psi_j multiple times
            DM_io = csrDM[IO:IO + ia_atom.no, :].tolil()
            DM_pj = _a.zerosd([ia_atom.no, grid_xyz.shape[0]])

            # Now we perform the loop on the connections for this atom
            # Remark that we have removed the diagonal atom (it-self)
            # As that will be calculated in the end
            for ja in a_col[a_ptr[ia]:a_ptr[ia + 1]]:
                # Retrieve atom (which contains the orbitals)
                ja_atom = atom[ja % na]
                JO = a2o(ja)
                jR = ja_atom.maxR()
                # Get actual coordinate of the atom
                ja_xyz = axyz(ja) + cell_offset

                # Reduce the ia'th grid points to those that connects to the ja'th atom
                ja_idx, ja_r, ja_theta, ja_cos_phi = xyz2sphericalR(
                    grid_xyz, ja_xyz, jR)

                if len(ja_idx) == 0:
                    # Quick step
                    continue

                # Loop on orbitals on this atom
                for jo in range(ja_atom.no):
                    o = ja_atom.orbital[jo]
                    oR = o.R

                    # Downsize to the correct indices
                    if jR - oR < 1e-6:
                        ja_idx1 = ja_idx.view()
                        ja_r1 = ja_r.view()
                        ja_theta1 = ja_theta.view()
                        ja_cos_phi1 = ja_cos_phi.view()
                    else:
                        ja_idx1 = indices_le(ja_r, oR)
                        if len(ja_idx1) == 0:
                            # Quick step
                            continue

                        # Reduce arrays
                        ja_r1 = ja_r[ja_idx1]
                        ja_theta1 = ja_theta[ja_idx1]
                        ja_cos_phi1 = ja_cos_phi[ja_idx1]
                        ja_idx1 = ja_idx[ja_idx1]

                    # Calculate the psi_j component
                    psi = o.psi_spher(ja_r1,
                                      ja_theta1,
                                      ja_cos_phi1,
                                      cos_phi=True)

                    # Now add this orbital to all components
                    for io in IO_range:
                        DM_pj[io, ja_idx1] += DM_io[io, JO + jo] * psi

                # Temporary clean up
                del ja_idx, ja_r, ja_theta, ja_cos_phi
                del ja_idx1, ja_r1, ja_theta1, ja_cos_phi1, psi

            # Now we have all components for all orbitals connection to all orbitals on atom
            # ia. We simply need to add the diagonal components

            # Loop on the orbitals on this atom
            ia_r, ia_theta, ia_cos_phi = xyz2spherical(grid_xyz, ia_xyz)
            del grid_xyz
            for io in IO_range:
                # Only loop halve the range.
                # This is because: triu + tril(-1).transpose()
                # removes the lower half of the on-site matrix.
                for jo in range(io + 1, ia_atom.no):
                    DM = DM_io[io, off + IO + jo]

                    oj = ia_atom.orbital[jo]
                    ojR = oj.R

                    # Downsize to the correct indices
                    if R - ojR < 1e-6:
                        ja_idx1 = slice(None)
                        ja_r1 = ia_r.view()
                        ja_theta1 = ia_theta.view()
                        ja_cos_phi1 = ia_cos_phi.view()
                    else:
                        ja_idx1 = indices_le(ia_r, ojR)
                        if len(ja_idx1) == 0:
                            # Quick step
                            continue

                        # Reduce arrays
                        ja_r1 = ia_r[ja_idx1]
                        ja_theta1 = ia_theta[ja_idx1]
                        ja_cos_phi1 = ia_cos_phi[ja_idx1]

                    # Calculate the psi_j component
                    DM_pj[io, ja_idx1] += DM * oj.psi_spher(
                        ja_r1, ja_theta1, ja_cos_phi1, cos_phi=True)

                # Calculate the psi_i component
                # Note that this one *also* zeroes points outside the shell
                # I.e. this step is important because it "nullifies" all but points where
                # orbital io is defined.
                psi = ia_atom.orbital[io].psi_spher(ia_r,
                                                    ia_theta,
                                                    ia_cos_phi,
                                                    cos_phi=True)
                DM_pj[io, :] += DM_io[io, off + IO + io] * psi
                DM_pj[io, :] *= psi

            # Temporary clean up
            ja_idx1 = ja_r1 = ja_theta1 = ja_cos_phi1 = None
            del ia_r, ia_theta, ia_cos_phi, psi, DM_io

            # Now add the density
            grid.grid[idx[:, 0], idx[:, 1], idx[:, 2]] += DM_pj.sum(0)

            # Clean-up
            del DM_pj, idx

            eta.update()
        eta.close()

        # Reset the error code for division
        np.seterr(**old_err)
Exemplo n.º 8
0
 def setUp(self):
     self.s1 = SparseCSR((10, 100), dtype=np.int32)
     self.s2 = SparseCSR((10, 100, 2))
Exemplo n.º 9
0
class TestSparseCSR(object):
    def setUp(self):
        self.s1 = SparseCSR((10, 100), dtype=np.int32)
        self.s2 = SparseCSR((10, 100, 2))

    def test_init1(self):
        assert_equal(self.s1.dtype, np.int32)
        assert_equal(self.s2.dtype, np.float64)
        assert_true(np.allclose(self.s1.data, self.s1.data))
        assert_true(np.allclose(self.s2.data, self.s2.data))

    def test_init2(self):
        SparseCSR((10, 100))
        for d in [np.int32, np.float64, np.complex128]:
            s = SparseCSR((10, 100), dtype=d)
            assert_equal(s.shape, (10, 100, 1))
            assert_equal(s.dim, 1)
            assert_equal(s.dtype, d)
            for k in [1, 2]:
                s = SparseCSR((10, 100, k), dtype=d)
                assert_equal(s.shape, (10, 100, k))
                assert_equal(s.dim, k)
                s = SparseCSR((10, 100), dim=k, dtype=d)
                assert_equal(s.shape, (10, 100, k))
                assert_equal(s.dim, k)
                s = SparseCSR((10, 100, 3), dim=k, dtype=d)
                assert_equal(s.shape, (10, 100, 3))
                assert_equal(s.dim, 3)

    def test_init3(self):
        csr = sc.sparse.csr_matrix((10, 10), dtype=np.int32)
        csr[0, 1] = 1
        csr[0, 2] = 2
        sp = SparseCSR(csr)
        assert_equal(sp.dtype, np.int32)
        assert_equal(sp.shape, (10, 10, 1))
        assert_equal(len(sp), 2)
        assert_equal(sp[0, 1], 1)
        assert_equal(sp[0, 2], 2)
        sp = SparseCSR(csr, dtype=np.float64)
        assert_equal(sp.shape, (10, 10, 1))
        assert_equal(sp.dtype, np.float64)
        assert_equal(len(sp), 2)
        assert_equal(sp[0, 1], 1)
        assert_equal(sp[0, 2], 2)

    def test_init4(self):
        csr = sc.sparse.csr_matrix((10, 10), dtype=np.int32)
        csr[0, 1] = 1
        csr[0, 2] = 2
        print(csr.indices, csr.indptr)
        sp = SparseCSR((csr.data, csr.indices, csr.indptr))
        assert_equal(sp.dtype, np.int32)
        assert_equal(sp.shape, (10, 10, 1))
        assert_equal(len(sp), 2)
        assert_equal(sp[0, 1], 1)
        assert_equal(sp[0, 2], 2)
        sp = SparseCSR((csr.data, csr.indices, csr.indptr), dtype=np.float64)
        assert_equal(sp.shape, (10, 10, 1))
        assert_equal(sp.dtype, np.float64)
        assert_equal(len(sp), 2)
        assert_equal(sp[0, 1], 1)
        assert_equal(sp[0, 2], 2)

    def test_create1(self):
        self.s1[0, [1, 2, 3]] = 1
        assert_equal(self.s1.nnz, 3)
        self.s1[2, [1, 2, 3]] = 1
        assert_equal(self.s1.nnz, 6)
        self.s1.empty(keep=True)
        assert_equal(self.s1.nnz, 6)
        self.s1.empty()
        assert_equal(self.s1.nnz, 0)

    def test_create2(self):
        for i in range(10):
            j = range(i * 4, i * 4 + 3)
            self.s1[0, j] = i
            assert_equal(len(self.s1), (i + 1) * 3)
            for jj in j:
                assert_equal(self.s1[0, jj], i)
                assert_equal(self.s1[1, jj], 0)
        self.s1.empty()

    def test_create3(self):
        for i in range(10):
            j = range(i * 4, i * 4 + 3)
            self.s1[0, j] = i
            assert_equal(len(self.s1), (i + 1) * 3)
            self.s1[0, range((i + 1) * 4, (i + 1) * 4 + 3)] = None
            assert_equal(len(self.s1), (i + 1) * 3)
            for jj in j:
                assert_equal(self.s1[0, jj], i)
                assert_equal(self.s1[1, jj], 0)
        self.s1.empty()

    def test_finalize1(self):
        self.s1[0, [1, 2, 3]] = 1
        self.s1[2, [1, 2, 3]] = 1.
        assert_false(self.s1.finalized)
        self.s1.finalize()
        assert_true(self.s1.finalized)
        self.s1.empty(keep=True)
        assert_true(self.s1.finalized)
        self.s1.empty()
        assert_false(self.s1.finalized)

    def test_delitem1(self):
        self.s1[0, [1, 2, 3]] = 1
        assert_equal(len(self.s1), 3)
        del self.s1[0, 1]
        assert_equal(len(self.s1), 2)
        assert_equal(self.s1[0, 1], 0)
        assert_equal(self.s1[0, 2], 1)
        assert_equal(self.s1[0, 3], 1)
        self.s1[0, [1, 2, 3]] = 1
        del self.s1[0, [1, 3]]
        assert_equal(len(self.s1), 1)
        assert_equal(self.s1[0, 1], 0)
        assert_equal(self.s1[0, 2], 1)
        assert_equal(self.s1[0, 3], 0)
        self.s1.empty()

    def test_op1(self):
        for i in range(10):
            j = range(i * 4, i * 4 + 3)
            self.s1[0, j] = i

            # i+
            self.s1 += 1
            for jj in j:
                assert_equal(self.s1[0, jj], i + 1)
                assert_equal(self.s1[1, jj], 0)

            # i-
            self.s1 -= 1
            for jj in j:
                assert_equal(self.s1[0, jj], i)
                assert_equal(self.s1[1, jj], 0)

            # i*
            self.s1 *= 2
            for jj in j:
                assert_equal(self.s1[0, jj], i * 2)
                assert_equal(self.s1[1, jj], 0)

            # //
            self.s1 //= 2
            for jj in j:
                assert_equal(self.s1[0, jj], i)
                assert_equal(self.s1[1, jj], 0)

            # i**
            self.s1 **= 2
            for jj in j:
                assert_equal(self.s1[0, jj], i**2)
                assert_equal(self.s1[1, jj], 0)

    def test_op2(self):
        for i in range(10):
            j = range(i * 4, i * 4 + 3)
            self.s1[0, j] = i

            # +
            s = self.s1 + 1
            for jj in j:
                assert_equal(s[0, jj], i + 1)
                assert_equal(self.s1[0, jj], i)
                assert_equal(s[1, jj], 0)

            # -
            s = self.s1 - 1
            for jj in j:
                assert_equal(s[0, jj], i - 1)
                assert_equal(self.s1[0, jj], i)
                assert_equal(s[1, jj], 0)

            # - (r)
            s = 1 - self.s1
            for jj in j:
                assert_equal(s[0, jj], 1 - i)
                assert_equal(self.s1[0, jj], i)
                assert_equal(s[1, jj], 0)

            # *
            s = self.s1 * 2
            for jj in j:
                assert_equal(s[0, jj], i * 2)
                assert_equal(self.s1[0, jj], i)
                assert_equal(s[1, jj], 0)

            # //
            s = s // 2
            for jj in j:
                assert_equal(s[0, jj], i)
                assert_equal(self.s1[0, jj], i)
                assert_equal(s[1, jj], 0)

            # **
            s = self.s1**2
            for jj in j:
                assert_equal(s[0, jj], i**2)
                assert_equal(self.s1[0, jj], i)
                assert_equal(s[1, jj], 0)

            # ** (r)
            s = 2**self.s1
            for jj in j:
                assert_equal(s[0, jj], 2**self.s1[0, jj])
                assert_equal(self.s1[0, jj], i)
                assert_equal(s[1, jj], 0)

    def test_op3(self):
        S = SparseCSR((10, 100), dtype=np.int32)
        # Create initial stuff
        for i in range(10):
            j = range(i * 4, i * 4 + 3)
            S[0, j] = i

        for op in ['add', 'sub', 'mul', 'pow']:
            func = getattr(S, '__{}__'.format(op))
            s = func(1)
            assert_equal(s.dtype, np.int32)
            s = func(1.)
            assert_equal(s.dtype, np.float64)
            if op != 'pow':
                s = func(1.j)
                assert_equal(s.dtype, np.complex128)

        S = S.copy(dtype=np.float64)
        for op in ['add', 'sub', 'mul', 'pow']:
            func = getattr(S, '__{}__'.format(op))
            s = func(1)
            assert_equal(s.dtype, np.float64)
            s = func(1.)
            assert_equal(s.dtype, np.float64)
            if op != 'pow':
                s = func(1.j)
                assert_equal(s.dtype, np.complex128)

        S = S.copy(dtype=np.complex128)
        for op in ['add', 'sub', 'mul', 'pow']:
            func = getattr(S, '__{}__'.format(op))
            s = func(1)
            assert_equal(s.dtype, np.complex128)
            s = func(1.)
            assert_equal(s.dtype, np.complex128)
            if op != 'pow':
                s = func(1.j)
                assert_equal(s.dtype, np.complex128)

    def test_op4(self):
        S = SparseCSR((10, 100), dtype=np.int32)
        # Create initial stuff
        for i in range(10):
            j = range(i * 4, i * 4 + 3)
            S[0, j] = i

        s = 1 + S
        assert_equal(s.dtype, np.int32)
        s = 1. + S
        assert_equal(s.dtype, np.float64)
        s = 1.j + S
        assert_equal(s.dtype, np.complex128)

        s = 1 - S
        assert_equal(s.dtype, np.int32)
        s = 1. - S
        assert_equal(s.dtype, np.float64)
        s = 1.j - S
        assert_equal(s.dtype, np.complex128)

        s = 1 * S
        assert_equal(s.dtype, np.int32)
        s = 1. * S
        assert_equal(s.dtype, np.float64)
        s = 1.j * S
        assert_equal(s.dtype, np.complex128)

        s = 1**S
        assert_equal(s.dtype, np.int32)
        s = 1.**S
        assert_equal(s.dtype, np.float64)
        s = 1.j**S
        assert_equal(s.dtype, np.complex128)
Exemplo n.º 10
0
    def reset(self, nnzpr=None, orthogonal=True, spin=1, dtype=None):
        """
        The sparsity pattern is cleaned and every thing
        is reset.

        The object will be the same as if it had been
        initialized with the same geometry as it were
        created with.
        
        Parameters
        ----------
        nnzpr: int
           number of non-zero elements per row
        orthogonal: boolean, True
           if there is an overlap matrix associated with the
           Hamiltonian
        spin: int, 1
           number of spin-components
        dtype: ``numpy.dtype``, `numpy.float64`
           the datatype of the Hamiltonian
        """
        # I know that this is not the most efficient way to
        # access a C-array, however, for constructing a
        # sparse pattern, it should be faster if memory elements
        # are closer...
        # Hence, this choice of having H and S like this

        # We check the first atom and its neighbours, we then
        # select max(5,len(nc) * 4)
        if nnzpr is None:
            nnzpr = self.geom.close(0)
            if nnzpr is None: nnzpr = [0,0]
            nnzpr = max(5, len(nnzpr) * 4)

        self._orthogonal = orthogonal

        # Reset the sparsity pattern
        if not orthogonal:
            self._data = SparseCSR((self.no, self.no_s, spin+1), nnzpr=nnzpr, dtype=dtype)
        else:
            self._data = SparseCSR((self.no, self.no_s, spin), nnzpr=nnzpr, dtype=dtype)


        self._spin = spin

        if spin == 1:
            self.UP = 0
            self.DOWN = 0
            self.S_idx = 1
        elif spin == 2:
            self.UP = 0
            self.DOWN = 1
            self.S_idx = 2
        else:
            raise ValueError("Currently the Hamiltonian has only been implemented with up to collinear spin.")

        if orthogonal:
            # There is no overlap matrix
            self.S_idx = -1

        # Denote that one *must* specify all details of the elements
        self._def_dim = -1
Exemplo n.º 11
0
class Hamiltonian(object):
    """ Hamiltonian object containing the coupling constants between orbitals.

    The Hamiltonian object contains information regarding the 
     - geometry
     - coupling constants between orbitals

    It contains an intrinsic sparse matrix of the Hamiltonian elements.
    
    Assigning or changing Hamiltonian elements is as easy as with
    standard ``numpy`` assignments:
      
    >>> ham = Hamiltonian(...)
    >>> ham.H[1,2] = 0.1
    
    which assigns 0.1 as the coupling constant between orbital 2 and 3.
    (remember that Python is 0-based elements).
    """

    # The order of the Energy
    # I.e. whether energy should be in other units than Ry
    # This conversion is made: [eV] ** _E_order
    _E_order = 1

    def __init__(self, geom, nnzpr=None, orthogonal=True, spin=1,
                 dtype=None, *args, **kwargs):
        """Create tight-binding model from geometry

        Initializes a tight-binding model using the :code:`geom` object
        as the underlying geometry for the tight-binding parameters.
        """
        self._geom = geom

        # Initialize the sparsity pattern
        self.reset(nnzpr=nnzpr, orthogonal=orthogonal, spin=spin, dtype=dtype)

    def reset(self, nnzpr=None, orthogonal=True, spin=1, dtype=None):
        """
        The sparsity pattern is cleaned and every thing
        is reset.

        The object will be the same as if it had been
        initialized with the same geometry as it were
        created with.
        
        Parameters
        ----------
        nnzpr: int
           number of non-zero elements per row
        orthogonal: boolean, True
           if there is an overlap matrix associated with the
           Hamiltonian
        spin: int, 1
           number of spin-components
        dtype: ``numpy.dtype``, `numpy.float64`
           the datatype of the Hamiltonian
        """
        # I know that this is not the most efficient way to
        # access a C-array, however, for constructing a
        # sparse pattern, it should be faster if memory elements
        # are closer...
        # Hence, this choice of having H and S like this

        # We check the first atom and its neighbours, we then
        # select max(5,len(nc) * 4)
        if nnzpr is None:
            nnzpr = self.geom.close(0)
            if nnzpr is None: nnzpr = [0,0]
            nnzpr = max(5, len(nnzpr) * 4)

        self._orthogonal = orthogonal

        # Reset the sparsity pattern
        if not orthogonal:
            self._data = SparseCSR((self.no, self.no_s, spin+1), nnzpr=nnzpr, dtype=dtype)
        else:
            self._data = SparseCSR((self.no, self.no_s, spin), nnzpr=nnzpr, dtype=dtype)


        self._spin = spin

        if spin == 1:
            self.UP = 0
            self.DOWN = 0
            self.S_idx = 1
        elif spin == 2:
            self.UP = 0
            self.DOWN = 1
            self.S_idx = 2
        else:
            raise ValueError("Currently the Hamiltonian has only been implemented with up to collinear spin.")

        if orthogonal:
            # There is no overlap matrix
            self.S_idx = -1

        # Denote that one *must* specify all details of the elements
        self._def_dim = -1


    def empty(self, keep=False):
        """ See `SparseCSR.empty` for specifics """
        self._data.empty(keep)

    def copy(self, dtype=None):
        """ Return a copy of the `Hamiltonian` object """
        if dtype is None:
            dtype = self.dtype
        H = self.__class__(self.geom, orthogonal=self.orthogonal,
                           spin=self.spin, dtype=dtype)
        # Be sure to copy the content of the SparseCSR object
        H._data = self._data.copy(dtype=dtype)
        return H

    ######### Definitions of overrides ############
    @property
    def geom(self):
        """ Return the attached geometry """
        return self._geom

    @property
    def spin(self):
        """ Return number of spin-components in Hamiltonian """
        return self._spin

    @property
    def dtype(self):
        """ Return data type of Hamiltonian (and overlap matrix) """
        return self._data.dtype

    @property
    def orthogonal(self):
        """ Return whether the Hamiltonian is orthogonal """
        return self._orthogonal

    def __len__(self):
        """ Returns number of rows in the Hamiltonian """
        return self.geom.no

    def __repr__(self):
        """ Representation of the tight-binding model """
        s = self.geom.__repr__()
        return s + '\nNumber of non-zero elements {0}'.format(self.nnz)

    def __getattr__(self, attr):
        """ Returns the attributes from the underlying geometry

        Any attribute not found in the tight-binding model will
        be looked up in the underlying geometry.
        """
        return getattr(self.geom, attr)


    def __getitem__(self, key):
        """ Return Hamiltonian coupling elements for the index(s) """
        dd = self._def_dim
        if dd >= 0:
            key = tuple(key) + (dd,)
            self._def_dim = -1
        d = self._data[key]
        return d


    def __setitem__(self, key, val):
        """ Set or create couplings between orbitals in the Hamiltonian

        Override set item for slicing operations and enables easy
        setting of tight-binding parameters in a sparse matrix
        """
        dd = self._def_dim
        if dd >= 0:
            key = tuple(key) + (dd,)
            self._def_dim = -1
        self._data[key] = val


    def __get_H(self):
        self._def_dim = self.UP
        return self
    _get_H = __get_H

    def __set_H(self, key, value):
        if len(key) == 2:
            self._def_dim = self.UP
        self[key] = value
    _set_H = __set_H

    H = property(__get_H, __set_H)

    def __get_S(self):
        if self.orthogonal:
            return None
        self._def_dim = self.S_idx
        return self
    _get_S = __get_S

    def __set_S(self, key, value):
        if self.orthogonal:
            return None
        self._def_dim = self.S_idx
        self[key] = value
    _set_S = __set_S

    S = property(__get_S, __set_S)


    # Create iterations module
    def iter_linear(self):
        """ Iterations of the orbital space, two indices from loop

        An iterator returning the current atomic index and the corresponding
        orbital index.

        >>> for ia, io in self:

        In the above case `io` always belongs to atom `ia` and `ia` may be
        repeated according to the number of orbitals associated with
        the atom `ia`.
        """
        for ia in self.geom:
            ia1, ia2 = self.geom.lasto[ia], self.geom.lasto[ia + 1]
            for io in range(ia1, ia2):
                yield ia, io

    __iter__ = iter_linear


    def construct(self, dR, param, eta=False):
        """ Automatically construct the Hamiltonian model based on ``dR`` and associated hopping integrals ``param``.

        Parameters
        ----------
        dR : array_like
           radii parameters for tight-binding parameters.
           Must have same length as ``param`` or one less.
           If one less it will be extended with ``dR[0]/100``
        param : array_like
           coupling constants corresponding to the ``dR``
           ranges. ``param[0,:]`` are the tight-binding parameter
           for the all atoms within ``dR[0]`` of each atom.
        eta : `bool` (`False`)
           whether an ETA will be printed...
        """
        # Ensure that we are dealing with a numpy array
        param = np.array(param)

        if len(dR) + 1 == len(param):
            R = np.hstack((dR[0] / 100, np.asarray(dR)))
        elif len(dR) == len(param):
            R = np.asarray(dR).copy()
        else:
            raise ValueError("Length of dR and param must be the same "
                             "or dR one shorter than param. "
                             "One tight-binding parameter for each radii.")

        if not self.orthogonal:
            if len(param[0]) != 2:
                raise ValueError("Number of parameters "
                                 "for each element is not 2. "
                                 "You must make len(param[0] == 2) for non-orthogonal Hamiltonians.")

        if np.any(np.diff(self.geom.lasto) > 1):
            warnings.warn("Automatically setting a tight-binding model "
                          "for systems with atoms having more than 1 "
                          "orbital is not adviced. Please do it your-self.")

        eq_atoms = []
        def print_equal(eq_atoms):
            if len(eq_atoms) > 0:
                s = ("The geometry has one or more atoms having the same "
                     "atomic position. "
                     "The atoms are within {} Ang "
                     "of each other.\n".format(R[0]))
                
                for ia, ja in eq_atoms:
                    s += "  {0:7d} -- {1:7d}\n".format(ia, ja)
                warnings.warn(s)

        def append_equal(eq_atoms, ia, idx):
            # Append to the list of equal atoms the atomic indices
            if len(idx) > 1:
                tmp = list(idx)
                # only add in "one" direction
                for ja in tmp:
                    if ja > ia:
                        eq_atoms.append( (ia,ja) )


        if len(self.geom) < 1501:
            # there is no need to do anything complex
            # for small systems
            for ia in self.geom:
                # Find atoms close to 'ia'
                idx = self.geom.close(ia, dR=R)
                append_equal(eq_atoms, ia, idx[0])

                for ix, h in zip(idx, param):
                    # Set the tight-binding parameters
                    self[ia, ix] = h

            print_equal(eq_atoms)
            return self

        # check how many atoms are within the standard 10 dR
        # range of some random atom.
        ia = np.random.randint(len(self.geom) - 1)

        # default block iterator
        d = self.geom.dR
        na = len(self.geom.close(ia, dR=d * 10))

        # Convert to 1000 atoms spherical radii
        iR = int(4 / 3 * np.pi * d ** 3 / na * 1000)

        # Get number of atoms
        na = len(self.geom)
        na_run = 0

        from time import time
        from sys import stdout
        t0 = time()

        # Do the loop
        for ias, idxs in self.geom.iter_block(iR=iR):
            # Loop the atoms inside
            for ia in ias:
                # Find atoms close to 'ia'
                idx = self.geom.close(ia, dR=R, idx=idxs)
                append_equal(eq_atoms, ia, idx[0])

                for ix, h in zip(idx, param):
                    # Set the tight-binding parameters
                    self[ia, ix] = h

            if eta:
                na_run += len(ias)
                na -= len(ias)
                t1 = time()
                # calculate hours, minutes, seconds
                m, s = divmod( float(t1-t0)/na_run * na, 60)
                h, m = divmod(m, 60)
                stdout.write("Hamiltonian.construct() ETA = {0:d}h {1:d}m {2:.2f}s\r".format(int(h), int(m), s))
                stdout.flush()

        print_equal(eq_atoms)


    @property
    def finalized(self):
        """ Whether the contained data is finalized and non-used elements have been removed """
        return self._data.finalized


    def finalize(self):
        """ Finalizes the tight-binding model

        Finalizes the tight-binding model so that no new sparse
        elements can be added.

        Sparse elements can still be changed.
        """
        self._data.finalize()

        # Get the folded Hamiltonian at the Gamma point
        Hk = self.Hk()

        nzs = Hk.nnz
        
        if nzs != (Hk + Hk.T).nnz:
            warnings.warn(
                'Hamiltonian does not retain symmetric couplings, this might be problematic.')


    @property
    def nnz(self):
        """ Returns number of non-zero elements in the tight-binding model """
        return self._data.nnz


    @property
    def no(self):
        """ Returns number of orbitals as used when the object was created """
        return self._data.nr


    def tocsr(self, index):
        """ Return a ``scipy.sparse.csr_matrix`` from the specified index
        """
        return self._data.tocsr(index)
        

    def Hk(self, k=(0, 0, 0), spin=0):
        """ Return the Hamiltonian in a ``scipy.sparse.csr_matrix`` at `k`.

        Parameters
        ----------
        k: float*3
           k-point 
        spin: int, 0
           the spin-index of the Hamiltonian
        """
        # Create csr sparse formats.
        # We import here as the user might not want to
        # rely on this feature.
        from scipy.sparse import csr_matrix

        dot = np.dot

        k = np.asarray(k, np.float64)
        k.shape = (-1,)
        
        # Setup the Hamiltonian for this k-point
        Hf = self.tocsr(spin)

        no = self.no
        s = (no, no)
        H = csr_matrix(s, dtype=np.complex128)

        # Get the reciprocal lattice vectors dotted with k
        kr = dot(self.rcell, k)
        for si in range(self.sc.n_s):
            isc = self.sc_off[si, :]
            phase = np.exp(-1j * dot(kr, dot(self.cell, isc)))
            H += Hf[:, si * no:(si + 1) * no] * phase
            
        del Hf
        
        return H


    def Sk(self, k=(0, 0, 0), spin=0):
        """ Return the overlap matrix in a ``scipy.sparse.csr_matrix`` at `k`.

        Parameters
        ----------
        k: float*3
           k-point 
        """
        if self.orthogonal:
            return None

        # Create csr sparse formats.
        # We import here as the user might not want to
        # rely on this feature.
        from scipy.sparse import csr_matrix

        dot = np.dot

        k = np.asarray(k, np.float64)
        k.shape = (-1,)
        
        # Setup the Hamiltonian for this k-point
        Sf = self.tocsr(self.S_idx)

        no = self.no
        s = (no, no)
        S = csr_matrix(s, dtype=np.complex128)

        # Get the reciprocal lattice vectors dotted with k
        kr = dot(self.rcell, k)
        for si in range(self.sc.n_s):
            isc = self.sc_off[si, :]
            phase = np.exp(-1j * dot(kr, dot(self.cell, isc)))
            S += Sf[:, si * no:(si + 1) * no] * phase
            
        del Sf
        
        return S


    def eigh(self,k=(0,0,0),
            atoms=None, eigvals_only=True,
            overwrite_a=True, overwrite_b=True,
            *args,
            **kwargs):
        """ Returns the eigenvalues of the tight-binding model

        Setup the Hamiltonian and overlap matrix with respect to
        the given k-point, then reduce the space to the specified atoms
        and calculate the eigenvalues.

        All subsequent arguments gets passed directly to :code:`scipy.linalg.eigh`
        """
        H = self.Hk(k=k)
        if not self.orthogonal:
            S = self.Sk(k=k)
        # Reduce sparsity pattern
        if not atoms is None:
            orbs = self.a2o(atoms)
            # Reduce space
            H = H[orbs, orbs]
            if not self.orthogonal:
                S = S[orbs, orbs]
        if self.orthogonal:
            return sli.eigh(H.todense(),
                *args,
                eigvals_only=eigvals_only,
                overwrite_a=overwrite_a,
                **kwargs)
        
        return sli.eigh(H.todense(), S.todense(),
            *args,
            eigvals_only=eigvals_only,
            overwrite_a=overwrite_a,
            overwrite_b=overwrite_b,
            **kwargs)


    def eigsh(self, k=(0,0,0), n=10,
            atoms=None, eigvals_only=True,
            *args,
            **kwargs):
        """ Returns the eigenvalues of the tight-binding model

        Setup the Hamiltonian and overlap matrix with respect to
        the given k-point, then reduce the space to the specified atoms
        and calculate the eigenvalues.

        All subsequent arguments gets passed directly to :code:`scipy.linalg.eigh`
        """
        
        # We always request the smallest eigenvalues... 
        kwargs.update({'which':kwargs.get('which', 'SM')})
        
        H = self.Hk(k=k)
        if not self.orthogonal:
            raise ValueError("The sparsity pattern is non-orthogonal, you cannot use the Arnoldi procedure with scipy")
        
        # Reduce sparsity pattern
        if not atoms is None:
            orbs = self.a2o(atoms)
            # Reduce space
            H = H[orbs, orbs]

        return ssli.eigsh(H, k=n,
                          *args,
                          return_eigenvectors=not eigvals_only,
                          **kwargs)

    def cut(self, seps, axis, *args, **kwargs):
        """ Cuts the tight-binding model into different parts.

        Creates a tight-binding model by retaining the parameters
        for the cut-out region, possibly creating a super-cell.

        Parameters
        ----------
        seps  : integer, optional
           number of times the structure will be cut.
        axis  : integer
           the axis that will be cut
        """
        new_w = None
        # Create new geometry
        with warnings.catch_warnings(record=True) as w:
            # Cause all warnings to always be triggered.
            warnings.simplefilter("always")
            # Create new cut geometry
            geom = self.geom.cut(seps, axis, *args, **kwargs)
            # Check whether the warning exists
            if len(w) > 0:
                if issubclass(w[-1].category, UserWarning):
                    new_w = str(w[-1].message)
                    new_w += ("\n---\n"
                              "The tight-binding model cannot be cut as the structure "
                              "cannot be tiled accordingly. ANY use of the model has been "
                              "relieved from sisl.")
        if new_w:
            warnings.warn(new_w, UserWarning)

        # Now we need to re-create the tight-binding model
        H = self.Hk()
        S = self.Sk()
        # they are created similarly, hence the following
        # should keep their order

        # First we need to figure out how long the interaction range is
        # in the cut-direction
        # We initialize to be the same as the parent direction
        nsc = np.copy(self.nsc) // 2
        nsc[axis] = 0  # we count the new direction
        isc = np.zeros([3], np.int32)
        isc[axis] -= 1
        out = False
        while not out:
            # Get supercell index
            isc[axis] += 1
            try:
                idx = self.sc_index(isc)
            except:
                break

            # Figure out if the Hamiltonian has interactions
            # to 'isc'
            sub = H[0:geom.no, idx * self.no:(idx + 1) * self.no].indices[:]
            sub = np.unique(np.hstack(
                (sub, S[0:geom.no, idx * self.no:(idx + 1) * self.no].indices[:])))
            if len(sub) == 0:
                break

            c_max = np.amax(sub)
            # Count the number of cells it interacts with
            i = (c_max % self.no) // geom.no
            ic = idx * self.no
            for j in range(i):
                idx = ic + geom.no * j
                # We need to ensure that every "in between" index exists
                # if it does not we discard those indices
                if len(np.where(
                        np.logical_and(idx <= sub,
                                       sub < idx + geom.no)
                )[0]) == 0:
                    i = j - 1
                    out = True
                    break
            nsc[axis] = isc[axis] * seps + i

            if out:
                warnings.warn(
                    'Cut the connection at nsc={0} in direction {1}.'.format(
                        nsc[axis], axis), UserWarning)

        # Update number of super-cells
        nsc[:] = nsc[:] * 2 + 1
        geom.sc.set_nsc(nsc)

        # Now we have a correct geometry, and
        # we are now ready to create the sparsity pattern
        # Reduce the sparsity pattern, first create the new one
        ham = self.__class__(geom, nc=np.amax(self.ncol), spin=self.spin)

        def sco2sco(M, o, m, seps, axis):
            # Converts an o from M to m
            isc = np.copy(M.o2isc(o))
            isc[axis] = isc[axis] * seps
            # Correct for cell-offset
            isc[axis] = isc[axis] + (o % M.no) // m.no
            # find the equivalent cell in m
            try:
                # If a fail happens it is due to a discarded
                # interaction across a non-interacting region
                return (o % m.no,
                        m.sc_index(isc) * m.no,
                        m.sc_index(-isc) * m.no)
            except:
                return None, None, None

        # Copy elements
        for jo in range(geom.no):

            # make smaller cut
            sH = H[jo, :]
            sS = S[jo, :]

            for io, iH in zip(sH.indices, sH.data):
                # Get the equivalent orbital in the smaller cell
                o, ofp, ofm = sco2sco(self.geom, io, ham.geom, seps, axis)
                if o is None:
                    continue
                ham[jo, o + ofp] = iH, S[jo, io]
                ham[o, jo + ofm] = iH, S[jo, io]

            if np.any(sH.indices != sS.indices):

                # Ensure that S is also cut
                for io, iS in zip(sS.indices, sS.data):
                    # Get the equivalent orbital in the smaller cell
                    o, ofp, ofm = sco2sco(self.geom, io, ham.geom, seps, axis)
                    if o is None:
                        continue
                    ham[jo, o + ofp] = H[jo, io], iS
                    ham[o, jo + ofm] = H[jo, io], iS

        return ham

    def tile(self, reps, axis):
        """ Returns a repeated tight-binding model for this, much like the `Geometry`

        The already existing tight-binding parameters are extrapolated
        to the new supercell by repeating them in blocks like the coordinates.

        Parameters
        ----------
        reps : number of tiles (repetitions)
        axis : direction of tiling
            0, 1, 2 according to the cell-direction
        """

        # Create the new geometry
        g = self.geom.tile(reps, axis)

        raise NotImplementedError(('tiling a Hamiltonian model has not been '
                              'fully implemented yet.'))

    def repeat(self, reps, axis):
        """ Refer to `tile` instead """
        # Create the new geometry
        g = self.geom.repeat(reps, axis)
        
        raise NotImplementedError(('repeating a Hamiltonian model has not been '
                              'fully implemented yet, use tile instead.'))

    @classmethod
    def sp2HS(cls, geom, H, S=None):
        """ Returns a tight-binding model from a preset H, S and Geometry
        """
        # Calculate number of connections
        nc = 0

        has_S = not S is None
        
        # Ensure csr format
        H = H.tocsr()
        if has_S:
            S = S.tocsr()
        for i in range(geom.no):
            nc = max(nc, H[i, :].getnnz())
            if has_S:
                nc = max(nc, S[i, :].getnnz())

        if has_S:
            ham = cls(geom, nnzpr=nc,
                      orthogonal=False, dtype=H.dtype)
        else:
            ham = cls(geom, nnzpr=nc, dtype=H.dtype)

        # Copy data to the model
        H = H.tocoo()
        if has_S:
            for jo, io, h in zip(H.row, H.col, H.data):
                ham[jo, io] = (h, S[jo, io])

            # Convert S to coo matrix
            S = S.tocoo()
            # If the Hamiltonian for one reason or the other
            # is zero in the diagonal, then we *must* account for
            # this as it isn't captured in the above loop.
            skip_S = np.all(H.row == S.row)
            skip_S = skip_S and np.all(H.col == S.col)
            if not skip_S:
                # Re-convert back to allow index retrieval
                H = H.tocsr()
                for jo, io, s in zip(S.row, S.col, S.data):
                    ham[jo, io] = (H[jo, io], s)

        else:
            for jo, io, h in zip(H.row, H.col, H.data):
                ham[jo, io] = h

        return ham


    @staticmethod
    def read(sile, *args, **kwargs):
        """ Reads Hamiltonian from `Sile` using `read_H`.

        Parameters
        ----------
        sile : `Sile`, str
            a `Sile` object which will be used to read the Hamiltonian
            and the overlap matrix (if any)
            if it is a string it will create a new sile using `get_sile`.
        * : args passed directly to ``read_es(,**)``
        """
        # This only works because, they *must*
        # have been imported previously
        from sisl.io import get_sile, BaseSile
        if isinstance(sile, BaseSile):
            return sile.read_es(*args, **kwargs)
        else:
            return get_sile(sile).read_es(*args, **kwargs)


    def write(self, sile, *args, **kwargs):
        """ Writes a tight-binding model to the `Sile` as implemented in the :code:`Sile.write_es` method """
        self.finalize()

        # This only works because, they *must*
        # have been imported previously
        from sisl.io import get_sile, BaseSile
        if isinstance(sile, BaseSile):
            sile.write_es(self, *args, **kwargs)
        else:
            get_sile(sile, 'w').write_es(self, *args, **kwargs)

    ###############################
    # Overload of math operations #
    ###############################
    def __add__(a, b):
        if isinstance(b, Hamiltonian):
            raise NotImplementedError
        c = a.copy(dtype=get_dtype(b, other=a.dtype))
        c += b
        return c
    __radd__ = __add__

    def __iadd__(a, b):
        if isinstance(b, Hamiltonian):
            raise NotImplementedError
        a._data += b
        return a

    def __sub__(a, b):
        if isinstance(b, Hamiltonian):
            raise NotImplementedError
        c = a.copy(dtype=get_dtype(b, other=a.dtype))
        c -= b
        return c

    def __rsub__(a, b):
        if isinstance(b, Hamiltonian):
            raise NotImplementedError
        c = b + (-1) * a
        return c

    def __isub__(a, b):
        if isinstance(b, Hamiltonian):
            raise NotImplementedError
        a._data -= b
        return a

    def __mul__(a, b):
        if isinstance(b, Hamiltonian):
            raise NotImplementedError
        c = a.copy(dtype=get_dtype(b, other=a.dtype))
        c *= b
        return c
    __rmul__ = __mul__

    def __imul__(a, b):
        if isinstance(b, Hamiltonian):
            raise NotImplementedError
        a._data *= b
        return a

    def __div__(a, b):
        if isinstance(a, Hamiltonian):
            if isinstance(b, Hamiltonian):
                raise NotImplementedError
            c = a.copy(dtype=get_dtype(b, other=a.dtype))
            c /= b
        elif isinstance(b, Hamiltonian):
            c = b.copy(dtype=get_dtype(a, other=b.dtype))
            c._data = a / c._data
        return c

    def __idiv__(a, b):
        if isinstance(b, Hamiltonian):
            raise NotImplementedError
        a._data /= b
        return a

    def __floordiv__(a, b):
        if isinstance(b, Hamiltonian):
            raise NotImplementedError
        c = a.copy(dtype=get_dtype(b, other=a.dtype))
        c //= b
        return c

    def __ifloordiv__(a, b):
        if isinstance(b, Hamiltonian):
            raise NotImplementedError
        a._data //= b
        return a

    def __truediv__(a, b):
        if isinstance(b, Hamiltonian):
            raise NotImplementedError
        c = a.copy(dtype=get_dtype(b, other=a.dtype))
        c /= b
        return c

    def __itruediv__(a, b):
        if isinstance(b, Hamiltonian):
            raise NotImplementedError
        a._data /= b
        return a

    def __pow__(a, b):
        if isinstance(b, Hamiltonian):
            raise NotImplementedError
        c = a.copy(dtype=get_dtype(b, other=a.dtype))
        c **= b
        return c

    def __rpow__(a, b):
        if isinstance(b, SparseCSR):
            raise NotImplementedError
        c = a.copy(dtype=get_dtype(b, other=a.dtype))
        c._data = b ** c._data
        return c

    def __ipow__(a, b):
        if isinstance(b, Hamiltonian):
            raise NotImplementedError
        a._data **= b
        return a
Exemplo n.º 12
0
    def reset(self, nnzpr=None, orthogonal=True, spin=1, dtype=None):
        """
        The sparsity pattern is cleaned and every thing
        is reset.

        The object will be the same as if it had been
        initialized with the same geometry as it were
        created with.

        Parameters
        ----------
        nnzpr: int
           number of non-zero elements per row
        orthogonal: boolean, True
           if there is an overlap matrix associated with the
           Hamiltonian
        spin: int, 1
           number of spin-components
        dtype: ``numpy.dtype``, `numpy.float64`
           the datatype of the Hamiltonian
        """
        # I know that this is not the most efficient way to
        # access a C-array, however, for constructing a
        # sparse pattern, it should be faster if memory elements
        # are closer...
        # Hence, this choice of having H and S like this

        # We check the first atom and its neighbours, we then
        # select max(5,len(nc) * 4)
        if nnzpr is None:
            nnzpr = self.geom.close(0)
            if nnzpr is None:
                nnzpr = 8
            else:
                nnzpr = max(5, len(nnzpr) * 4)

        self._orthogonal = orthogonal

        # Reset the sparsity pattern
        if not orthogonal:
            self._data = SparseCSR((self.no, self.no_s, spin + 1),
                                   nnzpr=nnzpr,
                                   dtype=dtype)
        else:
            self._data = SparseCSR((self.no, self.no_s, spin),
                                   nnzpr=nnzpr,
                                   dtype=dtype)

        self._spin = spin

        if spin == 1:
            self.UP = 0
            self.DOWN = 0
            self.S_idx = 1
            self.Hk = self._Hk_unpolarized
            self.Sk = self._Sk
        elif spin == 2:
            self.UP = 0
            self.DOWN = 1
            self.S_idx = 2
            self.Hk = self._Hk_polarized
            self.Sk = self._Sk
        elif spin == 4:
            self.Hk = self._Hk_non_collinear
            self.Sk = self._Sk_non_collinear
            self.S_idx = 4
        elif spin == 8:
            self.Hk = self._Hk_spin_orbit
            self.Sk = self._Sk_non_collinear
            self.S_idx = 8
            raise ValueError(
                "Currently the Hamiltonian has only been implemented with up to non-collinear spin."
            )

        if orthogonal:
            # There is no overlap matrix
            self.S_idx = -1

            def diagonal_Sk(self, k, dtype=None):
                """ For an orthogonal case we always return the identity matrix """
                if dtype is None:
                    dtype = np.float64
                no = self.no
                S = csr_matrix((no, no), dtype=dtype)
                S.setdiag(1.)
                return S

            self.Sk = diagonal_Sk

        # Denote that one *must* specify all details of the elements
        self._def_dim = -1
Exemplo n.º 13
0
class Hamiltonian(object):
    """ Hamiltonian object containing the coupling constants between orbitals.

    The Hamiltonian object contains information regarding the 
     - geometry
     - coupling constants between orbitals

    It contains an intrinsic sparse matrix of the Hamiltonian elements.

    Assigning or changing Hamiltonian elements is as easy as with
    standard ``numpy`` assignments:

    >>> ham = Hamiltonian(...)
    >>> ham.H[1,2] = 0.1

    which assigns 0.1 as the coupling constant between orbital 2 and 3.
    (remember that Python is 0-based elements).
    """

    # The order of the Energy
    # I.e. whether energy should be in other units than Ry
    # This conversion is made: [eV] ** _E_order
    _E_order = 1

    def __init__(self,
                 geom,
                 nnzpr=None,
                 orthogonal=True,
                 spin=1,
                 dtype=None,
                 *args,
                 **kwargs):
        """Create tight-binding model from geometry

        Initializes a tight-binding model using the :code:`geom` object
        as the underlying geometry for the tight-binding parameters.
        """
        self._geom = geom

        # Initialize the sparsity pattern
        self.reset(nnzpr=nnzpr, orthogonal=orthogonal, spin=spin, dtype=dtype)

    def reset(self, nnzpr=None, orthogonal=True, spin=1, dtype=None):
        """
        The sparsity pattern is cleaned and every thing
        is reset.

        The object will be the same as if it had been
        initialized with the same geometry as it were
        created with.

        Parameters
        ----------
        nnzpr: int
           number of non-zero elements per row
        orthogonal: boolean, True
           if there is an overlap matrix associated with the
           Hamiltonian
        spin: int, 1
           number of spin-components
        dtype: ``numpy.dtype``, `numpy.float64`
           the datatype of the Hamiltonian
        """
        # I know that this is not the most efficient way to
        # access a C-array, however, for constructing a
        # sparse pattern, it should be faster if memory elements
        # are closer...
        # Hence, this choice of having H and S like this

        # We check the first atom and its neighbours, we then
        # select max(5,len(nc) * 4)
        if nnzpr is None:
            nnzpr = self.geom.close(0)
            if nnzpr is None:
                nnzpr = 8
            else:
                nnzpr = max(5, len(nnzpr) * 4)

        self._orthogonal = orthogonal

        # Reset the sparsity pattern
        if not orthogonal:
            self._data = SparseCSR((self.no, self.no_s, spin + 1),
                                   nnzpr=nnzpr,
                                   dtype=dtype)
        else:
            self._data = SparseCSR((self.no, self.no_s, spin),
                                   nnzpr=nnzpr,
                                   dtype=dtype)

        self._spin = spin

        if spin == 1:
            self.UP = 0
            self.DOWN = 0
            self.S_idx = 1
            self.Hk = self._Hk_unpolarized
            self.Sk = self._Sk
        elif spin == 2:
            self.UP = 0
            self.DOWN = 1
            self.S_idx = 2
            self.Hk = self._Hk_polarized
            self.Sk = self._Sk
        elif spin == 4:
            self.Hk = self._Hk_non_collinear
            self.Sk = self._Sk_non_collinear
            self.S_idx = 4
        elif spin == 8:
            self.Hk = self._Hk_spin_orbit
            self.Sk = self._Sk_non_collinear
            self.S_idx = 8
            raise ValueError(
                "Currently the Hamiltonian has only been implemented with up to non-collinear spin."
            )

        if orthogonal:
            # There is no overlap matrix
            self.S_idx = -1

            def diagonal_Sk(self, k, dtype=None):
                """ For an orthogonal case we always return the identity matrix """
                if dtype is None:
                    dtype = np.float64
                no = self.no
                S = csr_matrix((no, no), dtype=dtype)
                S.setdiag(1.)
                return S

            self.Sk = diagonal_Sk

        # Denote that one *must* specify all details of the elements
        self._def_dim = -1

    def empty(self, keep=False):
        """ See `SparseCSR.empty` for details """
        self._data.empty(keep)

    def copy(self, dtype=None):
        """ Return a copy of the ``Hamiltonian`` object """
        if dtype is None:
            dtype = self.dtype
        H = self.__class__(self.geom,
                           orthogonal=self.orthogonal,
                           spin=self.spin,
                           dtype=dtype)
        # Be sure to copy the content of the SparseCSR object
        H._data = self._data.copy(dtype=dtype)
        return H

    ######### Definitions of overrides ############
    @property
    def geometry(self):
        """ Return the attached geometry """
        return self._geom

    geom = geometry

    @property
    def spin(self):
        """ Return number of spin-components in Hamiltonian """
        return self._spin

    @property
    def dtype(self):
        """ Return data type of Hamiltonian (and overlap matrix) """
        return self._data.dtype

    @property
    def orthogonal(self):
        """ Return whether the Hamiltonian is orthogonal """
        return self._orthogonal

    def __len__(self):
        """ Returns number of rows in the Hamiltonian """
        return self.geom.no

    def __repr__(self):
        """ Representation of the tight-binding model """
        s = self.geom.__repr__()
        return s + '\nNumber of spin / non-zero elements {0} / {1} '.format(
            self.spin, self.nnz)

    def __getattr__(self, attr):
        """ Returns the attributes from the underlying geometry

        Any attribute not found in the Hamiltonian class will
        be looked up in the underlying geometry.
        """
        return getattr(self.geom, attr)

    def __getitem__(self, key):
        """ Return Hamiltonian coupling elements for the index(s) """
        dd = self._def_dim
        if dd >= 0:
            key = tuple(key) + (dd, )
            self._def_dim = -1
        d = self._data[key]
        return d

    def __setitem__(self, key, val):
        """ Set or create couplings between orbitals in the Hamiltonian

        Override set item for slicing operations and enables easy
        setting of tight-binding parameters in a sparse matrix
        """
        dd = self._def_dim
        if dd >= 0:
            key = tuple(key) + (dd, )
            self._def_dim = -1
        self._data[key] = val

        if dd < 0 and not self.orthogonal:
            warnings.warn((
                'Hamiltonian specification of both H and S simultaneously is deprecated. '
                'This functionality will be removed in a future release.'))

    def __get_H(self):
        self._def_dim = self.UP
        return self

    _get_H = __get_H

    def __set_H(self, key, value):
        if len(key) == 2:
            self._def_dim = self.UP
        self[key] = value

    _set_H = __set_H

    H = property(__get_H, __set_H)

    def __get_S(self):
        if self.orthogonal:
            return None
        self._def_dim = self.S_idx
        return self

    _get_S = __get_S

    def __set_S(self, key, value):
        if self.orthogonal:
            return None
        self._def_dim = self.S_idx
        self[key] = value

    _set_S = __set_S

    S = property(__get_S, __set_S)

    # Create iterations on entire set of orbitals
    def iter(self, local=False):
        """ Iterations of the orbital space in the geometry, two indices from loop

        An iterator returning the current atomic index and the corresponding
        orbital index.

        >>> for ia, io in self:

        In the above case `io` always belongs to atom `ia` and `ia` may be
        repeated according to the number of orbitals associated with
        the atom `ia`.

        Parameters
        ----------
        local : `bool=False`
           whether the orbital index is the global index, or the local index relative to 
           the atom it resides on.
        """
        for ia, io in self.geom.iter_orbitals(local=local):
            yield ia, io

    __iter__ = iter

    # Create iterations on the non-zero elements
    def iter_nnz(self, atom=None, orbital=None):
        """ Iterations of the non-zero elements, returns a tuple of orbital and coupling orbital

        An iterator returning the current orbital index and the corresponding
        connected orbital where a non-zero is defined

        >>> for io, jo in self.iter_nnz():

        In the above case `io` and `jo` are orbitals such that:

        >>> self.H[io,jo] 

        returns the non-zero element of the Hamiltonian.

        One may reduce the iterated space by either requesting a specific set of atoms,
        or orbitals, _not_ both simultaneously.

        Examples
        --------
        Looping only on one or more atoms:

        >>> for io, jo in self.iter_nnz(atom=[2, 3]):
        >>>     # loop on all orbitals on atom 3 and 4 (0 indexing)

        >>> for io, jo in self.iter_nnz(orbital=[2, 3]):
        >>>     # loop on orbitals 3 and 4 (0 indexing)


        Parameters
        ----------
        atom : ``int``/``array_like``
           iterate on couplings to the set of atoms (not compatible with `orbital`)
        orbital : ``int``/``array_like``
           iterate on couplings to the set of orbitals (not compatible with `atom`)
        """
        if atom is not None and orbital is not None:
            raise ValueError(
                "iter_nnz: both atom and orbital has been passed, only one allowed."
            )

        if atom is not None:
            orbs = self.geom.a2o(atom, all=True)
            for io, jo in self._data.iter_nnz(orbs):
                yield io, jo
        elif orbital is not None:
            for io, jo in self._data.iter_nnz(orbital):
                yield io, jo
        else:
            for io, jo in self._data:
                yield io, jo

    def create_construct(self, dR, param):
        """ Returns a simple function for passing to the `construct` function.

        This is simply to leviate the creation of simplistic
        functions needed for setting up the Hamiltonian.

        Basically this returns a function:
        >>> def func(self, ia, idxs, idxs_xyz=None):
        >>>     idx = self.geom.close(ia, dR=dR, idx=idxs)
        >>>     for ix, p in zip(idx, param):
        >>>         self[ia, ix] = p

        Note
        ----
        This function only works for geometries with one orbital
        per atom.
        If you have more than one orbital on any atom, you should 
        define your own function.

        Parameters
        ----------
        dR : array_like
           radii parameters for tight-binding parameters.
           Must have same length as ``param`` or one less.
           If one less it will be extended with ``dR[0]/100``
        param : array_like
           coupling constants corresponding to the ``dR``
           ranges. ``param[0,:]`` are the tight-binding parameter
           for the all atoms within ``dR[0]`` of each atom.
        """

        if self.orthogonal:

            def func(self, ia, idxs, idxs_xyz=None):
                idx = self.geom.close(ia, dR=dR, idx=idxs, idx_xyz=idxs_xyz)
                for ix, p in zip(idx, param):
                    self[ia, ix] = p
        else:

            def func(self, ia, idxs, idxs_xyz=None):
                idx = self.geom.close(ia, dR=dR, idx=idxs, idx_xyz=idxs_xyz)
                for ix, p in zip(idx, param):
                    self.H[ia, ix] = p[:-1]
                    self.S[ia, ix] = p[-1]

        return func

    def construct(self, func, na_iR=1000, method='rand', eta=False):
        """ Automatically construct the Hamiltonian model based on a function that does the setting up of the Hamiltonian

        This may be called in two variants.

        1. Pass a function (``func``), see e.g. ``create_construct`` 
           which does the setting up.
        2. Pass a tuple/list in ``func`` which consists of two 
           elements, one is ``dR`` the radii parameters for
           the corresponding tight-binding parameters.
           The second is the tight-binding parameters
           corresponding to the ``dR[i]`` elements.
           In this second case all atoms must only have
           one orbital.

        Parameters
        ----------
        func: callable or array_like
           this function *must* take 4 arguments.
           1. Is the Hamiltonian object it-self (`self`)
           2. Is the currently examined atom (`ia`)
           3. Is the currently bounded indices (`idxs`)
           4. Is the currently bounded indices atomic coordinates (`idxs_xyz`)
           An example `func` could be:

           >>> def func(self, ia, idxs, idxs_xyz=None):
           >>>     idx = self.geom.close(ia, dR=[0.1, 1.44], idx=idxs, idx_xyz=idxs_xyz)
           >>>     self.H[ia, idx[0]] = 0.   # on-site
           >>>     self.H[ia, idx[1]] = -2.7 # nearest-neighbour
        na_iR : int, 1000
           number of atoms within the sphere for speeding
           up the `iter_block` loop.
        method : str, 'rand'
           method used in `Geometry.iter_block`, see there for details
        eta: bool, False
           whether an ETA will be printed
        """

        if not callable(func):
            if not isinstance(func, (tuple, list)):
                raise ValueError(
                    'Passed `func` which is not a function, nor tuple/list of `dR, param`'
                )

            if np.any(np.diff(self.geom.lasto) > 1):
                raise ValueError(
                    "Automatically setting a tight-binding model "
                    "for systems with atoms having more than 1 "
                    "orbital *must* be done by your-self. You have to define a corresponding `func`."
                )

            # Convert to a proper function
            func = self.create_construct(func[0], func[1])

        iR = self.geom.iR(na_iR)

        # Get number of atoms
        na = len(self.geom)
        na_run = 0

        from time import time
        from sys import stdout
        t0 = time()

        # Do the loop
        for ias, idxs in self.geom.iter_block(iR=iR, method=method):

            # Get all the indexed atoms...
            # This speeds up the searching for
            # coordinates...
            idxs_xyz = self.geom[idxs, :]

            # Loop the atoms inside
            for ia in ias:
                func(self, ia, idxs, idxs_xyz)

            if eta:
                # calculate the remaining atoms to process
                na_run += len(ias)
                na -= len(ias)
                t1 = time()
                # calculate hours, minutes, seconds
                m, s = divmod(float(t1 - t0) / na_run * na, 60)
                h, m = divmod(m, 60)
                stdout.write(
                    "Hamiltonian.construct() ETA = {0:5d}h {1:2d}m {2:5.2f}s\r"
                    .format(int(h), int(m), s))
                stdout.flush()

        if eta:
            stdout.write("Hamiltonian.construct() {0:23s}\n".format('DONE'))
            stdout.flush()

    @property
    def finalized(self):
        """ Whether the contained data is finalized and non-used elements have been removed """
        return self._data.finalized

    def finalize(self):
        """ Finalizes the tight-binding model

        Finalizes the tight-binding model so that no new sparse
        elements can be added.

        Sparse elements can still be changed.
        """
        self._data.finalize()

        # Get the folded Hamiltonian at the Gamma point
        Hk = self.Hk()

        nzs = Hk.nnz

        if nzs != (Hk + Hk.T).nnz:
            warnings.warn(
                'Hamiltonian does not retain symmetric couplings, this might be problematic.'
            )

    @property
    def nnz(self):
        """ Returns number of non-zero elements in the tight-binding model """
        return self._data.nnz

    @property
    def no(self):
        """ Returns number of orbitals as used when the object was created """
        return self._data.nr

    def tocsr(self, index, isc=None):
        """ Return a ``scipy.sparse.csr_matrix`` from the specified index

        Parameters
        ----------
        index : ``int``
           the index in the sparse matrix (for non-orthogonal cases the last
           dimension is the overlap matrix)
        isc : ``int``, `None`
           the supercell index (or all)
        """
        if isc is not None:
            raise NotImplementedError(
                "Requesting sub-Hamiltonian has not been implemented yet")
        return self._data.tocsr(index)

    def _Hk_unpolarized(self, k=(0, 0, 0), dtype=None):
        """ Return the Hamiltonian in a ``scipy.sparse.csr_matrix`` at `k`.

        Parameters
        ----------
        k: ``array_like``, `[0,0,0]`
           k-point 
        dtype : ``numpy.dtype``
           default to `numpy.complex128`
        """
        return self._Hk_polarized(k, dtype=dtype)

    def _Hk_polarized(self, k=(0, 0, 0), spin=0, dtype=None):
        """ Return the Hamiltonian in a ``scipy.sparse.csr_matrix`` at `k` for a polarized calculation

        Parameters
        ----------
        k: ``array_like``, `[0,0,0]`
           k-point 
        spin: ``int``, `0`
           the spin-index of the Hamiltonian
        dtype : ``numpy.dtype``
           default to `numpy.complex128`
        """
        if dtype is None:
            dtype = np.complex128

        exp = np.exp
        dot = np.dot

        k = np.asarray(k, np.float64)
        k.shape = (-1, )

        if not np.allclose(k, 0.):
            if np.dtype(dtype).kind != 'c':
                raise ValueError(
                    "Hamiltonian setup at k different from Gamma requires a complex matrix"
                )

        # Setup the Hamiltonian for this k-point
        Hf = self.tocsr(spin)

        no = self.no
        s = (no, no)
        H = csr_matrix(s, dtype=dtype)

        # Get the reciprocal lattice vectors dotted with k
        kr = dot(self.rcell, k)
        for si in range(self.sc.n_s):
            isc = self.sc_off[si, :]
            phase = exp(-1j * dot(kr, dot(self.cell, isc)))
            H += Hf[:, si * no:(si + 1) * no] * phase

        del Hf

        return H

    def _Hk_non_collinear(self, k=(0, 0, 0), dtype=None):
        """ Return the Hamiltonian in a ``scipy.sparse.csr_matrix`` at `k` for a non-collinear
        Hamiltonian.

        Parameters
        ----------
        k: ``array_like``, `[0,0,0]`
           k-point 
        dtype : ``numpy.dtype``
           default to `numpy.complex128`
        """
        if dtype is None:
            dtype = np.complex128

        if np.dtype(dtype).kind != 'c':
            raise ValueError(
                "Non-collinear Hamiltonian setup requires a complex matrix")

        exp = np.exp
        dot = np.dot

        k = np.asarray(k, np.float64)
        k.shape = (-1, )

        no = self.no * 2
        s = (no, no)
        H = csr_matrix(s, dtype=dtype)

        # get back-dimension of the intrinsic sparse matrix
        no = self.no

        # Get the reciprocal lattice vectors dotted with k
        kr = dot(self.rcell, k)
        for si in range(self.sc.n_s):
            isc = self.sc_off[si, :]
            phase = exp(-1j * dot(kr, dot(self.cell, isc)))

            # diagonal elements
            Hf1 = self.tocsr(0)[:, si * no:(si + 1) * no] * phase
            for i, j, h in ispmatrixd(Hf1):
                H[i * 2, j * 2] += h
            Hf1 = self.tocsr(1)[:, si * no:(si + 1) * no] * phase
            for i, j, h in ispmatrixd(Hf1):
                H[1 + i * 2, 1 + j * 2] += h

            # off-diagonal elements
            Hf1 = self.tocsr(2)[:, si * no:(si + 1) * no]
            Hf2 = self.tocsr(3)[:, si * no:(si + 1) * no]
            # We expect Hf1 and Hf2 to be aligned equivalently!
            # TODO CHECK
            for i, j, hr in ispmatrixd(Hf1):
                # get value for the imaginary part
                hi = Hf2[i, j]
                H[i * 2, 1 + j * 2] += (hr - 1j * hi) * phase
                H[1 + i * 2, j * 2] += (hr + 1j * hi) * phase

        del Hf1, Hf2

        return H

    def _Sk(self, k=(0, 0, 0), dtype=None):
        """ Return the Hamiltonian in a ``scipy.sparse.csr_matrix`` at `k`.

        Parameters
        ----------
        k: ``array_like``, `[0,0,0]`
           k-point 
        dtype : ``numpy.dtype``
           default to `numpy.complex128`
        """
        # we forward it to Hk_polarized (same thing for S)
        return self._Hk_polarized(k, spin=self.S_idx, dtype=dtype)

    def _Sk_non_collinear(self, k=(0, 0, 0), dtype=None):
        """ Return the Hamiltonian in a ``scipy.sparse.csr_matrix`` at `k`.

        Parameters
        ----------
        k: ``array_like``, `[0,0,0]`
           k-point 
        dtype : ``numpy.dtype``
           default to `numpy.complex128`
        """
        if dtype is None:
            dtype = np.complex128

        if not np.allclose(k, 0.):
            if np.dtype(dtype).kind != 'c':
                raise ValueError(
                    "Hamiltonian setup at k different from Gamma requires a complex matrix"
                )

        exp = np.exp
        dot = np.dot

        k = np.asarray(k, np.float64)
        k.shape = (-1, )

        # Get the overlap matrix
        Sf = self.tocsr(self.S_idx)

        no = self.no * 2
        s = (no, no)
        S = csr_matrix(s, dtype=dtype)

        # Get back dimensionality of the intrinsic orbitals
        no = self.no

        # Get the reciprocal lattice vectors dotted with k
        kr = dot(self.rcell, k)
        for si in range(self.sc.n_s):
            isc = self.sc_off[si, :]
            phase = exp(-1j * dot(kr, dot(self.cell, isc)))
            # Setup the overlap for this k-point
            sf = Sf[:, si * no:(si + 1) * no]
            for i, j, s in ispmatrixd(sf):
                S[i * 2, j * 2] += s
                S[1 + i * 2, 1 + j * 2] += s

        del Sf

        return S

    def eigh(self,
             k=(0, 0, 0),
             atoms=None,
             eigvals_only=True,
             overwrite_a=True,
             overwrite_b=True,
             *args,
             **kwargs):
        """ Returns the eigenvalues of the Hamiltonian

        Setup the Hamiltonian and overlap matrix with respect to
        the given k-point, then reduce the space to the specified atoms
        and calculate the eigenvalues.

        All subsequent arguments gets passed directly to :code:`scipy.linalg.eigh`
        """
        H = self.Hk(k=k)
        if not self.orthogonal:
            S = self.Sk(k=k)
        # Reduce sparsity pattern
        if not atoms is None:
            orbs = self.a2o(atoms)
            # Reduce space
            H = H[orbs, orbs]
            if not self.orthogonal:
                S = S[orbs, orbs]
        if self.orthogonal:
            return sli.eigh(H.todense(),
                            *args,
                            eigvals_only=eigvals_only,
                            overwrite_a=overwrite_a,
                            **kwargs)

        return sli.eigh(H.todense(),
                        S.todense(),
                        *args,
                        eigvals_only=eigvals_only,
                        overwrite_a=overwrite_a,
                        overwrite_b=overwrite_b,
                        **kwargs)

    def eigsh(self,
              k=(0, 0, 0),
              n=10,
              atoms=None,
              eigvals_only=True,
              *args,
              **kwargs):
        """ Returns the eigenvalues of the Hamiltonian

        Setup the Hamiltonian and overlap matrix with respect to
        the given k-point, then reduce the space to the specified atoms
        and calculate the eigenvalues.

        All subsequent arguments gets passed directly to :code:`scipy.linalg.eigh`
        """

        # We always request the smallest eigenvalues...
        kwargs.update({'which': kwargs.get('which', 'SM')})

        H = self.Hk(k=k)
        if not self.orthogonal:
            raise ValueError(
                "The sparsity pattern is non-orthogonal, you cannot use the Arnoldi procedure with scipy"
            )

        # Reduce sparsity pattern
        if not atoms is None:
            orbs = self.a2o(atoms)
            # Reduce space
            H = H[orbs, orbs]

        return ssli.eigsh(H,
                          k=n,
                          *args,
                          return_eigenvectors=not eigvals_only,
                          **kwargs)

    def cut(self, seps, axis, *args, **kwargs):
        """ Cuts the tight-binding model into different parts.

        Creates a tight-binding model by retaining the parameters
        for the cut-out region, possibly creating a super-cell.

        Parameters
        ----------
        seps  : integer, optional
           number of times the structure will be cut.
        axis  : integer
           the axis that will be cut
        """
        new_w = None
        # Create new geometry
        with warnings.catch_warnings(record=True) as w:
            # Cause all warnings to always be triggered.
            warnings.simplefilter("always")
            # Create new cut geometry
            geom = self.geom.cut(seps, axis, *args, **kwargs)
            # Check whether the warning exists
            if len(w) > 0:
                if issubclass(w[-1].category, UserWarning):
                    new_w = str(w[-1].message)
                    new_w += (
                        "\n---\n"
                        "The tight-binding model cannot be cut as the structure "
                        "cannot be tiled accordingly. ANY use of the model has been "
                        "relieved from sisl.")
        if new_w:
            warnings.warn(new_w, UserWarning)

        # Now we need to re-create the tight-binding model
        H = self.tocsr(0)
        has_S = self.S_idx > 0
        if has_S:
            S = self.tocsr(self.S_idx)
        # they are created similarly, hence the following
        # should keep their order

        # First we need to figure out how long the interaction range is
        # in the cut-direction
        # We initialize to be the same as the parent direction
        nsc = np.copy(self.nsc) // 2
        nsc[axis] = 0  # we count the new direction
        isc = np.zeros([3], np.int32)
        isc[axis] -= 1
        out = False
        while not out:
            # Get supercell index
            isc[axis] += 1
            try:
                idx = self.sc_index(isc)
            except:
                break

            # Figure out if the Hamiltonian has interactions
            # to 'isc'
            sub = H[0:geom.no, idx * self.no:(idx + 1) * self.no].indices[:]
            if has_S:
                sub = np.unique(
                    np.concatenate(
                        (sub, S[0:geom.no,
                                idx * self.no:(idx + 1) * self.no].indices[:]),
                        axis=0))
            if len(sub) == 0:
                break

            c_max = np.amax(sub)
            # Count the number of cells it interacts with
            i = (c_max % self.no) // geom.no
            ic = idx * self.no
            for j in range(i):
                idx = ic + geom.no * j
                # We need to ensure that every "in between" index exists
                # if it does not we discard those indices
                if len(
                        np.where(
                            np.logical_and(idx <= sub,
                                           sub < idx + geom.no))[0]) == 0:
                    i = j - 1
                    out = True
                    break
            nsc[axis] = isc[axis] * seps + i

            if out:
                warnings.warn(
                    'Cut the connection at nsc={0} in direction {1}.'.format(
                        nsc[axis], axis), UserWarning)

        # Update number of super-cells
        nsc[:] = nsc[:] * 2 + 1
        geom.sc.set_nsc(nsc)

        # Now we have a correct geometry, and
        # we are now ready to create the sparsity pattern
        # Reduce the sparsity pattern, first create the new one
        ham = self.__class__(geom,
                             nnzpr=np.amax(self._data.ncol),
                             spin=self.spin,
                             orthogonal=self.orthogonal)

        def sco2sco(M, o, m, seps, axis):
            # Converts an o from M to m
            isc = np.copy(M.o2isc(o))
            isc[axis] = isc[axis] * seps
            # Correct for cell-offset
            isc[axis] = isc[axis] + (o % M.no) // m.no
            # find the equivalent cell in m
            try:
                # If a fail happens it is due to a discarded
                # interaction across a non-interacting region
                return (o % m.no, m.sc_index(isc) * m.no,
                        m.sc_index(-isc) * m.no)
            except:
                return None, None, None

        # Copy elements
        if has_S:
            for jo in range(geom.no):

                # make smaller cut
                sH = H[jo, :]
                sS = S[jo, :]

                for io, iH in zip(sH.indices, sH.data):
                    # Get the equivalent orbital in the smaller cell
                    o, ofp, ofm = sco2sco(self.geom, io, ham.geom, seps, axis)
                    if o is None:
                        continue
                    ham.H[jo, o + ofp] = iH
                    ham.S[jo, o + ofp] = S[jo, io]
                    ham.H[o, jo + ofm] = iH
                    ham.S[o, jo + ofm] = S[jo, io]

                if np.any(sH.indices != sS.indices):

                    # Ensure that S is also cut
                    for io, iS in zip(sS.indices, sS.data):
                        # Get the equivalent orbital in the smaller cell
                        o, ofp, ofm = sco2sco(self.geom, io, ham.geom, seps,
                                              axis)
                        if o is None:
                            continue
                        ham.H[jo, o + ofp] = H[jo, io]
                        ham.S[jo, o + ofp] = iS
                        ham.H[o, jo + ofm] = H[jo, io]
                        ham.S[o, jo + ofm] = iS

        else:
            for jo in range(geom.no):
                sH = H[jo, :]

                for io, iH in zip(sH.indices, sH.data):
                    # Get the equivalent orbital in the smaller cell
                    o, ofp, ofm = sco2sco(self.geom, io, ham.geom, seps, axis)
                    if o is None:
                        continue
                    ham[jo, o + ofp] = iH
                    ham[o, jo + ofm] = iH

        return ham

    def tile(self, reps, axis):
        """ Returns a repeated tight-binding model for this, much like the `Geometry`

        The already existing tight-binding parameters are extrapolated
        to the new supercell by repeating them in blocks like the coordinates.

        Parameters
        ----------
        reps : number of tiles (repetitions)
        axis : direction of tiling
            0, 1, 2 according to the cell-direction
        """

        # Create the new geometry
        g = self.geom.tile(reps, axis)

        raise NotImplementedError(('tiling a Hamiltonian model has not been '
                                   'fully implemented yet.'))

    def repeat(self, reps, axis):
        """ Refer to `tile` instead """
        # Create the new geometry
        g = self.geom.repeat(reps, axis)

        raise NotImplementedError(
            ('repeating a Hamiltonian model has not been '
             'fully implemented yet, use tile instead.'))

    @classmethod
    def sp2HS(cls, geom, H, S=None):
        """ Returns a tight-binding model from a preset H, S and Geometry
        """
        # Calculate number of connections
        nc = 0

        has_S = not S is None

        # Ensure csr format
        H = H.tocsr()
        if has_S:
            S = S.tocsr()
        for i in range(geom.no):
            nc = max(nc, H[i, :].getnnz())
            if has_S:
                nc = max(nc, S[i, :].getnnz())

        # Create the Hamiltonian
        ham = cls(geom, nnzpr=nc, orthogonal=not has_S, dtype=H.dtype)

        # Copy data to the model
        if has_S:
            for jo, io in ispmatrix(H):
                ham.S[jo, io] = S[jo, io]

            # If the Hamiltonian for one reason or the other
            # is zero in the diagonal, then we *must* account for
            # this as it isn't captured in the above loop.
            skip_S = np.all(H.row == S.row)
            skip_S = skip_S and np.all(H.col == S.col)
            skip_S = False
            if not skip_S:
                # Re-convert back to allow index retrieval
                H = H.tocsr()
                for jo, io, s in ispmatrixd(S):
                    ham[jo, io] = (H[jo, io], s)

        else:
            for jo, io, h in ispmatrixd(H):
                ham[jo, io] = h

        return ham

    @staticmethod
    def read(sile, *args, **kwargs):
        """ Reads Hamiltonian from `Sile` using `read_H`.

        Parameters
        ----------
        sile : `Sile`, str
            a `Sile` object which will be used to read the Hamiltonian
            and the overlap matrix (if any)
            if it is a string it will create a new sile using `get_sile`.
        * : args passed directly to ``read_es(,**)``
        """
        # This only works because, they *must*
        # have been imported previously
        from sisl.io import get_sile, BaseSile
        if isinstance(sile, BaseSile):
            return sile.read_es(*args, **kwargs)
        else:
            return get_sile(sile).read_es(*args, **kwargs)

    def write(self, sile, *args, **kwargs):
        """ Writes a tight-binding model to the `Sile` as implemented in the :code:`Sile.write_es` method """
        self.finalize()

        # This only works because, they *must*
        # have been imported previously
        from sisl.io import get_sile, BaseSile
        if isinstance(sile, BaseSile):
            sile.write_es(self, *args, **kwargs)
        else:
            get_sile(sile, 'w').write_es(self, *args, **kwargs)

    ###############################
    # Overload of math operations #
    ###############################
    def __add__(a, b):
        c = a.copy(dtype=get_dtype(b, other=a.dtype))
        c += b
        return c

    __radd__ = __add__

    def __iadd__(a, b):
        if isinstance(b, Hamiltonian):
            a._data += b._data
        else:
            a._data += b
        return a

    def __sub__(a, b):
        c = a.copy(dtype=get_dtype(b, other=a.dtype))
        c -= b
        return c

    def __rsub__(a, b):
        if isinstance(b, Hamiltonian):
            c = b.copy(dtype=get_dtype(a, other=b.dtype))
            c._data += -1 * a._data
        else:
            c = b + (-1) * a
        return c

    def __isub__(a, b):
        if isinstance(b, Hamiltonian):
            a._data -= b._data
        else:
            a._data -= b
        return a

    def __mul__(a, b):
        c = a.copy(dtype=get_dtype(b, other=a.dtype))
        c *= b
        return c

    __rmul__ = __mul__

    def __imul__(a, b):
        if isinstance(b, Hamiltonian):
            a._data *= b._data
        else:
            a._data *= b
        return a

    def __div__(a, b):
        c = a.copy(dtype=get_dtype(b, other=a.dtype))
        c /= b
        return c

    def __rdiv__(a, b):
        c = b.copy(dtype=get_dtype(a, other=b.dtype))
        c /= a
        return c

    def __idiv__(a, b):
        if isinstance(b, Hamiltonian):
            a._data /= b._data
        else:
            a._data /= b
        return a

    def __floordiv__(a, b):
        if isinstance(b, Hamiltonian):
            raise NotImplementedError
        c = a.copy(dtype=get_dtype(b, other=a.dtype))
        c //= b
        return c

    def __ifloordiv__(a, b):
        if isinstance(b, Hamiltonian):
            raise NotImplementedError
        a._data //= b
        return a

    def __truediv__(a, b):
        if isinstance(b, Hamiltonian):
            raise NotImplementedError
        c = a.copy(dtype=get_dtype(b, other=a.dtype))
        c /= b
        return c

    def __itruediv__(a, b):
        if isinstance(b, Hamiltonian):
            raise NotImplementedError
        a._data /= b
        return a

    def __pow__(a, b):
        c = a.copy(dtype=get_dtype(b, other=a.dtype))
        c **= b
        return c

    def __rpow__(a, b):
        c = a.copy(dtype=get_dtype(b, other=a.dtype))
        c._data = b**c._data
        return c

    def __ipow__(a, b):
        if isinstance(b, Hamiltonian):
            a._data **= b._data
        else:
            a._data **= b
        return a
Exemplo n.º 14
0
 def setUp(self):
     self.s1 = SparseCSR((10,100), dtype=np.int32)
     self.s2 = SparseCSR((10,100,2))
Exemplo n.º 15
0
class TestSparseCSR(object):

    def setUp(self):
        self.s1 = SparseCSR((10,100), dtype=np.int32)
        self.s2 = SparseCSR((10,100,2))

    def test_init1(self):
        assert_equal(self.s1.dtype, np.int32)
        assert_equal(self.s2.dtype, np.float64)
        assert_true(np.allclose(self.s1.data, self.s1.data))
        assert_true(np.allclose(self.s2.data, self.s2.data))

    def test_init2(self):
        SparseCSR((10,100))
        for d in [np.int32, np.float64, np.complex128]:
            s = SparseCSR((10,100), dtype=d)
            assert_equal(s.shape, (10, 100, 1))
            assert_equal(s.dim, 1)
            assert_equal(s.dtype, d)
            for k in [1, 2]:
                s = SparseCSR((10,100,k), dtype=d)
                assert_equal(s.shape, (10, 100, k))
                assert_equal(s.dim, k)
                s = SparseCSR((10,100), dim=k, dtype=d)
                assert_equal(s.shape, (10, 100, k))
                assert_equal(s.dim, k)
                s = SparseCSR((10,100, 3), dim=k, dtype=d)
                assert_equal(s.shape, (10, 100, 3))
                assert_equal(s.dim, 3)

    def test_init3(self):
        csr = sc.sparse.csr_matrix( (10,10), dtype=np.int32)
        csr[0, 1] = 1
        csr[0, 2] = 2
        sp = SparseCSR(csr)
        assert_equal(sp.dtype, np.int32)
        assert_equal(sp.shape, (10, 10, 1))
        assert_equal(len(sp), 2)
        assert_equal(sp[0, 1], 1)
        assert_equal(sp[0, 2], 2)
        sp = SparseCSR(csr, dtype=np.float64)
        assert_equal(sp.shape, (10, 10, 1))
        assert_equal(sp.dtype, np.float64)
        assert_equal(len(sp), 2)
        assert_equal(sp[0, 1], 1)
        assert_equal(sp[0, 2], 2)

    def test_init4(self):
        csr = sc.sparse.csr_matrix( (10,10), dtype=np.int32)
        csr[0, 1] = 1
        csr[0, 2] = 2
        print(csr.indices, csr.indptr)
        sp = SparseCSR((csr.data, csr.indices, csr.indptr))
        assert_equal(sp.dtype, np.int32)
        assert_equal(sp.shape, (10, 10, 1))
        assert_equal(len(sp), 2)
        assert_equal(sp[0, 1], 1)
        assert_equal(sp[0, 2], 2)
        sp = SparseCSR((csr.data, csr.indices, csr.indptr), dtype=np.float64)
        assert_equal(sp.shape, (10, 10, 1))
        assert_equal(sp.dtype, np.float64)
        assert_equal(len(sp), 2)
        assert_equal(sp[0, 1], 1)
        assert_equal(sp[0, 2], 2)

    def test_create1(self):
        self.s1[0,[1,2,3]] = 1
        assert_equal(self.s1.nnz, 3)
        self.s1[2,[1,2,3]] = 1
        assert_equal(self.s1.nnz, 6)
        self.s1.empty(keep=True)
        assert_equal(self.s1.nnz, 6)
        self.s1.empty()
        assert_equal(self.s1.nnz, 0)
        
    def test_create2(self):
        for i in range(10):
            j = range(i*4, i*4+3)
            self.s1[0, j] = i
            assert_equal(len(self.s1), (i+1)*3)
            for jj in j:
                assert_equal(self.s1[0, jj], i)
                assert_equal(self.s1[1, jj], 0)
        self.s1.empty()

    def test_create3(self):
        for i in range(10):
            j = range(i*4, i*4+3)
            self.s1[0, j] = i
            assert_equal(len(self.s1), (i+1)*3)
            self.s1[0, range((i+1)*4, (i+1)*4+3)] = None
            assert_equal(len(self.s1), (i+1)*3)
            for jj in j:
                assert_equal(self.s1[0, jj], i)
                assert_equal(self.s1[1, jj], 0)
        self.s1.empty()

    def test_finalize1(self):
        self.s1[0,[1,2,3]] = 1
        self.s1[2,[1,2,3]] = 1.
        assert_false(self.s1.finalized)
        self.s1.finalize()
        assert_true(self.s1.finalized)
        self.s1.empty(keep=True)
        assert_true(self.s1.finalized)
        self.s1.empty()
        assert_false(self.s1.finalized)

    def test_delitem1(self):
        self.s1[0,[1,2,3]] = 1
        assert_equal(len(self.s1), 3)
        del self.s1[0,1]
        assert_equal(len(self.s1), 2)
        assert_equal(self.s1[0,1], 0)
        assert_equal(self.s1[0,2], 1)
        assert_equal(self.s1[0,3], 1)
        self.s1[0,[1,2,3]] = 1
        del self.s1[0,[1,3]]
        assert_equal(len(self.s1), 1)
        assert_equal(self.s1[0,1], 0)
        assert_equal(self.s1[0,2], 1)
        assert_equal(self.s1[0,3], 0)
        self.s1.empty()

    def test_op1(self):
        for i in range(10):
            j = range(i*4, i*4+3)
            self.s1[0, j] = i

            # i+
            self.s1 += 1
            for jj in j:
                assert_equal(self.s1[0, jj], i+1)
                assert_equal(self.s1[1, jj], 0)

            # i-
            self.s1 -= 1
            for jj in j:
                assert_equal(self.s1[0, jj], i)
                assert_equal(self.s1[1, jj], 0)

            # i*
            self.s1 *= 2
            for jj in j:
                assert_equal(self.s1[0, jj], i*2)
                assert_equal(self.s1[1, jj], 0)

            # //
            self.s1 //= 2
            for jj in j:
                assert_equal(self.s1[0, jj], i)
                assert_equal(self.s1[1, jj], 0)

            # i**
            self.s1 **= 2
            for jj in j:
                assert_equal(self.s1[0, jj], i**2)
                assert_equal(self.s1[1, jj], 0)

    def test_op2(self):
        for i in range(10):
            j = range(i*4, i*4+3)
            self.s1[0, j] = i

            # +
            s = self.s1 + 1
            for jj in j:
                assert_equal(s[0, jj], i+1)
                assert_equal(self.s1[0, jj], i)
                assert_equal(s[1, jj], 0)

            # -
            s = self.s1 - 1
            for jj in j:
                assert_equal(s[0, jj], i-1)
                assert_equal(self.s1[0, jj], i)
                assert_equal(s[1, jj], 0)

            # - (r)
            s = 1 - self.s1
            for jj in j:
                assert_equal(s[0, jj], 1 - i)
                assert_equal(self.s1[0, jj], i)
                assert_equal(s[1, jj], 0)

            # *
            s = self.s1 * 2
            for jj in j:
                assert_equal(s[0, jj], i*2)
                assert_equal(self.s1[0, jj], i)
                assert_equal(s[1, jj], 0)

            # //
            s = s // 2
            for jj in j:
                assert_equal(s[0, jj], i)
                assert_equal(self.s1[0, jj], i)
                assert_equal(s[1, jj], 0)

            # **
            s = self.s1 ** 2
            for jj in j:
                assert_equal(s[0, jj], i**2)
                assert_equal(self.s1[0, jj], i)
                assert_equal(s[1, jj], 0)

            # ** (r)
            s = 2 ** self.s1
            for jj in j:
                assert_equal(s[0, jj], 2 ** self.s1[0, jj])
                assert_equal(self.s1[0, jj], i)
                assert_equal(s[1, jj], 0)

    def test_op3(self):
        S = SparseCSR((10,100), dtype=np.int32)
        # Create initial stuff
        for i in range(10):
            j = range(i*4, i*4+3)
            S[0, j] = i

        for op in ['add', 'sub', 'mul', 'pow']:
            func = getattr(S, '__{}__'.format(op))
            s = func(1)
            assert_equal(s.dtype, np.int32)
            s = func(1.)
            assert_equal(s.dtype, np.float64)
            if op != 'pow':
                s = func(1.j)
                assert_equal(s.dtype, np.complex128)

        S = S.copy(dtype=np.float64)
        for op in ['add', 'sub', 'mul', 'pow']:
            func = getattr(S, '__{}__'.format(op))
            s = func(1)
            assert_equal(s.dtype, np.float64)
            s = func(1.)
            assert_equal(s.dtype, np.float64)
            if op != 'pow':
                s = func(1.j)
                assert_equal(s.dtype, np.complex128)

        S = S.copy(dtype=np.complex128)
        for op in ['add', 'sub', 'mul', 'pow']:
            func = getattr(S, '__{}__'.format(op))
            s = func(1)
            assert_equal(s.dtype, np.complex128)
            s = func(1.)
            assert_equal(s.dtype, np.complex128)
            if op != 'pow':
                s = func(1.j)
                assert_equal(s.dtype, np.complex128)

    def test_op4(self):
        S = SparseCSR((10,100), dtype=np.int32)
        # Create initial stuff
        for i in range(10):
            j = range(i*4, i*4+3)
            S[0, j] = i

        s = 1 + S
        assert_equal(s.dtype, np.int32)
        s = 1. + S
        assert_equal(s.dtype, np.float64)
        s = 1.j + S
        assert_equal(s.dtype, np.complex128)

        s = 1 - S
        assert_equal(s.dtype, np.int32)
        s = 1. - S
        assert_equal(s.dtype, np.float64)
        s = 1.j - S
        assert_equal(s.dtype, np.complex128)
        
        s = 1 * S
        assert_equal(s.dtype, np.int32)
        s = 1. * S
        assert_equal(s.dtype, np.float64)
        s = 1.j * S
        assert_equal(s.dtype, np.complex128)

        s = 1 ** S
        assert_equal(s.dtype, np.int32)
        s = 1. ** S
        assert_equal(s.dtype, np.float64)
        s = 1.j ** S
        assert_equal(s.dtype, np.complex128)