TB2J 0.9.0.3__py3-none-any.whl → 0.9.0.4__py3-none-any.whl

TB2J/lawaf/lawaf_interface.py

@@ -0,0 +1,433 @@
+import os
+from TB2J.myTB import MyTB, merge_tbmodels_spin
+import numpy as np
+from TB2J.exchange import ExchangeNCL
+from TB2J.exchangeCL2 import ExchangeCL2
+from TB2J.exchange_qspace import ExchangeCLQspace
+from TB2J.utils import read_basis, auto_assign_basis_name
+from ase.io import read
+from TB2J.sisl_wrapper import SislWrapper
+from TB2J.gpaw_wrapper import GPAWWrapper
+from TB2J.wannier import parse_atoms
+
+
+def gen_exchange(
+    path,
+    colinear=True,
+    groupby="spin",
+    posfile=None,
+    prefix_up="wannier90.up",
+    prefix_dn="wannier90.dn",
+    prefix_SOC="wannier90",
+    min_hopping_norm=1e-4,
+    max_distance=None,
+    efermi=0,
+    magnetic_elements=[],
+    kmesh=[4, 4, 4],
+    emin=-12.0,
+    emax=0.0,
+    nz=100,
+    exclude_orbs=[],
+    Rcut=None,
+    ne=None,
+    use_cache=False,
+    np=1,
+    output_path="TB2J_results",
+    wannier_type="wannier90",
+    qspace=False,
+    orb_decomposition=False,
+    write_density_matrix=False,
+    description="",
+):
+    try:
+        fname = os.path.join(path, posfile)
+        print(f"Reading atomic structure from file {fname}.")
+        atoms = read(os.path.join(path, posfile))
+    except Exception:
+        print(
+            f"Cannot read atomic structure from file {fname}. Trying to read from Wannier input."
+        )
+        if colinear:
+            fname = os.path.join(path, f"{prefix_up}.win")
+        else:
+            fname = os.path.join(path, f"{prefix_SOC}.win")
+
+        print(f"Reading atomic structure from file {fname}.")
+        atoms = parse_atoms(fname)
+
+    basis_fname = os.path.join(path, "basis.txt")
+    if colinear:
+        if wannier_type.lower() == "wannier90":
+            print("Reading Wannier90 hamiltonian: spin up.")
+            tbmodel_up = MyTB.read_from_wannier_dir(
+                path=path, prefix=prefix_up, atoms=atoms, nls=False
+            )
+            print("Reading Wannier90 hamiltonian: spin down.")
+            tbmodel_dn = MyTB.read_from_wannier_dir(
+                path=path, prefix=prefix_dn, atoms=atoms, nls=False
+            )
+            if os.path.exists(basis_fname):
+                basis = read_basis(basis_fname)
+            else:
+                basis, _ = auto_assign_basis_name(
+                    tbmodel_up.xred,
+                    atoms,
+                    write_basis_file=os.path.join(output_path, "assigned_basis.txt"),
+                )
+        elif wannier_type.lower() == "banddownfolder":
+            print("Reading Banddownfolder hamiltonian: spin up.")
+            tbmodel_up = MyTB.load_banddownfolder(
+                path=path, prefix=prefix_up, atoms=atoms, nls=False
+            )
+            print("Reading Banddownfolder hamiltonian: spin down.")
+            tbmodel_dn = MyTB.load_banddownfolder(
+                path=path, prefix=prefix_dn, atoms=atoms, nls=False
+            )
+
+            basis, _ = auto_assign_basis_name(
+                tbmodel_up.xred,
+                atoms,
+                write_basis_file=os.path.join(output_path, "assigned_basis.txt"),
+            )
+        else:
+            raise ValueError("wannier_type should be Wannier90 or banddownfolder.")
+
+        print("Starting to calculate exchange.")
+        description = f""" Input from collinear Wannier90 data.
+ Tight binding data from {path}. 
+ Prefix of wannier function files:{prefix_up} and {prefix_dn}.
+Warning: Please check if the noise level of Wannier function Hamiltonian to make sure it is much smaller than the exchange values.
+\n"""
+        if not qspace:
+            exchange = ExchangeCL2(
+                tbmodels=(tbmodel_up, tbmodel_dn),
+                atoms=atoms,
+                basis=basis,
+                efermi=efermi,
+                magnetic_elements=magnetic_elements,
+                kmesh=kmesh,
+                emin=emin,
+                emax=emax,
+                nz=nz,
+                exclude_orbs=exclude_orbs,
+                Rcut=Rcut,
+                ne=ne,
+                np=np,
+                use_cache=use_cache,
+                output_path=output_path,
+                write_density_matrix=write_density_matrix,
+                description=description,
+            )
+        else:
+            exchange = ExchangeCLQspace(
+                tbmodels=(tbmodel_up, tbmodel_dn),
+                atoms=atoms,
+                basis=basis,
+                efermi=efermi,
+                magnetic_elements=magnetic_elements,
+                kmesh=kmesh,
+                emin=emin,
+                emax=emax,
+                nz=nz,
+                exclude_orbs=exclude_orbs,
+                Rcut=Rcut,
+                ne=ne,
+                np=np,
+                use_cache=use_cache,
+                output_path=output_path,
+                write_density_matrix=write_density_matrix,
+                description=description,
+            )
+
+        exchange.run(path=output_path)
+        print("All calculation finsihed. The results are in TB2J_results directory.")
+
+    elif colinear and wannier_type.lower() == "banddownfolder":
+        print("Reading Wannier90 hamiltonian: spin up.")
+        tbmodel_up = MyTB.read_from_wannier_dir(
+            path=path, prefix=prefix_up, atoms=atoms, groupby=None, nls=False
+        )
+        print("Reading Wannier90 hamiltonian: spin down.")
+        tbmodel_dn = MyTB.read_from_wannier_dir(
+            path=path, prefix=prefix_dn, atoms=atoms, groupby=None, nls=False
+        )
+        tbmodel = merge_tbmodels_spin(tbmodel_up, tbmodel_dn)
+        basis, _ = auto_assign_basis_name(
+            tbmodel.xred,
+            atoms,
+            write_basis_file=os.path.join(output_path, "assigned_basis.txt"),
+        )
+        description = f""" Input from collinear BandDownfolder data.
+ Tight binding data from {path}. 
+ Prefix of wannier function files:{prefix_up} and {prefix_dn}.
+Warning: Please check if the noise level of Wannier function Hamiltonian to make sure it is much smaller than the exchange values.
+\n"""
+        print("Starting to calculate exchange.")
+        exchange = ExchangeCL2(
+            tbmodels=tbmodel,
+            atoms=atoms,
+            basis=basis,
+            efermi=efermi,
+            magnetic_elements=magnetic_elements,
+            kmesh=kmesh,
+            emin=emin,
+            emax=emax,
+            nz=nz,
+            exclude_orbs=exclude_orbs,
+            Rcut=Rcut,
+            ne=ne,
+            np=np,
+            use_cache=use_cache,
+            output_path=output_path,
+            write_density_matrix=write_density_matrix,
+            description=description,
+        )
+        exchange.run(path=output_path)
+        print("All calculation finsihed. The results are in TB2J_results directory.")
+    else:
+        print("Reading Wannier90 hamiltonian: non-colinear spin.")
+        groupby = groupby.lower().strip()
+        if groupby not in ["spin", "orbital"]:
+            raise ValueError("groupby can only be spin or orbital.")
+        tbmodel = MyTB.read_from_wannier_dir(
+            path=path, prefix=prefix_SOC, atoms=atoms, groupby=groupby, nls=True
+        )
+        if os.path.exists(basis_fname):
+            print("The use of basis file is deprecated. It will be ignored.")
+            # basis = read_basis(basis_fname)
+        else:
+            basis, _ = auto_assign_basis_name(
+                tbmodel.xred,
+                atoms,
+                write_basis_file=os.path.join(output_path, "assigned_basis.txt"),
+            )
+        description = f""" Input from non-collinear Wannier90 data.
+ Tight binding data from {path}. 
+ Prefix of wannier function files:{prefix_SOC}.
+Warning: Please check if the noise level of Wannier function Hamiltonian to make sure it is much smaller than the exchange values.
+ The DMI component parallel to the spin orientation, the Jani which has the component of that orientation should be disregarded
+ e.g. if the spins are along z, the xz, yz, zz, zx, zy components and the z component of DMI.
+ If you need these component, try to do three calculations with spin along x, y, z,  or use structure with z rotated to x, y and z. And then use TB2J_merge.py to get the full set of parameters.
+
+\n"""
+        print("Starting to calculate exchange.")
+        exchange = ExchangeNCL(
+            tbmodels=tbmodel,
+            atoms=atoms,
+            basis=basis,
+            efermi=efermi,
+            magnetic_elements=magnetic_elements,
+            kmesh=kmesh,
+            emin=emin,
+            emax=emax,
+            nz=nz,
+            exclude_orbs=exclude_orbs,
+            Rcut=Rcut,
+            ne=ne,
+            np=np,
+            use_cache=use_cache,
+            description=description,
+            output_path=output_path,
+            write_density_matrix=write_density_matrix,
+            orb_decomposition=orb_decomposition,
+        )
+        print("\n")
+        exchange.run(path=output_path)
+        print(f"All calculation finsihed. The results are in {output_path} directory.")
+
+
+def gen_exchange_siesta(
+    fdf_fname,
+    magnetic_elements=[],
+    include_orbs=None,
+    kmesh=[5, 5, 5],
+    emin=-12.0,
+    emax=0.0,
+    nz=100,
+    exclude_orbs=[],
+    Rcut=None,
+    ne=None,
+    np=1,
+    use_cache=False,
+    output_path="TB2J_results",
+    orb_decomposition=False,
+    description="",
+):
+    try:
+        import sisl
+    except:
+        raise ImportError("sisl cannot be imported. Please install sisl first.")
+
+    from packaging import version
+
+    if version.parse(sisl.__version__) <= version.parse("0.10.0"):
+        raise ImportError(
+            f"sisl version is {sisl.__version__}, but should be larger than 0.10.0."
+        )
+
+    include_orbs = {}
+    if isinstance(magnetic_elements, str):
+        magnetic_elements = [magnetic_elements]
+    for element in magnetic_elements:
+        if "_" in element:
+            elem = element.split("_")[0]
+            orb = element.split("_")[1:]
+            include_orbs[elem] = orb
+        else:
+            include_orbs[element] = None
+    magnetic_elements = list(include_orbs.keys())
+
+    fdf = sisl.get_sile(fdf_fname)
+    # geom = fdf.read_geometry()
+    H = fdf.read_hamiltonian()
+    geom = H.geometry
+    if H.spin.is_colinear:
+        print("Reading Siesta hamiltonian: colinear spin.")
+        tbmodel_up = SislWrapper(H, spin=0, geom=geom)
+        tbmodel_dn = SislWrapper(H, spin=1, geom=geom)
+        basis = dict(zip(tbmodel_up.orbs, list(range(tbmodel_up.norb))))
+        print("Starting to calculate exchange.")
+        description = f""" Input from collinear Siesta data.
+ working directory: {os.getcwd()}
+ fdf_fname: {fdf_fname}.
+\n"""
+        exchange = ExchangeCL2(
+            tbmodels=(tbmodel_up, tbmodel_dn),
+            atoms=tbmodel_up.atoms,
+            basis=basis,
+            efermi=0.0,
+            magnetic_elements=magnetic_elements,
+            include_orbs=include_orbs,
+            kmesh=kmesh,
+            emin=emin,
+            emax=emax,
+            nz=nz,
+            exclude_orbs=exclude_orbs,
+            Rcut=Rcut,
+            ne=ne,
+            np=np,
+            use_cache=use_cache,
+            output_path=output_path,
+            description=description,
+        )
+        exchange.run(path=output_path)
+        print("\n")
+        print(f"All calculation finsihed. The results are in {output_path} directory.")
+
+    elif H.spin.is_colinear and False:
+        print(
+            "Reading Siesta hamiltonian: colinear spin. Treat as non-colinear. For testing only."
+        )
+        tbmodel = SislWrapper(H, spin="merge", geom=geom)
+        basis = dict(zip(tbmodel.orbs, list(range(tbmodel.nbasis))))
+        print("Starting to calculate exchange.")
+        description = f""" Input from collinear Siesta data.
+ working directory: {os.getcwd()}
+ fdf_fname: {fdf_fname}.
+\n"""
+        exchange = ExchangeNCL(
+            tbmodels=tbmodel,
+            atoms=tbmodel.atoms,
+            basis=basis,
+            efermi=0.0,
+            magnetic_elements=magnetic_elements,
+            include_orbs=include_orbs,
+            kmesh=kmesh,
+            emin=emin,
+            emax=emax,
+            nz=nz,
+            exclude_orbs=exclude_orbs,
+            Rcut=Rcut,
+            ne=ne,
+            np=np,
+            use_cache=use_cache,
+            description=description,
+            output_path=output_path,
+            orb_decomposition=orb_decomposition,
+        )
+        exchange.run(path=output_path)
+        print("\n")
+        print(f"All calculation finsihed. The results are in {output_path} directory.")
+
+    elif H.spin.is_spinorbit or H.spin.is_noncolinear:
+        print("Reading Siesta hamiltonian: non-colinear spin.")
+        tbmodel = SislWrapper(H, spin=None, geom=geom)
+        basis = dict(zip(tbmodel.orbs, list(range(tbmodel.nbasis))))
+        print("Starting to calculate exchange.")
+        description = f""" Input from non-collinear Siesta data.
+ working directory: {os.getcwd()}
+ fdf_fname: {fdf_fname}.
+Warning: The DMI component parallel to the spin orientation, the Jani which has the component of that orientation should be disregarded
+ e.g. if the spins are along z, the xz, yz, zz, zx, zy components and the z component of DMI.
+ If you need these component, try to do three calculations with spin along x, y, z,  or use structure with z rotated to x, y and z. And then use TB2J_merge.py to get the full set of parameters.
+\n"""
+        exchange = ExchangeNCL(
+            tbmodels=tbmodel,
+            atoms=tbmodel.atoms,
+            basis=basis,
+            efermi=0.0,
+            magnetic_elements=magnetic_elements,
+            include_orbs=include_orbs,
+            kmesh=kmesh,
+            emin=emin,
+            emax=emax,
+            nz=nz,
+            exclude_orbs=exclude_orbs,
+            Rcut=Rcut,
+            ne=ne,
+            np=np,
+            use_cache=use_cache,
+            description=description,
+            output_path=output_path,
+            orb_decomposition=orb_decomposition,
+        )
+        exchange.run(path=output_path)
+        print("\n")
+        print(f"All calculation finsihed. The results are in {output_path} directory.")
+
+
+def gen_exchange_gpaw(
+    gpw_fname,
+    magnetic_elements=[],
+    kmesh=[3, 3, 3],
+    emin=-12.0,
+    emax=0.0,
+    nz=50,
+    exclude_orbs=[],
+    Rcut=None,
+    use_cache=False,
+    output_path="TB2J_results",
+    description="",
+):
+    print("Reading from GPAW data and calculate electronic structure.")
+    model = GPAWWrapper(gpw_fname=gpw_fname)
+    efermi = model.calc.get_fermi_level()
+    print(f"Fermi Energy: {efermi}")
+    poses = np.vstack([model.positions, model.positions])
+    basis, _ = auto_assign_basis_name(
+        poses,
+        model.atoms,
+        write_basis_file=os.path.join(output_path, "assigned_basis.txt"),
+    )
+
+    if model.calc.get_spin_polarized():
+        print("Starting to calculate exchange.")
+        exchange = ExchangeNCL(
+            tbmodels=model,
+            atoms=model.atoms,
+            efermi=efermi,
+            basis=basis,
+            magnetic_elements=magnetic_elements,
+            kmesh=kmesh,
+            emin=emin,
+            emax=emax,
+            nz=nz,
+            exclude_orbs=exclude_orbs,
+            Rcut=Rcut,
+            use_cache=use_cache,
+            output_path=output_path,
+            description=description,
+        )
+        exchange.run(path=output_path)
+        print("\n")
+        print(f"All calculation finsihed. The results are in {output_path} directory.")

TB2J/manager.py

@@ -9,6 +9,29 @@ from ase.io import read
 from TB2J.sisl_wrapper import SislWrapper
 from TB2J.gpaw_wrapper import GPAWWrapper
 from TB2J.wannier import parse_atoms
+from dataclasses import dataclass
+
+
+@dataclass
+class TB2JParams:
+    colinear: bool = True
+    efermi: float = 0
+    magnetic_elements = []
+    kmesh = [4, 4, 4]
+    emin: float = -12.0
+    emax: float = 0.0
+    nz: int = 100
+    exclude_orbs = []
+    Rcut: float = None
+    ne: int = None
+    use_cache: bool = False
+    np: int = 1
+    output_path: str = "TB2J_results"
+    wannier_type: str = "wannier90"
+    qspace: bool = False
+    orb_decomposition: bool = False
+    write_density_matrix: bool = False
+    description: str = ""
 
 
 def gen_exchange(
@@ -94,7 +117,7 @@ def gen_exchange(
 
         print("Starting to calculate exchange.")
         description = f""" Input from collinear Wannier90 data.
- Tight binding data from {path}. 
+ Tight binding data from {path}.
  Prefix of wannier function files:{prefix_up} and {prefix_dn}.
 Warning: Please check if the noise level of Wannier function Hamiltonian to make sure it is much smaller than the exchange values.
 \n"""
@@ -158,7 +181,7 @@ Warning: Please check if the noise level of Wannier function Hamiltonian to make
             write_basis_file=os.path.join(output_path, "assigned_basis.txt"),
         )
         description = f""" Input from collinear BandDownfolder data.
- Tight binding data from {path}. 
+ Tight binding data from {path}.
  Prefix of wannier function files:{prefix_up} and {prefix_dn}.
 Warning: Please check if the noise level of Wannier function Hamiltonian to make sure it is much smaller than the exchange values.
 \n"""
@@ -202,7 +225,7 @@ Warning: Please check if the noise level of Wannier function Hamiltonian to make
                 write_basis_file=os.path.join(output_path, "assigned_basis.txt"),
             )
         description = f""" Input from non-collinear Wannier90 data.
- Tight binding data from {path}. 
+ Tight binding data from {path}.
  Prefix of wannier function files:{prefix_SOC}.
 Warning: Please check if the noise level of Wannier function Hamiltonian to make sure it is much smaller than the exchange values.
  The DMI component parallel to the spin orientation, the Jani which has the component of that orientation should be disregarded

TB2J/patch.py

@@ -0,0 +1,50 @@
+import types
+from unittest import mock
+
+class A(object):#but seems to work for old style objects too
+    def funcx(self,x):
+        print("x1=",x)
+        print("called from", self)
+
+    def method(self,x):
+        print("xmethod=",x)
+        print("called from", self)
+
+def patch_me(target):
+    def method(target,x):
+        print("x=",x)
+        print("called from", target)
+        target.method = types.MethodType(method,target)
+
+def method(self,x):
+    print("x=",x)
+    print("called from", self)
+
+@mock.patch("__main__.A")
+def funcx(self,x):
+    print("new x=",x)
+    print("called from", self)
+        
+A.method=method
+#add more if needed
+a = A()
+print(A.__dict__)
+print(a)
+#out: <__main__.A object at 0x2b73ac88bfd0>  
+
+@mock.patch("__main__.a")
+def funcx(self,x):
+    print("x=",x)
+    print("called from", self)
+ 
+a.funcx(3)
+patch_me(a)    #patch instance
+a.method=method
+#a.method(5)
+#out: x= 5
+#out: called from <__main__.A object at 0x2b73ac88bfd0>
+patch_me(A)
+
+a.method(6)        #can patch class too
+#out: x= 6
+#out: called from <class '__main__.A'>

TB2J/spinham/h_matrix.py

@@ -0,0 +1,68 @@
+import numpy as np
+import matplotlib.pyplot as plt
+import aiida
+from aiida_tb2j.data import ExchangeData
+
+def plot_dispersion(bands, kpoint_labels, color='blue', title=None):
+
+    fig, axs = plt.subplots(1, 1, constrained_layout=True)
+    fig.set_size_inches(6, 6/1.618)
+
+    '''
+    Plot the bands
+    '''
+    kpoints = np.arange(len(bands))
+    axs.plot(kpoints, bands, color=color, linewidth=1.5)
+
+    '''
+    Plot the symmetry points
+    '''
+    bmin = bands.min(); bmax = bands.max()
+    ymin = bmin - 0.05*np.abs(bmin-bmax); ymax = bmax + 0.05*np.abs(bmax-bmin);
+    axs.set_xticks(kpoint_labels[0], kpoint_labels[1], fontsize=10)
+    axs.vlines(x=kpoint_labels[0], ymin=ymin, ymax=ymax, color='black', linewidth=0.5)
+    axs.set_xlim([0, len(kpoints)])
+    axs.set_ylim([ymin, ymax])
+
+    if title is not None:
+        plt.title(title, fontsize=10)
+
+    plt.show()
+
+
+if __name__ == "__main__":
+    
+    from argparse import ArgumentParser
+
+    parser = ArgumentParser()
+    parser.add_argument('-f', '--pickle_filename', type=str, help="Path of the 'TB2J.pickle' file.", required=True)
+    args = parser.parse_args()
+    
+    '''
+    Right now the implementation depends on AiiDA and so we must create and load an AiiDA profile,
+    even if we do not store any information on a data base.
+    '''
+    aiida.load_profile()
+    '''
+    Create an ExchangeData object with the informations from the TB2J.pickle file
+    '''
+    exchange = ExchangeData.load_tb2j(pickle_file=args.pickle_filename, isotropic=False, pbc=(True, True, True))
+    '''
+    Compute the magnon band structure along a high symmetry path generated with
+    the ASE package. The informations is stored in an AiiDA BandsData object.
+    Here tol is the symmetry tolerance to determine the space group of the system.
+    They are in units of eV
+    '''
+    magnon_data = exchange.get_magnon_bands(npoints=300, tol=1e-1, with_DMI=True, with_Jani=True)
+    magnon_bands = 1000*magnon_data.get_bands() # Convert to meV
+    raw_labels = [(k, '$\Gamma$') if s == 'GAMMA' else (k, s) for k, s in magnon_data.labels]
+    kpoint_labels = list( zip( *raw_labels ) )
+    plot_dispersion(magnon_bands, kpoint_labels, color='blue', title='Magnon Bands')
+    '''
+    We can also obtain the dynamical matrix h instead of the actual magnon bands. The result
+    is stored in a numpy array with shape (number of kpoints, 2*natoms, 2*natoms)
+    '''
+    kpoints = magnon_data.get_kpoints() # The shape of the kpoints must be (nkpoints, 3)
+    h_matrix = 1000*exchange._H_matrix(kpoints, with_DMI=True, with_Jani=True) # Convert to meV
+    h_dispersion = np.linalg.eigvalsh(h_matrix) # We can also get the eigenvectors with np.linalg.eigh
+    plot_dispersion(h_dispersion, kpoint_labels, color='red', title='h matrix dispersion')

TB2J/spinham/obtain_J.py

@@ -0,0 +1,79 @@
+import numpy as np
+from aiida_tb2j.data import ExchangeData
+from aiida_tb2j.data.exchange import get_rotation_arrays
+from itertools import combinations_with_replacement
+
+ux, uy, uz = np.eye(3).reshape((3, 1, 3))
+
+def combine_arrays(u, v):
+
+    return np.concatenate([u*v, np.roll(u, -1, axis=-1)*v, np.roll(v, -1, axis=-1)*u], axis=-1)
+
+def get_coefficients(magmoms, indices):
+
+    i, j = indices
+
+    U, V = zip(*[get_rotation_arrays(magmoms, u=u) for u in [ux, uy, uz]])
+    U = np.stack(U).swapaxes(0, 1)
+    V = np.stack(V).swapaxes(0, 1)
+
+    uc = combine_arrays(U[i], U[j].conj())
+    ur = combine_arrays(U[i], U[j])
+    uc2 = combine_arrays(U[i].conj(), U[j])
+    u = np.concatenate([uc, ur, uc2], axis=1)
+
+    return u, V
+
+def get_C(H0, u, V):
+
+    n = int(H0.shape[-1]/2)
+    upi = np.triu_indices(n)
+    dig = np.diag_indices(n)
+    
+    i, j = upi
+    AB0 = H0[:, [i, i, i+n], [j, j+n, j+n]]
+    AB0 = np.swapaxes(AB0, 0, 2).reshape(len(i), 9)
+    J0_flat = np.linalg.solve(u, AB0)
+
+    J0 = np.empty((n, n, 3, 3), dtype=complex)
+    J0[*upi] = J0_flat[:, [0, 6, 5, 3, 1, 7, 8, 4, 2]].reshape(-1, 3, 3)
+    J0 += J0.swapaxes(0, 1)
+    J0[*dig] = 0.0
+
+    C = np.array([np.diag(a) for a in np.einsum('imx,ijxy,jmy->mi', V, 2*J0, V)]) 
+
+    return C
+
+def get_J(H, kpoints, exchange):
+
+    n = int(H.shape[-1]/2)
+    upi = np.triu_indices(n)
+    dig = np.diag_indices(n)
+    i, j = upi
+
+    magmoms = exchange.magmoms()[np.unique(exchange.pairs)]
+    magmoms /= np.linalg.norm(magmoms, axis=-1).reshape(-1, 1)
+    u, V = get_coefficients(magmoms, indices=upi)
+
+    H0 = np.stack([1000*exchange._H_matrix(kpoints=np.zeros((1, 3)), with_DMI=True, with_Jani=True, u=u) for u in [ux, uy, uz]])[:, 0, :, :]
+    C = get_C(H0, u, V)
+    H[:, :, :n, :n] += C.reshape(3, 1, n, n)
+    H[:, :, n:, n:] += C.reshape(3, 1, n, n)
+
+    AB = H[:, :, [i, i, i+n], [j, j+n, j+n]]
+    AB[:, :, 2, :] = AB[:, ::-1, 2, :]
+    AB = np.moveaxis(AB, [2, 3], [1, 0]).reshape(len(i), 9, -1)
+
+    vectors = exchange.get_vectors()
+    exp_summand = np.exp( -2j*np.pi*vectors @ kpoints.T )
+    nAB = np.einsum('nik,ndk->nid', AB, exp_summand) / len(kpoints)
+    
+    ii = np.where(i == j)
+    i0 = np.where(np.linalg.norm(vectors, axis=-1) == 0.0)
+    J = np.linalg.solve(u, nAB).swapaxes(1, 2)
+    J = J[:, :, [0, 6, 5, 3, 1, 7, 8, 4, 2]].reshape(len(i), -1, 3, 3)
+    J *= -1
+    J[ii] *= 2
+    J[i0] *= 0
+
+    return J