TB2J 0.9.12.13__py3-none-any.whl → 0.9.12.15__py3-none-any.whl

TB2J/exchangeCL2.py

@@ -4,6 +4,7 @@ Exchange from Green's function
 
 import os
 from collections import defaultdict
+from itertools import product
 
 import numpy as np
 from tqdm import tqdm
@@ -76,8 +77,8 @@ class ExchangeCL2(ExchangeCL):
         del self.Gdn.tbmodel
 
     def _adjust_emin(self):
-        emin_up = self.Gup.find_energy_ingap(rbound=self.efermi - 5.0) - self.efermi
-        emin_dn = self.Gdn.find_energy_ingap(rbound=self.efermi - 5.0) - self.efermi
+        emin_up = self.Gup.adjusted_emin
+        emin_dn = self.Gdn.adjusted_emin
         self.emin = min(emin_up, emin_dn)
         print(f"A gap is found at {self.emin}, set emin to it.")
 
@@ -108,11 +109,15 @@ class ExchangeCL2(ExchangeCL):
         Rij_done = set()
         Jorb_list = dict()
         JJ_list = dict()
-        for R, ijpairs in self.R_ijatom_dict.items():
-            if (iatom, jatom) in ijpairs and (R, iatom, jatom) not in Rij_done:
-                Gij_up = self.GR_atom(Gup[R], iatom, jatom)
-                Rm = tuple(-x for x in R)
-                Gji_dn = self.GR_atom(Gdn[Rm], jatom, iatom)
+        for iR, ijpairs in self.R_ijatom_dict.items():
+            if (iatom, jatom) in ijpairs and (iR, iatom, jatom) not in Rij_done:
+                Gij_up = self.GR_atom(Gup[iR], iatom, jatom)
+                iRm = self.R_negative_index[iR]
+                if iRm is None:
+                    R_vec = self.R_vectors[iR]
+                    Rm_vec = tuple(-x for x in R_vec)
+                    raise KeyError(f"Negative R vector {Rm_vec} not found in R_vectors")
+                Gji_dn = self.GR_atom(Gdn[iRm], jatom, iatom)
                 tmp = 0.0j
                 Deltai = self.get_Delta(iatom)
                 Deltaj = self.get_Delta(jatom)
@@ -132,13 +137,13 @@ class ExchangeCL2(ExchangeCL):
                 #        np.matmul(Deltaj, Gji_down),
                 #    )
                 tmp = np.sum(t)
-                self.Jorb_list[(R, iatom, jatom)].append(t / (4.0 * np.pi))
-                self.JJ_list[(R, iatom, jatom)].append(tmp / (4.0 * np.pi))
-                Rij_done.add((R, iatom, jatom))
-                if (Rm, jatom, iatom) not in Rij_done:
-                    Jorb_list[(Rm, jatom, iatom)] = t.T / (4.0 * np.pi)
-                    JJ_list[(Rm, jatom, iatom)] = tmp / (4.0 * np.pi)
-                    Rij_done.add((Rm, jatom, iatom))
+                self.Jorb_list[(iR, iatom, jatom)].append(t / (4.0 * np.pi))
+                self.JJ_list[(iR, iatom, jatom)].append(tmp / (4.0 * np.pi))
+                Rij_done.add((iR, iatom, jatom))
+                if (iRm, jatom, iatom) not in Rij_done:
+                    Jorb_list[(iRm, jatom, iatom)] = t.T / (4.0 * np.pi)
+                    JJ_list[(iRm, jatom, iatom)] = tmp / (4.0 * np.pi)
+                    Rij_done.add((iRm, jatom, iatom))
         return Jorb_list, JJ_list
 
     def get_all_A(self, Gup, Gdn):
@@ -151,16 +156,23 @@ class ExchangeCL2(ExchangeCL):
         Rij_done = set()
         Jorb_list = dict()
         JJ_list = dict()
-        for R, ijpairs in self.R_ijatom_dict.items():
-            for iatom, jatom in ijpairs:
-                if (R, iatom, jatom) not in Rij_done:
-                    Rm = tuple(-x for x in R)
-                    if (Rm, jatom, iatom) in Rij_done:
+        for iR in self.R_ijatom_dict:
+            for iatom, jatom in self.R_ijatom_dict[iR]:
+                if (iR, iatom, jatom) not in Rij_done:
+                    iRm = self.R_negative_index[iR]
+                    if iRm is None:
+                        R_vec = self.short_Rlist[iR]
+                        Rm_vec = tuple(-x for x in R_vec)
                         raise KeyError(
-                            f"Strange (Rm, jatom, iatom) has already been calculated! {(Rm, jatom, iatom)}"
+                            f"Negative R vector {Rm_vec} not found in short_Rlist"
                         )
-                    Gij_up = self.GR_atom(Gup[R], iatom, jatom)
-                    Gji_dn = self.GR_atom(Gdn[Rm], jatom, iatom)
+
+                    if (iRm, jatom, iatom) in Rij_done:
+                        raise KeyError(
+                            f"Strange (iRm, jatom, iatom) has already been calculated! {(iRm, jatom, iatom)}"
+                        )
+                    Gij_up = self.GR_atom(Gup[iR], iatom, jatom)
+                    Gji_dn = self.GR_atom(Gdn[iRm], jatom, iatom)
                     tmp = 0.0j
                     # t = self.get_Delta(iatom) @ Gij_up @ self.get_Delta(jatom) @ Gji_dn
                     t = np.einsum(
@@ -169,15 +181,78 @@ class ExchangeCL2(ExchangeCL):
                         np.matmul(self.get_Delta(jatom), Gji_dn),
                     )
                     tmp = np.sum(t)
-                    Jorb_list[(R, iatom, jatom)] = t / (4.0 * np.pi)
-                    JJ_list[(R, iatom, jatom)] = tmp / (4.0 * np.pi)
-                    Rij_done.add((R, iatom, jatom))
-                    if (Rm, jatom, iatom) not in Rij_done:
-                        Jorb_list[(Rm, jatom, iatom)] = t.T / (4.0 * np.pi)
-                        JJ_list[(Rm, jatom, iatom)] = tmp / (4.0 * np.pi)
-                        Rij_done.add((Rm, jatom, iatom))
+                    # Use R vectors for keys for I/O compatibility
+                    R_vec = self.short_Rlist[iR]
+                    Rm_vec = self.short_Rlist[iRm]
+                    Jorb_list[(R_vec, iatom, jatom)] = t / (4.0 * np.pi)
+                    JJ_list[(R_vec, iatom, jatom)] = tmp / (4.0 * np.pi)
+                    Rij_done.add((iR, iatom, jatom))
+                    if (iRm, jatom, iatom) not in Rij_done:
+                        Jorb_list[(Rm_vec, jatom, iatom)] = t.T / (4.0 * np.pi)
+                        JJ_list[(Rm_vec, jatom, iatom)] = tmp / (4.0 * np.pi)
+                        Rij_done.add((iRm, jatom, iatom))
         return Jorb_list, JJ_list
 
+    def get_all_A_vectorized(self, Gup, Gdn):
+        """
+        Vectorized calculation of all A matrix elements for collinear exchange.
+        Follows the pattern from TB2J/exchange.py get_all_A_vectorized method.
+
+        :param Gup: Green's function array for spin up, shape (nR, nbasis, nbasis)
+        :param Gdn: Green's function array for spin down, shape (nR, nbasis, nbasis)
+        :returns: tuple of (Jorb_list, JJ_list) with R vector keys
+        """
+        # Get magnetic sites and their orbital indices
+        magnetic_sites = self.ind_mag_atoms
+        iorbs = [self.iorb(site) for site in magnetic_sites]
+
+        # Build the Delta matrices for all magnetic sites
+        Delta = [self.get_Delta(site) for site in magnetic_sites]
+
+        # Initialize results dictionaries
+        Jorb_list = {}
+        JJ_list = {}
+
+        # Batch compute all exchange tensors following the vectorized approach
+        for i, j in product(range(len(magnetic_sites)), repeat=2):
+            idx, jdx = iorbs[i], iorbs[j]
+
+            # Extract Gij and Gji for all R vectors at once
+            Gij = Gup[:, idx][:, :, jdx]  # Shape: (nR, ni, nj)
+            # Since short_Rlist is properly ordered, we can flip Gji along R axis
+            # to get Gji(-R) for Gij(R)
+            Gji = np.flip(Gdn[:, jdx][:, :, idx], axis=0)  # Shape: (nR, nj, ni)
+
+            # Compute exchange tensors for all R vectors at once
+            # Following collinear formula: Delta_i @ Gij @ Delta_j @ Gji
+            t_tensor = np.einsum(
+                "ab,rbc,cd,rda->rac", Delta[i], Gij, Delta[j], Gji, optimize="optimal"
+            ) / (4.0 * np.pi)
+            tmp_tensor = np.sum(t_tensor, axis=(1, 2))  # Shape: (nR,)
+
+            mi, mj = (magnetic_sites[i], magnetic_sites[j])
+
+            # Store results for each R vector
+            for iR, R_vec in enumerate(self.short_Rlist):
+                # Store with R vector key for compatibility
+                Jorb_list[(R_vec, mi, mj)] = t_tensor[iR]  # Shape: (ni, nj)
+                JJ_list[(R_vec, mi, mj)] = tmp_tensor[iR]  # Scalar
+
+        # Filter results to only include atom pairs within cutoff distance
+        # This matches the behavior of the original get_all_A method
+        filtered_Jorb_list = {}
+        filtered_JJ_list = {}
+
+        for iR, atom_pairs in self.R_ijatom_dict.items():
+            R_vec = self.short_Rlist[iR]
+            for iatom, jatom in atom_pairs:
+                key = (R_vec, iatom, jatom)
+                if key in Jorb_list:
+                    filtered_Jorb_list[key] = Jorb_list[key]
+                    filtered_JJ_list[key] = JJ_list[key]
+
+        return filtered_Jorb_list, filtered_JJ_list
+
     def A_to_Jtensor(self):
         for key, val in self.JJ.items():
             # key:(R, iatom, jatom)
@@ -230,27 +305,182 @@ class ExchangeCL2(ExchangeCL):
         #    integrate = trapezoidal_nonuniform
         # elif method == "simpson":
         #    integrate = simpson_nonuniform
-        for R, ijpairs in self.R_ijatom_dict.items():
-            for iatom, jatom in ijpairs:
-                # self.Jorb[(R, iatom, jatom)] = integrate(
-                #    self.contour.path, self.Jorb_list[(R, iatom, jatom)]
+        for iR in self.R_ijatom_dict:
+            for iatom, jatom in self.R_ijatom_dict[iR]:
+                # Use R vector key consistently
+                R_vec = self.short_Rlist[iR]
+                key = (R_vec, iatom, jatom)
+                # self.Jorb[key] = integrate(
+                #    self.contour.path, self.Jorb_list[key]
                 # )
-                # self.JJ[(R, iatom, jatom)] = integrate(
-                #    self.contour.path, self.JJ_list[(R, iatom, jatom)]
+                # self.JJ[key] = integrate(
+                #    self.contour.path, self.JJ_list[key]
                 # )
-                self.Jorb[(R, iatom, jatom)] = self.contour.integrate_values(
-                    self.Jorb_list[(R, iatom, jatom)]
-                )
-                self.JJ[(R, iatom, jatom)] = self.contour.integrate_values(
-                    self.JJ_list[(R, iatom, jatom)]
-                )
+                self.Jorb[key] = self.contour.integrate_values(self.Jorb_list[key])
+                self.JJ[key] = self.contour.integrate_values(self.JJ_list[key])
 
     def get_quantities_per_e(self, e):
-        GR_up = self.Gup.get_GR(self.short_Rlist, energy=e, get_rho=False)
-        GR_dn = self.Gdn.get_GR(self.short_Rlist, energy=e, get_rho=False)
-        Jorb_list, JJ_list = self.get_all_A(GR_up, GR_dn)
+        # short_Rlist now contains actual R vectors
+        GR_up = self.Gup.get_GR(self.short_Rlist, energy=e)
+        GR_dn = self.Gdn.get_GR(self.short_Rlist, energy=e)
+
+        # Save diagonal elements of Green's functions for charge and magnetic moment calculation
+        # Only if debug option is enabled
+        if self.debug_options.get("compute_charge_moments", False):
+            self.save_greens_function_diagonals_collinear(GR_up, GR_dn, e)
+
+        # Use vectorized method with fallback to original method
+        try:
+            Jorb_list, JJ_list = self.get_all_A_vectorized(GR_up, GR_dn)
+            # Jorb_list, JJ_list = self.get_all_A(GR_up, GR_dn)
+        except Exception as ex:
+            print(f"Vectorized method failed: {ex}, falling back to original method")
+            Jorb_list, JJ_list = self.get_all_A(GR_up, GR_dn)
         return dict(Jorb_list=Jorb_list, JJ_list=JJ_list)
 
+    def save_greens_function_diagonals_collinear(self, GR_up, GR_dn, energy):
+        """
+        Save diagonal elements of spin-resolved Green's functions for each atom.
+        These will be used to compute charge and magnetic moments in collinear case.
+
+        :param GR_up: Spin-up Green's function array of shape (nR, nbasis, nbasis)
+        :param GR_dn: Spin-down Green's function array of shape (nR, nbasis, nbasis)
+        :param energy: Current energy value
+        """
+        # For proper charge and magnetic moment calculation, we need to sum over k-points
+        # with weights: Σ_k S(k)·G(k)·w(k)
+
+        # Initialize the k-summed SG matrices for this energy
+        nbasis = GR_up.shape[1]
+        SG_up_ksum = np.zeros((nbasis, nbasis), dtype=complex)
+        SG_dn_ksum = np.zeros((nbasis, nbasis), dtype=complex)
+
+        # Get k-points and weights from Green's function object
+        kpts = self.Gup.kpts
+        kweights = self.Gup.kweights
+
+        # Use the passed energy parameter
+        current_energy = energy
+
+        # Sum over all k-points
+        for ik, kpt in enumerate(kpts):
+            # Get G(k) for current energy
+            Gk_up = self.Gup.get_Gk(ik, energy=current_energy)
+            Gk_dn = self.Gdn.get_Gk(ik, energy=current_energy)
+
+            if not self.Gup.is_orthogonal:
+                Sk = self.Gup.get_Sk(ik)  # Same overlap for both spins
+                SG_up_ksum += Sk @ Gk_up * kweights[ik]
+                SG_dn_ksum += Sk @ Gk_dn * kweights[ik]
+            else:
+                # For orthogonal case, S is identity
+                SG_up_ksum += Gk_up * kweights[ik]
+                SG_dn_ksum += Gk_dn * kweights[ik]
+
+        # Store both spin channels
+        if not hasattr(self, "G_diagonal_up"):
+            self.G_diagonal_up = {iatom: [] for iatom in range(len(self.atoms))}
+            self.G_diagonal_dn = {iatom: [] for iatom in range(len(self.atoms))}
+
+        for iatom in self.orb_dict:
+            # Get orbital indices for this atom
+            orbi = self.iorb(iatom)
+            # Extract diagonal elements for this atom
+            G_up_diag = np.diag(SG_up_ksum[np.ix_(orbi, orbi)])
+            G_dn_diag = np.diag(SG_dn_ksum[np.ix_(orbi, orbi)])
+
+            self.G_diagonal_up[iatom].append(G_up_diag)
+            self.G_diagonal_dn[iatom].append(G_dn_diag)
+
+    def compute_charge_and_magnetic_moments(self):
+        """
+        Compute charge and magnetic moments from stored spin-resolved Green's function diagonals.
+        Uses the relation for collinear case:
+        - Charge: n_i = -1/π ∫ (Im[Tr(S·G_ii^↑(E))] + Im[Tr(S·G_ii^↓(E))]) dE
+        - Magnetic moment: m_i = -1/π ∫ (Im[Tr(S·G_ii^↑(E))] - Im[Tr(S·G_ii^↓(E))]) dE
+        where S is the overlap matrix.
+        """
+        # Only run if debug option is enabled
+        if not self.debug_options.get("compute_charge_moments", False):
+            # Just use density matrix method directly
+            self.get_rho_atom()
+            return
+
+        if not hasattr(self, "G_diagonal_up") or not self.G_diagonal_up:
+            print(
+                "Warning: No Green's function diagonals stored. Cannot compute charge and magnetic moments."
+            )
+            return
+
+        self.charges = np.zeros(len(self.atoms))
+        self.spinat = np.zeros((len(self.atoms), 3))
+
+        # Arrays to store Green's function results for comparison
+        gf_charges = np.zeros(len(self.atoms))
+        gf_spinat = np.zeros((len(self.atoms), 3))
+
+        for iatom in range(len(self.atoms)):
+            if not self.G_diagonal_up[iatom] or not self.G_diagonal_dn[iatom]:
+                continue
+
+            # Stack all diagonal elements for this atom
+            G_up_diags = np.array(
+                self.G_diagonal_up[iatom]
+            )  # shape: (n_energies, n_orbitals)
+            G_dn_diags = np.array(
+                self.G_diagonal_dn[iatom]
+            )  # shape: (n_energies, n_orbitals)
+
+            # Integrate over energy using the same contour as exchange calculation
+            # Charge: -1/π ∫ Im[G_up + G_dn] dE
+            integrated_up = -np.imag(self.contour.integrate_values(G_up_diags)) / np.pi
+            integrated_dn = -np.imag(self.contour.integrate_values(G_dn_diags)) / np.pi
+
+            # Sum over orbitals and spin channels to get total charge
+            gf_charge = np.sum(integrated_up) + np.sum(integrated_dn)
+
+            # Magnetic moment (assuming z-direction for collinear case)
+            # m_z = -1/π ∫ Im[G_up - G_dn] dE
+            gf_spin_z = np.sum(integrated_up) - np.sum(integrated_dn)
+
+            # Store Green's function results
+            gf_charges[iatom] = gf_charge
+            gf_spinat[iatom, 2] = gf_spin_z
+
+            # Compute using density matrix method for comparison
+            self.get_rho_atom()  # This computes charges and spinat using density matrix
+            dm_charge = self.charges[iatom]
+            dm_spin_z = self.spinat[iatom, 2]
+
+            # Compare methods if difference is above threshold
+            charge_diff = abs(gf_charge - dm_charge)
+            spin_diff = abs(gf_spin_z - dm_spin_z)
+            threshold = self.debug_options.get("charge_moment_threshold", 1e-4)
+
+            if charge_diff > threshold or spin_diff > threshold:
+                print(f"Atom {iatom}:")
+                print(f"  Green's function charge: {gf_charge:.6f}")
+                print(f"  Density matrix charge: {dm_charge:.6f}")
+                if charge_diff > threshold:
+                    print(
+                        f"  Charge difference: {charge_diff:.6f} (threshold: {threshold})"
+                    )
+                print(f"  Green's function magnetic moment (z): {gf_spin_z:.6f}")
+                print(f"  Density matrix magnetic moment (z): {dm_spin_z:.6f}")
+                if spin_diff > threshold:
+                    print(
+                        f"  Spin difference: {spin_diff:.6f} (threshold: {threshold})"
+                    )
+
+            # By default, use density matrix output unless debug option says otherwise
+            if not self.debug_options.get("use_density_matrix_output", True):
+                # Override with Green's function values (not recommended)
+                self.charges[iatom] = gf_charge
+                self.spinat[iatom, 2] = gf_spin_z
+
+        # Restore the final values from density matrix method
+        self.get_rho_atom()
+
     def calculate_all(self):
         """
         The top level.
@@ -271,13 +501,19 @@ class ExchangeCL2(ExchangeCL):
         for i, result in enumerate(results):
             Jorb_list = result["Jorb_list"]
             JJ_list = result["JJ_list"]
-            for iR, R in enumerate(self.R_ijatom_dict):
-                for iatom, jatom in self.R_ijatom_dict[R]:
-                    key = (R, iatom, jatom)
+            for iR in self.R_ijatom_dict:
+                for iatom, jatom in self.R_ijatom_dict[iR]:
+                    # Use R vector key consistently across all dictionaries
+                    R_vec = self.short_Rlist[iR]
+                    key = (R_vec, iatom, jatom)
                     self.Jorb_list[key].append(Jorb_list[key])
                     self.JJ_list[key].append(JJ_list[key])
         self.integrate()
         self.get_rho_atom()
+
+        # Compute charge and magnetic moments from Green's function diagonals
+        self.compute_charge_and_magnetic_moments()
+
         self.A_to_Jtensor()
 
     def write_output(self, path="TB2J_results"):

TB2J/exchange_params.py

@@ -33,6 +33,12 @@ class ExchangeParams:
     mae_angles = None
     orth = False
     ibz = False
+    # Debug options
+    debug_options = {
+        "compute_charge_moments": False,  # Whether to compute charge and magnetic moments with Green's function method
+        "use_density_matrix_output": True,  # Whether to use density matrix results as default output
+        "charge_moment_threshold": 1e-4,  # Threshold for printing comparison when difference > threshold
+    }
 
     def __init__(
         self,
@@ -57,6 +63,7 @@ class ExchangeParams:
         orth=False,
         ibz=False,
         index_magnetic_atoms=None,
+        debug_options=None,
     ):
         self.efermi = efermi
         self.basis = basis
@@ -83,6 +90,23 @@ class ExchangeParams:
         self.ibz = ibz
         self.index_magnetic_atoms = index_magnetic_atoms
 
+        # Initialize debug options
+        if debug_options is None:
+            self.debug_options = {
+                "compute_charge_moments": False,
+                "use_density_matrix_output": True,
+                "charge_moment_threshold": 1e-4,
+            }
+        else:
+            # Update default debug options with user-provided ones
+            default_debug = {
+                "compute_charge_moments": False,
+                "use_density_matrix_output": True,
+                "charge_moment_threshold": 1e-4,
+            }
+            default_debug.update(debug_options)
+            self.debug_options = default_debug
+
     def set_params(self, **kwargs):
         for key, val in kwargs.items():
             setattr(self, key, val)

TB2J/green.py

@@ -56,13 +56,26 @@ def find_energy_ingap(evals, rbound, gap=2.0):
     find a energy inside a gap below rbound (right bound),
     return the energy gap top - 0.5.
     """
+    # print("Finding energy in gap...")
+    # print(f"Right bound: {rbound}, min gap size: {gap}")
     m0 = np.sort(evals.flatten())
-    m = m0[m0 < rbound]
+    # print(f"Min eigenvalue: {m0[0]}, Max eigenvalue: {m0[-1]}")
+    m = m0[m0 <= rbound]
+    # append the next state above rbound to m
+    if len(m0) > len(m):
+        m = np.append(m, m0[len(m)])
+    # print(f"Number of states below right bound: {len(m)}")
+    # print(f"Max eigenvalue below right bound: {m[-1]}")
     ind = np.where(np.diff(m) > gap)[0]
+    # print(f"Number of gaps found: {len(ind)}")
+    # print("ind[-1]: ", ind[-1] if len(ind) > 0 else "N/A")
+    emin = 0.0
     if len(ind) == 0:
-        return m0[0] - 0.5
+        emin = m0[0] - 0.5
     else:
-        return m[ind[-1] + 1] - 0.5
+        emin = m[ind[-1] + 1] - 0.5
+    # print("emin:", emin)
+    return emin
 
 
 class TBGreen:
@@ -79,12 +92,22 @@ class TBGreen:
         use_cache=False,
         cache_path=None,
         nproc=1,
+        initial_emin=-25,
     ):
         """
         :param tbmodel: A tight binding model
         :param kmesh: size of monkhorst pack. e.g [6,6,6]
         :param efermi: fermi energy.
+        :param gamma: whether to include gamma point in monkhorst pack
+        :param kpts: user defined kpoints
+        :param kweights: weights for user defined kpoints
+        :param k_sym: whether the kpoints are symmetrized
+        :param use_cache: whether to use cache to store wavefunctions
+        :param cache_path: path to store cache
+        :param nproc: number of processes to use
+        :param emin: minimum energy relative to fermi level to consider
         """
+        self.initial_emin = initial_emin
         self.tbmodel = tbmodel
         self.is_orthogonal = tbmodel.is_orthogonal
         self.R2kfactor = tbmodel.R2kfactor
@@ -235,9 +258,21 @@ class TBGreen:
             self.efermi = occ.efermi(copy.deepcopy(self.evals))
             print(f"Fermi energy found: {self.efermi}")
 
-        # self.evals, self.evecs = self._reduce_eigens(
-        #    self.evals, self.evecs, emin=self.efermi - 15.0, emax=self.efermi + 10.1
-        # )
+        self.adjusted_emin = (
+            find_energy_ingap(
+                self.evals, rbound=self.efermi + self.initial_emin, gap=2.0
+            )
+            - self.efermi
+        )
+        # print(f"Adjusted emin relative to Fermi level: {self.adjusted_emin}")
+        self.evals, self.evecs = self._reduce_eigens(
+            self.evals,
+            self.evecs,
+            emin=self.efermi + self.adjusted_emin,
+            emax=self.efermi + 5.1,
+            # emin=self.efermi -10,
+            # emax=self.efermi + 10,
+        )
         if self._use_cache:
             evecs = self.evecs
             self.evecs_shape = self.evecs.shape
@@ -376,38 +411,28 @@ class TBGreen:
             Gk_all[ik] = self.get_Gk(ik, energy)
         return Gk_all
 
-    def get_GR(self, Rpts, energy, get_rho=False, Gk_all=None):
+    def compute_GR(self, Rpts, kpts, Gks):
+        Rvecs = np.array(Rpts)
+        phase = np.exp(self.k2Rfactor * np.einsum("ni,mi->nm", Rvecs, kpts))
+        phase *= self.kweights[None]
+        GR = np.einsum("kij,rk->rij", Gks, phase, optimize="optimal")
+        return GR
+
+    def get_GR(self, Rpts, energy, Gk_all=None):
         """calculate real space Green's function for one energy, all R points.
         G(R, epsilon) = G(k, epsilon) exp(-2\pi i R.dot. k)
         :param Rpts: R points
-        :param energy:
-        :returns:  real space green's function for one energy for a list of R.
-        :rtype:  dictionary, the keys are tuple of R, values are matrices of nbasis*nbasis
+        :param energy: energy value
+        :param Gk_all: optional pre-computed Gk for all k-points
+        :returns: real space green's function for one energy for a list of R.
+        :rtype: numpy array of shape (len(Rpts), nbasis, nbasis)
         """
-        Rpts = [tuple(R) for R in Rpts]
-        GR = defaultdict(lambda: 0.0j)
-        rhoR = defaultdict(lambda: 0.0j)
-        for ik, kpt in enumerate(self.kpts):
-            if Gk_all is None:
-                Gk = self.get_Gk(ik, energy)
-            else:
-                Gk = Gk_all[ik]
-            if get_rho:
-                if self.is_orthogonal:
-                    rhok = Gk
-                else:
-                    rhok = self.get_Sk(ik) @ Gk
-            for iR, R in enumerate(Rpts):
-                phase = np.exp(self.k2Rfactor * np.dot(R, kpt))
-                tmp = Gk * (phase * self.kweights[ik])
-                GR[R] += tmp
-                # change this if need full rho
-                if get_rho and R == (0, 0, 0):
-                    rhoR[R] += rhok * (phase * self.kweights[ik])
-        if get_rho:
-            return GR, rhoR
+        if Gk_all is None:
+            Gks = self.get_Gk_all(energy)
         else:
-            return GR
+            Gks = Gk_all
+
+        return self.compute_GR(Rpts, self.kpts, Gks)
 
     def get_GR_and_dGRdx1(self, Rpts, energy, dHdx):
         """

TB2J/myTB.py

@@ -1,15 +1,17 @@
-import os
-import numpy as np
 import copy
-from scipy.linalg import eigh
-from scipy.sparse import csr_matrix
-from scipy.io import netcdf_file
+import os
 from collections import defaultdict
 
+import numpy as np
+
 # from tbmodels import Model
 from ase.atoms import Atoms
+from scipy.io import netcdf_file
+from scipy.linalg import eigh
+from scipy.sparse import csr_matrix
+
 from TB2J.utils import auto_assign_basis_name
-from TB2J.wannier import parse_ham, parse_xyz, parse_atoms, parse_tb
+from TB2J.wannier import parse_atoms, parse_ham, parse_tb, parse_xyz
 
 
 class AbstractTB:
@@ -187,9 +189,7 @@ class MyTB(AbstractTB):
         if os.path.exists(tb_fname):
             xcart, nbasis, data, R_degens = parse_tb(fname=tb_fname)
         else:
-            nbasis, data, R_degens = parse_ham(
-                fname=os.path.join(path, prefix + "_hr.dat")
-            )
+            nbasis, data, R_degens = parse_ham(fname=hr_fname)
             xcart, _, _ = parse_xyz(fname=os.path.join(path, prefix + "_centres.xyz"))
 
         if atoms is None:
@@ -237,7 +237,7 @@ class MyTB(AbstractTB):
         return m
 
     def gen_ham(self, k, convention=2):
-        """
+        r"""
         generate hamiltonian matrix at k point.
         H_k( i, j)=\sum_R H_R(i, j)^phase.
         There are two conventions,
@@ -451,7 +451,7 @@ class MyTB(AbstractTB):
         nbasis = root.dimensions["nbasis"]
         nspin = root.dimensions["nspin"]
         ndim = root.dimensions["ndim"]
-        natom = root.dimensions["natom"]
+        _natom = root.dimensions["natom"]
         Rlist = root.variables["R"][:]
         mdata_real = root.variables["data_real"][:]
         mdata_imag = root.variables["data_imag"][:]