TB2J 0.9.9.4__py3-none-any.whl → 0.9.9.6__py3-none-any.whl

TB2J/magnon/magnon3.py

@@ -0,0 +1,334 @@
+from dataclasses import dataclass
+
+import numpy as np
+from scipy.spatial.transform import Rotation
+
+from ..io_exchange import SpinIO
+from ..mathutils import Hermitize, get_rotation_arrays
+from .plot import BandsPlot
+
+
+@dataclass
+class Magnon:
+    """
+    Magnon calculator implementation using dataclass
+    """
+
+    nspin: int
+    # ind_atoms: list
+    magmom: np.ndarray
+    Rlist: np.ndarray
+    JR: np.ndarray
+    _Q: np.ndarray = np.array([0.0, 0.0, 0.0], dtype=float)
+    _uz: np.ndarray = np.array([[0.0, 0.0, 1.0]], dtype=float)
+    _n: np.ndarray = np.array([0, 0, 1], dtype=float)
+
+    def set_reference(self, Q, uz, n):
+        """
+        Set reference propagation vector and quantization axis
+
+        Parameters
+        ----------
+        Q : array_like
+            Propagation vector
+        uz : array_like
+            Quantization axis
+        n : array_like
+            Normal vector for rotation
+        """
+        self.set_propagation_vector(Q)
+        self._uz = np.array(uz, dtype=float)
+        self._n = np.array(n, dtype=float)
+
+    def set_propagation_vector(self, Q):
+        """Set propagation vector"""
+        self._Q = np.array(Q)
+
+    @property
+    def Q(self):
+        """Get propagation vector"""
+        if self._Q is None:
+            raise ValueError("Propagation vector Q is not set.")
+        return self._Q
+
+    @Q.setter
+    def Q(self, value):
+        if not isinstance(value, (list, np.ndarray)):
+            raise TypeError("Propagation vector Q must be a list or numpy array.")
+        if len(value) != 3:
+            raise ValueError("Propagation vector Q must have three components.")
+        self._Q = np.array(value)
+
+    def Jq(self, kpoints):
+        """
+        Compute the exchange interactions in reciprocal space.
+
+        The exchange interactions J(q) are computed using the Fourier transform:
+        J(q) = ∑_R J(R) exp(iq·R)
+
+        Array shapes and indices:
+        - kpoints: (nkpt, 3) array of k-points
+        - Rlist: (nR, 3) array of real-space lattice vectors
+        - JR: (nR, nspin, nspin, 3, 3) array of exchange tensors in real space
+            where nspin is number of magnetic atoms
+        - Output Jq: (nkpt, nspin, nspin, 3, 3) array of exchange tensors in q-space
+
+        If propagation vector Q is set, each J(R) is rotated before the Fourier transform:
+        J'_mn(R) = R_m(ϕ)^T J(R) R_n(ϕ)
+        where ϕ = 2π R·Q and R(ϕ) is rotation matrix around axis n by angle ϕ
+
+        Parameters
+        ----------
+        kpoints : array_like (nkpt, 3)
+            k-points at which to evaluate the exchange interactions
+
+        Returns
+        -------
+        numpy.ndarray (nkpt, nspin, nspin, 3, 3)
+            Exchange interaction tensors J(q) at each k-point
+            First two indices are for magnetic atom pairs
+            Last two indices are for 3x3 tensor components
+        """
+        Rlist = np.array(self.Rlist)
+        JR = self.JR
+        JRprime = JR.copy()
+
+        for iR, R in enumerate(Rlist):
+            if self._Q is not None:
+                # Rotate exchange tensors based on propagation vector
+                phi = 2 * np.pi * R @ self._Q  # angle ϕ = 2π R·Q
+                rv = phi * self._n  # rotation vector
+                Rmat = Rotation.from_rotvec(rv).as_matrix()
+                # J'_mn(R) = R_m(ϕ)^T J(R) R_n(ϕ) using Einstein summation.
+                # Here m is always in the R=0, thus the rotation is only applied on the
+                # n , so only on the right.
+                JRprime[iR] = np.einsum(" rijxy, yb -> rijab", JR[iR], Rmat)
+
+        nkpt = kpoints.shape[0]
+        Jq = np.zeros((nkpt, self.nspin, self.nspin, 3, 3), dtype=complex)
+
+        for iR, R in enumerate(Rlist):
+            for iqpt, qpt in enumerate(kpoints):
+                # Fourier transform of exchange tensors
+                phase = 2 * np.pi * R @ qpt
+                Jq[iqpt] += np.exp(1j * phase) * JRprime[iR]
+
+        # Ensure Hermiticity: J(q) = J(-q)†
+        for iqpt in range(nkpt):
+            Jq[iqpt, :, :, :, :] += np.conj(
+                np.moveaxis(Jq[iqpt, :, :, :, :], [1, 3], [2, 4])
+            )
+            Jq[iqpt, :, :, :, :] /= 2.0
+        return Jq
+
+    def Hq(self, kpoints, anisotropic=True):
+        """
+        Compute the magnon Hamiltonian in reciprocal space.
+
+        Parameters
+        ----------
+        kpoints : array_like
+            k-points at which to evaluate the Hamiltonian
+        anisotropic : bool, optional
+            Whether to include anisotropic interactions, default True
+
+        Returns
+        -------
+        numpy.ndarray
+            Magnon Hamiltonian matrix at each k-point
+        """
+        magmoms = self.magmom.copy()
+        magmoms /= np.linalg.norm(magmoms, axis=-1)[:, None]
+
+        U, V = get_rotation_arrays(magmoms, u=self._uz)
+
+        J0 = self.Jq(np.zeros((1, 3)), anisotropic=anisotropic)
+        J0 = -Hermitize(J0)[:, :, 0]
+        Jq = -Hermitize(self.Jq(kpoints, anisotropic=anisotropic))
+
+        C = np.diag(np.einsum("ix,ijxy,jy->i", V, 2 * J0, V))
+        B = np.einsum("ix,ijkxy,jy->kij", U, Jq, U)
+        A1 = np.einsum("ix,ijkxy,jy->kij", U, Jq, U.conj())
+        A2 = np.einsum("ix,ijkxy,jy->kij", U.conj(), Jq, U)
+
+        return np.block([[A1 - C, B], [B.swapaxes(-1, -2).conj(), A2 - C]])
+
+    def _magnon_energies(self, kpoints, anisotropic=True, u=None):
+        """Calculate magnon energies"""
+        H = self.Hq(kpoints, anisotropic=anisotropic)
+        n = H.shape[-1] // 2
+        I = np.eye(n)
+
+        min_eig = 0.0
+        try:
+            K = np.linalg.cholesky(H)
+        except np.linalg.LinAlgError:
+            try:
+                K = np.linalg.cholesky(H + 1e-6 * np.eye(2 * n))
+            except np.linalg.LinAlgError:
+                from warnings import warn
+
+                min_eig = np.min(np.linalg.eigvalsh(H))
+                K = np.linalg.cholesky(H - (min_eig - 1e-6) * np.eye(2 * n))
+                warn(
+                    f"WARNING: The system may be far from the magnetic ground-state. Minimum eigenvalue: {min_eig}. The magnon energies might be unphysical."
+                )
+
+        g = np.block([[1 * I, 0 * I], [0 * I, -1 * I]])
+        KH = K.swapaxes(-1, -2).conj()
+
+        return np.linalg.eigvalsh(KH @ g @ K)[:, n:] + min_eig
+
+    def get_magnon_bands(
+        self,
+        kpoints: np.array = np.array([]),
+        path: str = None,
+        npoints: int = 300,
+        special_points: dict = None,
+        tol: float = 2e-4,
+        pbc: tuple = None,
+        cartesian: bool = False,
+        labels: list = None,
+        anisotropic: bool = True,
+        u: np.array = None,
+    ):
+        """Get magnon band structure"""
+        pbc = self._pbc if pbc is None else pbc
+        u = self._uz if u is None else u
+        if kpoints.size == 0:
+            from ase.cell import Cell
+
+            bandpath = Cell(self._cell).bandpath(
+                path=path,
+                npoints=npoints,
+                special_points=special_points,
+                eps=tol,
+                pbc=pbc,
+            )
+            kpoints = bandpath.kpts
+            spk = bandpath.special_points
+            spk[r"$\Gamma$"] = spk.pop("G", np.zeros(3))
+            labels = [
+                (i, symbol)
+                for symbol in spk
+                for i in np.where((kpoints == spk[symbol]).all(axis=1))[0]
+            ]
+        elif cartesian:
+            kpoints = np.linalg.solve(self._cell.T, kpoints.T).T
+
+        bands = self._magnon_energies(kpoints, anisotropic=anisotropic)
+
+        return labels, bands
+
+    def plot_magnon_bands(self, **kwargs):
+        """
+        Plot magnon band structure.
+
+        Parameters
+        ----------
+        **kwargs
+            Additional keyword arguments passed to get_magnon_bands and plotting functions
+        """
+        filename = kwargs.pop("filename", None)
+        kpath, bands = self.get_magnon_bands(**kwargs)
+        bands_plot = BandsPlot(bands, kpath)
+        bands_plot.plot(filename=filename)
+
+    @classmethod
+    def load_from_io(cls, exc: SpinIO, **kwargs):
+        """
+        Create Magnon instance from SpinIO
+
+        Parameters
+        ----------
+        exc : SpinIO
+            SpinIO instance with exchange parameters
+        **kwargs : dict
+            Additional arguments passed to get_full_Jtensor_for_Rlist
+
+        Returns
+        -------
+        Magnon
+            Initialized Magnon instance
+        """
+        return cls(
+            nspin=exc.nspin,
+            magmom=exc.magmoms,
+            Rlist=exc.Rlist,
+            JR=exc.get_full_Jtensor_for_Rlist(order="ij33", **kwargs),
+        )
+
+    @classmethod
+    def from_TB2J_results(cls, path=None, fname="TB2J.pickle", **kwargs):
+        """
+        Create Magnon instance from TB2J results.
+
+        Parameters
+        ----------
+        path : str, optional
+            Path to the TB2J results file
+        fname : str, optional
+            Filename of the TB2J results file, default "TB2J.pickle"
+        **kwargs : dict
+            Additional arguments passed to load_from_io
+
+        Returns
+        -------
+        Magnon
+            Initialized Magnon instance
+        """
+        exc = SpinIO.load_pickle(path=path, fname=fname)
+        return cls.load_from_io(exc, **kwargs)
+
+
+def test_magnon(path="TB2J_results"):
+    """Test the magnon calculator by loading from TB2J_results and computing at high-symmetry points."""
+    from pathlib import Path
+
+    import numpy as np
+
+    # Check if TB2J_results exists
+    results_path = Path(path)
+    if not results_path.exists():
+        raise FileNotFoundError(f"TB2J_results directory not found at {path}")
+
+    # Load magnon calculator from TB2J results
+    print(f"Loading exchange parameters from {path}...")
+    magnon = Magnon.from_TB2J_results(path=path, iso_only=True)
+
+    # Define high-symmetry points for a cube
+    kpoints = np.array(
+        [
+            [0.0, 0.0, 0.0],  # Γ (Gamma)
+            [0.5, 0.0, 0.0],  # X
+            [0.5, 0.5, 0.0],  # M
+            [0.5, 0.5, 0.5],  # R
+        ]
+    )
+    klabels = ["Gamma", "X", "M", "R"]
+
+    print("\nComputing exchange interactions at high-symmetry points...")
+    Jq = magnon.Jq(kpoints)
+
+    print(f"\nResults for {len(kpoints)} k-points:")
+    print("-" * 50)
+    print("Exchange interactions J(q):")
+    print(f"Shape of Jq tensor: {Jq.shape}")
+    print(
+        f"Dimensions: (n_kpoints={Jq.shape[0]}, n_spin={Jq.shape[1]}, n_spin={Jq.shape[2]}, xyz={Jq.shape[3]}, xyz={Jq.shape[4]})"
+    )
+
+    print("\nComputing magnon energies...")
+    energies = magnon._magnon_energies(kpoints)
+
+    print("\nMagnon energies at high-symmetry points (in meV):")
+    print("-" * 50)
+    for i, (k, label) in enumerate(zip(kpoints, klabels)):
+        print(f"\n{label}-point k={k}:")
+        print(f"Energies: {energies[i] * 1000:.3f} meV")  # Convert to meV
+
+    return magnon, Jq, energies
+
+
+if __name__ == "__main__":
+    test_magnon()

TB2J/magnon/magnon_io.py

@@ -0,0 +1,48 @@
+import numpy as np
+
+from TB2J.io_exchange import SpinIO
+
+
+class MagnonIO:
+    """Handle IO operations for magnon calculations"""
+
+    def __init__(self, exc: SpinIO):
+        """
+        Initialize MagnonIO with a SpinIO instance
+
+        Parameters
+        ----------
+        exc : SpinIO
+            SpinIO instance containing exchange information
+        """
+        self.exc = exc
+
+    def get_nspin(self):
+        """Get number of spins"""
+        return self.exc.get_nspin()
+
+    def get_ind_atoms(self):
+        """Get atom indices"""
+        return self.exc.ind_atoms
+
+    def get_magmom(self):
+        """Get magnetic moments"""
+        nspin = self.get_nspin()
+        return np.array([self.exc.spinat[self.exc.iatom(i)] for i in range(nspin)])
+
+    def get_rlist(self):
+        """Get R-vectors list"""
+        return self.exc.Rlist
+
+    def get_jtensor(self, asr=False, iso_only=False):
+        """
+        Get full J tensor for R-list
+
+        Parameters
+        ----------
+        asr : bool, optional
+            Acoustic sum rule, default False
+        iso_only : bool, optional
+            Only isotropic interactions, default False
+        """
+        return self.exc.get_full_Jtensor_for_Rlist(asr=asr, iso_only=iso_only)

TB2J/magnon/magnon_math.py

@@ -0,0 +1,53 @@
+import numpy as np
+
+__all__ = ["generate_grid", "get_rotation_arrays", "round_to_precision", "uz", "I"]
+
+I = np.eye(3)
+uz = np.array([[0.0, 0.0, 1.0]])
+
+
+def generate_grid(kmesh, sort=True):
+    half_grid = [int(n / 2) for n in kmesh]
+    grid = np.stack(
+        np.meshgrid(*[np.arange(-n, n + 1) for n in half_grid]), axis=-1
+    ).reshape(-1, 3)
+
+    if sort:
+        idx = np.linalg.norm(grid, axis=-1).argsort()
+        grid = grid[idx]
+
+    return grid
+
+
+# def JR_to_Jq(JR, Rlist, qpt, vecn):
+#    for iR, R in enumerate(Rlist):
+#        phase = 2 * np.pi * R @ qpt
+#        rv = phase * vecn
+#        Rot = Rotation.from_rotvec(rv.reshape(-1, 3)).as_matrix().reshape(R.shape[0], 3, 3)
+#
+
+
+def get_rotation_arrays(magmoms, u=uz):
+    dim = magmoms.shape[0]
+    v = magmoms
+    n = np.cross(u, v)
+    n /= np.linalg.norm(n, axis=-1).reshape(dim, 1)
+    z = np.repeat(u, dim, axis=0)
+    A = np.stack([z, np.cross(n, z), n], axis=1)
+    B = np.stack([v, np.cross(n, v), n], axis=1)
+    R = np.einsum("nki,nkj->nij", A, B)
+
+    Rnan = np.isnan(R)
+    if Rnan.any():
+        nanidx = np.where(Rnan)[0]
+        R[nanidx] = I
+        R[nanidx, 2] = v[nanidx]
+
+    U = R[:, 0] + 1j * R[:, 1]
+    V = R[:, 2]
+
+    return U, V
+
+
+def round_to_precision(array, precision):
+    return precision * np.round(array / precision)

TB2J/magnon/plot.py

@@ -0,0 +1,58 @@
+import matplotlib.pyplot as plt
+import numpy as np
+
+
+class BandsPlot:
+    _UNITS = "meV"
+    _NSYSTEMS = 1
+
+    def __init__(self, bands, kpath, **kwargs):
+        self.bands = bands
+        self.kpath = kpath
+        self.bands *= 1000
+
+        plot_options = kwargs
+        self.linewidth = plot_options.pop("linewidth", 1.5)
+        self.color = plot_options.pop("color", "blue")
+        self.fontsize = plot_options.pop("fontsize", 12)
+        self.ticksize = plot_options.pop("ticksize", 10)
+        self.plot_options = plot_options
+
+    def plot(self, filename=None):
+        fig, axs = plt.subplots(1, self._NSYSTEMS, constrained_layout=True)
+
+        kdata = np.arange(self.bands.shape[0])
+        for band in self.bands.T:
+            axs.plot(
+                kdata,
+                band,
+                linewidth=self.linewidth,
+                color=self.color,
+                **self.plot_options,
+            )
+
+        bmin, bmax = self.bands.min(), self.bands.max()
+        ymin, ymax = (
+            bmin - 0.05 * np.abs(bmin - bmax),
+            bmax + 0.05 * np.abs(bmax - bmin),
+        )
+
+        axs.set_ylim([ymin, ymax])
+        axs.set_xlim([0, kdata[-1]])
+
+        kpoint_labels = list(zip(*self.kpath))
+        axs.set_xticks(*kpoint_labels, fontsize=self.ticksize)
+        axs.vlines(
+            x=kpoint_labels[0],
+            ymin=ymin,
+            ymax=ymax,
+            color="black",
+            linewidth=self.linewidth / 5,
+        )
+
+        axs.set_ylabel(f"Energy ({self._UNITS})", fontsize=self.fontsize)
+
+        if filename is None:
+            plt.show()
+        else:
+            fig.save(filename, dpi=300, bbox_inches="tight")