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)