TB2J 0.9.12.13__py3-none-any.whl → 0.9.12.15__py3-none-any.whl

TB2J/MAEGreen.py

@@ -51,7 +51,7 @@ class MAEGreen(ExchangeNCL):
 
         nangles = len(self.thetas)
         self.es = np.zeros(nangles, dtype=complex)
-        self.es_atom = np.zeros((nangles, self.natoms), dtype=complex)
+        self.es_matrix = np.zeros((nangles, self.natoms, self.natoms), dtype=complex)
         self.es_atom_orb = DefaultDict(lambda: 0)
 
     def set_angles_xyz(self):
@@ -153,7 +153,7 @@ class MAEGreen(ExchangeNCL):
         Hsoc_k = self.tbmodel.get_Hk_soc(self.G.kpts)
         na = len(thetas)
         dE_angle = np.zeros(na, dtype=complex)
-        dE_angle_atom = np.zeros((na, self.natoms), dtype=complex)
+        dE_angle_matrix = np.zeros((na, self.natoms, self.natoms), dtype=complex)
         # dE_angle_orbitals = np.zeros((na, self.natoms, self.norb, self.norb), dtype=complex)
         dE_angle_atom_orb = DefaultDict(lambda: 0)
         for iangle, (theta, phi) in enumerate(zip(thetas, phis)):
@@ -178,23 +178,30 @@ class MAEGreen(ExchangeNCL):
                 # dE_angle[iangle] += np.trace(GdH@GdH) * self.G.kweights[ik]
                 # dE_angle[iangle] += np.trace(GdH@G0K[ik].T.conj()@dHi ) * self.G.kweights[ik]
                 dE_angle[iangle] += dG2sum * self.G.kweights[ik]
+
+                # Calculate atom-atom matrix interactions
                 for iatom in range(self.natoms):
                     iorb = self.iorb(iatom)
-                    # dG2= dG2[::2, ::2] + dG2[1::2, 1::2] + dG2[1::2, ::2] + dG2[::2, 1::2]
-                    dE_atom_orb = dG2[np.ix_(iorb, iorb)] * self.G.kweights[ik]
-                    dE_atom_orb = (
-                        dE_atom_orb[::2, ::2]
-                        + dE_atom_orb[1::2, 1::2]
-                        + 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_angle_atom[iangle, iatom] += dE_atom
-        return dE_angle, dE_angle_atom, dE_angle_atom_orb
+                    for jatom in range(self.natoms):
+                        jorb = self.iorb(jatom)
+                        # Calculate cross terms between atoms i and j
+                        dE_ij_orb = dG2[np.ix_(iorb, jorb)] * self.G.kweights[ik]
+                        dE_ij_orb = (
+                            dE_ij_orb[::2, ::2]
+                            + dE_ij_orb[1::2, 1::2]
+                            + dE_ij_orb[1::2, ::2]
+                            + dE_ij_orb[::2, 1::2]
+                        )
+                        dE_ij = np.sum(dE_ij_orb)
+                        # Transform to local orbital basis
+                        mmat_i = self.mmats[iatom]
+                        mmat_j = self.mmats[jatom]
+                        dE_ij_orb = mmat_i.T @ dE_ij_orb @ mmat_j
+                        dE_angle_matrix[iangle, iatom, jatom] += dE_ij
+                        # Store orbital-resolved data for diagonal terms
+                        if iatom == jatom:
+                            dE_angle_atom_orb[(iangle, iatom)] += dE_ij_orb
+        return dE_angle, dE_angle_matrix, dE_angle_atom_orb
 
     def get_perturbed_R(self, e, thetas, phis):
         self.tbmodel.set_so_strength(0.0)
@@ -232,16 +239,16 @@ class MAEGreen(ExchangeNCL):
             npole = len(self.contour.path)
             results = map(func, tqdm.tqdm(self.contour.path, total=npole))
         for i, result in enumerate(results):
-            dE_angle, dE_angle_atom, dE_angle_atom_orb = result
+            dE_angle, dE_angle_matrix, dE_angle_atom_orb = result
             self.es += dE_angle * self.contour.weights[i]
-            self.es_atom += dE_angle_atom * self.contour.weights[i]
+            self.es_matrix += dE_angle_matrix * self.contour.weights[i]
             for key, value in dE_angle_atom_orb.items():
                 self.es_atom_orb[key] += (
                     dE_angle_atom_orb[key] * self.contour.weights[i]
                 )
 
         self.es = -np.imag(self.es) / (2 * np.pi)
-        self.es_atom = -np.imag(self.es_atom) / (2 * np.pi)
+        self.es_matrix = -np.imag(self.es_matrix) / (2 * np.pi)
         for key, value in self.es_atom_orb.items():
             self.es_atom_orb[key] = -np.imag(value) / (2 * np.pi)
 
@@ -259,6 +266,7 @@ class MAEGreen(ExchangeNCL):
         Path(output_path).mkdir(exist_ok=True)
         fname = f"{output_path}/MAE.dat"
         fname_orb = f"{output_path}/MAE_orb.dat"
+        fname_matrix = f"{output_path}/MAE_matrix.dat"
         # fname_tensor = f"{output_path}/MAE_tensor.dat"
         # if figure3d is not None:
         #    fname_fig3d = f"{output_path}/{figure3d}"
@@ -269,24 +277,32 @@ class MAEGreen(ExchangeNCL):
         if with_eigen:
             fname_eigen = f"{output_path}/MAE_eigen.dat"
             with open(fname_eigen, "w") as f:
-                f.write("# theta, phi, MAE(total), MAE(atom-wise) Unit: meV\n")
-                for i, (theta, phi, e, es) in enumerate(
-                    zip(self.thetas, self.phis, self.es2, self.es_atom)
+                f.write("# theta, phi, MAE(total) Unit: meV\n")
+                for i, (theta, phi, e) in enumerate(
+                    zip(self.thetas, self.phis, self.es2)
                 ):
-                    f.write(f"{theta:.5f} {phi:.5f} {e*1e3:.8f} ")
-                    for ea in es:
-                        f.write(f"{ea*1e3:.8f} ")
-                    f.write("\n")
+                    f.write(f"{theta:.5f} {phi:.5f} {e*1e3:.8f}\n")
 
         with open(fname, "w") as f:
-            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%np.pi:.5f} {phi%(2*np.pi):.5f} {e*1e3:.8f} ")
-                for ea in es:
-                    f.write(f"{ea*1e3:.8f} ")
-                f.write("\n")
+            f.write("# theta (rad), phi(rad), MAE(total) Unit: meV\n")
+            for i, (theta, phi, e) in enumerate(zip(self.thetas, self.phis, self.es)):
+                f.write(f"{theta%np.pi:.5f} {phi%(2*np.pi):.5f} {e*1e3:.8f}\n")
+
+        # Write matrix data to MAE_matrix.dat
+        with open(fname_matrix, "w") as fmat:
+            fmat.write("# MAE atom-atom interaction matrices\n")
+            fmat.write("# Format: angle_index theta phi atom_i atom_j MAE_ij(meV)\n")
+            fmat.write("# Units: theta and phi in radians, MAE in meV\n")
+            for iangle, (theta, phi) in enumerate(zip(self.thetas, self.phis)):
+                for iatom in range(self.natoms):
+                    for jatom in range(self.natoms):
+                        mae_ij = (
+                            self.es_matrix[iangle, iatom, jatom] * 1e3
+                        )  # Convert to meV
+                        fmat.write(
+                            f"{iangle:4d} {theta:.5f} {phi:.5f} {iatom:4d} {jatom:4d} {mae_ij:.8f}\n"
+                        )
+                fmat.write("\n")  # Empty line between angles for readability
 
         # self.ani = self.fit_anisotropy_tensor()
         # with open(fname_tensor, "w") as f:
@@ -312,34 +328,37 @@ class MAEGreen(ExchangeNCL):
                 for orb in self.orbital_names[iatom]:
                     f.write(f"{orb} ")
                 f.write("\n")
-            for i, (theta, phi, e, eatom) in enumerate(
-                zip(self.thetas, self.phis, self.es, self.es_atom)
-            ):
+            for i, (theta, phi, e) in enumerate(zip(self.thetas, self.phis, self.es)):
                 f.write("-" * 60 + "\n")
                 f.write(f"Angle {i:03d}:   theta={theta:.5f} phi={phi:.5f} \n ")
                 f.write(f"E: {e*1e3:.8f} \n")
-                for iatom, ea in enumerate(eatom):
-                    f.write(f"Atom {iatom:03d}: {ea*1e3:.8f} \n")
-                    f.write("Orbital: ")
-                    eorb = self.es_atom_orb[(i, iatom)]
-
-                    # write numpy matrix to file
-                    f.write(
-                        np.array2string(
-                            eorb * 1e3, precision=4, separator=",", suppress_small=True
-                        )
-                    )
-
-                    eorb_diff = eorb - self.es_atom_orb[(0, iatom)]
-                    f.write("Diference to the first angle: ")
-                    f.write(
-                        np.array2string(
-                            eorb_diff * 1e3,
-                            precision=4,
-                            separator=",",
-                            suppress_small=True,
+                for iatom in range(self.natoms):
+                    f.write(f"Atom {iatom:03d} orbital matrix:\n")
+                    if (i, iatom) in self.es_atom_orb:
+                        eorb = self.es_atom_orb[(i, iatom)]
+                        # write numpy matrix to file
+                        f.write(
+                            np.array2string(
+                                eorb * 1e3,
+                                precision=4,
+                                separator=",",
+                                suppress_small=True,
+                            )
                         )
-                    )
+                        f.write("\n")
+
+                        if (0, iatom) in self.es_atom_orb:
+                            eorb_diff = eorb - self.es_atom_orb[(0, iatom)]
+                            f.write("Difference to the first angle: ")
+                            f.write(
+                                np.array2string(
+                                    eorb_diff * 1e3,
+                                    precision=4,
+                                    separator=",",
+                                    suppress_small=True,
+                                )
+                            )
+                            f.write("\n")
                 f.write("\n")
 
     def run(self, output_path="TB2J_anisotropy", with_eigen=False):

TB2J/contour.py

@@ -27,8 +27,9 @@ class Contour:
         integrate f along the path
         """
         ret = 0
-        for i in range(len(values)):
-            ret += values[i] * self.weights[i]
+        ret = np.einsum("i...,i->...", values, self.weights)
+        # for i in range(len(values)):
+        #    ret += values[i] * self.weights[i]
         return ret
 
     def build_path_semicircle(self, npoints, endpoint=True):