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)
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_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_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_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)
def setUp(self):
self.s1 = SparseCSR((10, 100), dtype=np.int32)
self.s2 = SparseCSR((10, 100, 2))
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)
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
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
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
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
def setUp(self):
self.s1 = SparseCSR((10,100), dtype=np.int32)
self.s2 = SparseCSR((10,100,2))
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)
