TB2J 0.9.12.9__py3-none-any.whl → 0.9.12.22__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
@@ -28,6 +29,7 @@ class ExchangeCL2(ExchangeCL):
             efermi=self.efermi,
             use_cache=self._use_cache,
             nproc=self.nproc,
+            smearing_width=self.smearing,
         )
         self.Gdn = TBGreen(
             tbmodel=self.tbmodel_dn,
@@ -35,6 +37,7 @@ class ExchangeCL2(ExchangeCL):
             efermi=self.efermi,
             use_cache=self._use_cache,
             nproc=self.nproc,
+            smearing_width=self.smearing,
         )
         if self.write_density_matrix:
             self.Gup.write_rho_R(
@@ -76,8 +79,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 +111,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 +139,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 +158,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 +183,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 +307,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 +503,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

@@ -1,6 +1,7 @@
 import argparse
 from dataclasses import dataclass
 
+import ase.units
 import yaml
 
 __all__ = ["ExchangeParams", "add_exchange_args_to_parser", "parser_argument_to_dict"]
@@ -13,6 +14,7 @@ class ExchangeParams:
     """
 
     efermi: float
+    smearing: float = 600.0 * ase.units.kB
     basis = []
     magnetic_elements = []
     index_magnetic_atoms = None
@@ -33,10 +35,17 @@ 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,
         efermi=-10.0,
+        smearing=600.0 * ase.units.kB,
         basis=None,
         magnetic_elements=None,
         include_orbs=None,
@@ -57,8 +66,10 @@ class ExchangeParams:
         orth=False,
         ibz=False,
         index_magnetic_atoms=None,
+        debug_options=None,
     ):
         self.efermi = efermi
+        self.smearing = smearing
         self.basis = basis
         # self.magnetic_elements = magnetic_elements
         # self.include_orbs = include_orbs
@@ -83,6 +94,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)
@@ -133,6 +161,12 @@ def add_exchange_args_to_parser(parser: argparse.ArgumentParser):
         type=float,
     )
     parser.add_argument("--efermi", help="Fermi energy in eV", default=None, type=float)
+    parser.add_argument(
+        "--smearing",
+        help="Fermi smearing width for density computation and CFR. Units: K (default) or eV. Example: '600K', '0.2eV'. Default: 600K.",
+        default="600K",
+        type=str,
+    )
     parser.add_argument(
         "--ne",
         help="number of electrons in the unit cell. If not given, TB2J will use the fermi energy to compute it.",
@@ -255,8 +289,22 @@ def parser_argument_to_dict(args) -> dict:
         ind_mag_atoms = [int(i) - 1 for i in ind_mag_atoms]
     else:
         ind_mag_atoms = None
+
+    smearing_str = args.smearing
+    if smearing_str is None:
+        smearing = 600.0 * ase.units.kB
+    else:
+        s = smearing_str.strip().lower()
+        if s.endswith("ev"):
+            smearing = float(s[:-2])
+        elif s.endswith("k"):
+            smearing = float(s[:-1]) * ase.units.kB
+        else:
+            smearing = float(s) * ase.units.kB
+
     return {
         "efermi": args.efermi,
+        "smearing": smearing,
         "magnetic_elements": args.elements,
         "kmesh": args.kmesh,
         "emin": args.emin,