merge upstraeam

TB2J/Jdownfolder.py

@@ -1,7 +1,9 @@
 import os
 from collections import defaultdict
+
 import numpy as np
 from ase.dft.kpoints import monkhorst_pack
+
 from TB2J.io_exchange import SpinIO
 from TB2J.Jtensor import decompose_J_tensor
 
@@ -73,8 +75,8 @@ def downfold_oneq(self, J):
 class PWFDownfolder:
     def __init__(self, JR, Rlist, iM, iL, qmesh, atoms=None, iso_only=False, **kwargs):
         from lawaf.interfaces.magnon.magnon_downfolder import (
-            MagnonWrapper,
             MagnonDownfolder,
+            MagnonWrapper,
         )
 
         model = MagnonWrapper(JR, Rlist, atoms)
@@ -92,13 +94,13 @@ def __init__(self, JR, Rlist, iM, iL, qmesh, atoms=None, iso_only=False, **kwarg
             # anchors={(0, 0, 0): (-1, -2, -3, -4)},
             # anchors={(0, 0, 0): ()},
             # use_proj=True,
-            enhance_Amn=2.0,
+            enhance_Amn=0.0,
         )
         params.update(kwargs)
         wann.set_parameters(**params)
         print("begin downfold")
         ewf = wann.downfold()
-        ewf.save_hr_pickle("downfolded_JR.pickle")
+        # ewf.save_hr_pickle("downfolded_JR.pickle")
 
         # Plot the band structure.
         wann.plot_band_fitting(
@@ -135,7 +137,7 @@ def __init__(
         qmesh=[7, 7, 7],
         iso_only=False,
         method="pwf",
-        **kwargs
+        **kwargs,
     ):
         self.exc = SpinIO.load_pickle(path=inpath, fname="TB2J.pickle")
 
@@ -193,7 +195,7 @@ def _downfold(self, **kwargs):
                 qmesh=self.qmesh,
                 atoms=self.atoms,
                 iso_only=self.iso_only,
-                **kwargs
+                **kwargs,
             )
             Jd, Rlist = d.get_JR()
         return Jd, Rlist
@@ -274,10 +276,10 @@ def _prepare_index_spin(self):
 def test():
     # pass
     # inpath = "/home/hexu/projects/NiCl2/vasp_inputs/TB2J_results"
-    # inpath = "/home/hexu/projects/TB2J_example/CrI3/TB2J_results"
-    inpath = "/home/hexu/projects/TB2J_projects/NiCl2/TB2J_NiCl/TB2J_results"
-    fname = os.path.join(inpath, "TB2J.pickle")
-    p = JDownfolder_pickle(
+    inpath = "/home/hexu/projects/TB2J_example/CrI3/TB2J_results"
+    # inpath = "/home/hexu/projects/TB2J_projects/NiCl2/TB2J_NiCl/TB2J_results"
+    _fname = os.path.join(inpath, "TB2J.pickle")
+    _p = JDownfolder_pickle(
         inpath=inpath, metals=["Ni"], ligands=["Cl"], outpath="TB2J_results_downfolded"
     )
 

TB2J/MAEGreen.py

@@ -1,13 +1,17 @@
+import gc
 from pathlib import Path
 
+import matplotlib.pyplot as plt
 import numpy as np
 import tqdm
 from HamiltonIO.abacus.abacus_wrapper import AbacusSplitSOCParser
 from HamiltonIO.model.occupations import GaussOccupations
 from typing_extensions import DefaultDict
 
+from TB2J.anisotropy import Anisotropy
 from TB2J.exchange import ExchangeNCL
 from TB2J.external import p_map
+from TB2J.mathutils.fibonacci_sphere import fibonacci_semisphere
 
 # from HamiltonIO.model.rotate_spin import rotate_Matrix_from_z_to_axis, rotate_Matrix_from_z_to_sperical
 # from TB2J.abacus.abacus_wrapper import AbacusWrapper, AbacusParser
@@ -23,12 +27,21 @@ def get_occupation(evals, kweights, nel, width=0.1):
 
 class MAEGreen(ExchangeNCL):
     def __init__(self, angles=None, **kwargs):
+        """
+        angles are defined as theta, phi pairs, where theta is the angle between the z-axis and the magnetization direction, and phi is the angle between the x-axis and the projection of the magnetization direction on the x-y plane.
+        """
         super().__init__(**kwargs)
         self.natoms = len(self.atoms)
         if angles is None or angles == "axis":
             self.set_angles_axis()
         elif angles == "scan":
             self.set_angles_scan()
+        elif angles == "fib":
+            self.set_angles_fib()
+        elif angles == "random":
+            self.set_angles_random()
+        elif angles == "miller":
+            self.set_angles_miller()
         else:
             self.thetas = angles[0]
             self.phis = angles[1]
@@ -39,9 +52,32 @@ def __init__(self, angles=None, **kwargs):
         self.es_atom_orb = DefaultDict(lambda: 0)
 
     def set_angles_axis(self):
-        """theta and phi are defined as the x, y, z, axis."""
-        self.thetas = [0, np.pi / 2, np.pi / 2, np.pi / 2, np.pi]
-        self.phis = [0, 0, np.pi / 2, np.pi / 4, 0]
+        """theta and phi are defined as the x, y, z, xy, yz, xz, xyz, x-yz, -xyz, -x-yz axis."""
+        self.thetas = [0, np.pi / 2, np.pi / 2, np.pi / 2, np.pi, 0, np.pi / 2, 0, 0, 0]
+        self.phis = [0, 0, np.pi / 2, np.pi / 4, 0, 0, 0, np.pi]
+
+    def set_angles_miller(self, nmax=2):
+        """theta and angles corresponding to the miller index. remove duplicates.
+        e.g. 002 and 001 are the same.
+        """
+        thetas = []
+        phis = []
+        for k in range(0, nmax + 1):
+            for j in range(-nmax, nmax + 1):
+                for i in range(-nmax, nmax + 1):
+                    if i == 0 and j == 0 and k == 0:
+                        continue
+                    thetas.append(np.arccos(k / np.sqrt(i**2 + j**2 + k**2)))
+                    if i == 0 and j == 0:
+                        phis.append(0)
+                    else:
+                        phis.append(np.arctan2(j, i))
+        self.thetas = thetas
+        self.phis = phis
+        self.angle_pairs = list(zip(thetas, phis))
+        # remove duplicates of angles using sets.
+        self.angle_pairs = list(set(self.angle_pairs))
+        self.thetas, self.phis = zip(*self.angle_pairs)
 
     def set_angles_scan(self, step=15):
         self.thetas = []
@@ -51,6 +87,21 @@ def set_angles_scan(self, step=15):
                 self.thetas.append(i * np.pi / 180)
                 self.phis.append(j * np.pi / 180)
 
+    def set_angles_random(self, n=16):
+        # n random pairs of theta, phi
+        self.thetas = np.random.random(n) * np.pi
+        self.phis = np.random.random(n) * 2 * np.pi
+
+    def set_angles_fib(self, n=35):
+        self.thetas, self.phis = fibonacci_semisphere(n)
+        # thetas and phis are np.array
+        # add (theta, phi): (pi/2, pi/2) and (pi/2, pi/4)
+        # self.thetas += [np.pi / 2, np.pi / 2]
+        # self.phis += [np.pi / 2, np.pi / 4]
+        for i in range(8):
+            self.thetas = np.concatenate([self.thetas, [np.pi / 2]])
+            self.phis = np.concatenate([self.phis, [np.pi * i / 8]])
+
     def get_band_energy_vs_angles_from_eigen(
         self,
         thetas,
@@ -79,16 +130,12 @@ def get_band_energy(self):
 
     def get_perturbed(self, e, thetas, phis):
         self.tbmodel.set_so_strength(0.0)
-        maxsoc = self.tbmodel.get_max_Hsoc_abs()
-        maxH0 = self.tbmodel.get_max_H0_spin_abs()
-        if maxsoc > maxH0 * 0.01:
-            print(f"""Warning: The SOC of the Hamiltonian has a maximum of {maxsoc} eV,
-                  comparing to the maximum of {maxH0} eV of the spin part of the Hamiltonian.
-                  The SOC is too strong, the perturbation theory may not be valid.""")
-
-        print(f"""Warning: The SOC of the Hamiltonian has a maximum of {maxsoc} eV,
-                  comparing to the maximum of {maxH0} eV of the spin part of the Hamiltonian.
-                  The SOC is too strong, the perturbation theory may not be valid.""")
+        # maxsoc = self.tbmodel.get_max_Hsoc_abs()
+        # maxH0 = self.tbmodel.get_max_H0_spin_abs()
+        # if maxsoc > maxH0 * 0.01:
+        #    print(f"""Warning: The SOC of the Hamiltonian has a maximum of {maxsoc} eV,
+        #          comparing to the maximum of {maxH0} eV of the spin part of the Hamiltonian.
+        #          The SOC is too strong, the perturbation theory may not be valid.""")
 
         G0K = self.G.get_Gk_all(e)
         Hsoc_k = self.tbmodel.get_Hk_soc(self.G.kpts)
@@ -104,12 +151,12 @@ def get_perturbed(self, e, thetas, phis):
                 GdH = G0K[ik] @ dHi
                 # dE += np.trace(GdH @ G0K[i].T.conj() @ dHi) * self.kweights[i]
                 # diagonal of second order perturbation.
-                dG2diag = np.diag(GdH @ GdH)
+                # dG2diag = np.diag(GdH @ GdH)
                 # dG2 = np.einsum("ij,ji->ij", GdH,   GdH)
                 dG2 = GdH * GdH.T
                 dG2sum = np.sum(dG2)
                 # print(f"dG2sum-sum: {dG2sum}")
-                dG2sum = np.sum(dG2diag)
+                # dG2sum = np.sum(dG2diag)
 
                 # dG2sum = np.trace(GdH @ GdH)
                 # print(f"dG2sum-Tr: {dG2sum}")
@@ -130,13 +177,11 @@ def get_perturbed(self, e, thetas, phis):
                         + dE_atom_orb[1::2, ::2]
                         + dE_atom_orb[::2, 1::2]
                     )
+                    dE_atom = np.sum(dE_atom_orb)
                     mmat = self.mmats[iatom]
                     dE_atom_orb = mmat.T @ dE_atom_orb @ mmat
-
                     dE_angle_atom_orb[(iangle, iatom)] += dE_atom_orb
-
-                    dE_atom = np.sum(dG2diag[iorb]) * self.G.kweights[ik]
-                    # dE_atom = np.sum(dE_atom_orb)
+                    # dE_atom = np.sum(dG2diag[iorb]) * self.G.kweights[ik]
                     dE_angle_atom[iangle, iatom] += dE_atom
         return dE_angle, dE_angle_atom, dE_angle_atom_orb
 
@@ -189,10 +234,25 @@ def func(e):
         for key, value in self.es_atom_orb.items():
             self.es_atom_orb[key] = -np.imag(value) / (2 * np.pi)
 
-    def output(self, output_path="TB2J_anisotropy", with_eigen=True):
+    def fit_anisotropy_tensor(self):
+        self.ani = Anisotropy.fit_from_data(self.thetas, self.phis, self.es)
+        return self.ani
+
+    def output(
+        self,
+        output_path="TB2J_anisotropy",
+        with_eigen=True,
+        figure3d="MAE_3d.png",
+        figure_contourf="MAE_contourf.png",
+    ):
         Path(output_path).mkdir(exist_ok=True)
         fname = f"{output_path}/MAE.dat"
         fname_orb = f"{output_path}/MAE_orb.dat"
+        fname_tensor = f"{output_path}/MAE_tensor.dat"
+        if figure3d is not None:
+            fname_fig3d = f"{output_path}/{figure3d}"
+        if figure_contourf is not None:
+            fname_figcontourf = f"{output_path}/{figure_contourf}"
 
         # ouput with eigenvalues.
         if with_eigen:
@@ -208,15 +268,29 @@ def output(self, output_path="TB2J_anisotropy", with_eigen=True):
                     f.write("\n")
 
         with open(fname, "w") as f:
-            f.write("# theta, phi, MAE(total), MAE(atom-wise) Unit: meV\n")
+            f.write("# theta (rad), phi(rad), MAE(total), MAE(atom-wise) Unit: meV\n")
             for i, (theta, phi, e, es) in enumerate(
                 zip(self.thetas, self.phis, self.es, self.es_atom)
             ):
-                f.write(f"{theta:.5f} {phi:.5f} {e*1e3:.8f} ")
+                f.write(f"{theta%np.pi:.5f} {phi%(2*np.pi):.5f} {e*1e3:.8f} ")
                 for ea in es:
                     f.write(f"{ea*1e3:.8f} ")
                 f.write("\n")
 
+        self.ani = self.fit_anisotropy_tensor()
+        with open(fname_tensor, "w") as f:
+            f.write("# Anisotropy tensor in meV\n")
+            f.write(f"{self.ani.tensor_strings(include_isotropic=False)}\n")
+
+        if figure3d is not None:
+            self.ani.plot_3d(figname=fname_fig3d, show=False)
+
+        if figure_contourf is not None:
+            self.ani.plot_contourf(figname=fname_figcontourf, show=False)
+
+        plt.close()
+        gc.collect()
+
         with open(fname_orb, "w") as f:
             f.write("=" * 80 + "\n")
             f.write("Orbitals for each atom: \n")