TB2J 0.9.4rc0__py3-none-any.whl → 0.9.6rc0__py3-none-any.whl

TB2J/manager.py

@@ -1,441 +0,0 @@
-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
-try:
-    from HamiltonIO.siesta import SislWrapper
-except ImportError:
-    print(
-        "Cannot import SislWrapper from HamiltonIO.siesta. Please install HamiltonIO first."
-    )
-    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(fdf_fname=None, sisl_hamiltonian=H, spin=0, geom=geom)
-        tbmodel_dn = SislWrapper(fdf_fname=None, sisl_hamiltonian=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(fdf_fname=None, sisl_hamiltonian=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/mathutils.py

@@ -1,12 +0,0 @@
-import numpy as np
-from scipy.linalg import inv, eigh
-
-
-def Lowdin(S):
-    """
-    Calculate S^(-1/2).
-    Which is used in lowind's symmetric orthonormalization.
-    psi_prime = S^(-1/2) psi
-    """
-    eigval, eigvec = eigh(S)
-    return eigvec @ np.diag(np.sqrt(1.0 / eigval)) @ (eigvec.T.conj())

TB2J/patch.py

@@ -1,50 +0,0 @@
-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

@@ -1,68 +0,0 @@
-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

@@ -1,79 +0,0 @@
-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