def _call(self, *args, **kwargs):
func = getattr(self.parent, self.__attr)
wrap = kwargs.pop('wrap', _do_nothing)
eta = tqdm_eta(len(self), self.__class__.__name__ + '.asyield()',
'k', kwargs.pop('eta', False))
for k in self:
yield wrap(
func(*args, k=k, **kwargs).astype(dtype, copy=False))
eta.update()
eta.close()
def _call(self, *args, **kwargs):
func = getattr(self.parent, self.__attr)
wrap = kwargs.pop('wrap', _do_nothing)
eta = tqdm_eta(len(self), self.__class__.__name__ + '.aslist()',
'k', kwargs.pop('eta', False))
a = [None] * len(self)
for i, k in enumerate(self):
a[i] = wrap(func(*args, k=k, **kwargs))
eta.update()
eta.close()
return a
def _parse_kwargs(self, wrap, eta=False, eta_key=""):
""" Parse kwargs """
bz = self._obj
parent = bz.parent
if wrap is None:
# we always return a wrap
def wrap(v, parent=None, k=None, weight=None):
return v
else:
wrap = allow_kwargs("parent", "k", "weight")(wrap)
eta = tqdm_eta(len(bz), f"{bz.__class__.__name__}.{eta_key}", "k", eta)
return bz, parent, wrap, eta
def _call(self, *args, **kwargs):
func = getattr(self.parent, self._bz_attr)
wrap = allow_kwargs('parent', 'k', 'weight')(kwargs.pop('wrap', _do_nothing))
eta = tqdm_eta(len(self), self.__class__.__name__ + '.asyield',
'k', kwargs.pop('eta', False))
parent = self.parent
k = self.k
w = self.weight
for i in range(len(k)):
yield wrap(func(*args, k=k[i], **kwargs), parent=parent, k=k[i], weight=w[i])
eta.update()
eta.close()
def _call(self, *args, **kwargs):
func = getattr(self.parent, self.__attr)
wrap = kwargs.pop('wrap', _do_nothing)
eta = tqdm_eta(len(self), self.__class__.__name__ + '.asaverage()',
'k', kwargs.pop('eta', False))
w = self.weight.view()
for i, k in enumerate(self):
if i == 0:
v = wrap(func(*args, k=k, **kwargs)) * w[i]
else:
v += wrap(func(*args, k=k, **kwargs)) * w[i]
eta.update()
eta.close()
return v
def _call(self, *args, **kwargs):
func = getattr(self.parent, self._bz_attr)
wrap = allow_kwargs('parent', 'k', 'weight')(kwargs.pop('wrap', _do_nothing))
eta = tqdm_eta(len(self), self.__class__.__name__ + '.asarray',
'k', kwargs.pop('eta', False))
parent = self.parent
k = self.k
w = self.weight
v = wrap(func(*args, k=k[0], **kwargs), parent=parent, k=k[0], weight=w[0])
if v.ndim == 0:
a = np.empty([len(self)], dtype=v.dtype)
else:
a = np.empty((len(self), ) + v.shape, dtype=v.dtype)
a[0] = v
del v
for i in range(1, len(k)):
a[i] = wrap(func(*args, k=k[i], **kwargs), parent=parent, k=k[i], weight=w[i])
eta.update()
eta.close()
return a
def _call(self, *args, **kwargs):
func = getattr(self.parent, self.__attr)
wrap = kwargs.pop('wrap', _do_nothing)
eta = tqdm_eta(len(self), self.__class__.__name__ + '.asarray()',
'k', kwargs.pop('eta', False))
for i, k in enumerate(self):
if i == 0:
v = wrap(func(*args, k=k, **kwargs))
if len(self) == 1:
return v
shp = [len(self)]
shp.extend(v.shape)
a = np.empty(shp, dtype=dtype)
a[i, :] = v[:]
del v
else:
a[i, :] = wrap(func(*args, k=k, **kwargs))
eta.update()
eta.close()
return a
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 test_tqdm_false():
eta = sm.tqdm_eta(2, 'Hello', 'unit', False)
eta.update()
eta.update()
eta.close()
def test_tqdm_eta_true():
eta = sm.tqdm_eta(2, 'Hello', 'unit', True)
eta.update()
eta.update()
eta.close()
def wavefunction(v,
grid,
geometry=None,
k=None,
spinor=0,
spin=None,
eta=False):
r""" Add the wave-function (`Orbital.psi`) component of each orbital to the grid
This routine calculates the real-space wave-function components in the
specified grid.
This is an *in-place* operation that *adds* to the current values in the grid.
It may be instructive to check that an eigenstate is normalized:
>>> grid = Grid(...) # doctest: +SKIP
>>> psi(state, grid) # doctest: +SKIP
>>> (np.abs(grid.grid) ** 2).sum() * grid.dvolume == 1. # doctest: +SKIP
Note: To calculate :math:`\psi(\mathbf r)` in a unit-cell different from the
originating geometry, simply pass a grid with a unit-cell smaller than the originating
supercell.
The wavefunctions are calculated in real-space via:
.. math::
\psi(\mathbf r) = \sum_i\phi_i(\mathbf r) |\psi\rangle_i \exp(-i\mathbf k \mathbf R)
While for non-collinear/spin-orbit calculations the wavefunctions are determined from the
spinor component (`spinor`)
.. math::
\psi_{\alpha/\beta}(\mathbf r) = \sum_i\phi_i(\mathbf r) |\psi_{\alpha/\beta}\rangle_i \exp(-i\mathbf k \mathbf R)
where ``spinor in [0, 1]`` determines :math:`\alpha` or :math:`\beta`, respectively.
Notes
-----
Currently this method only works for :math:`\Gamma` states
Parameters
----------
v : array_like
coefficients for the orbital expansion on the real-space grid.
If `v` is a complex array then the `grid` *must* be complex as well.
grid : Grid
grid on which the wavefunction will be plotted.
If multiple eigenstates are in this object, they will be summed.
geometry : Geometry, optional
geometry where the orbitals are defined. This geometry's orbital count must match
the number of elements in `v`.
If this is ``None`` the geometry associated with `grid` will be used instead.
k : array_like, optional
k-point associated with wavefunction, by default the inherent k-point used
to calculate the eigenstate will be used (generally shouldn't be used unless the `EigenstateElectron` object
has not been created via `Hamiltonian.eigenstate`).
spinor : int, optional
the spinor for non-collinear/spin-orbit calculations. This is only used if the
eigenstate object has been created from a parent object with a `Spin` object
contained, *and* if the spin-configuration is non-collinear or spin-orbit coupling.
Default to the first spinor component.
spin : Spin, optional
specification of the spin configuration of the orbital coefficients. This only has
influence for non-collinear wavefunctions where `spinor` choice is important.
eta : bool, optional
Display a console progressbar.
"""
if geometry is None:
geometry = grid.geometry
warn(
'wavefunction was not passed a geometry associated, will use the geometry associated with the Grid.'
)
if geometry is None:
raise SislError(
'wavefunction did not find a usable Geometry through keywords or the Grid!'
)
# In case the user has passed several vectors we sum them to plot the summed state
if v.ndim == 2:
v = v.sum(0)
if spin is None:
if len(v) // 2 == geometry.no:
# We can see from the input that the vector *must* be a non-collinear calculation
v = v.reshape(-1, 2)[:, spinor]
info(
'wavefunction assumes the input wavefunction coefficients to originate from a non-collinear calculation!'
)
elif spin.kind > Spin.POLARIZED:
# For non-collinear cases the user selects the spinor component.
v = v.reshape(-1, 2)[:, spinor]
if len(v) != geometry.no:
raise ValueError(
"wavefunction require wavefunction coefficients corresponding to number of orbitals in the geometry."
)
# Check for k-points
k = _a.asarrayd(k)
kl = (k**2).sum()**0.5
has_k = kl > 0.000001
if has_k:
raise NotImplementedError(
'wavefunction for k != Gamma does not produce correct wavefunctions!'
)
# Check that input/grid makes sense.
# If the coefficients are complex valued, then the grid *has* to be
# complex valued.
# Likewise if a k-point has been passed.
is_complex = np.iscomplexobj(v) or has_k
if is_complex and not np.iscomplexobj(grid.grid):
raise SislError(
"wavefunction input coefficients are complex, while grid only contains real."
)
if is_complex:
psi_init = _a.zerosz
else:
psi_init = _a.zerosd
# Extract sub variables used throughout the loop
shape = _a.asarrayi(grid.shape)
dcell = grid.dcell
ic = grid.sc.icell * shape.reshape(1, -1)
geom_shape = dot(ic, geometry.cell.T).T
# 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')
addouter = add.outer
def idx2spherical(ix, iy, iz, offset, dc, R):
""" Calculate the spherical coordinates from indices """
rx = addouter(addouter(ix * dc[0, 0], iy * dc[1, 0]),
iz * dc[2, 0] - offset[0]).ravel()
ry = addouter(addouter(ix * dc[0, 1], iy * dc[1, 1]),
iz * dc[2, 1] - offset[1]).ravel()
rz = addouter(addouter(ix * dc[0, 2], iy * dc[1, 2]),
iz * dc[2, 2] - offset[2]).ravel()
# Total size of the indices
n = rx.shape[0]
# Reduce our arrays to where the radius is "fine"
idx = indices_le(rx**2 + ry**2 + rz**2, R**2)
rx = rx[idx]
ry = ry[idx]
rz = rz[idx]
xyz_to_spherical_cos_phi(rx, ry, rz)
return n, idx, rx, ry, rz
# Figure out the max-min indices with a spacing of 1 radian
rad1 = pi / 180
theta, phi = ogrid[-pi:pi:rad1, 0:pi:rad1]
cphi, sphi = cos(phi), sin(phi)
ctheta_sphi = cos(theta) * sphi
stheta_sphi = sin(theta) * sphi
del sphi
nrxyz = (theta.size, phi.size, 3)
del theta, phi, rad1
# First we calculate the min/max indices for all atoms
idx_mm = _a.emptyi([geometry.na, 2, 3])
rxyz = _a.emptyd(nrxyz)
rxyz[..., 0] = ctheta_sphi
rxyz[..., 1] = stheta_sphi
rxyz[..., 2] = cphi
# Reshape
rxyz.shape = (-1, 3)
idx = dot(ic, rxyz.T)
idxm = idx.min(1).reshape(1, 3)
idxM = idx.max(1).reshape(1, 3)
del ctheta_sphi, stheta_sphi, cphi, idx, rxyz, nrxyz
origo = grid.sc.origo.reshape(1, -1)
for atom, ia in geometry.atom.iter(True):
if len(ia) == 0:
continue
R = atom.maxR()
# Now do it for all the atoms to get indices of the middle of
# the atoms
# The coordinates are relative to origo, so we need to shift (when writing a grid
# it is with respect to origo)
xyz = geometry.xyz[ia, :] - origo
idx = dot(ic, xyz.T).T
# Get min-max for all atoms
idx_mm[ia, 0, :] = idxm * R + idx
idx_mm[ia, 1, :] = idxM * R + idx
# Now we have min-max for all atoms
# When we run the below loop all indices can be retrieved by looking
# up in the above table.
# Before continuing, we can easily clean up the temporary arrays
del origo, idx
aranged = _a.aranged
# In case this grid does not have a Geometry associated
# We can *perhaps* easily attach a geometry with the given
# atoms in the unit-cell
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)
r_k = dot(geometry.rcell, k)
r_k_cell = dot(r_k, geometry.cell)
phase = 1
# Retrieve progressbar
eta = tqdm_eta(len(IA), 'wavefunction', 'atom', eta)
# Loop over all atoms in the grid-cell
for ia, xyz, isc in zip(IA, XYZ - grid.origo.reshape(1, 3), ISC):
# Get current atom
atom = geometry.atom[ia]
# Extract maximum R
R = atom.maxR()
if R <= 0.:
warn("Atom '{}' does not have a wave-function, skipping atom.".
format(atom))
eta.update()
continue
# Get indices in the supercell grid
idx = (isc.reshape(3, 1) * geom_shape).sum(0)
idxm = floor(idx_mm[ia, 0, :] + idx).astype(int32)
idxM = ceil(idx_mm[ia, 1, :] + idx).astype(int32) + 1
# Fast check whether we can skip this point
if idxm[0] >= shape[0] or idxm[1] >= shape[1] or idxm[2] >= shape[2] or \
idxM[0] <= 0 or idxM[1] <= 0 or idxM[2] <= 0:
eta.update()
continue
# Truncate values
if idxm[0] < 0:
idxm[0] = 0
if idxM[0] > shape[0]:
idxM[0] = shape[0]
if idxm[1] < 0:
idxm[1] = 0
if idxM[1] > shape[1]:
idxM[1] = shape[1]
if idxm[2] < 0:
idxm[2] = 0
if idxM[2] > shape[2]:
idxM[2] = shape[2]
# Now idxm/M contains min/max indices used
# Convert to spherical coordinates
n, idx, r, theta, phi = idx2spherical(aranged(idxm[0], idxM[0]),
aranged(idxm[1], idxM[1]),
aranged(idxm[2], idxM[2]), xyz,
dcell, R)
# Get initial orbital
io = geometry.a2o(ia)
if has_k:
phase = np.exp(-1j * (dot(r_k_cell, isc)))
# TODO
# Possibly the phase should be an additional
# array for the position in the unit-cell!
# + np.exp(-1j * dot(r_k, spher2cart(r, theta, np.arccos(phi)).T) )
# Allocate a temporary array where we add the psi elements
psi = psi_init(n)
# Loop on orbitals on this atom, grouped by radius
for os in atom.iter(True):
# Get the radius of orbitals (os)
oR = os[0].R
if oR <= 0.:
warn(
"Orbital(s) '{}' does not have a wave-function, skipping orbital!"
.format(os))
# Skip these orbitals
io += len(os)
continue
# Downsize to the correct indices
if R - oR < 1e-6:
idx1 = idx.view()
r1 = r.view()
theta1 = theta.view()
phi1 = phi.view()
else:
idx1 = indices_le(r, oR)
# Reduce arrays
r1 = r[idx1]
theta1 = theta[idx1]
phi1 = phi[idx1]
idx1 = idx[idx1]
# Loop orbitals with the same radius
for o in os:
# Evaluate psi component of the wavefunction and add it for this atom
psi[idx1] += o.psi_spher(r1, theta1, phi1,
cos_phi=True) * (v[io] * phase)
io += 1
# Clean-up
del idx1, r1, theta1, phi1, idx, r, theta, phi
# Convert to correct shape and add the current atom contribution to the wavefunction
psi.shape = idxM - idxm
grid.grid[idxm[0]:idxM[0], idxm[1]:idxM[1], idxm[2]:idxM[2]] += psi
# Clean-up
del psi
# Step progressbar
eta.update()
eta.close()
# Reset the error code for division
np.seterr(**old_err)
def _call(self, *args, **kwargs):
data_axis = kwargs.pop('data_axis', None)
grid_unit = kwargs.pop('grid_unit', 'b')
func = getattr(self.parent, self._bz_attr)
wrap = allow_kwargs('parent', 'k', 'weight')(kwargs.pop('wrap', _do_nothing))
eta = tqdm_eta(len(self), self.__class__.__name__ + '.asgrid',
'k', kwargs.pop('eta', False))
parent = self.parent
k = self.k
w = self.weight
# Extract information from the MP grid, these values
# define the Grid size, etc.
diag = self._diag.copy()
if not np.all(self._displ == 0):
raise SislError(self.__class__.__name__ + '.{} requires the displacement to be 0 for all k-points.'.format(self._bz_attr))
displ = self._displ.copy()
size = self._size.copy()
steps = size / diag
if self._centered:
offset = np.where(diag % 2 == 0, steps, steps / 2)
else:
offset = np.where(diag % 2 == 0, steps / 2, steps)
# Instead of doing
# _in_primitive(k) + 0.5 - offset
# we can do it here
# _in_primitive(k) + offset'
offset -= 0.5
# Check the TRS direction
trs_axis = self._trs
_in_primitive = self.in_primitive
_rint = np.rint
_int32 = np.int32
def k2idx(k):
# In case TRS is applied two indices may be returned
return _rint((_in_primitive(k) - offset) / steps).astype(_int32)
# To find the opposite k-point, do this
# idx[i] = [diag[i] - idx[i] - 1, idx[i]
# with i in [0, 1, 2]
# Create cell from the reciprocal cell.
if grid_unit == 'b':
cell = np.diag(self._size)
else:
cell = parent.sc.rcell * self._size.reshape(1, -1) / units('Ang', grid_unit)
# Find the grid origo
origo = -(cell * 0.5).sum(0)
# Calculate first k-point (to get size and dtype)
v = wrap(func(*args, k=k[0], **kwargs), parent=parent, k=k[0], weight=w[0])
if data_axis is None:
if v.size != 1:
raise SislError(self.__class__.__name__ + '.{} requires one value per-kpoint because of the 3D grid values'.format(self._bz_attr))
else:
# Check the weights
weights = self.grid(diag[data_axis], displ[data_axis], size[data_axis],
centered=self._centered, trs=trs_axis == data_axis)[1]
# Correct the Grid size
diag[data_axis] = len(v)
# Create the orthogonal cell direction to ensure it is orthogonal
# Since array axis is cyclic for negative numbers, we simply do this
cell[data_axis, :] = cross(cell[data_axis-1, :], cell[data_axis-2, :])
# Check whether we should rotate it
if cart2spher(cell[data_axis, :])[2] > pi / 4:
cell[data_axis, :] *= -1
# Correct cell for the grid
if trs_axis >= 0:
origo[trs_axis] = 0.
# Correct offset since we only have the positive halve
if self._diag[trs_axis] % 2 == 0 and not self._centered:
offset[trs_axis] = steps[trs_axis] / 2
else:
offset[trs_axis] = 0.
# Find number of points
if trs_axis != data_axis:
diag[trs_axis] = len(self.grid(diag[trs_axis], displ[trs_axis], size[trs_axis],
centered=self._centered, trs=True)[1])
# Create the grid in the reciprocal cell
sc = SuperCell(cell, origo=origo)
grid = Grid(diag, sc=sc, dtype=v.dtype)
if data_axis is None:
grid[k2idx(k[0])] = v
else:
idx = k2idx(k[0]).tolist()
weight = weights[idx[data_axis]]
idx[data_axis] = slice(None)
grid[idx] = v * weight
del v
# Now perform calculation
if data_axis is None:
for i in range(1, len(k)):
grid[k2idx(k[i])] = wrap(func(*args, k=k[i], **kwargs),
parent=parent, k=k[i], weight=w[i])
eta.update()
else:
for i in range(1, len(k)):
idx = k2idx(k[i]).tolist()
weight = weights[idx[data_axis]]
idx[data_axis] = slice(None)
grid[idx] = wrap(func(*args, k=k[i], **kwargs),
parent=parent, k=k[i], weight=w[i]) * weight
eta.update()
eta.close()
return grid
