TB2J 0.9.12.9__py3-none-any.whl → 0.9.12.22__py3-none-any.whl

TB2J/io_exchange/io_txt.py

@@ -1,4 +1,5 @@
 import os
+import sys
 
 import numpy as np
 from numpy import array_str
@@ -9,6 +10,7 @@ from TB2J.utils import symbol_number
 def write_info_section(cls, myfile):
     myfile.write("=" * 90 + "\n")
     myfile.write("Information: \n")
+    myfile.write(f"Executed:  \n{' '.join(sys.argv)}\n")
     # myfile.write("Exchange parameters generated by TB2J %s\n" % (__version__))
     # myfile.write("Generation time: %s\n" % (now.strftime("%y/%m/%d %H:%M:%S")))
     myfile.write(cls.description)
@@ -82,6 +84,24 @@ def write_atom_section(cls, myfile):
             )
         )
 
+    # write single ion anisotropy
+    if cls.k1 is not None and cls.k1dir is not None:
+        myfile.write("\n")
+        myfile.write("-" * 90 + "\n")
+        myfile.write("Single Ion Anisotropy (meV): \n")
+        myfile.write("(k1 is the anisotropy constant, k1dir is the direction vector)\n")
+        myfile.write("{:^12s} {:^12s} {:^24s}\n".format("Atom number", "k1", "k1dir"))
+        for i, s in enumerate(symnum):
+            ispin = cls.index_spin[i]
+            if ispin >= 0:
+                k1 = cls.k1[ispin] * 1e3
+                k1dir = cls.k1dir[ispin]
+                myfile.write(
+                    "{:<12s} {:12.4f} ({:7.4f}, {:7.4f}, {:7.4f})\n".format(
+                        s, k1, k1dir[0], k1dir[1], k1dir[2]
+                    )
+                )
+
     myfile.write("\n")
 
 
@@ -159,17 +179,17 @@ def write_exchange_section(
 
         if cls.has_biquadratic and write_experimental:
             Jprime, B = cls.biquadratic_Jdict[ll]
-            myfile.write(f"[Testing!] Jprime: {Jprime*1e3:.3f},  B: {B*1e3:.3f}\n")
+            myfile.write(f"[Testing!] Jprime: {Jprime * 1e3:.3f},  B: {B * 1e3:.3f}\n")
 
         if cls.dJdx is not None:
             dJdx = cls.dJdx[ll]
-            myfile.write(f"dJ/dx: {dJdx*1e3:.3f}\n")
+            myfile.write(f"dJ/dx: {dJdx * 1e3:.3f}\n")
 
         if cls.dJdx2 is not None:
             dJdx2 = cls.dJdx2[ll]
-            myfile.write(f"d2J/dx2: {dJdx2*1e3:.3f}\n")
+            myfile.write(f"d2J/dx2: {dJdx2 * 1e3:.3f}\n")
 
-        if cls.dmi_ddict is not None:
+        if cls.dmi_ddict is not None and ll in cls.dmi_ddict:
             DMI = cls.dmi_ddict[ll] * 1e3
             myfile.write(
                 "[Testing!] DMI: ({:7.4f} {:7.4f} {:7.4f})\n".format(
@@ -188,13 +208,13 @@ def write_exchange_section(
             except Exception as e:
                 myfile.write(f"[Debug!] DMI2 not available: {e}\n")
 
-        if cls.Jani_dict is not None:
+        if cls.Jani_dict is not None and ll in cls.Jani_dict:
             J = cls.Jani_dict[ll] * 1e3
             myfile.write(
                 f"[Testing!]J_ani:\n{array_str(J, precision=3, suppress_small=True)}\n"
             )
 
-        if cls.NJT_ddict is not None:
+        if cls.NJT_ddict is not None and ll in cls.NJT_ddict:
             DMI = cls.NJT_ddict[ll] * 1e3
             myfile.write(
                 "[Experimental!] DMI_NJt: ({:7.4f} {:7.4f} {:7.4f})\n".format(
@@ -202,7 +222,7 @@ def write_exchange_section(
                 )
             )
 
-        if cls.NJT_Jdict is not None:
+        if cls.NJT_Jdict is not None and ll in cls.NJT_Jdict:
             J = cls.NJT_Jdict[ll] * 1e3
             myfile.write(
                 "[Testing!] Jani_NJt: ({:7.4f} {:7.4f} {:7.4f})\n".format(
@@ -305,6 +325,50 @@ def write_Jq_info(cls, kpts, evals, evecs, myfile, special_kpoints={}):
         print(f"{sns[cls.ind_atoms[i]]}, {v[0]: 8.3f}, {v[2]: 8.3f}, {v[2]: 8.3f}\n")
 
 
+def write_sia_section(cls, myfile):
+    """
+    Write single ion anisotropy tensor section.
+    """
+    if cls.has_sia_tensor and cls.sia_tensor is not None:
+        myfile.write("\n")
+        myfile.write("=" * 90 + "\n")
+        myfile.write("Single Ion Anisotropy: \n")
+        symnum = symbol_number(cls.atoms)
+        sns = list(symnum.keys())
+        myfile.write(
+            "{:^12s} {:^13s} {:^13s} {:^13s}\n".format(
+                "Atom_number", "A_xx", "A_xy", "A_xz"
+            )
+        )
+        myfile.write(
+            "{:^12s} {:^13s} {:^13s} {:^13s}\n".format("", "A_yx", "A_yy", "A_yz")
+        )
+        myfile.write(
+            "{:^12s} {:^13s} {:^13s} {:^13s}\n".format("", "A_zx", "A_zy", "A_zz")
+        )
+        for i, tensor in cls.sia_tensor.items():
+            iatom = cls.ind_atoms[i]
+            myfile.write(
+                "{:<12s} {:13.8f} {:13.8f} {:13.8f}\n".format(
+                    sns[iatom],
+                    tensor[0, 0] * 1e3,
+                    tensor[0, 1] * 1e3,
+                    tensor[0, 2] * 1e3,
+                )
+            )
+            myfile.write(
+                "{:<12s} {:13.8f} {:13.8f} {:13.8f}\n".format(
+                    "", tensor[1, 0] * 1e3, tensor[1, 1] * 1e3, tensor[1, 2] * 1e3
+                )
+            )
+            myfile.write(
+                "{:<12s} {:13.8f} {:13.8f} {:13.8f}\n".format(
+                    "", tensor[2, 0] * 1e3, tensor[2, 1] * 1e3, tensor[2, 2] * 1e3
+                )
+            )
+            myfile.write("-" * 88 + "\n")
+
+
 def write_txt(
     cls,
     path="TB2J_results",
@@ -321,6 +385,7 @@ def write_txt(
         write_info_section(cls, myfile)
         write_atom_section(cls, myfile)
         write_orbital_section(cls, myfile)
+        write_sia_section(cls, myfile)
         write_exchange_section(
             cls,
             myfile,

TB2J/io_exchange/io_vampire.py

@@ -50,7 +50,7 @@ def write_vampire_unitcell_file(cls, fname):
         counter = -1
         for key in cls.exchange_Jdict:
             R, ispin, jspin = key
-            Jtensor = cls.get_J_tensor(ispin, jspin, R)
+            Jtensor = cls.get_J_tensor(ispin, jspin, R, Jani=True, DMI=True)
             counter += 1  # starts at 0
             myfile.write(
                 f"{counter:5d} {ispin:3d} {jspin:3d} {R[0]:3d} {R[1]:3d} {R[2]:3d} "
@@ -66,7 +66,7 @@ def write_vampire_unitcell_file(cls, fname):
 
 def write_vampire_mat_file(cls, fname):
     mat_tmpl = """#---------------------------------------------------
-# Material {id} 
+# Material {id}
 #---------------------------------------------------
 material[{id}]:material-name={name}
 material[{id}]:damping-constant={damping}
@@ -75,6 +75,8 @@ material[{id}]:uniaxial-anisotropy-constant={k1}
 material[{id}]:material-element={name}
 material[{id}]:initial-spin-direction = {spinat}
 material[{id}]:uniaxial-anisotropy-direction = {k1dir}
+# The following line is required for vampire 5 and later. 
+material[{id}]:unit-cell-category =  {id}
 #---------------------------------------------------
 """
     with open(fname, "w") as myfile:

TB2J/io_merge.py

@@ -107,22 +107,23 @@ class Merger:
         projv = {}
         coeff_matrix = {}
         projectors = {}
-        for key in self.main_dat.projv.keys():
-            vectors = [obj.projv[key] for obj in self.dat]
-            coefficients, u, v = zip(
-                *[get_Jani_coefficients(vectors[i], R=R[i]) for i in indices]
-            )
-            projectors[key] = np.vstack([u[i] @ R[i].T for i in indices])
-            coeff_matrix[key] = np.vstack(coefficients)
-            proju[key] = np.stack(u)
-            projv[key] = np.stack(v)
-            if np.linalg.matrix_rank(coeff_matrix[key], tol=1e-2) < 6:
-                warnings.warn("""
-                    WARNING: The matrix of equations to reconstruct the exchange tensors is
-                    close to being singular. This happens when the magnetic moments between
-                    different configurations are cloes to being parallel. You need to consider 
-                    more rotated spin configurations, otherwise the results might have a large 
-                    error.""")
+        if self.main_dat.projv is not None:
+            for key in self.main_dat.projv.keys():
+                vectors = [obj.projv[key] for obj in self.dat]
+                coefficients, u, v = zip(
+                    *[get_Jani_coefficients(vectors[i], R=R[i]) for i in indices]
+                )
+                projectors[key] = np.vstack([u[i] @ R[i].T for i in indices])
+                coeff_matrix[key] = np.vstack(coefficients)
+                proju[key] = np.stack(u)
+                projv[key] = np.stack(v)
+                if np.linalg.matrix_rank(coeff_matrix[key], tol=1e-2) < 6:
+                    warnings.warn("""
+                        WARNING: The matrix of equations to reconstruct the exchange tensors is
+                        close to being singular. This happens when the magnetic moments between
+                        different configurations are cloes to being parallel. You need to consider 
+                        more rotated spin configurations, otherwise the results might have a large 
+                        error.""")
 
         self.proju = proju
         self.projv = projv
@@ -134,13 +135,23 @@ class Merger:
         proju = self.proju
         projv = self.projv
         coeff_matrix = self.coeff_matrix
-        for key in self.main_dat.Jani_dict.keys():
+        for key in self.main_dat.exchange_Jdict.keys():
             try:
                 R, i, j = key
                 u = proju[i, j]
                 v = projv[i, j]
-                Jani = np.stack([sio.Jani_dict[key] for sio in self.dat])
-                projections = np.einsum("nmi,nij,nmj->nm", u, Jani, v).flatten()
+                projections = []
+                for idx, sio in enumerate(self.dat):
+                    if sio.has_bilinear and key in sio.Jani_dict:
+                        Jmat = sio.get_J_tensor(
+                            i, j, R, Jiso=False, Jani=True, DMI=False
+                        )
+                    else:
+                        #J_scalar = sio.exchange_Jdict.get(key, 0.0)
+                        Jmat = np.eye(3) * 0.0
+                    proj = np.einsum("mi,ij,mj->m", u[idx], Jmat, v[idx])
+                    projections.append(proj)
+                projections = np.concatenate(projections)
             except KeyError as err:
                 raise KeyError(
                     "Can not find key: %s, Please make sure the three calculations use the same k-mesh and same Rcut."
@@ -158,28 +169,36 @@ class Merger:
 
     def merge_Jiso(self):
         Jdict = {}
-        for key in self.main_dat.exchange_Jdict.keys():
-            try:
-                J = np.mean([obj.exchange_Jdict[key] for obj in self.dat])
-            except KeyError as err:
-                raise KeyError(
-                    "Can not find key: %s, Please make sure the three calculations use the same k-mesh and same Rcut."
-                    % err
-                )
-            Jdict[key] = J
-        self.main_dat.exchange_Jdict = Jdict
+        if self.main_dat.exchange_Jdict is not None:
+            for key in self.main_dat.exchange_Jdict.keys():
+                try:
+                    J = np.mean([obj.exchange_Jdict[key] for obj in self.dat])
+                except KeyError as err:
+                    raise KeyError(
+                        "Can not find key: %s, Please make sure the three calculations use the same k-mesh and same Rcut."
+                        % err
+                    )
+                Jdict[key] = J
+            self.main_dat.exchange_Jdict = Jdict
 
     def merge_DMI(self):
         dmi_ddict = {}
-        if all(obj.has_dmi for obj in self.dat):
+        if any(obj.has_dmi for obj in self.dat):
             projectors = self.projectors
             proju = self.proju
-            for key in self.main_dat.dmi_ddict.keys():
+            for key in self.main_dat.exchange_Jdict.keys():
                 try:
                     R, i, j = key
                     u = proju[i, j]
-                    DMI = np.stack([sio.dmi_ddict[key] for sio in self.dat])
-                    projections = np.einsum("nmi,ni->nm", u, DMI).flatten()
+                    projections = []
+                    for idx, sio in enumerate(self.dat):
+                        if sio.has_dmi and key in sio.dmi_ddict:
+                            DMI = sio.dmi_ddict[key]
+                        else:
+                            DMI = np.zeros(3)
+                        proj = np.einsum("mi,i->m", u[idx], DMI)
+                        projections.append(proj)
+                    projections = np.concatenate(projections)
                 except KeyError as err:
                     raise KeyError(
                         "Can not find key: %s, Please make sure the three calculations use the same k-mesh and same Rcut."
@@ -193,13 +212,14 @@ class Merger:
         # make sure that the Jani has the trace of zero
         Jdict = self.main_dat.exchange_Jdict
         Jani_dict = self.main_dat.Jani_dict
-        for key in self.main_dat.Jani_dict.keys():
-            Jani = self.main_dat.Jani_dict[key]
-            shift = np.trace(Jani) / 3.0
-            Jani_dict[key] -= shift * np.eye(3)
-            Jdict[key] += shift
-        self.main_dat.Jani_dict = Jani_dict
-        self.main_dat.exchange_Jdict = Jdict
+        if Jani_dict is not None:
+            for key in Jani_dict.keys():
+                Jani = Jani_dict[key]
+                shift = np.trace(Jani) / 3.0
+                Jani_dict[key] -= shift * np.eye(3)
+                Jdict[key] += shift
+            self.main_dat.Jani_dict = Jani_dict
+            self.main_dat.exchange_Jdict = Jdict
 
 
 def merge(*paths, main_path=None, save=True, write_path="TB2J_results"):

TB2J/magnon/magnon3.py

@@ -28,6 +28,7 @@ class MagnonParameters:
     Jiso: bool = True
     Jani: bool = False
     DMI: bool = False
+    SIA: bool = True
     Q: Optional[List[float]] = None
     uz_file: Optional[str] = None
     n: Optional[List[float]] = None
@@ -398,11 +399,18 @@ class Magnon:
         cell = exc.atoms.get_cell()
         pbc = exc.atoms.get_pbc()
 
+        # Extract SIA from kwargs, default True if not present
+        include_SIA = kwargs.pop("SIA", True)
+
+        JR = exc.get_full_Jtensor_for_Rlist(
+            order="ij33", asr=False, SIA=include_SIA, **kwargs
+        )
+
         return cls(
             nspin=exc.nspin,
             magmom=magmoms,
             Rlist=exc.Rlist,
-            JR=exc.get_full_Jtensor_for_Rlist(order="ij33", asr=False, **kwargs),
+            JR=JR,
             cell=cell,
             _Q=np.zeros(3),  # Default propagation vector
             _uz=np.array([[0.0, 0.0, 1.0]]),  # Default quantization axis
@@ -654,7 +662,11 @@ def plot_magnon_bands_from_TB2J(
     # Load magnon calculator from TB2J results
     print(f"Loading exchange parameters from {params.path}...")
     magnon = Magnon.from_TB2J_results(
-        path=params.path, Jiso=params.Jiso, Jani=params.Jani, DMI=params.DMI
+        path=params.path,
+        Jiso=params.Jiso,
+        Jani=params.Jani,
+        DMI=params.DMI,
+        SIA=params.SIA,
     )
 
     # Set reference vectors if provided
@@ -815,6 +827,18 @@ def plot_magnon_bands_cli():
         default=False,
         help="Include anisotropic exchange interactions (default: False)",
     )
+    parser.add_argument(
+        "--SIA",
+        action="store_true",
+        default=True,
+        help="Include single ion anisotropy (default: True)",
+    )
+    parser.add_argument(
+        "--no-SIA",
+        action="store_false",
+        dest="SIA",
+        help="Exclude single ion anisotropy",
+    )
     parser.add_argument(
         "-d",
         "--DMI",
@@ -881,6 +905,7 @@ def plot_magnon_bands_cli():
             filename=args.output,
             Jiso=args.Jiso,
             Jani=args.Jani,
+            SIA=args.SIA,
             DMI=args.DMI,
             Q=args.Q if args.Q is not None else None,
             uz_file=args.uz_file,

TB2J/mathutils/rotate_spin.py

@@ -92,7 +92,7 @@ def rotate_spinor_matrix_einsum(M, theta, phi):
     Rotate the spinor matrix M by theta and phi,
     """
     shape = M.shape
-    n1 = np.product(shape[:-1]) // 2
+    n1 = np.prod(shape[:-1]) // 2
     n2 = M.shape[-1] // 2
     Mnew = np.reshape(M, (n1, 2, n2, 2))  # .swapaxes(1, 2)
     U = rotation_matrix(theta, phi)
@@ -203,15 +203,15 @@ def test_rotate_spinor_M():
     Mrot4 = rotate_spinor_matrix_kron(M, np.pi / 2, np.pi / 2)
     Mrot5 = rotate_spinor_matrix_spkron(M, np.pi / 2, np.pi / 2)
     print(f"Rotated M with jit:\n {Mrot1}")
-    print(f"Rotated M with einsum:\n {Mrot2-Mrot1}")
-    print(f"Rotated M with reshape:\n {Mrot3-Mrot1}")
-    print(f"Rotated M with kron:\n {Mrot4-Mrot1}")
-    print(f"Rotated M with spkron:\n {Mrot5-Mrot1}")
+    print(f"Rotated M with einsum:\n {Mrot2 - Mrot1}")
+    print(f"Rotated M with reshape:\n {Mrot3 - Mrot1}")
+    print(f"Rotated M with kron:\n {Mrot4 - Mrot1}")
+    print(f"Rotated M with spkron:\n {Mrot5 - Mrot1}")
 
     M_rot00 = rotate_spinor_matrix(M, 0, 0)
     M_rot00_sph = rotate_Matrix_from_z_to_spherical(M, 0, 0)
-    print(f"Rotated M with theta=0, phi=0 compared with M:\n {M_rot00-M}")
-    print(f"Rotated M with theta=0, phi=0 compared with M:\n {M_rot00_sph-M}")
+    print(f"Rotated M with theta=0, phi=0 compared with M:\n {M_rot00 - M}")
+    print(f"Rotated M with theta=0, phi=0 compared with M:\n {M_rot00_sph - M}")
 
 
 def test_rotate_spinor_oneblock():

TB2J/myTB.py

@@ -1,15 +1,17 @@
-import os
-import numpy as np
 import copy
-from scipy.linalg import eigh
-from scipy.sparse import csr_matrix
-from scipy.io import netcdf_file
+import os
 from collections import defaultdict
 
+import numpy as np
+
 # from tbmodels import Model
 from ase.atoms import Atoms
+from scipy.io import netcdf_file
+from scipy.linalg import eigh
+from scipy.sparse import csr_matrix
+
 from TB2J.utils import auto_assign_basis_name
-from TB2J.wannier import parse_ham, parse_xyz, parse_atoms, parse_tb
+from TB2J.wannier import parse_atoms, parse_ham, parse_tb, parse_xyz
 
 
 class AbstractTB:
@@ -187,9 +189,7 @@ class MyTB(AbstractTB):
         if os.path.exists(tb_fname):
             xcart, nbasis, data, R_degens = parse_tb(fname=tb_fname)
         else:
-            nbasis, data, R_degens = parse_ham(
-                fname=os.path.join(path, prefix + "_hr.dat")
-            )
+            nbasis, data, R_degens = parse_ham(fname=hr_fname)
             xcart, _, _ = parse_xyz(fname=os.path.join(path, prefix + "_centres.xyz"))
 
         if atoms is None:
@@ -237,7 +237,7 @@ class MyTB(AbstractTB):
         return m
 
     def gen_ham(self, k, convention=2):
-        """
+        r"""
         generate hamiltonian matrix at k point.
         H_k( i, j)=\sum_R H_R(i, j)^phase.
         There are two conventions,
@@ -451,7 +451,7 @@ class MyTB(AbstractTB):
         nbasis = root.dimensions["nbasis"]
         nspin = root.dimensions["nspin"]
         ndim = root.dimensions["ndim"]
-        natom = root.dimensions["natom"]
+        _natom = root.dimensions["natom"]
         Rlist = root.variables["R"][:]
         mdata_real = root.variables["data_real"][:]
         mdata_imag = root.variables["data_imag"][:]

TB2J/mycfr.py

@@ -5,21 +5,18 @@
 Continued fraction representation.
 """
 
+import ase.units.kB as kb
 import numpy as np
 from scipy.linalg import eig
 
-kb = 8.61733e-5  # Boltzmann constant in eV
-
 
 class CFR:
-    """
-    Integration with the continued fraction representation.
-    """
+    """Integration with the continued fraction representation."""
 
-    def __init__(self, nz=50, T=60):
+    def __init__(self, nz: int = 50, T: float = 200.0):
         self.nz = nz
-        self.T = 600
-        self.beta = 1 / (kb * self.T)
+        self.T = float(T)
+        self.beta = 1.0 / (kb * self.T)
         self.Rinf = 1e10
         if nz <= 0:
             raise ValueError("nz should be larger than 0.")
@@ -65,8 +62,8 @@ class CFR:
         # Rp = 0.25 * np.diag(eigv)**2 * zp **2
 
         # print the poles and the weights
-        for i in range(len(self.poles)):
-            print("Pole: ", self.poles[i], "Weight: ", self.weights[i])
+        # for i in range(len(self.poles)):
+        #    print("Pole: ", self.poles[i], "Weight: ", self.weights[i])
 
         # add a point to the poles: 1e10j
         self.path = np.concatenate((self.path, [self.Rinf * 1j]))
@@ -110,7 +107,10 @@ def test_cfr():
 
     r = cfr.integrate_func(test_gf, ef=2)
     r = -np.imag(r) / np.pi * 2
-    print(r)
+    if not np.close(r, -1.0):
+        raise AssertionError(
+            f"Test failed, the integration should be close to -1.0 but get {r}"
+        )
     return r
 
 

TB2J/pauli.py

@@ -59,7 +59,7 @@ def pauli_decomp2(M):
     )
 
 
-def pauli_sigma_norm(M):
+def pauli_sigma_norm_old(M):
     MI, Mx, My, Mz = pauli_decomp2(M)
     return np.linalg.norm([Mx, My, Mz])
 
@@ -134,7 +134,7 @@ def pauli_block(M, idim):
     return tmp
 
 
-def pauli_block_all(M):
+def pauli_block_all_old(M):
     MI = (M[::2, ::2] + M[1::2, 1::2]) / 2.0
     Mx = (M[::2, 1::2] + M[1::2, ::2]) / 2.0
     # Note that this is not element wise product with sigma_y but dot product
@@ -143,6 +143,28 @@ def pauli_block_all(M):
     return MI, Mx, My, Mz
 
 
+def pauli_block_all(array):
+    A00 = array[..., ::2, ::2]
+    A01 = array[..., ::2, 1::2]
+    A10 = array[..., 1::2, ::2]
+    A11 = array[..., 1::2, 1::2]
+    n2 = array.shape[-1] // 2
+
+    out_dtype = np.result_type(array.dtype, np.complex64)
+    block = np.empty((*array.shape[:-2], 4, n2, n2), dtype=out_dtype)
+
+    np.add(A00, A11, out=block[..., 0, :, :])
+    block[..., 0, :, :] *= 0.5
+    np.add(A01, A10, out=block[..., 1, :, :])
+    block[..., 1, :, :] *= 0.5
+    np.subtract(A01, A10, out=block[..., 2, :, :])
+    block[..., 2, :, :] *= 0.5j
+    np.subtract(A00, A11, out=block[..., 3, :, :])
+    block[..., 3, :, :] *= 0.5
+
+    return block
+
+
 def gather_pauli_blocks(MI, Mx, My, Mz, coeffs=[1.0, 1.0, 1.0, 1.0]):
     """
     Gather the I, x, y, z component of a matrix.
@@ -212,3 +234,11 @@ def pauli_block_sigma_norm(M):
     evec = np.array((ex, ey, ez))
     evec = evec / np.linalg.norm(evec)
     return Mx * evec[0] + My * evec[1] + Mz * evec[2]
+
+
+def pauli_sigma_norm(array):
+    block = pauli_block_all(array)[..., 1:, :, :]
+    E = np.trace(block, axis1=-2, axis2=-1)
+    E /= np.linalg.norm(E, axis=-1, keepdims=True)
+    np.multiply(block, E[..., None, None], out=block)
+    return block.sum(axis=-3)

TB2J/plot.py

@@ -14,6 +14,32 @@ def write_eigen(qmesh, gamma=True, path="./", output_fname="EigenJq.txt", **kwar
     m.write_Jq(kmesh=qmesh, path=path, gamma=gamma, output_fname=output_fname, **kwargs)
 
 
+def plot_exchange(
+    input_data,
+    ax=None,
+    marker="o",
+    fname=None,
+    show=False,
+    by_species=True,
+    **kwargs,
+):
+    """Plot exchange parameters vs distance grouped by species.
+    input_data: SpinIO object or path to the pickle file/directory.
+    """
+    if isinstance(input_data, str):
+        if os.path.isdir(input_data):
+            m = SpinIO.load_pickle(input_data)
+        else:
+            m = SpinIO.load_pickle(
+                os.path.dirname(input_data), os.path.basename(input_data)
+            )
+    else:
+        m = input_data
+    return m.plot_JvsR(
+        ax=ax, marker=marker, fname=fname, show=show, by_species=by_species, **kwargs
+    )
+
+
 def plot_magnon_band(
     fname="exchange.xml",
     path="./",

TB2J/rotate_atoms.py

@@ -1,7 +1,8 @@
 #!/usr/bin/env python3
 import copy
-from ase.io import read, write
+
 import numpy as np
+from ase.io import read, write
 from TB2J.tensor_rotate import Rzx, Rzy, Rzz
 
 
@@ -15,14 +16,16 @@ def rotate_atom_xyz(atoms, noncollinear=False):
     will be generated.
     """
 
+    yield atoms
     rotation_axes = [(1, 0, 0), (0, 1, 0)]
     if noncollinear:
         rotation_axes += [(1, 1, 0), (1, 0, 1), (0, 1, 1)]
-    
+
     for axis in rotation_axes:
         rotated_atoms = copy.deepcopy(atoms)
         rotated_atoms.rotate(90, axis, rotate_cell=True)
         yield rotated_atoms
+    yield atoms
 
 
 def rotate_atom_spin_one_rotation(atoms, Rotation):
@@ -109,7 +112,7 @@ def rotate_xyz(fname, ftype="xyz", noncollinear=False):
     rotated = rotate_atom_xyz(atoms, noncollinear=noncollinear)
 
     for i, rotated_atoms in enumerate(rotated):
-        write(f"atoms_{i+1}.{ftype}", rotated_atoms)
-    write(f"atoms_0.{ftype}", atoms)
-
-    print(f"The output has been written to the atoms_i.{ftype} files. atoms_0.{ftype} contains the reference structure.")
+        write(f"atoms_{i}.{ftype}", rotated_atoms)
+    print(
+        f"The output has been written to the atoms_i.{ftype} files. atoms_0.{ftype} contains the reference structure."
+    )