TB2J 0.9.0.1__py3-none-any.whl → 0.9.2rc0__py3-none-any.whl

TB2J/abacus/occupations.py

@@ -0,0 +1,278 @@
+"""
+This file is stolen from the hotbit programm, with some modification.
+"""
+
+import numpy as np
+from scipy.optimize import brentq
+import sys
+
+from ase.dft.dos import DOS
+from scipy import integrate
+
+# import numba
+
+# from numba import float64, int32
+
+MAX_EXP_ARGUMENT = np.log(sys.float_info.max)
+
+# @numba.vectorize(nopython=True)
+# def myfermi(e, mu, width, nspin):
+#    x = (e - mu) / width
+#    if x < -10:
+#        ret = 2.0 / nspin
+#    elif x > 10:
+#        ret = 0.0
+#    else:
+#        ret = 2.0 / nspin / (math.exp(x) + 1)
+#    return ret
+
+
+def myfermi(e, mu, width, nspin):
+    x = (e - mu) / width
+    return np.where(x < 10, 2.0 / (nspin * (np.exp(x) + 1.0)), 0.0)
+
+
+class Occupations(object):
+    def __init__(self, nel, width, wk, nspin=1):
+        """
+        Initialize parameters for occupations.
+        :param nel: Number of electrons
+        :param width: Fermi-broadening
+        :param wk: k-point weights. eg. If only gamma, [1.0]
+        :param nspin(optional): number of spin, if spin=1 multiplicity=2 else, multiplicity=1.
+        """
+        self.nel = nel
+        self.width = width
+        self.wk = wk
+        self.nk = len(wk)
+        self.nspin = nspin
+
+    def get_mu(self):
+        """Return the Fermi-level (or chemical potential)."""
+        return self.mu
+
+    def fermi(self, mu):
+        """
+        Occupy states with given chemical potential.
+        Occupations are 0...2; without k-point weights
+        """
+        return myfermi(self.e, mu, self.width, self.nspin)
+
+    def root_function(self, mu):
+        """This function is exactly zero when mu is right."""
+        f = self.fermi(mu)
+        return np.einsum("i, ij->", self.wk, f) - self.nel
+
+    def occupy(self, e, xtol=1e-11):
+        """
+        Calculate occupation numbers with given Fermi-broadening.
+
+        @param e: e[ind_k,ind_orb] energy of k-point, state a
+            Note added by hexu:  With spin=2,e[k,a,sigma], it also work. only the *2 should be removed.
+        @param wk: wk[:] weights for k-points
+        @param width: The Fermi-broadening
+
+        Returns: fermi[ind_k, ind_orb]
+        """
+        self.e = e
+        eflat = e.flatten()
+        ind = np.argsort(eflat)
+        e_sorted = eflat[ind]
+        if self.nspin == 1:
+            m = 2
+        elif self.nspin == 2:
+            m = 1
+        n_sorted = (self.wk[:, None, None] * np.ones_like(e) * m).flatten()[ind]
+
+        sum = n_sorted.cumsum()
+        if self.nel < sum[0]:
+            ifermi = 0
+        elif self.nel > sum[-1]:
+            raise ("number of electrons larger than number of orbital*spin")
+        else:
+            ifermi = np.searchsorted(sum, self.nel)
+        try:
+            if ifermi == 0:
+                elo = e_sorted[0]
+            else:
+                elo = e_sorted[ifermi - 1]
+            if ifermi == len(e_sorted) - 1:
+                ehi = e_sorted[-1]
+            else:
+                ehi = e_sorted[ifermi + 1]
+            guess = e_sorted[ifermi]
+            dmu = np.max((self.width, guess - elo, ehi - guess))
+            mu = brentq(self.root_function, guess - dmu, guess + dmu, xtol=xtol)
+            # mu = brent(
+            #    self.root_function,
+            #    brack=(guess - elo, guess, guess + dmu),
+            #    tol=xtol)
+        except Exception as E:
+            # probably a bad guess
+            print("Error in finding Fermi level: ", E)
+            dmu = self.width
+            if self.nel < 1e-3:
+                mu = min(e_sorted) - dmu * 20
+            elif self.nel - sum[-1] > -1e-3:
+                mu = max(e_sorted) + dmu * 20
+            else:
+                # mu = brent(
+                #         self.root_function,
+                #         brack=(e_sorted[0] - dmu * 10,
+                #                guess,
+                #                e_sorted[-1] + dmu * 10),
+                #         tol=xtol)
+                mu = brentq(
+                    self.root_function,
+                    e_sorted[0] - dmu * 20,
+                    e_sorted[-1] + dmu * 20,
+                    xtol=xtol,
+                )
+
+        if np.abs(self.root_function(mu)) > xtol * 1e4:
+            # raise RuntimeError(
+            #    'Fermi level could not be assigned reliably. Has the system fragmented?'
+            # )
+            print(
+                "Fermi level could not be assigned reliably. Has the system fragmented?"
+            )
+
+        f = self.fermi(mu)
+        # rho=(self.eigenvecs*f).dot(self.eigenvecs.transpose())
+
+        self.mu, self.f = mu, f
+        return f
+
+    def plot(self):
+        import pylab as pl
+
+        for ik in range(self.nk):
+            pl.plot(self.e[ik, :], self.f[ik, :])
+            pl.scatter(self.e[ik, :], self.f[ik, :])
+        pl.title("occupations")
+        pl.xlabel("energy (Ha)")
+        pl.ylabel("occupation")
+        pl.show()
+
+
+class GaussOccupations(Occupations):
+    def get_mu(self):
+        return self.mu
+
+    def delta(self, energy):
+        """Return a delta-function centered at 'energy'."""
+        x = -(((self.e - energy) / self.width) ** 2)
+        return np.exp(x) / (np.sqrt(np.pi) * self.width)
+
+    def get_dos(self, npts=500):
+        eflat = self.e.flatten()
+        ind = np.argsort(eflat)
+        ##e_sorted = eflat[ind]
+        if self.nspin == 1:
+            m = 2
+        elif self.nspin == 2:
+            m = 1
+        # n_sorted = (self.wk * np.ones_like(self.e) * m).flatten()[ind]
+        dos = np.zeros(npts)
+        for w, e_n in zip(self.w_k, self.e_skn[0]):
+            for e in e_n:
+                dos += w * self.delta(e)
+
+    def root_function(self, mu):
+        pass
+
+    # @profile
+    def occupy(self, e, xtol=1e-8, guess=0.0):
+        self.e = e
+        dos = myDOS(kweights=self.wk, eigenvalues=e, width=self.width, npts=501)
+        edos = dos.get_energies()
+        d = dos.get_dos()
+        idos = integrate.cumtrapz(d, edos, initial=0) - self.nel
+        # f_idos = interpolate.interp1d(edos, idos)
+        # ret = optimize.fmin(f_idos, x0=edos[400], xtol=xtol, disp=True)
+        ifermi = np.searchsorted(idos, 0.0)
+        # self.mu = ret[0]
+        self.mu = edos[ifermi]
+        self.f = self.fermi(self.mu)
+        return self.f
+
+
+class myDOS(DOS):
+    def __init__(
+        self, kweights, eigenvalues, nspin=1, width=0.1, window=None, npts=1001
+    ):
+        """Electronic Density Of States object.
+
+        calc: calculator object
+            Any ASE compliant calculator object.
+        width: float
+            Width of guassian smearing.  Use width=0.0 for linear tetrahedron
+            interpolation.
+        window: tuple of two float
+            Use ``window=(emin, emax)``.  If not specified, a window
+            big enough to hold all the eigenvalues will be used.
+        npts: int
+            Number of points.
+
+        """
+        self.npts = npts
+        self.width = width
+        # self.w_k = calc.get_k_point_weights()
+        self.w_k = kweights
+        self.nspins = nspin
+        # self.e_skn = np.array([[calc.get_eigenvalues(kpt=k, spin=s)
+        #                        for k in range(len(self.w_k))]
+        #                       for s in range(self.nspins)])
+        # self.e_skn -= calc.get_fermi_level()
+        self.e_skn = np.array([eigenvalues.T])  # eigenvalues: iband, ikpt
+
+        if window is None:
+            emin = None
+            emax = None
+        else:
+            emin, emax = window
+
+        if emin is None:
+            emin = self.e_skn.min() - 10 * self.width
+        if emax is None:
+            emax = self.e_skn.max() + 10 * self.width
+
+        self.energies = np.linspace(emin, emax, npts)
+
+        # if width == 0.0: # To use tetrahedron method
+        #    bzkpts = calc.get_bz_k_points()
+        #    size, offset = get_monkhorst_pack_size_and_offset(bzkpts)
+        #    bz2ibz = calc.get_bz_to_ibz_map()
+        #    shape = (self.nspins,) + tuple(size) + (-1,)
+        #    self.e_skn = self.e_skn[:, bz2ibz].reshape(shape)
+        #    self.cell = calc.atoms.cell
+
+    def get_idos(self):
+        e, d = self.get_dos()
+        return np.trapz(d, e)
+
+    def delta(self, energy):
+        """Return a delta-function centered at 'energy'."""
+        x = -(((self.energies - energy) / self.width) ** 2)
+        return np.exp(x) / (np.sqrt(np.pi) * self.width)
+
+    def get_dos(self, spin=None):
+        """Get array of DOS values.
+
+        The *spin* argument can be 0 or 1 (spin up or down) - if not
+        specified, the total DOS is returned.
+        """
+
+        if spin is None:
+            if self.nspins == 2:
+                # Spin-polarized calculation, but no spin specified -
+                # return the total DOS:
+                return self.get_dos(spin=0) + self.get_dos(spin=1)
+            else:
+                spin = 0
+
+        dos = np.zeros(self.npts)
+        for w, e_n in zip(self.w_k, self.e_skn[spin]):
+            for e in e_n:
+                dos += w * self.delta(e)
+        return dos

TB2J/abacus/test_density_matrix.py

@@ -0,0 +1,38 @@
+from scipy.linalg import eigh
+import numpy as np
+
+
+def gen_random_hermitean_matrix(n):
+    A = np.random.rand(n, n) + 1j * np.random.rand(n, n)
+    return A + A.conj().T
+
+
+def gen_overlap_matrix(n):
+    A = np.random.rand(n, n) + 1j * np.random.rand(n, n)
+    return np.dot(A, A.conj().T)
+
+
+def fermi_function(x, ef, beta):
+    return 1.0 / (np.exp(beta * (x - ef)) + 1)
+
+
+def test():
+    n = 10
+    A = gen_random_hermitean_matrix(n)
+    S = gen_overlap_matrix(n)
+    beta = 0.1
+    ef = 0
+
+    evals, evecs = eigh(A, S)
+
+    etot = np.sum(evals * fermi_function(evals, ef, beta))
+
+    rho = np.einsum("ib,b,jb->ij", evecs, fermi_function(evals, ef, beta), evecs.conj())
+
+    etot2 = np.trace(np.dot(A, rho))
+
+    print(etot, etot2)
+
+
+if __name__ == "__main__":
+    test()

TB2J/cut_cell.py

@@ -0,0 +1,82 @@
+import numpy as np
+from TB2J.supercell import find_primitive_cell, map_to_primitive
+from TB2J.io_exchange.io_exchange import SpinIO, gen_distance_dict
+
+
+def cut_cell(path, output_path, sc_matrix, origin_atom_id, thr=1e-5):
+    """
+    Cut the exchange parameters
+    :param path: the original TB2J_results path
+    :param output_path: the output path.
+    :param sc_matrix: the matrix which maps the primitive cell to supercell.
+    :param origin: the origin of the primitive cell.
+    :param thr: the atoms which the reduced position is within -thr to 1.0+thr are considered as inside the primitive atoms
+    :returns:
+    """
+    sc_matrix = np.asarray(sc_matrix, dtype=int)
+    sc_excparams = SpinIO.load_pickle(path=path, fname='TB2J.pickle')
+    sc_atoms = sc_excparams.atoms
+    uc_atoms, ids = find_primitive_cell(sc_atoms,
+                                        sc_matrix,
+                                        origin_atom_id=origin_atom_id,
+                                        thr=thr,
+                                        perfect=False)
+    uc_charges = sc_excparams.charges[ids]
+    uc_spinat = sc_excparams.spinat[ids]
+    indmap, Rmap = map_to_primitive(sc_atoms, uc_atoms)
+
+    # TODO index_spin: {iatom: ispin}
+
+    # list of iatom for each spin index.
+    uc_index_spin = [sc_excparams.index_spin[i] for i in ids]
+
+    uc_ind_atoms = {}
+    for iatom, ispin in enumerate(uc_index_spin):
+        if ispin >= 0:
+            uc_ind_atoms[ispin] = iatom
+
+    uc_Jdict = {}
+    uc_Rset = set()
+    for key, val in sc_excparams.exchange_Jdict.items():
+        R, ispin, jspin = key
+        iatom = sc_excparams.ind_atoms[ispin]
+        jatom = sc_excparams.ind_atoms[jspin]
+        uc_ispin = uc_index_spin[indmap[iatom]]
+        uc_jspin = uc_index_spin[indmap[jatom]]
+        uc_R = R @ sc_matrix + Rmap[jatom] - Rmap[iatom]
+        uc_R = tuple(uc_R)
+        if iatom in ids:
+            #print(f"{iatom=}, {indmap[iatom]=},{uc_ispin=}, {uc_jspin=}, {uc_R=} ")
+            uc_Jdict[(uc_R, uc_ispin, uc_jspin)] = val
+            uc_Rset.add(uc_R)
+
+    uc_distance_dict = gen_distance_dict(uc_ind_atoms, uc_atoms, list(uc_Rset))
+
+    assert sc_excparams.colinear, "Cut supercell for non-collinear spin is not yet implemented."
+
+    uc_exc = SpinIO(uc_atoms,
+                    spinat=uc_spinat,
+                    charges=uc_charges,
+                    index_spin=uc_index_spin,
+                    colinear=sc_excparams.colinear,
+                    distance_dict=uc_distance_dict,
+                    exchange_Jdict=uc_Jdict,
+                    description="Cutted")
+
+    uc_exc.write_all(output_path)
+
+
+def run_cut_cell(
+    path="TB2J_results",
+    output_path="./TB2J_cutted",
+    sc_matrix=np.array([[1, 1, 0], [-1, 1, 0], [0, 0, 2]]),
+):
+    cut_cell(path,
+             output_path,
+             np.array(sc_matrix),
+             origin_atom_id=0,
+             thr=1e-19)
+
+
+if __name__ == "__main__":
+    run_cut_cell()

TB2J/exchange.py

@@ -1,6 +1,9 @@
 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
@@ -10,25 +13,47 @@ from TB2J.external import p_map
 from TB2J.contour import Contour
 from TB2J.utils import simpson_nonuniform, trapezoidal_nonuniform, split_symbol_number
 from TB2J.orbmap import map_orbs_matrix
-import pickle
+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,
+        efermi=-10.0,
         basis=None,
-        magnetic_elements=[],
-        include_orbs={},
+        magnetic_elements=None,
+        include_orbs=None,
         kmesh=[4, 4, 4],
-        emin=-15,  # integration lower bound, relative to fermi energy
-        # integration upper bound. Should be 0 (fermi energy). But DFT codes define Fermi energy in various ways.
+        emin=-15,
         emax=0.05,
         nz=100,
-        # the delta in the (i delta) in green's function to prevent divergence
-        exclude_orbs=[],  #
-        ne=None,  # number of electrons in Wannier function.
-        Rcut=None,  # Rcut.
+        exclude_orbs=[],
+        ne=None,
+        Rcut=None,
         use_cache=False,
         np=1,
         description="",
@@ -37,70 +62,41 @@ class ExchangeParams:
         output_path="TB2J_results",
     ):
         self.efermi = efermi
-        self.emin = emin
-        self.emax = emax
-        self.nz = nz
-        self.Rcut = Rcut
         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._kmesh = kmesh
-        self.orb_decomposition = orb_decomposition
-        self.write_density_matrix = write_density_matrix
         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)
+
 
 class Exchange(ExchangeParams):
     def __init__(
         self,
         tbmodels,
         atoms,
-        efermi,
-        basis=None,
-        magnetic_elements=[],
-        include_orbs={},
-        kmesh=[4, 4, 4],
-        emin=-15,  # integration lower bound, relative to fermi energy
-        # integration upper bound. Should be 0 (fermi energy). But DFT codes define Fermi energy in various ways.
-        emax=0.05,
-        nz=100,
-        # the delta in the (i delta) in green's function to prevent divergence
-        exclude_orbs=[],  #
-        ne=None,  # number of electrons in Wannier function.
-        Rcut=None,  # Rcut.
-        use_cache=False,
-        np=1,
-        description="",
-        write_density_matrix=False,
-        output_path="TB2J_results",
-        orb_decomposition=False,
+        **params: ExchangeParams,
     ):
         self.atoms = atoms
-        super().__init__(
-            efermi=efermi,
-            basis=basis,
-            magnetic_elements=magnetic_elements,
-            include_orbs=include_orbs,
-            kmesh=kmesh,
-            emin=emin,
-            emax=emax,
-            nz=nz,
-            exclude_orbs=exclude_orbs,
-            ne=ne,
-            Rcut=Rcut,
-            use_cache=use_cache,
-            np=np,
-            description=description,
-            write_density_matrix=write_density_matrix,
-            orb_decomposition=orb_decomposition,
-            output_path=output_path,
-        )
+        super().__init__(**params)
         self._prepare_kmesh(self._kmesh)
         self._prepare_Rlist()
         self.set_tbmodels(tbmodels)
@@ -334,7 +330,6 @@ class ExchangeNCL(Exchange):
         """
         self.tbmodel = tbmodels
         self.backend_name = self.tbmodel.name
-        # TODO: check if tbmodels are really a tbmodel with SOC.
         self.G = TBGreen(
             self.tbmodel,
             self.kmesh,
@@ -355,6 +350,12 @@ class ExchangeNCL(Exchange):
         if self.write_density_matrix:
             self.G.write_rho_R()
 
+    def get_MAE(self, thetas, phis):
+        """
+        Calculate the magnetic anisotropy energy.
+        """
+        pass
+
     def _prepare_NijR(self):
         self.N = {}
         for R in self.Rlist:
@@ -615,7 +616,7 @@ class ExchangeNCL(Exchange):
     def calculate_DMI_NJT(self):
         """
         calculate exchange and DMI with the
-        D(i,j) =
+        This did not work and should be removed.
         """
         Ddict_NJT = {}
         Jdict_NJT = {}
@@ -679,10 +680,36 @@ class ExchangeNCL(Exchange):
                         self.contour.path, AijRs_orb[(R, iatom, jatom)]
                     )
 
-    def get_AijR(self, e):
-        GR = self.G.get_GR(self.short_Rlist, energy=e, get_rho=False)
+    def get_quantities_per_e(self, e):
+        Gk_all = self.G.get_Gk_all(e)
+        mae = self.get_mae_kspace(Gk_all)
+        # 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 AijR, AijR_orb
+        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)
@@ -693,6 +720,7 @@ class ExchangeNCL(Exchange):
         """
         Do some sanity check before proceding.
         """
+        pass
 
     def calculate_all(self):
         """
@@ -701,40 +729,38 @@ class ExchangeNCL(Exchange):
         print("Green's function Calculation started.")
 
         AijRs = {}
-
         AijRs_orb = {}
 
         self.validate()
 
         npole = len(self.contour.path)
         if self.np > 1:
-            results = p_map(self.get_AijR, self.contour.path, num_cpus=self.np)
+            results = p_map(
+                self.get_quantities_per_e, self.contour.path, num_cpus=self.np
+            )
         else:
-            results = map(self.get_AijR, tqdm(self.contour.path, total=npole))
+            results = map(
+                self.get_quantities_per_e, tqdm(self.contour.path, total=npole)
+            )
 
         for i, result in enumerate(results):
             for iR, R in enumerate(self.R_ijatom_dict):
                 for iatom, jatom in self.R_ijatom_dict[R]:
                     if (R, iatom, jatom) in AijRs:
-                        AijRs[(R, iatom, jatom)].append(result[0][R, iatom, jatom])
+                        AijRs[(R, iatom, jatom)].append(result["AijR"][R, iatom, jatom])
                         if self.orb_decomposition:
                             AijRs_orb[(R, iatom, jatom)].append(
-                                result[1][R, iatom, jatom]
+                                result["AijR_orb"][R, iatom, jatom]
                             )
 
                     else:
                         AijRs[(R, iatom, jatom)] = []
-                        AijRs[(R, iatom, jatom)].append(result[0][R, iatom, jatom])
+                        AijRs[(R, iatom, jatom)].append(result["AijR"][R, iatom, jatom])
                         if self.orb_decomposition:
                             AijRs_orb[(R, iatom, jatom)] = []
                             AijRs_orb[(R, iatom, jatom)].append(
-                                result[1][R, iatom, jatom]
+                                result["AijR_orb"][R, iatom, jatom]
                             )
-        if self.np > 1:
-            # executor.close()
-            # executor.join()
-            # executor.clear()
-            pass
 
         # self.save_AijRs(AijRs)
         self.integrate(AijRs, AijRs_orb)