TB2J 0.9.5rc0__py3-none-any.whl → 0.9.7rc0__py3-none-any.whl

TB2J/exchange.py

@@ -1,91 +1,23 @@
-from collections import defaultdict, OrderedDict
 import os
-import numpy as np
 import pickle
-from dataclasses import dataclass
-import yaml
-from TB2J.green import TBGreen
-from TB2J.pauli import pauli_block_all, pauli_block_sigma_norm, pauli_mat
-from TB2J.utils import symbol_number, kmesh_to_R
-from TB2J.io_exchange import SpinIO
+from collections import OrderedDict, defaultdict
+
+import numpy as np
 from tqdm import tqdm
-from TB2J.external import p_map
+
 from TB2J.contour import Contour
-from TB2J.utils import simpson_nonuniform, trapezoidal_nonuniform, split_symbol_number
+from TB2J.exchange_params import ExchangeParams
+from TB2J.external import p_map
+from TB2J.green import TBGreen
+from TB2J.io_exchange import SpinIO
 from TB2J.orbmap import map_orbs_matrix
-from typing import List, Tuple
-
-
-@dataclass
-class ExchangeParams:
-    """
-    A class to store the parameters for exchange calculation.
-    """
-
-    efermi: float
-    basis: list = None
-    magnetic_elements: list = None
-    include_orbs = {}
-    _kmesh = [4, 4, 4]
-    emin: float = -15
-    emax: float = 0.05
-    nz: int = 100
-    exclude_orbs = []
-    ne: int = None
-    Rcut: float = None
-    _use_cache: bool = False
-    np: int = 1
-    description: str = ""
-    write_density_matrix: bool = False
-    orb_decomposition: bool = False
-    output_path: str = "TB2J_results"
-    mae_angles = None
-
-    def __init__(
-        self,
-        efermi=-10.0,
-        basis=None,
-        magnetic_elements=None,
-        include_orbs=None,
-        kmesh=[4, 4, 4],
-        emin=-15,
-        emax=0.05,
-        nz=100,
-        exclude_orbs=[],
-        ne=None,
-        Rcut=None,
-        use_cache=False,
-        np=1,
-        description="",
-        write_density_matrix=False,
-        orb_decomposition=False,
-        output_path="TB2J_results",
-    ):
-        self.efermi = efermi
-        self.basis = basis
-        self.magnetic_elements = magnetic_elements
-        self.include_orbs = include_orbs
-        self._kmesh = kmesh
-        self.emin = emin
-        self.emax = emax
-        self.nz = nz
-        self.exclude_orbs = exclude_orbs
-        self.ne = ne
-        self.Rcut = Rcut
-        self._use_cache = use_cache
-        self.np = np
-        self.description = description
-        self.write_density_matrix = write_density_matrix
-        self.orb_decomposition = orb_decomposition
-        self.output_path = output_path
-
-    def set_params(self, **kwargs):
-        for key, val in kwargs.items():
-            setattr(self, key, val)
-
-    def save_to_yaml(self, fname):
-        with open(fname, "w") as myfile:
-            yaml.dump(self.__dict__, myfile)
+from TB2J.pauli import pauli_block_all, pauli_block_sigma_norm
+from TB2J.utils import (
+    kmesh_to_R,
+    simpson_nonuniform,
+    symbol_number,
+    trapezoidal_nonuniform,
+)
 
 
 class Exchange(ExchangeParams):
@@ -107,10 +39,7 @@ class Exchange(ExchangeParams):
         self._prepare_distance()
 
         # whether to calculate J and DMI with NJt method.
-        self.calc_NJt = True
         # self._prepare_NijR()
-        self.Ddict_NJT = None
-        self.Jdict_NJT = None
         self._is_collinear = True
         self.has_elistc = False
         self._clean_tbmodels()
@@ -128,8 +57,9 @@ class Exchange(ExchangeParams):
         pass
 
     def _clean_tbmodels(self):
-        del self.tbmodel
-        del self.G.tbmodel
+        # del self.tbmodel
+        # del self.G.tbmodel
+        pass
 
     def _prepare_kmesh(self, kmesh):
         for k in kmesh:
@@ -219,9 +149,9 @@ class Exchange(ExchangeParams):
                     raise ValueError(
                         f"""The number of spin-orbitals for atom {iatom} is not even,
 {nsorb} spin-orbitals are found near this atom.
-which means the spin up/down does not have same number of orbitals. 
+which means the spin up/down does not have same number of orbitals.
 This is often because the Wannier functions are wrongly defined,
-or badly localized. Please check the Wannier centers in the Wannier90 output file. 
+or badly localized. Please check the Wannier centers in the Wannier90 output file.
 """
                     )
         self._spin_dict = {}
@@ -330,13 +260,17 @@ class ExchangeNCL(Exchange):
         """
         self.tbmodel = tbmodels
         self.backend_name = self.tbmodel.name
+        print(self.kmesh)
         self.G = TBGreen(
-            self.tbmodel,
-            self.kmesh,
-            self.efermi,
+            tbmodel=self.tbmodel,
+            kmesh=self.kmesh,
+            gamma=True,
+            efermi=self.efermi,
             use_cache=self._use_cache,
-            nproc=self.np,
+            nproc=self.nproc,
         )
+        if self.efermi is None:
+            self.efermi = self.G.efermi
         self.norb = self.G.norb
         self.nbasis = self.G.nbasis
         # self.rho = np.zeros((self.nbasis, self.nbasis), dtype=complex)
@@ -344,7 +278,11 @@ class ExchangeNCL(Exchange):
         self.A_ijR_list = defaultdict(lambda: [])
         self.A_ijR = defaultdict(lambda: np.zeros((4, 4), dtype=complex))
         self.A_ijR_orb = dict()
-        self.HR0 = self.G.H0
+        # self.HR0 = self.tbmodel.get_H0()
+        if hasattr(self.tbmodel, "get_H0"):
+            self.HR0 = self.tbmodel.get_H0()
+        else:
+            self.HR0 = self.G.H0
         self._is_collinear = False
         self.Pdict = {}
         if self.write_density_matrix:
@@ -477,7 +415,7 @@ class ExchangeNCL(Exchange):
         if self.orb_decomposition:
             for key, val in self.A_ijR_orb.items():
                 R, iatom, jatom = key
-                Rm = tuple(-x for x in R)
+                # Rm = tuple(-x for x in R)
                 # valm = self.A_ijR_orb[(Rm, jatom, iatom)]
                 ni = self.norb_reduced[iatom]
                 nj = self.norb_reduced[jatom]
@@ -613,53 +551,6 @@ class ExchangeNCL(Exchange):
         self.rho_dict = rho
         return self.rho_dict
 
-    def calculate_DMI_NJT(self):
-        """
-        calculate exchange and DMI with the
-        This did not work and should be removed.
-        """
-        Ddict_NJT = {}
-        Jdict_NJT = {}
-        for R in self.short_Rlist:
-            N = self.N[tuple(-np.array(R))]  # density matrix
-            t = self.tbmodel.get_hamR(R)  # hopping parameter
-            for iatom in self.ind_mag_atoms:
-                orbi = self.iorb(iatom)
-                ni = len(orbi)
-                for jatom in self.ind_mag_atoms:
-                    orbj = self.iorb(jatom)
-                    nj = len(orbj)
-                    Nji = N[np.ix_(orbj, orbi)]
-                    tij = t[np.ix_(orbi, orbj)]
-                    D = np.zeros(3, dtype=float)
-                    J = np.zeros(3, dtype=float)
-                    for dim in range(3):
-                        # S_i = pauli_mat(ni, dim +
-                        #                1)  #*self.rho[np.ix_(orbi, orbi)]
-                        # S_j = pauli_mat(nj, dim +
-                        #                1)  #*self.rho[np.ix_(orbj, orbj)]
-                        # TODO: Note that rho is complex, not the imaginary part
-                        S_i = pauli_mat(ni, dim + 1) * self.rho[np.ix_(orbi, orbi)]
-                        S_j = pauli_mat(nj, dim + 1) * self.rho[np.ix_(orbj, orbj)]
-
-                        # [S, t]+  = Si tij + tij Sj, where
-                        # Si and Sj are the spin operator
-                        # Here we do not have L operator, so J-> S
-                        Jt = np.matmul(S_i, tij) + np.matmul(tij, S_j)
-
-                        Jtminus = np.matmul(S_i, tij) - np.matmul(tij, S_j)
-                        # D = -1/2 Tr Nji [J, tij]
-                        # Trace over spin and orb
-                        D[dim] = -0.5 * np.imag(np.trace(np.matmul(Nji, Jt)))
-                        J[dim] = -0.5 * np.imag(np.trace(np.matmul(Nji, Jtminus)))
-                    ispin = self.ispin(iatom)
-                    jspin = self.ispin(jatom)
-                    Ddict_NJT[(R, ispin, jspin)] = D
-                    Jdict_NJT[(R, ispin, jspin)] = J
-        self.Jdict_NJT = Jdict_NJT
-        self.Ddict_NJT = Ddict_NJT
-        return Ddict_NJT
-
     def integrate(self, AijRs, AijRs_orb=None, method="simpson"):
         """
         AijRs: a list of AijR,
@@ -682,35 +573,14 @@ class ExchangeNCL(Exchange):
 
     def get_quantities_per_e(self, e):
         Gk_all = self.G.get_Gk_all(e)
-        mae = self.get_mae_kspace(Gk_all)
+        # mae = self.get_mae_kspace(Gk_all)
+        mae = None
         # TODO: get the MAE from Gk_all
         GR = self.G.get_GR(self.short_Rlist, energy=e, get_rho=False, Gk_all=Gk_all)
         # TODO: define the quantities for one energy.
         AijR, AijR_orb = self.get_all_A(GR)
         return dict(AijR=AijR, AijR_orb=AijR_orb, mae=mae)
 
-    def get_mae_kspace(self, Gk_all):
-        """
-        get the MAE from Gk_all
-        TODO: This is only a place holder.
-        """
-        return None
-        Hso_k = self.model.get_Hso_k(k)
-        mae_e = np.zeros(len(self.mae_angles), dtype=complex)
-        # rotate the Hso_k to angles
-        for ik, k in enumerate(self.G.kpoints):
-            Gk = Gk_all[ik]
-            for i_angle, angle in enumerate(self.mae_angles):
-                # TODO: implementa rotate_H
-                Hso_k_rot = rotate_H(Hso_k, angle)
-                mae_e[i_angle] = Gk @ Hso_k @ Gk @ Hso_k
-        return mae_e
-
-    def get_mae_rspace(self, GR):
-        """
-        get the MAE from GR
-        """
-
     def save_AijR(self, AijRs, fname):
         result = dict(path=self.contour.path, AijRs=AijRs)
         with open(fname, "wb") as myfile:
@@ -734,9 +604,9 @@ class ExchangeNCL(Exchange):
         self.validate()
 
         npole = len(self.contour.path)
-        if self.np > 1:
+        if self.nproc > 1:
             results = p_map(
-                self.get_quantities_per_e, self.contour.path, num_cpus=self.np
+                self.get_quantities_per_e, self.contour.path, num_cpus=self.nproc
             )
         else:
             results = map(
@@ -799,8 +669,6 @@ class ExchangeNCL(Exchange):
             Jani_dict=self.Jani,
             DMI_orb=self.DMI_orb,
             Jani_orb=self.Jani_orb,
-            NJT_Jdict=self.Jdict_NJT,
-            NJT_ddict=self.Ddict_NJT,
             biquadratic_Jdict=self.B,
             debug_dict=self.debug_dict,
             description=self.description,
@@ -839,8 +707,6 @@ class ExchangeCL(ExchangeNCL):
             distance_dict=self.distance_dict,
             exchange_Jdict=self.exchange_Jdict,
             dmi_ddict=None,
-            NJT_Jdict=None,
-            NJT_ddict=None,
             biquadratic_Jdict=self.B,
             description=self.description,
         )

TB2J/exchangeCL2.py

@@ -20,18 +20,18 @@ class ExchangeCL2(ExchangeCL):
         self.tbmodel_up, self.tbmodel_dn = tbmodels
         self.backend_name = self.tbmodel_up.name
         self.Gup = TBGreen(
-            self.tbmodel_up,
-            self.kmesh,
-            self.efermi,
+            tbmodel=self.tbmodel_up,
+            kmesh=self.kmesh,
+            efermi=self.efermi,
             use_cache=self._use_cache,
-            nproc=self.np,
+            nproc=self.nproc,
         )
         self.Gdn = TBGreen(
-            self.tbmodel_dn,
-            self.kmesh,
-            self.efermi,
+            tbmodel=self.tbmodel_dn,
+            kmesh=self.kmesh,
+            efermi=self.efermi,
             use_cache=self._use_cache,
-            nproc=self.np,
+            nproc=self.nproc,
         )
         if self.write_density_matrix:
             self.Gup.write_rho_R(
@@ -247,15 +247,15 @@ class ExchangeCL2(ExchangeCL):
         print("Green's function Calculation started.")
 
         npole = len(self.contour.path)
-        if self.np == 1:
+        if self.nproc == 1:
             results = map(
                 self.get_quantities_per_e, tqdm(self.contour.path, total=npole)
             )
         else:
-            # pool = ProcessPool(nodes=self.np)
+            # pool = ProcessPool(nodes=self.nproc)
             # results = pool.map(self.get_AijR_rhoR ,self.contour.path)
             results = p_map(
-                self.get_quantities_per_e, self.contour.path, num_cpus=self.np
+                self.get_quantities_per_e, self.contour.path, num_cpus=self.nproc
             )
         for i, result in enumerate(results):
             Jorb_list = result["Jorb_list"]
@@ -265,8 +265,6 @@ class ExchangeCL2(ExchangeCL):
                     key = (R, iatom, jatom)
                     self.Jorb_list[key].append(Jorb_list[key])
                     self.JJ_list[key].append(JJ_list[key])
-        if self.np > 1:
-            pass
         self.integrate()
         self.get_rho_atom()
         self.A_to_Jtensor()

TB2J/exchange_params.py

@@ -0,0 +1,213 @@
+import argparse
+from dataclasses import dataclass
+
+import yaml
+
+__all__ = ["ExchangeParams", "add_exchange_args_to_parser", "parser_argument_to_dict"]
+
+
+@dataclass
+class ExchangeParams:
+    """
+    A class to store the parameters for exchange calculation.
+    """
+
+    efermi: float
+    basis = []
+    magnetic_elements = []
+    include_orbs = {}
+    _kmesh = [4, 4, 4]
+    emin: float = -15
+    emax: float = 0.05
+    nz: int = 100
+    exclude_orbs = []
+    ne: int = 0
+    Rcut: float = None
+    _use_cache: bool = False
+    nproc: int = 1
+    description: str = ""
+    write_density_matrix: bool = False
+    orb_decomposition: bool = False
+    output_path: str = "TB2J_results"
+    mae_angles = None
+    orth = False
+
+    def __init__(
+        self,
+        efermi=-10.0,
+        basis=None,
+        magnetic_elements=None,
+        include_orbs=None,
+        kmesh=[4, 4, 4],
+        emin=-15,
+        emax=0.05,
+        nz=100,
+        ne=None,
+        Rcut=None,
+        use_cache=False,
+        nproc=1,
+        description="",
+        write_density_matrix=False,
+        orb_decomposition=False,
+        output_path="TB2J_results",
+        exclude_orbs=[],
+        mae_angles=None,
+        orth=False,
+    ):
+        self.efermi = efermi
+        self.basis = basis
+        self.magnetic_elements = magnetic_elements
+        self.include_orbs = include_orbs
+        self._kmesh = kmesh
+        self.emin = emin
+        self.emax = emax
+        self.nz = nz
+        self.exclude_orbs = exclude_orbs
+        self.ne = ne
+        self.Rcut = Rcut
+        self._use_cache = use_cache
+        self.nproc = nproc
+        self.description = description
+        self.write_density_matrix = write_density_matrix
+        self.orb_decomposition = orb_decomposition
+        self.output_path = output_path
+        self.mae_angles = mae_angles
+        self.orth = orth
+
+    def set_params(self, **kwargs):
+        for key, val in kwargs.items():
+            setattr(self, key, val)
+
+    def save_to_yaml(self, fname):
+        with open(fname, "w") as myfile:
+            yaml.dump(self.__dict__, myfile)
+
+
+def add_exchange_args_to_parser(parser: argparse.ArgumentParser):
+    parser.add_argument(
+        "--elements",
+        help="elements to be considered in Heisenberg model",
+        default=None,
+        type=str,
+        nargs="*",
+    )
+    parser.add_argument(
+        "--rcut",
+        help="cutoff of spin pair distance. The default is to calculate all commensurate R point to the k mesh.",
+        default=None,
+        type=float,
+    )
+    parser.add_argument("--efermi", help="Fermi energy in eV", default=None, type=float)
+    parser.add_argument(
+        "--ne",
+        help="number of electrons in the unit cell. If not given, TB2J will use the fermi energy to compute it.",
+    )
+    parser.add_argument(
+        "--kmesh",
+        help="kmesh in the format of kx ky kz",
+        type=int,
+        nargs="*",
+        default=[5, 5, 5],
+    )
+    parser.add_argument(
+        "--emin",
+        help="energy minimum below efermi, default -14 eV",
+        type=float,
+        default=-14.0,
+    )
+    parser.add_argument(
+        "--emax",
+        help="energy maximum above efermi, default 0.0 eV",
+        type=float,
+        default=0.0,
+    )
+    parser.add_argument(
+        "--nz",
+        help="number of steps for semicircle contour, default: 100",
+        default=100,
+        type=int,
+    )
+    parser.add_argument(
+        "--cutoff",
+        help="The minimum of J amplitude to write, (in eV), default is 1e-5 eV",
+        default=1e-5,
+        type=float,
+    )
+    parser.add_argument(
+        "--exclude_orbs",
+        help="the indices of wannier functions to be excluded from magnetic site. counting start from 0",
+        default=[],
+        type=int,
+        nargs="+",
+    )
+
+    parser.add_argument(
+        "--np",
+        "--nproc",
+        help="number of cpu cores to use in parallel, default: 1",
+        default=1,
+        type=int,
+    )
+
+    parser.add_argument(
+        "--use_cache",
+        help="whether to use disk file for temporary storing wavefunctions and hamiltonian to reduce memory usage. Default: False",
+        action="store_true",
+        default=False,
+    )
+
+    parser.add_argument(
+        "--description",
+        help="add description of the calculatiion to the xml file. Essential information, like the xc functional, U values, magnetic state should be given.",
+        type=str,
+        default="Calculated with TB2J.",
+    )
+
+    parser.add_argument(
+        "--orb_decomposition",
+        default=False,
+        action="store_true",
+        help="whether to do orbital decomposition in the non-collinear mode.",
+    )
+
+    parser.add_argument(
+        "--output_path",
+        help="The path of the output directory, default is TB2J_results",
+        type=str,
+        default="TB2J_results",
+    )
+    parser.add_argument(
+        "--write_dm",
+        help="whether to write density matrix",
+        action="store_true",
+        default=False,
+    )
+
+    parser.add_argument(
+        "--orth",
+        help="whether to use lowdin orthogonalization before diagonalization (for testing only)",
+        action="store_true",
+        default=False,
+    )
+    return parser
+
+
+def parser_argument_to_dict(args) -> dict:
+    return {
+        "efermi": args.efermi,
+        "magnetic_elements": args.elements,
+        "kmesh": args.kmesh,
+        "emin": args.emin,
+        "emax": args.emax,
+        "nz": args.nz,
+        "exclude_orbs": args.exclude_orbs,
+        "ne": args.ne,
+        "Rcut": args.rcut,
+        "use_cache": args.use_cache,
+        "nproc": args.np,
+        "description": args.description,
+        "write_density_matrix": args.write_dm,
+        "orb_decomposition": args.orb_decomposition,
+        "output_path": args.output_path,
+        "orth": args.orth,
+    }