def read_energy_density_matrix(self, **kwargs): """ Returns the energy density matrix from the siesta.DM file """ # Now read the sizes used... spin, no, nsc, nnz = _siesta.read_tsde_sizes(self.file) _bin_check(self, 'read_energy_density_matrix', 'could not read energy density matrix sizes.') ncol, col, dEDM = _siesta.read_tsde_edm(self.file, spin, no, nsc, nnz) _bin_check(self, 'read_energy_density_matrix', 'could not read energy density matrix.') # Try and immediately attach a geometry geom = kwargs.get('geometry', kwargs.get('geom', None)) if geom is None: # We truly, have no clue, # Just generate a boxed system xyz = [[x, 0, 0] for x in range(no)] sc = SuperCell([no, 1, 1], nsc=nsc) geom = Geometry(xyz, Atom(1), sc=sc) if nsc[0] != 0 and np.any(geom.nsc != nsc): # We have to update the number of supercells! geom.set_nsc(nsc) if geom.no != no: raise SileError( str(self) + '.read_energy_density_matrix could ' 'not use the passed geometry as the number of atoms or orbitals ' 'is inconsistent with DM file.') # Create the energy density matrix container EDM = EnergyDensityMatrix(geom, spin, nnzpr=1, dtype=np.float64, orthogonal=False) # Create the new sparse matrix EDM._csr.ncol = ncol.astype(np.int32, copy=False) EDM._csr.ptr = np.insert(np.cumsum(ncol, dtype=np.int32), 0, 0) # Correct fortran indices EDM._csr.col = col.astype(np.int32, copy=False) - 1 EDM._csr._nnz = len(col) EDM._csr._D = _a.emptyd([nnz, spin + 1]) EDM._csr._D[:, :spin] = dEDM[:, :] # EDM file does not contain overlap matrix... so neglect it for now. EDM._csr._D[:, spin] = 0. _mat_spin_convert(EDM) # Convert the supercells to sisl supercells if nsc[0] != 0 or geom.no_s >= col.max(): _csr_from_siesta(geom, EDM._csr) else: warn( str(self) + '.read_energy_density_matrix may result in a wrong sparse pattern!' ) return EDM
def _pushfile(self, f): if self.dir_file(f).is_file(): self._parent_fh.append(self.fh) self.fh = self.dir_file(f).open(self._mode) else: warn(str(self) + f' is trying to include file: {f} but the file seems not to exist? Will disregard file!')
def _geometry_align(geom_b, geom_u, cls, method): """ Routine used to align two geometries There are a few twists in this since the fdf-reads will automatically try and pass a geometry from the output files. In cases where the *.ion* files are non-existing this will result in a twist. This routine will select and return a merged Geometry which fulfills the correct number of atoms and orbitals. However, if the input geometries have mis-matching number of atoms a SislError will be raised. Parameters ---------- geom_b : Geometry from binary file geom_u : Geometry supplied by user Raises ------ SislError : if the geometries have non-equal atom count """ if geom_b is None: return geom_u elif geom_u is None: return geom_b # Default to use the users geometry geom = geom_u is_copy = False def get_copy(geom, is_copy): if is_copy: return geom, True return geom.copy(), True if geom_b.na != geom.na: # we have no way of solving this issue... raise SileError( "{cls}.{method} could not use the passed geometry as the " "of atoms is not consistent, user-atoms={u_na}, file-atoms={b_na}." .format(cls=cls.__name__, method=method, b_na=geom_b.na, u_na=geom_u.na)) # Try and figure out what to do if not np.allclose(geom_b.xyz, geom.xyz): warn( f"{cls.__name__}.{method} has mismatched atomic coordinates, will copy geometry and use file XYZ." ) geom, is_copy = get_copy(geom, is_copy) geom.xyz[:, :] = geom_b.xyz[:, :] if not np.allclose(geom_b.sc.cell, geom.sc.cell): warn( f"{cls.__name__}.{method} has non-equal lattice vectors, will copy geometry and use file lattice." ) geom, is_copy = get_copy(geom, is_copy) geom.sc.cell[:, :] = geom_b.sc.cell[:, :] if not np.array_equal(geom_b.nsc, geom.nsc): warn( f"{cls.__name__}.{method} has non-equal number of supercells, will copy geometry and use file supercell count." ) geom, is_copy = get_copy(geom, is_copy) geom.set_nsc(geom_b.nsc) # Now for the difficult part. # If there is a mismatch in the number of orbitals we will # prefer to use the user-supplied atomic species, but fill with # *random* orbitals if not np.array_equal(geom_b.atoms.orbitals, geom.atoms.orbitals): warn( f"{cls.__name__}.{method} has non-equal number of orbitals per atom, will correct with *empty* orbitals." ) geom, is_copy = get_copy(geom, is_copy) # Now create a new atom specie with the correct number of orbitals norbs = geom_b.atoms.orbitals[:] atoms = Atoms([ geom.atoms[i].copy(orbital=[-1] * norbs[i]) for i in range(geom.na) ]) geom._atoms = atoms return geom
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( f"{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() 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( f"{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[:, 0, 1] = csr._D[idx, 2] + 1j * csr._D[idx, 3] DM[:, 1, 0] = np.conj(DM[:, 0, 1]) DM[:, 1, 1] = csr._D[idx, 1] else: # spin-orbit DM[:, 0, 0] = csr._D[idx, 0] + 1j * csr._D[idx, 4] DM[:, 0, 1] = csr._D[idx, 2] + 1j * csr._D[idx, 3] DM[:, 1, 0] = csr._D[idx, 6] + 1j * csr._D[idx, 7] DM[:, 1, 1] = csr._D[idx, 1] + 1j * csr._D[idx, 5] # 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( f"{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], _ncol_to_indptr(csr.ncol)), 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() # Find the periodic directions pbc = [ bc == grid.PERIODIC or geometry.nsc[i] > 1 for i, bc in enumerate(grid.bc[:, 0]) ] if grid.geometry is None: # Create the actual geometry that encompass the grid ia, xyz, _ = geometry.within_inf(sc, periodic=pbc) if len(ia) > 0: grid.set_geometry(Geometry(xyz, geometry.atoms[ia], sc=sc)) # Instead of looping all atoms in the supercell we find the exact atoms # and their supercell indices. add_R = _a.fulld(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 + np.diag(2 * add_R), origo=o.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, periodic=pbc) XYZ -= grid.sc.origo.reshape(1, 3) # Retrieve progressbar eta = tqdm_eta(len(IA), f"{self.__class__.__name__}.density", "atom", eta) cell = geometry.cell atoms = geometry.atoms 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 = _ncol_to_indptr(csr.ncol) 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, ISC): # Get current atom ia_atom = atoms[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( f"Atom '{ia_atom}' does not have a wave-function, skipping 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 = atoms[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.orbitals[jo] oR = o.R # Downsize to the correct indices if jR - oR < 1e-6: ja_idx1 = ja_idx ja_r1 = ja_r ja_theta1 = ja_theta ja_cos_phi1 = ja_cos_phi 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.orbitals[jo] ojR = oj.R # Downsize to the correct indices if R - ojR < 1e-6: ja_idx1 = slice(None) ja_r1 = ia_r ja_theta1 = ia_theta ja_cos_phi1 = ia_cos_phi 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.orbitals[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 _r_geometry_fdf(self, *args, **kwargs): """ Returns Geometry object from the FDF file NOTE: Interaction range of the Atoms are currently not read. """ sc = self.read_supercell(order=['fdf']) # No fractional coordinates is_frac = False # Read atom scaling lc = self.get('AtomicCoordinatesFormat', default='Bohr').lower() if 'ang' in lc or 'notscaledcartesianang' in lc: s = 1. elif 'bohr' in lc or 'notscaledcartesianbohr' in lc: s = Bohr2Ang elif 'scaledcartesian' in lc: # the same scaling as the lattice-vectors s = self.get('LatticeConstant', 'Ang') elif 'fractional' in lc or 'scaledbylatticevectors' in lc: # no scaling of coordinates as that is entirely # done by the latticevectors s = 1. is_frac = True # If the user requests a shifted geometry # we correct for this origo = np.zeros([3], np.float64) lor = self.get('AtomicCoordinatesOrigin') if lor: if kwargs.get('origin', True): origo = _a.asarrayd(map(float, lor[0].split()[:3])) * s # Origo cannot be interpreted with fractional coordinates # hence, it is not transformed. # Read atom block atms = self.get('AtomicCoordinatesAndAtomicSpecies') if atms is None: raise SileError( 'AtomicCoordinatesAndAtomicSpecies block could not be found') # Read number of atoms and block # We default to the number of elements in the # AtomicCoordinatesAndAtomicSpecies block na = self.get('NumberOfAtoms', default=len(atms)) # Reduce space if number of atoms specified if na < len(atms): # align number of atoms and atms array atms = atms[:na] elif na > len(atms): raise SileError( 'NumberOfAtoms is larger than the atoms defined in the blocks') elif na == 0: raise SileError( 'NumberOfAtoms has been determined to be zero, no atoms.') # Create array xyz = np.empty([na, 3], np.float64) species = np.empty([na], np.int32) for ia in range(na): l = atms[ia].split() xyz[ia, :] = [float(k) for k in l[:3]] species[ia] = int(l[3]) - 1 if is_frac: xyz = np.dot(xyz, sc.cell) xyz *= s xyz += origo # Read the block (not strictly needed, if so we simply set all atoms to H) atom = self.read_basis() if atom is None: warn( SileWarning( 'Block ChemicalSpeciesLabel does not exist, cannot determine the basis (all Hydrogen).' )) # Default atom (hydrogen) atom = Atom(1) else: atom = [atom[i] for i in species] # Create and return geometry object return Geometry(xyz, atom=atom, sc=sc)
def write_hamiltonian(self, ham, hermitian=True, **kwargs): """ Writes the Hamiltonian model to the file Writes a Hamiltonian model to the intrinsic Hamiltonian file format. The file can be constructed by the implict force of Hermiticity, or without. Utilizing the Hermiticity we reduce the file-size by approximately 50%. Parameters ---------- ham : `Hamiltonian` model hermitian : boolean=True whether the stored data is halved using the Hermitian property """ # We use the upper-triangular form of the Hamiltonian # and the overlap matrix for hermitian problems geom = ham.geometry # First write the geometry self.write_geometry(geom, **kwargs) # We default to the advanced layuot if we have more than one # orbital on any one atom advanced = kwargs.get( 'advanced', np.any(np.array([a.no for a in geom.atom.atom], np.int32) > 1)) fmt = kwargs.get('fmt', 'g') if advanced: fmt1_str = ' {{0:d}}[{{1:d}}] {{2:d}}[{{3:d}}] {{4:{0}}}\n'.format( fmt) fmt2_str = ' {{0:d}}[{{1:d}}] {{2:d}}[{{3:d}}] {{4:{0}}} {{5:{0}}}\n'.format( fmt) else: fmt1_str = ' {{0:d}} {{1:d}} {{2:{0}}}\n'.format(fmt) fmt2_str = ' {{0:d}} {{1:d}} {{2:{0}}} {{3:{0}}}\n'.format(fmt) # We currently force the model to be finalized # before we can write it # This should be easily circumvented H = ham.tocsr(0) if not ham.orthogonal: S = ham.tocsr(ham.S_idx) # If the model is Hermitian we can # do with writing out half the entries if hermitian: herm_acc = kwargs.get('herm_acc', 1e-6) # We check whether it is Hermitian (not S) for i, isc in enumerate(geom.sc.sc_off): oi = i * geom.no oj = geom.sc_index(-isc) * geom.no # get the difference between the ^\dagger elements diff = H[:, oi:oi + geom.no] - \ H[:, oj:oj + geom.no].transpose() diff.eliminate_zeros() if np.any(np.abs(diff.data) > herm_acc): amax = np.amax(np.abs(diff.data)) warn( SileWarning( 'The model could not be asserted to be Hermitian ' 'within the accuracy required ({0}).'.format( amax))) hermitian = False del diff if hermitian: # Remove all double stuff for i, isc in enumerate(geom.sc.sc_off): if np.any(isc < 0): # We have ^\dagger element, remove it o = i * geom.no # Ensure that we remove all nullified quantities # (setting elements to zero will add them internally # :(, hence this actually constructs the full matrix # Therefore we do it on a row basis, to limit memory # requirements for j in range(geom.no): H[j, o:o + geom.no] = 0. H.eliminate_zeros() if not ham.orthogonal: S[j, o:o + geom.no] = 0. S.eliminate_zeros() o = geom.sc_index(np.zeros([3], np.int32)) # Get upper-triangular matrix of the unit-cell H and S ut = triu(H[:, o:o + geom.no], k=0).tocsr() for j in range(geom.no): H[j, o:o + geom.no] = 0. H[j, o:o + geom.no] = ut[j, :] H.eliminate_zeros() if not ham.orthogonal: ut = triu(S[:, o:o + geom.no], k=0).tocsr() for j in range(geom.no): S[j, o:o + geom.no] = 0. S[j, o:o + geom.no] = ut[j, :] S.eliminate_zeros() # Ensure that S and H have the same sparsity pattern for jo, io in ispmatrix(S): H[jo, io] = H[jo, io] del ut # Start writing of the model # We loop on all super-cells for i, isc in enumerate(geom.sc.sc_off): # Check that we have any contributions in this # sub-section Hsub = H[:, i * geom.no:(i + 1) * geom.no] if not ham.orthogonal: Ssub = S[:, i * geom.no:(i + 1) * geom.no] if Hsub.getnnz() == 0: continue # We have a contribution, write out the information self._write('\nbegin matrix {0:d} {1:d} {2:d}\n'.format(*isc)) if advanced: for jo, io, h in ispmatrixd(Hsub): o = np.array([jo, io], np.int32) a = geom.o2a(o) o = o - geom.a2o(a) if not ham.orthogonal: s = Ssub[jo, io] elif jo == io: s = 1. else: s = 0. if s == 0.: self._write(fmt1_str.format(a[0], o[0], a[1], o[1], h)) else: self._write( fmt2_str.format(a[0], o[0], a[1], o[1], h, s)) else: for jo, io, h in ispmatrixd(Hsub): if not ham.orthogonal: s = Ssub[jo, io] elif jo == io: s = 1. else: s = 0. if s == 0.: self._write(fmt1_str.format(jo, io, h)) else: self._write(fmt2_str.format(jo, io, h, s)) self._write('end matrix {0:d} {1:d} {2:d}\n'.format(*isc))
def read_data(self, as_dataarray=False): r""" Returns data associated with the PDOS file For spin-polarized calculations the returned values are up/down, orbitals, energy. For non-collinear calculations the returned values are sum/x/y/z, orbitals, energy. Parameters ---------- as_dataarray: bool, optional If True the returned PDOS is a `xarray.DataArray` with energy, spin and orbital information as coordinates in the data. The geometry, unit and Fermi level are stored as attributes in the DataArray. Returns ------- geom : Geometry instance with positions, atoms and orbitals. E : the energies at which the PDOS has been evaluated at (if Fermi-level present in file energies are shifted to :math:`E - E_F = 0`). PDOS : an array of DOS with dimensions ``(nspin, geom.no, len(E))`` (with different spin-components) or ``(geom.no, len(E))`` (spin-symmetric). DataArray : if `as_dataarray` is True, only this data array is returned, in this case all data can be post-processed using the `xarray` selection routines. """ # Get the element-tree root = xml_parse(self.file).getroot() # Get number of orbitals nspin = int(root.find('nspin').text) # Try and find the fermi-level Ef = root.find('fermi_energy') E = arrayd(list(map(float, root.find('energy_values').text.split()))) if Ef is None: warn( str(self) + '.read_data could not locate the Fermi-level in the XML tree, using E_F = 0. eV' ) else: Ef = float(Ef.text) E -= Ef ne = len(E) # All coordinate, atoms and species data xyz = [] atoms = [] atom_species = [] def ensure_size(ia): while len(atom_species) <= ia: atom_species.append(None) xyz.append(None) def ensure_size_orb(ia, i): while len(atoms) <= ia: atoms.append([]) while len(atoms[ia]) <= i: atoms[ia].append(None) if nspin == 4: def process(D): tmp = np.empty(D.shape[0], D.dtype) tmp[:] = D[:, 3] D[:, 3] = D[:, 0] - D[:, 1] D[:, 0] = D[:, 0] + D[:, 1] D[:, 1] = D[:, 2] D[:, 2] = tmp[:] return D else: def process(D): return D if as_dataarray: import xarray as xr if nspin == 1: spin = ['sum'] elif nspin == 2: spin = ['up', 'down'] elif nspin == 4: spin = ['sum', 'x', 'y' 'z'] # Dimensions of the PDOS data-array dims = ['E', 'spin', 'n', 'l', 'm', 'zeta', 'polarization'] shape = (ne, nspin, 1, 1, 1, 1, 1) def to(o, DOS): # Coordinates for this dataarray coords = [E, spin, [o.n], [o.l], [o.m], [o.Z], [o.P]] return xr.DataArray(data=process(DOS).reshape(shape), dims=dims, coords=coords, name='PDOS') D = xr.DataArray([]) else: def to(o, DOS): return process(DOS) D = [] for orb in root.findall('orbital'): # Short-hand function to retrieve integers for the attributes def oi(name): return int(orb.get(name)) # Get indices ia = oi('atom_index') - 1 i = oi('index') - 1 species = orb.get('species') # Create the atomic orbital try: Z = oi('Z') except: try: Z = PeriodicTable().Z(species) except: # Unknown Z = -1 try: P = orb.get('P') == 'true' except: P = False ensure_size(ia) xyz[ia] = list(map(float, orb.get('position').split())) atom_species[ia] = Z # Construct the atomic orbital O = AtomicOrbital(n=oi('n'), l=oi('l'), m=oi('m'), Z=oi('z'), P=P) # We know that the index is far too high. However, # this ensures a consecutive orbital ensure_size_orb(ia, i) atoms[ia][i] = O # it is formed like : spin-1, spin-2 (however already in eV) DOS = arrayd(list(map(float, orb.find('data').text.split()))).reshape( -1, nspin) if as_dataarray: D = D.combine_first(to(O, DOS)) else: D.append(process(DOS)) # Now we need to parse the data # First reduce the atom atoms = [[o for o in a if o] for a in atoms] atoms = Atoms([Atom(Z, os) for Z, os in zip(atom_species, atoms)]) geom = Geometry(arrayd(xyz) * Bohr2Ang, atoms) if as_dataarray: # Add attributes D.attrs['geometry'] = geom D.attrs['units'] = '1/eV' if Ef is None: D.attrs['Ef'] = 'Unknown' else: D.attrs['Ef'] = Ef return D D = np.moveaxis(np.stack(D, axis=0), 2, 0) if nspin == 1: return geom, E, D[0] return geom, E, D
def test_warn_category(): with pytest.warns(sm.SislWarning): sm.warn('Warning', sm.SislWarning)
def write_delta(self, delta, **kwargs): r""" Writes a :math:`\delta` Hamiltonian to the file This term may be of - level-1: no E or k dependence - level-2: k-dependent - level-3: E-dependent - level-4: k- and E-dependent Parameters ---------- delta : SparseOrbitalBZSpin the model to be saved in the NC file k : array_like, optional a specific k-point :math:`\delta` term. I.e. only save the :math:`\delta` term for the given k-point. May be combined with `E` for a specific k and energy point. E : float, optional an energy dependent :math:`\delta` term. I.e. only save the :math:`\delta` term for the given energy. May be combined with `k` for a specific k and energy point. """ csr = delta._csr.copy() if csr.nnz == 0: raise SileError(f"{self!s}.write_overlap cannot write a zero element sparse matrix!") # convert to siesta thing and store _csr_to_siesta(delta.geometry, csr) # delta should always write sorted matrices csr.finalize(sort=True) _mat_spin_convert(csr, delta.spin) # Ensure that the geometry is written self.write_geometry(delta.geometry) self._crt_dim(self, 'spin', len(delta.spin)) # Determine the type of delta we are storing... k = kwargs.get('k', None) E = kwargs.get('E', None) ilvl, ik, iE = self._get_lvl_k_E(**kwargs) lvl = self._add_lvl(ilvl) # Append the sparsity pattern # Create basis group if 'n_col' in lvl.variables: if len(lvl.dimensions['nnzs']) != csr.nnz: raise ValueError("The sparsity pattern stored in delta *MUST* be equivalent for " "all delta entries [nnz].") if np.any(lvl.variables['n_col'][:] != csr.ncol[:]): raise ValueError("The sparsity pattern stored in delta *MUST* be equivalent for " "all delta entries [n_col].") if np.any(lvl.variables['list_col'][:] != csr.col[:]+1): raise ValueError("The sparsity pattern stored in delta *MUST* be equivalent for " "all delta entries [list_col].") if np.any(lvl.variables['isc_off'][:] != siesta_sc_off(*delta.geometry.sc.nsc).T): raise ValueError("The sparsity pattern stored in delta *MUST* be equivalent for " "all delta entries [sc_off].") else: self._crt_dim(lvl, 'nnzs', csr.nnz) v = self._crt_var(lvl, 'n_col', 'i4', ('no_u',)) v.info = "Number of non-zero elements per row" v[:] = csr.ncol[:] v = self._crt_var(lvl, 'list_col', 'i4', ('nnzs',), chunksizes=(csr.nnz,), **self._cmp_args) v.info = "Supercell column indices in the sparse format" v[:] = csr.col[:] + 1 # correct for fortran indices v = self._crt_var(lvl, 'isc_off', 'i4', ('n_s', 'xyz')) v.info = "Index of supercell coordinates" v[:] = siesta_sc_off(*delta.geometry.sc.nsc).T warn_E = True if ilvl in [3, 4]: if iE < 0: # We need to add the new value iE = lvl.variables['E'].shape[0] lvl.variables['E'][iE] = E * eV2Ry warn_E = False warn_k = True if ilvl in [2, 4]: if ik < 0: ik = lvl.variables['kpt'].shape[0] lvl.variables['kpt'][ik, :] = k warn_k = False if ilvl == 4 and warn_k and warn_E and False: # As soon as we have put the second k-point and the first energy # point, this warning will proceed... # I.e. even though the variable has not been set, it will WARN # Hence we out-comment this for now... #warn(f"Overwriting k-point {ik} and energy point {iE} correction.") pass elif ilvl == 3 and warn_E: warn(f"Overwriting energy point {iE} correction.") elif ilvl == 2 and warn_k: warn(f"Overwriting k-point {ik} correction.") if ilvl == 1: dim = ('spin', 'nnzs') sl = [slice(None)] * 2 csize = [1] * 2 elif ilvl == 2: dim = ('nkpt', 'spin', 'nnzs') sl = [slice(None)] * 3 sl[0] = ik csize = [1] * 3 elif ilvl == 3: dim = ('ne', 'spin', 'nnzs') sl = [slice(None)] * 3 sl[0] = iE csize = [1] * 3 elif ilvl == 4: dim = ('nkpt', 'ne', 'spin', 'nnzs') sl = [slice(None)] * 4 sl[0] = ik sl[1] = iE csize = [1] * 4 # Number of non-zero elements csize[-1] = csr.nnz if delta.spin.kind > delta.spin.POLARIZED: print(delta.spin) raise ValueError(f"{self.__class__.__name__}.write_delta only allows spin-polarized delta values") if delta.dtype.kind == 'c': v1 = self._crt_var(lvl, 'Redelta', 'f8', dim, chunksizes=csize, attrs={'info': "Real part of delta", 'unit': "Ry"}, **self._cmp_args) v2 = self._crt_var(lvl, 'Imdelta', 'f8', dim, chunksizes=csize, attrs={'info': "Imaginary part of delta", 'unit': "Ry"}, **self._cmp_args) for i in range(len(delta.spin)): sl[-2] = i v1[sl] = csr._D[:, i].real * eV2Ry v2[sl] = csr._D[:, i].imag * eV2Ry else: v = self._crt_var(lvl, 'delta', 'f8', dim, chunksizes=csize, attrs={'info': "delta", 'unit': "Ry"}, **self._cmp_args) for i in range(len(delta.spin)): sl[-2] = i v[sl] = csr._D[:, i] * eV2Ry
def test_warn_method(): with pytest.warns(sm.SislWarning): sm.warn('Warning')
def test_warn_specific(): with pytest.warns(sm.SislWarning): sm.warn(sm.SislWarning('Warning'))
def write_hamiltonian(self, H, **kwargs): """ Writes Hamiltonian model to file Parameters ---------- H : Hamiltonian the model to be saved in the NC file spin : int, optional the spin-index of the Hamiltonian object that is stored. Default is the first index. """ # Ensure finalization H.finalize() # Ensure that the geometry is written self.write_geometry(H.geom) self._crt_dim(self, 'spin', len(H.spin)) # Determine the type of dH we are storing... k = kwargs.get('k', None) E = kwargs.get('E', None) ilvl, ik, iE = self._get_lvl_k_E(**kwargs) lvl = self._add_lvl(ilvl) # Append the sparsity pattern # Create basis group if 'n_col' in lvl.variables: if len(lvl.dimensions['nnzs']) != H.nnz: raise ValueError("The sparsity pattern stored in dH *MUST* be equivalent for " "all dH entries [nnz].") if np.any(lvl.variables['n_col'][:] != H._csr.ncol[:]): raise ValueError("The sparsity pattern stored in dH *MUST* be equivalent for " "all dH entries [n_col].") if np.any(lvl.variables['list_col'][:] != H._csr.col[:]+1): raise ValueError("The sparsity pattern stored in dH *MUST* be equivalent for " "all dH entries [list_col].") if np.any(lvl.variables['isc_off'][:] != H.geom.sc.sc_off): raise ValueError("The sparsity pattern stored in dH *MUST* be equivalent for " "all dH entries [sc_off].") else: self._crt_dim(lvl, 'nnzs', H._csr.col.shape[0]) v = self._crt_var(lvl, 'n_col', 'i4', ('no_u',)) v.info = "Number of non-zero elements per row" v[:] = H._csr.ncol[:] v = self._crt_var(lvl, 'list_col', 'i4', ('nnzs',), chunksizes=(len(H._csr.col),), **self._cmp_args) v.info = "Supercell column indices in the sparse format" v[:] = H._csr.col[:] + 1 # correct for fortran indices v = self._crt_var(lvl, 'isc_off', 'i4', ('n_s', 'xyz')) v.info = "Index of supercell coordinates" v[:] = H.geom.sc.sc_off[:, :] warn_E = True if ilvl in [3, 4]: if iE < 0: # We need to add the new value iE = len(lvl.variables['E']) lvl.variables['E'][iE] = E * eV2Ry warn_E = False warn_k = True if ilvl in [2, 4]: if ik < 0: ik = len(lvl.variables['kpt']) lvl.variables['kpt'][ik, :] = k warn_k = False if ilvl == 4 and warn_k and warn_E and False: # As soon as we have put the second k-point and the first energy # point, this warning will proceed... # I.e. even though the variable has not been set, it will WARN # Hence we out-comment this for now... warn(SileWarning('Overwriting k-point {0} and energy point {1} correction.'.format(ik, iE))) elif ilvl == 3 and warn_E: warn(SileWarning('Overwriting energy point {0} correction.'.format(iE))) elif ilvl == 2 and warn_k: warn(SileWarning('Overwriting k-point {0} correction.'.format(ik))) if ilvl == 1: dim = ('spin', 'nnzs') sl = [slice(None)] * 2 csize = [1] * 2 elif ilvl == 2: dim = ('nkpt', 'spin', 'nnzs') sl = [slice(None)] * 3 sl[0] = ik csize = [1] * 3 elif ilvl == 3: dim = ('ne', 'spin', 'nnzs') sl = [slice(None)] * 3 sl[0] = iE csize = [1] * 3 elif ilvl == 4: dim = ('nkpt', 'ne', 'spin', 'nnzs') sl = [slice(None)] * 4 sl[0] = ik sl[1] = iE csize = [1] * 4 # Number of non-zero elements csize[-1] = H.nnz if H.dtype.kind == 'c': v1 = self._crt_var(lvl, 'RedH', 'f8', dim, chunksizes=csize, attr = {'info': "Real part of dH", 'unit': "Ry"}, **self._cmp_args) for i in range(len(H.spin)): sl[-2] = i v1[sl] = H._csr._D[:, i].real * eV2Ry v2 = self._crt_var(lvl, 'ImdH', 'f8', dim, chunksizes=csize, attr = {'info': "Imaginary part of dH", 'unit': "Ry"}, **self._cmp_args) for i in range(len(H.spin)): sl[-2] = i v2[sl] = H._csr._D[:, i].imag * eV2Ry else: v = self._crt_var(lvl, 'dH', 'f8', dim, chunksizes=csize, attr = {'info': "dH", 'unit': "Ry"}, **self._cmp_args) for i in range(len(H.spin)): sl[-2] = i v[sl] = H._csr._D[:, i] * eV2Ry
def inner(self, ket=None, matrix=None, diag=True): r""" Calculate the inner product as :math:`\mathbf A_{ij} = \langle\psi_i|\mathbf M|\psi'_j\rangle` Parameters ---------- ket : State, optional the ket object to calculate the inner product with, if not passed it will do the inner product with itself. The object itself will always be the bra :math:`\langle\psi_i|` matrix : array_like, optional whether a matrix is sandwiched between the bra and ket, default to the identity matrix diag : bool, optional only return the diagonal matrix :math:`\mathbf A_{ii}`. Notes ----- This does *not* take into account a possible overlap matrix when non-orthogonal basis sets are used. Raises ------ ValueError if the number of state coefficients are different for the bra and ket Returns ------- numpy.ndarray a matrix with the sum of inner state products """ if matrix is None: M = _FakeMatrix(self.shape[-1]) else: M = matrix ndim = M.ndim bra = self.state # decide on the ket if ket is None: ket = self.state elif isinstance(ket, State): # check whether this, and ket are both originating from # non-orthogonal basis. That would be non-ideal ket = ket.state if len(ket.shape) == 1: ket.shape = (1, -1) # They *must* have same number of basis points per state if self.shape[-1] != ket.shape[-1]: raise ValueError(f"{self.__class__.__name__}.inner requires the objects to have the same number of coefficients per vector {self.shape[-1]} != {ket.shape[-1]}") if diag: if len(bra) != len(ket): warn(f"{self.__class__.__name__}.inner matrix product is non-square, only the first {min(len(bra), len(ket))} diagonal elements will be returned.") if len(bra) < len(ket): ket = ket[:len(bra)] else: bra = bra[:len(ket)] if ndim == 2: Aij = einsum('ij,ji->i', _conj(bra), M.dot(ket.T)) elif ndim == 1: Aij = einsum('ij,j,ij->i', _conj(bra), M, ket) elif ndim == 2: Aij = _conj(bra) @ M.dot(ket.T) elif ndim == 1: Aij = einsum('ij,j,kj->ik', _conj(bra), M, ket) return Aij
def create_construct(self, R, param): r""" Create a simple function for passing to the `construct` function. This is to relieve the creation of simplistic functions needed for setting up sparse elements. For simple matrices this returns a function: >>> def func(self, ia, atoms, atoms_xyz=None): ... idx = self.geometry.close(ia, R=R, atoms=atoms, atoms_xyz=atoms_xyz) ... for ix, p in zip(idx, param): ... self[ia, ix] = p In the non-colinear case the matrix element :math:`M_{ij}` will be set to input values `param` if :math:`i \le j` and the Hermitian conjugated values for :math:`j < i`. Notes ----- This function only works for geometry sparse matrices (i.e. one element per atom). If you have more than one element per atom you have to implement the function your-self. This method issues warnings if the on-site terms are not Hermitian for spin-orbit systems. Do note that it *still* creates the matrices based on the input. Parameters ---------- R : array_like radii parameters for different shells. Must have same length as `param` or one less. If one less it will be extended with ``R[0]/100`` param : array_like coupling constants corresponding to the `R` ranges. ``param[0,:]`` are the elements for the all atoms within ``R[0]`` of each atom. See Also -------- construct : routine to create the sparse matrix from a generic function (as returned from `create_construct`) """ if len(R) != len(param): raise ValueError( f"{self.__class__.__name__}.create_construct got different lengths of `R` and `param`" ) if self.spin.has_noncolinear: is_complex = self.dkind == 'c' if self.spin.is_spinorbit: if is_complex: nv = 4 # Hermitian parameters paramH = [[ p[0].conj(), p[1].conj(), p[3].conj(), p[2].conj(), *p[4:] ] for p in param] else: nv = 8 # Hermitian parameters paramH = [[ p[0], p[1], p[6], -p[7], -p[4], -p[5], p[2], -p[3], *p[8:] ] for p in param] if not self.orthogonal: nv += 1 # ensure we have correct number of values assert all(len(p) == nv for p in param) if R[0] <= 0.1001: # no atom closer than 0.1001 Ang! # We check that the the parameters here is Hermitian p = param[0] if is_complex: onsite = np.array([[p[0], p[2]], [p[3], p[1]]], self.dtype) else: onsite = np.array( [[p[0] + 1j * p[4], p[2] + 1j * p[3]], [p[6] + 1j * p[7], p[1] + 1j * p[5]]], np.complex128) if not np.allclose(onsite, onsite.T.conj()): warn( f"{self.__class__.__name__}.create_construct is NOT Hermitian for on-site terms. This is your responsibility!" ) elif self.spin.is_noncolinear: if is_complex: nv = 3 # Hermitian parameters paramH = [[p[0].conj(), p[1].conj(), p[2], *p[3:]] for p in param] else: nv = 4 # Hermitian parameters # Note that we don't need to do anything here. # H_ij = [[0, 2 + 1j 3], # [2 - 1j 3, 1]] # H_ji = [[0, 2 + 1j 3], # [2 - 1j 3, 1]] # H_ij^H == H_ji^H paramH = param if not self.orthogonal: nv += 1 # we don't need to check hermiticity for NC # Since the values are ensured Hermitian in the on-site case anyways. # ensure we have correct number of values assert all(len(p) == nv for p in param) na = self.geometry.na # Now create the function that returns the assignment function def func(self, ia, atoms, atoms_xyz=None): idx = self.geometry.close(ia, R=R, atoms=atoms, atoms_xyz=atoms_xyz) for ix, p, pc in zip(idx, param, paramH): ix_ge = (ix % na) >= ia self[ia, ix[ix_ge]] = p self[ia, ix[~ix_ge]] = pc return func return super().create_construct(R, param)
def read_geometry(self, ret_dynamic=False): r""" Returns Geometry object from the CONTCAR/POSCAR file Possibly also return the dynamics (if present). Parameters ---------- ret_dynamic : bool, optional also return selective dynamics (if present), if not, None will be returned. """ sc = self.read_supercell() # The species labels are not always included in *CAR line1 = self.readline().split() opt = self.readline().split() try: species = line1 species_count = np.array(opt, np.int32) except: species_count = np.array(line1, np.int32) # We have no species... # We default to consecutive elements in the # periodic table. species = [i + 1 for i in range(len(species_count))] err = '\n'.join([ "POSCAR best format:", " <Specie-1> <Specie-2>", " <#Specie-1> <#Specie-2>", "Format not found, the species are defaulted to the first elements of the periodic table." ]) warn(err) # Create list of atoms to be used subsequently atom = [ Atom[spec] for spec, nsp in zip(species, species_count) for i in range(nsp) ] # Number of atoms na = len(atom) # check whether this is Selective Dynamics opt = self.readline() if opt[0] in 'Ss': dynamics = True # pre-create the dynamic list dynamic = np.empty([na, 3], dtype=np.bool_) opt = self.readline() else: dynamics = False dynamic = None # Check whether this is in fractional or direct # coordinates (Direct == fractional) cart = False if opt[0] in 'CcKk': cart = True xyz = _a.emptyd([na, 3]) for ia in range(na): line = self.readline().split() xyz[ia, :] = list(map(float, line[:3])) if dynamics: dynamic[ia] = list(map(lambda x: x.lower() == 't', line[3:6])) if cart: # The unit of the coordinates are cartesian xyz *= self._scale else: xyz = xyz.dot(sc.cell) # The POT/CONT-CAR does not contain information on the atomic species geom = Geometry(xyz=xyz, atom=atom, sc=sc) if ret_dynamic: return geom, dynamic return geom
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)