def mgrid(cls, *slices):
""" Return a list of indices corresponding to the slices
The returned values are equivalent to `numpy.mgrid` but they are returned
in a (:, 3) array.
Parameters
----------
*slices : slice or list of int or int
return a linear list of indices that points to the collective slice
made by the passed arguments
Returns
-------
indices : (:, 3), linear indices for each of the sliced values
"""
if len(slices) == 1:
g = np.mgrid[slices[0]]
else:
g = np.mgrid[slices]
indices = _a.emptyi(g.size).reshape(-1, 3)
indices[:, 0] = g[0].flatten()
indices[:, 1] = g[1].flatten()
indices[:, 2] = g[2].flatten()
del g
return indices
def align_norm(self, other, ret_index=False):
r""" Align `other.state` with the site-norms for this state, a copy of `other` is returned with re-ordered states
To determine the new ordering of `other` we first calculate the residual norm of the site-norms.
.. math::
\delta N_{\alpha\beta} = \sum_i \big(\langle \psi^\alpha_i | \psi^\alpha_i\rangle - \langle \psi^\beta_i | \psi^\beta_i\rangle\big)^2
where :math:`\alpha` and :math:`\beta` correspond to state indices in `self` and `other`, respectively.
The new states (from `other`) returned is then ordered such that the index
:math:`\alpha \equiv \beta'` where :math:`\delta N_{\alpha\beta}` is smallest.
Parameters
----------
other : State
the other state to align onto this state
ret_index : bool, optional
also return indices for the swapped indices
Returns
-------
other_swap : State
A swapped instance of `other`
index : array of int
the indices that swaps `other` to be ``other_swap``, i.e. ``other_swap = other.sub(index)``
Notes
-----
The input state and output state have the same states, but their ordering is not necessarily the same.
See Also
--------
align_phase : rotate states such that their phases align
"""
snorm = self.norm2(False)
onorm = other.norm2(False)
# Now find new orderings
show_warn = False
idx = _a.fulli(len(other), -1)
idxr = _a.emptyi(len(other))
for i in range(len(other)):
R = snorm - onorm[i, :].reshape(1, -1)
R = einsum('ij,ij->i', R, R)
# Figure out which band it should correspond to
# find closest largest one
for j in np.argsort(R):
if j not in idx[:i]:
idx[i] = j
idxr[j] = i
break
show_warn = True
if show_warn:
warn(self.__class__.__name__ + '.align_norm found multiple possible candidates with minimal residue, swapping not unique')
if ret_index:
return other.sub(idxr), idxr
return other.sub(idxr)
def write_geometry(self, geometry):
""" Creates the NetCDF file and writes the geometry information """
sile_raise_write(self)
# Create initial dimensions
self.write_supercell(geometry.sc)
self._crt_dim(self, 'no_s', np.prod(geometry.nsc) * geometry.no)
self._crt_dim(self, 'no_u', geometry.no)
self._crt_dim(self, 'na_u', geometry.na)
# Create initial geometry
v = self._crt_var(self, 'lasto', 'i4', ('na_u', ))
v.info = 'Last orbital of equivalent atom'
v = self._crt_var(self, 'xa', 'f8', ('na_u', 'xyz'))
v.info = 'Atomic coordinates'
v.unit = 'Bohr'
# Save stuff
self.variables['xa'][:] = geometry.xyz / Bohr2Ang
bs = self._crt_grp(self, 'BASIS')
b = self._crt_var(bs, 'basis', 'i4', ('na_u', ))
b.info = "Basis of each atom by ID"
orbs = _a.emptyi([geometry.na])
for ia, a, isp in geometry.iter_species():
b[ia] = isp + 1
orbs[ia] = a.no
if a.tag in bs.groups:
# Assert the file sizes
if bs.groups[a.tag].Number_of_orbitals != a.no:
raise ValueError(
('File {}'
' has erroneous data in regards of '
'of the alreay stored dimensions.').format(self.file))
else:
ba = bs.createGroup(a.tag)
ba.ID = np.int32(isp + 1)
ba.Atomic_number = np.int32(a.Z)
ba.Mass = a.mass
ba.Label = a.tag
ba.Element = a.symbol
ba.Number_of_orbitals = np.int32(a.no)
# Store the lasto variable as the remaining thing to do
self.variables['lasto'][:] = _a.cumsumi(orbs)
def _index_shape(self, shape):
""" Internal routine for shape-indices """
# First grab the sphere, subsequent indices will be reduced
# by the actual shape
cuboid = shape.toCuboid()
ellipsoid = shape.toEllipsoid()
if ellipsoid.volume() > cuboid.volume():
idx = self._index_shape_cuboid(cuboid)
else:
idx = self._index_shape_ellipsoid(ellipsoid)
# Get min/max
imin = idx.min(0)
imax = idx.max(0)
del idx
dc = self.dcell
# Now to find the actual points inside the shape
# First create all points in the square and then retrieve all indices
# within.
ix = _a.aranged(imin[0], imax[0] + 0.5)
iy = _a.aranged(imin[1], imax[1] + 0.5)
iz = _a.aranged(imin[2], imax[2] + 0.5)
output_shape = (ix.size, iy.size, iz.size, 3)
rxyz = _a.emptyd(output_shape)
ao = add.outer
ao(ao(ix * dc[0, 0], iy * dc[1, 0]), iz * dc[2, 0], out=rxyz[:, :, :, 0])
ao(ao(ix * dc[0, 1], iy * dc[1, 1]), iz * dc[2, 1], out=rxyz[:, :, :, 1])
ao(ao(ix * dc[0, 2], iy * dc[1, 2]), iz * dc[2, 2], out=rxyz[:, :, :, 2])
idx = shape.within_index(rxyz.reshape(-1, 3))
del rxyz
i = _a.emptyi(output_shape)
i[:, :, :, 0] = ix.reshape(-1, 1, 1)
i[:, :, :, 1] = iy.reshape(1, -1, 1)
i[:, :, :, 2] = iz.reshape(1, 1, -1)
del ix, iy, iz
i.shape = (-1, 3)
i = take(i, idx, axis=0)
del idx
return i
def _mulliken(self):
# Calculate the Mulliken elements
# First we re-create the sparse matrix as required for csr_matrix
ptr = self._csr.ptr
ncol = self._csr.ncol
# Indices of non-zero elements
idx = array_arange(ptr[:-1], n=ncol)
# Create the new pointer array
new_ptr = _a.emptyi(len(ptr))
new_ptr[0] = 0
col = self._csr.col[idx]
_a.cumsumi(ncol, out=new_ptr[1:])
# The shape of the matrices
shape = self.shape[:2]
# Create list of charges to be returned
Q = list()
if self.orthogonal:
# We only need the diagonal elements
S = csr_matrix(shape, dtype=self.dtype)
S.setdiag(1.)
for i in range(self.shape[2]):
DM = csr_matrix((self._csr._D[idx, i], col, new_ptr),
shape=shape)
Q.append(DM.multiply(S))
Q[-1].eliminate_zeros()
else:
# We now what S is and do it element-wise.
q = self._csr._D[idx, :-1] * self._csr._D[idx, self.S_idx].reshape(
-1, 1)
for i in range(q.shape[1]):
Q.append(csr_matrix((q[:, i], col, new_ptr), shape=shape))
Q[-1].eliminate_zeros()
return Q
def _geometry_group(geometry, ret_index=False):
r""" Order atoms in geometry according to species such that all of one specie is consecutive
When creating VASP input files (`poscarSileVASP` for instance) the equivalent
``POTCAR`` file needs to contain the pseudos for each specie as they are provided
in blocks.
I.e. for a geometry like this:
.. code::
[Atom(6), Atom(4), Atom(6)]
the resulting ``POTCAR`` needs to contain the pseudo for Carbon twice.
This method will re-order atoms according to the species"
Parameters
----------
geometry : Geometry
geometry to be re-ordered
ret_index : bool, optional
return sorted indices
Returns
-------
geometry: reordered geometry
"""
na = len(geometry)
idx = _a.emptyi(na)
ia = 0
for _, idx_s in geometry.atoms.iter(species=True):
idx[ia:ia + len(idx_s)] = idx_s
ia += len(idx_s)
assert ia == na
if ret_index:
return geometry.sub(idx), idx
return geometry.sub(idx)
def orbital_momentum(self, projection='orbital', method='onsite'):
r""" Calculate orbital angular momentum on either atoms or orbitals
Currently this implementation equals the Siesta implementation in that
the on-site approximation is enforced thus limiting the calculated quantities
to obey the following conditions:
1. Same atom
2. :math:`l>0`
3. :math:`l_\nu \equiv l_\mu`
4. :math:`m_\nu \neq m_\mu`
5. :math:`\zeta_\nu \equiv \zeta_\mu`
This allows one to sum the orbital angular moments on a per atom site.
Parameters
----------
projection : {'orbital', 'atom'}
whether the angular momentum is resolved per atom, or per orbital
method : {'onsite'}
method used to calculate the angular momentum
Returns
-------
numpy.ndarray
orbital angular momentum with the last dimension equalling the :math:`L_x`, :math:`L_y` and :math:`L_z` components
"""
# Check that the spin configuration is correct
if not self.spin.is_spinorbit:
raise ValueError(
f"{self.__class__.__name__}.orbital_momentum requires a spin-orbit matrix"
)
# First we calculate
orb_lmZ = _a.emptyi([self.no, 3])
for atom, idx in self.geometry.atoms.iter(True):
# convert to FIRST orbital index per atom
oidx = self.geometry.a2o(idx)
# loop orbitals
for io, orb in enumerate(atom):
orb_lmZ[oidx + io, :] = orb.l, orb.m, orb.Z
# Now we need to calculate the stuff
DM = self.copy()
# The Siesta convention *only* calculates contributions
# in the primary unit-cell.
DM.set_nsc([1] * 3)
geom = DM.geometry
csr = DM._csr
# The siesta moments are only *on-site* per atom.
# 1. create a logical index for the matrix elements
# that is true for ia-ia interaction and false
# otherwise
idx = repeat(_a.arangei(geom.no), csr.ncol)
aidx = geom.o2a(idx)
# Sparse matrix indices for data
sidx = array_arange(csr.ptr[:-1], n=csr.ncol, dtype=np.int32)
jdx = csr.col[sidx]
ajdx = geom.o2a(jdx)
# Now only take the elements that are *on-site* and which are *not*
# having the same m quantum numbers (if the orbital index is the same
# it means they have the same m quantum number)
#
# 1. on the same atom
# 2. l > 0
# 3. same quantum number l
# 4. different quantum number m
# 5. same zeta
onsite_idx = ((aidx == ajdx) & \
(orb_lmZ[idx, 0] > 0) & \
(orb_lmZ[idx, 0] == orb_lmZ[jdx, 0]) & \
(orb_lmZ[idx, 1] != orb_lmZ[jdx, 1]) & \
(orb_lmZ[idx, 2] == orb_lmZ[jdx, 2])).nonzero()[0]
# clean variables we don't need
del aidx, ajdx
# Now reduce arrays to the orbital connections that obey the
# above criteria
idx = idx[onsite_idx]
idx_l = orb_lmZ[idx, 0]
idx_m = orb_lmZ[idx, 1]
jdx = jdx[onsite_idx]
jdx_m = orb_lmZ[jdx, 1]
sidx = sidx[onsite_idx]
# Sum the spin-box diagonal imaginary parts
DM = csr._D[sidx][:, [4, 5]].sum(1)
# Define functions to calculate L projections
def La(idx_l, DM, sub):
if len(sub) == 0:
return []
return (idx_l[sub] * (idx_l[sub] + 1) * 0.5)**0.5 * DM[sub]
def Lb(idx_l, DM, sub):
if len(sub) == 0:
return
return (idx_l[sub] * (idx_l[sub] + 1) - 2)**0.5 * 0.5 * DM[sub]
def Lc(idx, idx_l, DM, sub):
if len(sub) == 0:
return [], []
sub = sub[idx_l[sub] >= 3]
if len(sub) == 0:
return [], []
return idx[sub], (idx_l[sub] *
(idx_l[sub] + 1) - 6)**0.5 * 0.5 * DM[sub]
# construct for different m
# in Siesta the spin orbital angular momentum
# is calculated by swapping i and j indices.
# This is somewhat confusing to me, so I reversed everything.
# This will probably add to the confusion when comparing the two
# Additionally Siesta calculates L for <i|L|j> and then does:
# L(:) = [L(3), -L(2), -L(1)]
# Here we *directly* store the quantities used.
# Pre-allocate the L_xyz quantity per orbital.
L = np.zeros([geom.no, 3])
L0 = L[:, 0]
L1 = L[:, 1]
L2 = L[:, 2]
# Pre-calculate all those which have m_i + m_j == 0
b = (idx_m + jdx_m == 0).nonzero()[0]
subtract.at(L2, idx[b], idx_m[b] * DM[b])
del b
# mi == 0
i_m = idx_m == 0
# mj == -1
sub = logical_and(i_m, jdx_m == -1).nonzero()[0]
subtract.at(L0, idx[sub], La(idx_l, DM, sub))
# mj == 1
sub = logical_and(i_m, jdx_m == 1).nonzero()[0]
add.at(L1, idx[sub], La(idx_l, DM, sub))
# mi == 1
i_m = idx_m == 1
# mj == -2
sub = logical_and(i_m, jdx_m == -2).nonzero()[0]
subtract.at(L0, idx[sub], Lb(idx_l, DM, sub))
# mj == 0
sub = logical_and(i_m, jdx_m == 0).nonzero()[0]
subtract.at(L1, idx[sub], La(idx_l, DM, sub))
# mj == 2
sub = logical_and(i_m, jdx_m == 2).nonzero()[0]
add.at(L1, idx[sub], Lb(idx_l, DM, sub))
# mi == -1
i_m = idx_m == -1
# mj == -2
sub = logical_and(i_m, jdx_m == -2).nonzero()[0]
add.at(L1, idx[sub], Lb(idx_l, DM, sub))
# mj == 0
sub = logical_and(i_m, jdx_m == 0).nonzero()[0]
add.at(L0, idx[sub], La(idx_l, DM, sub))
# mj == 2
sub = logical_and(i_m, jdx_m == 2).nonzero()[0]
add.at(L0, idx[sub], Lb(idx_l, DM, sub))
# mi == 2
i_m = idx_m == 2
# mj == -3
sub = logical_and(i_m, jdx_m == -3).nonzero()[0]
subtract.at(L0, *Lc(idx, idx_l, DM, sub))
# mj == -1
sub = logical_and(i_m, jdx_m == -1).nonzero()[0]
subtract.at(L0, idx[sub], Lb(idx_l, DM, sub))
# mj == 1
sub = logical_and(i_m, jdx_m == 1).nonzero()[0]
subtract.at(L1, idx[sub], Lb(idx_l, DM, sub))
# mj == 3
sub = logical_and(i_m, jdx_m == 3).nonzero()[0]
add.at(L1, *Lc(idx, idx_l, DM, sub))
# mi == -2
i_m = idx_m == -2
# mj == -3
sub = logical_and(i_m, jdx_m == -3).nonzero()[0]
add.at(L1, *Lc(idx, idx_l, DM, sub))
# mj == -1
sub = logical_and(i_m, jdx_m == -1).nonzero()[0]
subtract.at(L1, idx[sub], Lb(idx_l, DM, sub))
# mj == 1
sub = logical_and(i_m, jdx_m == 1).nonzero()[0]
add.at(L0, idx[sub], Lb(idx_l, DM, sub))
# mj == 3
sub = logical_and(i_m, jdx_m == 3).nonzero()[0]
add.at(L0, *Lc(idx, idx_l, DM, sub))
# mi == -3
i_m = idx_m == -3
# mj == -2
sub = logical_and(i_m, jdx_m == -2).nonzero()[0]
subtract.at(L1, *Lc(idx, idx_l, DM, sub))
# mj == 2
sub = logical_and(i_m, jdx_m == 2).nonzero()[0]
add.at(L0, *Lc(idx, idx_l, DM, sub))
# mi == 3
i_m = idx_m == 3
# mj == -2
sub = logical_and(i_m, jdx_m == -2).nonzero()[0]
subtract.at(L0, *Lc(idx, idx_l, DM, sub))
# mj == 2
sub = logical_and(i_m, jdx_m == 2).nonzero()[0]
subtract.at(L1, *Lc(idx, idx_l, DM, sub))
if "orbital" == projection:
return L
elif "atom" == projection:
# Now perform summation per atom
l = np.zeros([geom.na, 3], dtype=L.dtype)
add.at(l, geom.o2a(np.arange(geom.no)), L)
return l
raise ValueError(
f"{self.__class__.__name__}.orbital_momentum must define projection to be 'orbital' or 'atom'."
)
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)
