TB2J 0.9.5rc0__py3-none-any.whl → 0.9.7rc0__py3-none-any.whl

TB2J/interfaces/siesta_interface.py

@@ -0,0 +1,202 @@
+import os
+
+import numpy as np
+
+from TB2J.exchange import ExchangeNCL
+from TB2J.exchangeCL2 import ExchangeCL2
+from TB2J.MAEGreen import MAEGreen
+
+try:
+    from HamiltonIO.siesta import SislParser
+except ImportError:
+    print(
+        "Cannot import SislWrapper from HamiltonIO.siesta. Please install HamiltonIO first."
+    )
+
+
+def siesta_anisotropy(**kwargs):
+    pass
+
+
+def gen_exchange_siesta(fdf_fname, read_H_soc=False, **kwargs):
+    """
+    parameters:
+        fdf_fname: str
+            The fdf file for the calculation.
+        read_H_soc: bool
+            Whether to read the SOC Hamiltonian. Default is False.
+
+    parameters in **kwargs:
+        magnetic_elements: list of str
+            The magnetic elements to calculate the exchange. e.g. ["Fe", "Co", "Ni"]
+        include_orbs: dict
+            The included orbitals for each magnetic element. e.g. {"Fe": ["d"]}. Default is None.
+        kmesh: list of int
+            The k-point mesh for the calculation. e.g. [5, 5, 5]. Default is [5, 5, 5].
+        emin: float
+            The minimum energy for the calculation. Default is -12.0.
+        emax: float
+            The maximum energy for the calculation. Default is 0.0.
+        nz: int
+            The number of points for the energy mesh. Default is 100.
+        exclude_orbs: list of str
+            The excluded orbitals for the calculation. Default is [].
+        Rcut: float
+            The cutoff radius for the exchange calculation in angstrom. Default is None.
+        ne: int
+            The number of electrons for the calculation. Default is None.
+        nproc: int
+            The number of processors for the calculation. Default is 1.
+        use_cache: bool
+            Whether to use the cache for the calculation. Default is False.
+        output_path: str
+            The output path for the calculation. Default is "TB2J_results".
+        orb_decomposition: bool
+            Whether to calculate the orbital decomposition. Default is False.
+        orth: bool
+            Whether to orthogonalize the orbitals. Default is False.
+        description: str
+            The description for the calculation. Default is "".
+    """
+
+    exargs = dict(
+        magnetic_elements=[],
+        include_orbs=None,
+        kmesh=[5, 5, 5],
+        emin=-12.0,
+        emax=0.0,
+        nz=100,
+        exclude_orbs=[],
+        Rcut=None,
+        ne=None,
+        nproc=1,
+        use_cache=False,
+        output_path="TB2J_results",
+        orb_decomposition=False,
+        orth=False,
+        description="",
+    )
+    exargs.update(kwargs)
+    try:
+        import sisl
+    except ImportError:
+        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."
+        )
+
+    magnetic_elements = exargs.pop("magnetic_elements")
+    include_orbs = exargs.pop("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())
+    output_path = exargs.pop("output_path")
+
+    parser = SislParser(
+        fdf_fname=fdf_fname, ispin=None, read_H_soc=read_H_soc, orth=exargs["orth"]
+    )
+    if parser.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)
+        tbmodel_up, tbmodel_dn = parser.get_model()
+        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,
+            **exargs,
+        )
+        exchange.run(path=output_path)
+        print("\n")
+        print(f"All calculation finished. The results are in {output_path} directory.")
+
+    elif parser.spin.is_spinorbit or parser.spin.is_noncolinear:
+        print("Reading Siesta hamiltonian: non-colinear spin.")
+        model = parser.get_model()
+        basis = dict(zip(model.orbs, list(range(model.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"""
+        exargs["description"] = description
+        if not model.split_soc:
+            exchange = ExchangeNCL(
+                tbmodels=model,
+                atoms=model.atoms,
+                basis=basis,
+                efermi=0.0,
+                magnetic_elements=magnetic_elements,
+                include_orbs=include_orbs,
+                output_path=output_path,
+                **exargs,
+            )
+            exchange.run(path=output_path)
+            print("\n")
+            print(
+                f"All calculation finished. The results are in {output_path} directory."
+            )
+        else:
+            print("Starting to calculate MAE.")
+            model.set_so_strength(0.0)
+            MAE = MAEGreen(
+                tbmodels=model,
+                atoms=model.atoms,
+                basis=basis,
+                efermi=None,
+                magnetic_elements=magnetic_elements,
+                include_orbs=include_orbs,
+                **exargs,
+            )
+            # thetas = [0, np.pi / 2, np.pi, 3 * np.pi / 2]
+            # phis = [0, 0, 0, 0]
+            # MAE.set_angles(thetas=thetas, phis=phis)
+            MAE.run(output_path=output_path)
+            print(
+                f"MAE calculation finished. The results are in {output_path} directory."
+            )
+
+            angle = {"x": (np.pi / 2, 0), "y": (np.pi / 2, np.pi / 2), "z": (0, 0)}
+            for key, val in angle.items():
+                # model = parser.get_model()
+                theta, phi = val
+                model.set_Hsoc_rotation_angle([theta, phi])
+                basis = dict(zip(model.orbs, list(range(model.nbasis))))
+                output_path_full = f"{output_path}_{key}"
+                exchange = ExchangeNCL(
+                    tbmodels=model,
+                    atoms=model.atoms,
+                    basis=basis,
+                    efermi=None,  # set to None, compute from efermi.
+                    magnetic_elements=magnetic_elements,
+                    include_orbs=include_orbs,
+                    **exargs,
+                )
+                exchange.run(path=output_path_full)
+                print("\n")
+                print(
+                    f"All calculation finished. The results are in {output_path_full} directory."
+                )

TB2J/interfaces/wannier90_interface.py

@@ -0,0 +1,115 @@
+import os
+
+from ase.io import read
+
+from TB2J.myTB import MyTB
+from TB2J.utils import auto_assign_basis_name
+from TB2J.wannier import parse_atoms
+
+from .manager import Manager
+
+
+class WannierManager(Manager):
+    def __init__(
+        self,
+        path,
+        prefix_up,
+        prefix_dn,
+        prefix_SOC,
+        colinear=True,
+        groupby="spin",
+        posfile=None,
+        basis_fname=None,
+        qspace=False,
+        wannier_type="wannier90",
+        **kwargs,
+    ):
+        # atoms
+        atoms = self.prepare_atoms(path, posfile, prefix_up, prefix_SOC, colinear)
+        output_path = kwargs.get("output_path", "TB2J_results")
+        # models and basis
+        if colinear:
+            tbmodels, basis = self.prepare_model_colinear(
+                path, prefix_up, prefix_dn, atoms, output_path=output_path
+            )
+        else:
+            tbmodels, basis = self.prepare_model_ncl(
+                path, prefix_SOC, atoms, groupby, output_path=output_path
+            )
+
+        description = self.description(path, prefix_up, prefix_dn, prefix_SOC, colinear)
+        kwargs["description"] = description
+
+        super().__init__(atoms, tbmodels, basis, colinear=colinear, **kwargs)
+
+    def prepare_atoms(self, path, posfile, prefix_up, prefix_SOC, colinear=True):
+        fname = os.path.join(path, posfile)
+        try:
+            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)
+        return atoms
+
+    def prepare_model_colinear(self, path, prefix_up, prefix_dn, atoms, output_path):
+        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
+        )
+        basis, _ = auto_assign_basis_name(
+            tbmodel_up.xred,
+            atoms,
+            write_basis_file=os.path.join(output_path, "assigned_basis.txt"),
+        )
+        tbmodels = (tbmodel_up, tbmodel_dn)
+        return tbmodels, basis
+
+    def prepare_model_ncl(self, path, prefix_SOC, atoms, groupby, output_path):
+        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
+        )
+        basis, _ = auto_assign_basis_name(
+            tbmodel.xred,
+            atoms,
+            write_basis_file=os.path.join(output_path, "assigned_basis.txt"),
+        )
+        return tbmodel, basis
+
+    def description(self, path, prefix_up, prefix_dn, prefix_SOC, colinear=True):
+        if colinear:
+            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"""
+
+        else:
+            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"""
+
+        return description
+
+
+gen_exchange = WannierManager

TB2J/io_exchange/io_exchange.py

@@ -9,16 +9,18 @@ write not only xml output.
 """
 
 import os
+import pickle
 from collections.abc import Iterable
+from datetime import datetime
+
+import matplotlib.pyplot as plt
 import numpy as np
-from TB2J.kpoints import monkhorst_pack
-import pickle
+
 from TB2J import __version__
+from TB2J.io_exchange.io_txt import write_Jq_info
 from TB2J.Jtensor import combine_J_tensor
-from datetime import datetime
-import matplotlib.pyplot as plt
+from TB2J.kpoints import monkhorst_pack
 from TB2J.spinham.spin_api import SpinModel
-from TB2J.io_exchange.io_txt import write_Jq_info
 from TB2J.utils import symbol_number
 
 
@@ -238,6 +240,18 @@ Generation time: {now.strftime("%y/%m/%d %H:%M:%S")}
     def get_charge_iatom(self, iatom):
         return self.charges[iatom]
 
+    def ijR_index_spin_to_atom(self, i, j, R):
+        return (self.iatom(i), self.iatom(j), R)
+
+    def ijR_index_atom_to_spin(self, iatom, jatom, R):
+        return (self.index_spin[iatom], self.index_spin[jatom], R)
+
+    def ijR_list(self):
+        return [(i, j, R) for R, i, j in self.exchange_Jdict]
+
+    def ijR_list_index_atom(self):
+        return [self.ijR_index_spin_to_atom(i, j, R) for R, i, j in self.exchange_Jdict]
+
     def get_J(self, i, j, R, default=None):
         i = self.i_spin(i)
         j = self.i_spin(j)
@@ -380,6 +394,7 @@ Generation time: {now.strftime("%y/%m/%d %H:%M:%S")}
         self.write_multibinit(path=os.path.join(path, "Multibinit"))
         self.write_tom_format(path=os.path.join(path, "TomASD"))
         self.write_vampire(path=os.path.join(path, "Vampire"))
+
         self.plot_all(savefile=os.path.join(path, "JvsR.pdf"))
         # self.write_Jq(kmesh=[9, 9, 9], path=path)
 
@@ -562,8 +577,8 @@ def gen_distance_dict(ind_mag_atoms, atoms, Rlist):
 
 
 def test_spin_io():
-    from ase import Atoms
     import numpy as np
+    from ase import Atoms
 
     atoms = Atoms(
         "SrMnO3",
@@ -579,7 +594,6 @@ def test_spin_io():
     spinat = [[0, 0, x] for x in [0, 3, 0, 0, 0]]
     charges = [2, 4, 5, 5, 5]
     index_spin = [-1, 0, -1, -1, -1]
-    colinear = True
     Rlist = [[0, 0, 0], [0, 0, 1]]
     ind_mag_atoms = [1]
     distance_dict = gen_distance_dict(ind_mag_atoms, atoms, Rlist)

TB2J/io_merge.py

@@ -48,7 +48,8 @@ class SpinIO_merge(SpinIO):
 
     def _set_projection_vectors(self):
 
-        spinat = self.spinat
+        norm = np.linalg.norm(self.spinat, axis=-1).reshape(-1, 1)
+        spinat = self.spinat / norm
         idx = [self.ind_atoms[i] for i in self.index_spin if i >= 0]
         projv = {}
         for i, j in combinations_with_replacement(range(self.nspin), 2):

TB2J/mathutils/fermi.py

@@ -1,8 +1,9 @@
-import numpy as np
-import warnings
 import sys
 
-MAX_EXP_ARGUMENT = np.log(sys.float_info.max)
+import numpy as np
+
+MAX_EXP_ARGUMENT = np.log(sys.float_info.max / 100000)
+print(MAX_EXP_ARGUMENT)
 
 
 def fermi(e, mu, width=0.01):
@@ -15,8 +16,12 @@ def fermi(e, mu, width=0.01):
     """
     x = (e - mu) / width
     # disable overflow warning
-    with warnings.catch_warnings():
-        warnings.simplefilter("ignore")
-        ret = np.where(x < MAX_EXP_ARGUMENT, 1 / (1.0 + np.exp(x)), 0.0)
+    # with warnings.catch_warnings():
+    #    warnings.simplefilter("ignore")
+    #    ret = np.where(x < MAX_EXP_ARGUMENT, 1 / (1.0 + np.exp(x)), 0.0)
 
+    ret = np.zeros_like(x, dtype=float)
+    for i, xi in enumerate(x):
+        if xi < 700:
+            ret[i] = 1 / (1.0 + np.exp(xi))
     return ret

TB2J/mathutils/lowdin.py

@@ -1,5 +1,5 @@
 import numpy as np
-from scipy.linalg import inv, eigh
+from scipy.linalg import eigh
 
 
 def Lowdin(S):
@@ -9,4 +9,14 @@ def Lowdin(S):
     psi_prime = S^(-1/2) psi
     """
     eigval, eigvec = eigh(S)
-    return eigvec @ np.diag(np.sqrt(1.0 / eigval)) @ (eigvec.T.conj())
+    S_half = eigvec @ np.diag(np.sqrt(1.0 / eigval)) @ (eigvec.T.conj())
+    return S_half
+
+
+def Lowdin_symmetric_orthonormalization(H, S):
+    """
+    Lowdin's symmetric orthonormalization.
+    """
+    S_half = Lowdin(S)
+    H_prime = S_half @ H @ S_half
+    return H_prime

TB2J/mathutils/rotate_spin.py

@@ -1,5 +1,7 @@
 import numpy as np
-from TB2J.pauli import pauli_block_all, s0, s1, s2, s3, gather_pauli_blocks
+from scipy.sparse import eye_array, kron
+
+from TB2J.pauli import gather_pauli_blocks, pauli_block_all
 
 
 def rotate_Matrix_from_z_to_axis(M, axis, normalize=True):
@@ -27,12 +29,224 @@ def spherical_to_cartesian(theta, phi, normalize=True):
     return vec
 
 
+def rotation_matrix(theta, phi):
+    """
+    The unitray operator U, that U^dagger * s3 * U is the rotated s3 by theta and phi
+    """
+    U = np.array(
+        [
+            [np.cos(theta / 2), np.exp(-1j * phi) * np.sin(theta / 2)],
+            [-np.exp(1j * phi) * np.sin(theta / 2), np.cos(theta / 2)],
+        ]
+    )
+    return U
+
+
+def rotate_spinor_single_block(M, theta, phi):
+    """
+    Rotate the spinor matrix M by theta and phi
+    """
+    U = rotation_matrix(theta, phi)
+    Uinv = np.linalg.inv(U)
+    return Uinv @ M @ U
+
+
+def rotate_spinor_matrix(M, theta, phi, method="einsum"):
+    """
+    Rotate the spinor matrix M by theta and phi,
+    """
+    if method == "plain":
+        return rotate_spinor_matrix_plain(M, theta, phi)
+    elif method == "einsum":
+        return rotate_spinor_matrix_einsum(M, theta, phi)
+    elif method == "reshape":
+        return rotate_spinor_matrix_reshape(M, theta, phi)
+    elif method == "kron":
+        return rotate_spinor_matrix_kron(M, theta, phi)
+    elif method == "spkron":
+        return rotate_spinor_matrix_spkron(M, theta, phi)
+    else:
+        raise ValueError(f"Unknown method: {method}")
+
+
+def rotate_spinor_matrix_plain(M, theta, phi):
+    """
+    M is a matrix with shape (2N, 2N), where N is the number of sites, and each site has a 2x2 matrix
+    rotate each 2x2 block by theta and phi
+    """
+    Mnew = np.zeros_like(M)
+    U = rotation_matrix(theta, phi)
+    UT = U.conj().T
+    tmp = np.zeros((2, 2), dtype=np.complex128)
+    for i in range(M.shape[0] // 2):
+        for j in range(M.shape[0] // 2):
+            for k in range(2):
+                for l in range(2):
+                    tmp[k, l] = M[2 * i + k, 2 * j + l]
+            # tmp[:, :]=M[2*i:2*i+2, 2*j:2*j+2]
+            Mnew[2 * i : 2 * i + 2, 2 * j : 2 * j + 2] = UT @ tmp @ U
+    return Mnew
+
+
+def rotate_spinor_matrix_einsum(M, theta, phi):
+    """
+    Rotate the spinor matrix M by theta and phi,
+    """
+    N = M.shape[0] // 2
+    Mnew = np.reshape(M, (N, 2, N, 2))  # .swapaxes(1, 2)
+    # print("Mnew:", Mnew)
+    U = rotation_matrix(theta, phi)
+    UT = U.conj().T
+    Mnew = np.einsum(
+        "ij, rjsk, kl -> risl", UT, Mnew, U, optimize=True, dtype=np.complex128
+    )
+    Mnew = Mnew.reshape(2 * N, 2 * N)
+    return Mnew
+
+
+def rotate_spinor_matrix_einsum_R(M, theta, phi):
+    """
+    Rotate the spinor matrix M by theta and phi,
+    """
+    nR = M.shape[0]
+    N = M.shape[1] // 2
+    Mnew = np.reshape(M, (nR, N, 2, N, 2))  # .swapaxes(1, 2)
+    # print("Mnew:", Mnew)
+    U = rotation_matrix(theta, phi)
+    UT = U.conj().T
+    Mnew = np.einsum(
+        "ij, nrjsk, kl -> nrisl", UT, Mnew, U, optimize=True, dtype=np.complex128
+    )
+    Mnew = Mnew.reshape(nR, 2 * N, 2 * N)
+    return Mnew
+
+
+def rotate_spinor_Matrix_R(M, theta, phi):
+    return rotate_spinor_matrix_einsum_R(M, theta, phi)
+
+
+def rotate_spinor_matrix_reshape(M, theta, phi):
+    """
+    Rotate the spinor matrix M by theta and phi,
+    """
+    N = M.shape[0] // 2
+    Mnew = np.reshape(M, (N, 2, N, 2)).swapaxes(1, 2)
+    # print("Mnew:", Mnew)
+    U = rotation_matrix(theta, phi)
+    UT = U.conj().T
+    Mnew = UT @ Mnew @ U
+    Mnew = Mnew.swapaxes(1, 2).reshape(2 * N, 2 * N)
+    return Mnew
+
+
+def rotate_spinor_matrix_kron(M, theta, phi):
+    """ """
+    U = rotation_matrix(theta, phi)
+    # U = np.kron( U, np.eye(M.shape[0]//2))
+    # U = kron(eye_array(M.shape[0]//2), U)
+    U = np.kron(np.eye(M.shape[0] // 2), U)
+    M = U.conj().T @ M @ U
+    return M
+
+
+def rotate_spinor_matrix_spkron(M, theta, phi):
+    """ """
+    U = rotation_matrix(theta, phi)
+    # U = np.kron( U, np.eye(M.shape[0]//2))
+    U = kron(eye_array(M.shape[0] // 2), U)
+    # U = np.kron(np.eye(M.shape[0]//2), U)
+    M = U.conj().T @ M @ U
+    return M
+
+
+def test_rotate_spinor_M():
+    N = 2
+    M_re = np.random.rand(N * 2, N * 2)
+    M_im = np.random.rand(N * 2, N * 2)
+    M = M_re + 1j * M_im
+    M = M + M.T.conj()
+    # M=np.array([[1, 0], [0, 1]], dtype=np.complex128)
+    print(f"Original M: {M}")
+
+    import timeit
+
+    print("Time for rotate_spinor_matrix")
+    print(
+        timeit.timeit(lambda: rotate_spinor_matrix(M, np.pi / 2, np.pi / 2), number=10)
+    )
+    print("Time for rotate_spinor_matrix_einsum")
+    print(
+        timeit.timeit(
+            lambda: rotate_spinor_matrix_einsum(M, np.pi / 2, np.pi / 2), number=10
+        )
+    )
+    print("Time for rotate_spinor_matrix_reshape")
+    print(
+        timeit.timeit(
+            lambda: rotate_spinor_matrix_reshape(M, np.pi / 2, np.pi / 2), number=10
+        )
+    )
+    print("Time for rotate_spinor_matrix_kron")
+    print(
+        timeit.timeit(
+            lambda: rotate_spinor_matrix_kron(M, np.pi / 2, np.pi / 2), number=10
+        )
+    )
+    print("Time for rotate_spinor_matrix_spkron")
+    print(
+        timeit.timeit(
+            lambda: rotate_spinor_matrix_spkron(M, np.pi / 2, np.pi / 2), number=10
+        )
+    )
+
+    Mrot1 = rotate_spinor_matrix(M, np.pi / 2, np.pi / 2)
+    Mrot2 = rotate_spinor_matrix_einsum(M, np.pi / 2, np.pi / 2)
+    Mrot3 = rotate_spinor_matrix_reshape(M, np.pi / 2, np.pi / 2)
+    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}")
+
+    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}")
+
+
+def test_rotate_spinor_oneblock():
+    M = np.array([[1.1, 0], [0, 0.9]])
+    print(np.array(pauli_block_all(M)).ravel())
+    print("Rotate by pi/2, pi/2 (z to y)")
+    Mnew = rotate_spinor_matrix_einsum(M, np.pi / 2, np.pi / 2)
+    print(np.array(pauli_block_all(Mnew)).ravel())
+
+    print("Rotate by pi/2, 0 (z to x)")
+    Mnew = rotate_spinor_matrix_kron(M, np.pi / 2, 0)
+
+    print(np.array(pauli_block_all(Mnew)).ravel())
+
+    print(Mnew)
+
+
 def rotate_Matrix_from_z_to_spherical(M, theta, phi, normalize=True):
     """
     Given a spinor matrix M, rotate it from z-axis to spherical coordinates
     """
-    axis = spherical_to_cartesian(theta, phi, normalize)
-    return rotate_Matrix_from_z_to_axis(M, axis, normalize)
+    # axis = spherical_to_cartesian(theta, phi, normalize)
+    # return rotate_Matrix_from_z_to_axis(M, axis, normalize)
+    return rotate_spinor_matrix_einsum(M, theta, phi)
+
+
+def test_rotate_Matrix_from_z_to_spherical():
+    M_re = np.random.rand(2, 2)
+    M_im = np.random.rand(2, 2)
+    M = M_re + 1j * M_im
+    print(M)
+    M_rot = rotate_Matrix_from_z_to_spherical(M, 0, 0)
+    print(M_rot)
 
 
 def test_rotate_Matrix_from_z_to_axis():
@@ -53,4 +267,8 @@ def test_rotate_Matrix_from_z_to_axis():
 
 
 if __name__ == "__main__":
-    test_rotate_Matrix_from_z_to_axis()
+    # test_rotate_Matrix_from_z_to_axis()
+    # test_rotate_Matrix_from_z_to_spherical()
+    test_rotate_spinor_M()
+    # test_rotate_spinor_oneblock()
+    pass