TB2J 0.9.5rc0__py3-none-any.whl → 0.9.6rc0__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, pauli_mat
+from TB2J.utils import (
+    kmesh_to_R,
+    simpson_nonuniform,
+    symbol_number,
+    trapezoidal_nonuniform,
+)
 
 
 class Exchange(ExchangeParams):
@@ -219,9 +151,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 = {}
@@ -331,12 +263,14 @@ class ExchangeNCL(Exchange):
         self.tbmodel = tbmodels
         self.backend_name = self.tbmodel.name
         self.G = TBGreen(
-            self.tbmodel,
-            self.kmesh,
-            self.efermi,
+            tbmodel=self.tbmodel,
+            kmesh=self.kmesh,
+            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]
@@ -695,21 +633,23 @@ class ExchangeNCL(Exchange):
         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
+
+    #    mae_e = np.zeros(len(self.mae_angles), dtype=complex)
+    #    # rotate the Hso_k to angles
+    #    for ik, k in enumerate(self.G.kpoints):
+    #        Hso_k = self.model.get_Hso_k(k)
+    #        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
         """
+        raise NotImplementedError("Not yet implemented.")
 
     def save_AijR(self, AijRs, fname):
         result = dict(path=self.contour.path, AijRs=AijRs)
@@ -734,9 +674,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(

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,
+    }

TB2J/green.py

@@ -1,14 +1,19 @@
-import numpy as np
-from collections import defaultdict
-from ase.dft.kpoints import monkhorst_pack
-from shutil import rmtree
+import copy
 import os
+import pickle
+import sys
 import tempfile
+from collections import defaultdict
+from shutil import rmtree
+
+import numpy as np
+from HamiltonIO.model.occupations import Occupations
+from HamiltonIO.model.occupations import myfermi as fermi
 from pathos.multiprocessing import ProcessPool
-import sys
-import pickle
-import warnings
-from TB2J.mathutils.fermi import fermi
+
+from TB2J.kpoints import monkhorst_pack
+
+# from TB2J.mathutils.fermi import fermi
 
 MAX_EXP_ARGUMENT = np.log(sys.float_info.max)
 
@@ -22,9 +27,16 @@ def eigen_to_G(evals, evecs, efermi, energy):
     :returns: Green's function G,
     :rtype:  Matrix with same shape of the Hamiltonian (and eigenvector)
     """
-    return (
-        np.einsum("ij, j-> ij", evecs, 1.0 / (-evals + (energy + efermi)))
-        @ evecs.conj().T
+    # return (
+    #    np.einsum("ij, j-> ij", evecs, 1.0 / (-evals + (energy + efermi)))
+    #    @ evecs.conj().T
+    # )
+    return np.einsum(
+        "ib, b, jb-> ij",
+        evecs,
+        1.0 / (-evals + (energy + efermi)),
+        evecs.conj(),
+        optimize=True,
     )
 
 
@@ -46,8 +58,10 @@ class TBGreen:
     def __init__(
         self,
         tbmodel,
-        kmesh,  # [ikpt, 3]
-        efermi,  # efermi
+        kmesh=None,  # [ikpt, 3]
+        efermi=None,  # efermi
+        gamma=False,
+        kpts=None,
         k_sym=False,
         use_cache=False,
         cache_path=None,
@@ -67,12 +81,14 @@ class TBGreen:
         self.cache_path = cache_path
         if use_cache:
             self._prepare_cache()
-        if kmesh is not None:
-            self.kpts = monkhorst_pack(size=kmesh)
+        if kpts is not None:
+            self.kpts = kpts
+        elif kmesh is not None:
+            self.kpts = monkhorst_pack(kmesh, gamma_center=gamma)
         else:
             self.kpts = tbmodel.get_kpts()
         self.nkpts = len(self.kpts)
-        self.kweights = [1.0 / self.nkpts] * self.nkpts
+        self.kweights = np.array([1.0 / self.nkpts] * self.nkpts)
         self.norb = tbmodel.norb
         self.nbasis = tbmodel.nbasis
         self.k_sym = k_sym
@@ -145,9 +161,19 @@ class TBGreen:
         if not saveH:
             self.H = None
 
-        self.evals, self.evecs = self._reduce_eigens(
-            self.evals, self.evecs, emin=self.efermi - 10.0, emax=self.efermi + 10.1
-        )
+        # get efermi
+        if self.efermi is None:
+            print("Calculating Fermi energy from eigenvalues")
+            print(f"Number of electrons: {self.tbmodel.nel} ")
+            occ = Occupations(
+                nel=self.tbmodel.nel, width=0.05, wk=self.kweights, nspin=2
+            )
+            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
+        # )
         if self._use_cache:
             evecs = self.evecs
             self.evecs_shape = self.evecs.shape
@@ -216,20 +242,22 @@ class TBGreen:
         rho = np.zeros((self.nbasis, self.nbasis), dtype=complex)
         if self.is_orthogonal:
             for ik, _ in enumerate(self.kpts):
+                evecs_k = self.get_evecs(ik)
+                # chekc if any of the evecs element is nan
                 rho += (
-                    (self.get_evecs(ik) * fermi(self.evals[ik], self.efermi))
-                    @ self.get_evecs(ik).T.conj()
+                    (evecs_k * fermi(self.evals[ik], self.efermi, nspin=2))
+                    @ evecs_k.T.conj()
                     * self.kweights[ik]
                 )
         else:
             for ik, _ in enumerate(self.kpts):
                 rho += (
-                    (self.get_evecs(ik) * fermi(self.evals[ik], self.efermi))
+                    (self.get_evecs(ik) * fermi(self.evals[ik], self.efermi, nspin=2))
                     @ self.get_evecs(ik).T.conj()
                     @ self.get_Sk(ik)
                     * self.kweights[ik]
                 )
-
+        # check if rho has nan values
         return rho
 
     def get_rho_R(self, Rlist):
@@ -238,7 +266,10 @@ class TBGreen:
         for ik, kpt in enumerate(self.kpts):
             evec = self.get_evecs(ik)
             rhok = np.einsum(
-                "ib,b, bj-> ij", evec, fermi(self.evals[ik], self.efermi), evec.conj().T
+                "ib,b, bj-> ij",
+                evec,
+                fermi(self.evals[ik], self.efermi, nspin=2),
+                evec.conj().T,
             )
             for iR, R in enumerate(Rlist):
                 rho_R[iR] += rhok * np.exp(self.k2Rfactor * kpt @ R) * self.kweights[ik]
@@ -252,16 +283,20 @@ class TBGreen:
     def get_density(self):
         return np.real(np.diag(self.get_density_matrix()))
 
-    def get_Gk(self, ik, energy):
+    def get_Gk(self, ik, energy, evals=None, evecs=None):
         """Green's function G(k) for one energy
         G(\epsilon)= (\epsilon I- H)^{-1}
         :param ik: indices for kpoint
         :returns: Gk
         :rtype:  a matrix of indices (nbasis, nbasis)
         """
+        if evals is None:
+            evals = self.get_evalue(ik)
+        if evecs is None:
+            evecs = self.get_evecs(ik)
         Gk = eigen_to_G(
-            evals=self.get_evalue(ik),
-            evecs=self.get_evecs(ik),
+            evals=evals,
+            evecs=evecs,
             efermi=self.efermi,
             energy=energy,
         )

TB2J/interfaces/__init__.py

@@ -0,0 +1,12 @@
+from .abacus import gen_exchange_abacus
+from .manager import Manager
+from .siesta_interface import gen_exchange_siesta
+from .wannier90_interface import WannierManager, gen_exchange
+
+__all__ = [
+    "Manager",
+    "gen_exchange_siesta",
+    "WannierManager",
+    "gen_exchange",
+    "gen_exchange_abacus",
+]

TB2J/interfaces/abacus/__init__.py

@@ -0,0 +1,4 @@
+from .abacus_wrapper import AbacusParser, AbacusWrapper
+from .gen_exchange_abacus import gen_exchange_abacus
+
+__all__ = ["AbacusWrapper", "AbacusParser", "gen_exchange_abacus"]

TB2J/{abacus → interfaces/abacus}/abacus_api.py

@@ -1,10 +1,10 @@
 # -*- coding: utf-8 -*-
-import numpy as np
-from scipy import sparse
-from scipy.sparse import csr_matrix
 import re
 import struct
 
+import numpy as np
+from scipy.sparse import csr_matrix
+
 
 class XR_matrix:
     def __init__(self, nspin, XR_fileName):