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

TB2J/magnon/io_exchange2.py

@@ -0,0 +1,695 @@
+"""
+This module provides functionality for handling magnetic exchange interactions and computing magnon band structures.
+
+It includes classes and functions for:
+- Reading and writing exchange interaction data
+- Computing exchange tensors and magnon energies
+- Plotting magnon band structures
+- Converting between different magnetic structure representations
+"""
+
+import numpy as np
+from scipy.spatial.transform import Rotation
+
+from .magnon_math import generate_grid, get_rotation_arrays, round_to_precision, uz
+from .plot import BandsPlot
+from .structure import BaseMagneticStructure, get_attribute_array
+
+__all__ = [
+    "ExchangeIO",
+    "plot_tb2j_magnon_bands",
+]
+
+
+def branched_keys(tb2j_keys, npairs):
+    """
+    Organize TB2J keys into branches based on magnetic site pairs.
+
+    Parameters
+    ----------
+    tb2j_keys : list
+        List of TB2J dictionary keys containing interaction information
+    npairs : int
+        Number of magnetic site pairs
+
+    Returns
+    -------
+    list
+        List of branched keys organized by magnetic site pairs
+    """
+    msites = int((2 * npairs) ** 0.5)
+    branch_size = len(tb2j_keys) // msites**2
+    new_keys = sorted(tb2j_keys, key=lambda x: -x[1] + x[2])[
+        (npairs - msites) * branch_size :
+    ]
+    new_keys.sort(key=lambda x: x[1:])
+    bkeys = [
+        new_keys[i : i + branch_size] for i in range(0, len(new_keys), branch_size)
+    ]
+
+    return [sorted(branch, key=lambda x: np.linalg.norm(x[0])) for branch in bkeys]
+
+
+def correct_content(content, quadratic=False):
+    """
+    Ensure content dictionary has all required entries with proper initialization.
+
+    Parameters
+    ----------
+    content : dict
+        Dictionary containing exchange interaction data
+    quadratic : bool, optional
+        Whether to include biquadratic interactions, by default False
+    """
+    n = max(content["index_spin"]) + 1
+    data_shape = {"exchange_Jdict": ()}
+
+    if not content["colinear"]:
+        data_shape |= {"Jani_dict": (3, 3), "dmi_ddict": (3,)}
+    if quadratic:
+        data_shape["biquadratic_Jdict"] = (2,)
+
+    for label, shape in data_shape.items():
+        content[label] |= {((0, 0, 0), i, i): np.zeros(shape) for i in range(n)}
+
+
+def Hermitize(array):
+    """
+    Convert an array into its Hermitian form by constructing a Hermitian matrix.
+
+    A Hermitian matrix H has the property that H = H†, where H† is the conjugate transpose.
+    This means H[i,j] = conj(H[j,i]) for all indices i,j. The function takes an input array
+    representing the upper triangular part of the matrix and constructs the full Hermitian
+    matrix by:
+    1. Placing the input values in the upper triangular part
+    2. Computing the conjugate transpose of these values for the lower triangular part
+
+    This is commonly used in quantum mechanics and magnetic systems where Hamiltonians
+    must be Hermitian to ensure real eigenvalues.
+
+    Parameters
+    ----------
+    array : numpy.ndarray
+        Input array containing the upper triangular elements of the matrix.
+        Shape should be (n*(n+1)/2, ...) where n is the dimension of
+        the resulting square matrix.
+
+    Returns
+    -------
+    numpy.ndarray
+        Full Hermitian matrix with shape (n, n, ...), where n is computed
+        from the input array size. The result satisfies result[i,j] = conj(result[j,i])
+        for all indices i,j.
+
+    Example
+    -------
+    >>> arr = np.array([1+0j, 2+1j, 3+0j])  # Upper triangular elements for 2x2 matrix
+    >>> Hermitize(arr)
+    array([[1.+0.j, 2.+1.j],
+           [2.-1.j, 3.+0.j]])
+    """
+    n = int((2 * array.shape[0]) ** 0.5)
+    result = np.zeros((n, n) + array.shape[1:], dtype=complex)
+    u_indices = np.triu_indices(n)
+
+    # for python>=3.11
+    # result[*u_indices] = array
+    # result.swapaxes(0, 1)[*u_indices] = array.swapaxes(-1, -2).conj()
+    # for python<3.11
+    result[u_indices[0], u_indices[1]] = array
+    result.swapaxes(0, 1)[u_indices[0], u_indices[1]] = array.swapaxes(-1, -2).conj()
+
+    return result
+
+
+class ExchangeIO(BaseMagneticStructure):
+    """
+    Class for handling magnetic exchange interactions and computing magnon properties.
+
+    This class provides functionality for:
+    - Managing magnetic structure information
+    - Computing exchange tensors
+    - Calculating magnon band structures
+    - Reading TB2J format files
+    - Visualizing magnon bands
+
+    Parameters
+    ----------
+    atoms : ase.Atoms, optional
+        ASE atoms object containing the structure
+    cell : array_like, optional
+        3x3 matrix defining the unit cell
+    elements : list, optional
+        List of chemical symbols for atoms
+    positions : array_like, optional
+        Atomic positions
+    magmoms : array_like, optional
+        Magnetic moments for each atom
+    pbc : tuple, optional
+        Periodic boundary conditions, default (True, True, True)
+    magnetic_elements : list, optional
+        List of magnetic elements in the structure
+    kmesh : list, optional
+        k-point mesh dimensions, default [1, 1, 1]
+    collinear : bool, optional
+        Whether the magnetic structure is collinear, default True
+    """
+
+    def __init__(
+        self,
+        atoms=None,
+        cell=None,
+        elements=None,
+        positions=None,
+        magmoms=None,
+        pbc=(True, True, True),
+        magnetic_elements=[],
+        kmesh=[1, 1, 1],
+        collinear=True,
+    ):
+        super().__init__(
+            atoms=atoms,
+            cell=cell,
+            elements=elements,
+            positions=positions,
+            pbc=pbc,
+            magmoms=magmoms,
+            collinear=collinear,
+        )
+
+        self.magnetic_elements = magnetic_elements
+        self.kmesh = kmesh
+
+        num_terms = 4 if collinear else 18
+        self._exchange_values = np.empty((0, 0, num_terms), dtype=float)
+
+    @property
+    def magnetic_elements(self):
+        """List of magnetic elements in the structure."""
+        return self._magnetic_elements
+
+    @magnetic_elements.setter
+    def magnetic_elements(self, value):
+        from .structure import validate_symbols
+
+        symbols = validate_symbols(value)
+        for symbol in symbols:
+            if symbol not in self.elements:
+                raise ValueError(f"Symbol '{symbol}' not in 'elements'.")
+
+        self._magnetic_elements = symbols
+        self._set_index_pairs()
+
+    @property
+    def interacting_pairs(self):
+        """List of pairs of interacting magnetic sites."""
+        return self._pairs
+
+    def _set_index_pairs(self):
+        from itertools import combinations_with_replacement
+
+        magnetic_elements = self.magnetic_elements
+        elements = self.elements
+        indices = [
+            i for i, symbol in enumerate(elements) if symbol in magnetic_elements
+        ]
+        index_pairs = list(combinations_with_replacement(indices, 2))
+        index_spin = np.sort(np.unique(index_pairs))
+
+        self._pairs = index_pairs
+        self._index_spin = index_spin
+
+    @property
+    def kmesh(self):
+        """K-point mesh dimensions for sampling the Brillouin zone."""
+        return self._kmesh
+
+    @kmesh.setter
+    def kmesh(self, values):
+        try:
+            the_kmesh = [int(k) for k in values]
+        except (ValueError, TypeError):
+            raise ValueError("Argument must be an iterable with 'int' elements.")
+        if len(the_kmesh) != 3:
+            raise ValueError("Argument must be of length 3.")
+        if any(k < 1 for k in the_kmesh):
+            raise ValueError("Argument must contain only positive numbers.")
+
+        self._kmesh = the_kmesh
+
+    @property
+    def vectors(self):
+        """Array of interaction vectors between magnetic sites."""
+        return self._exchange_values[:, :, :3]
+
+    def set_vectors(self, values=None, cartesian=False):
+        """
+        Set the interaction vectors between magnetic sites.
+
+        Parameters
+        ----------
+        values : array_like, optional
+            Array of interaction vectors
+        cartesian : bool, optional
+            Whether the vectors are in Cartesian coordinates, default False
+        """
+        try:
+            pairs = self._pairs
+        except AttributeError:
+            raise AttributeError("'magnetic_elements' attribute has not been set yet.")
+        else:
+            n_pairs = len(pairs)
+
+        if values is None:
+            i, j = zip(*pairs)
+            positions = self.positions
+            base_vectors = positions[i, :] - positions[j, :]
+            grid = generate_grid(self.kmesh)
+            vectors = base_vectors[:, None, :] + grid[None, :, :]
+            m_interactions = np.prod(self.kmesh)
+        else:
+            vectors = get_attribute_array(values, "vectors", dtype=float)
+            if vectors.ndim != 3 or vectors.shape[::2] != (n_pairs, 3):
+                raise ValueError(
+                    f"'vectors' must have the shape (n, m, ), where n={n_pairs} is the number of\n"
+                    "pairs of interacting species."
+                )
+            if cartesian:
+                vectors = np.linalg.solve(self.cell.T, vectors.swapaxes(1, -1))
+                vectors = vectors.swapaxes(-1, 1)
+            m_interactions = vectors.shape[1]
+
+        shape = (
+            (n_pairs, m_interactions, 4)
+            if self.collinear
+            else (n_pairs, m_interactions, 18)
+        )
+        exchange_values = np.zeros(shape, dtype=float)
+        exchange_values[:, :, :3] = vectors
+        self._exchange_values = exchange_values
+
+    def _get_neighbor_indices(self, neighbors, tol=1e-4):
+        """
+        Get indices of neighbor pairs based on distance.
+
+        Parameters
+        ----------
+        neighbors : list
+            List of neighbor shells to consider
+        tol : float, optional
+            Distance tolerance for neighbor shell assignment, default 1e-4
+
+        Returns
+        -------
+        tuple
+            Indices corresponding to the specified neighbor shells
+        """
+        distance = np.linalg.norm(self.vectors @ self.cell, axis=-1)
+        distance = round_to_precision(distance, tol)
+        neighbors_distance = np.unique(np.sort(distance))
+        indices = np.where(
+            distance[:, :, None] == neighbors_distance[neighbors][None, None, :]
+        )
+
+        return indices
+
+    def set_exchange_array(self, name, values, neighbors=None, tol=1e-4):
+        """
+        Set exchange interaction values for specified neighbors.
+
+        Parameters
+        ----------
+        name : str
+            Type of exchange interaction ('Jiso', 'Biquad', 'DMI', or 'Jani')
+        values : array_like
+            Exchange interaction values
+        neighbors : list, optional
+            List of neighbor shells to assign values to
+        tol : float, optional
+            Distance tolerance for neighbor shell assignment, default 1e-4
+        """
+        if self.vectors.size == 0:
+            raise AttributeError("The intraction vectors must be set first.")
+
+        array = get_attribute_array(values, name, dtype=float)
+
+        if neighbors is not None:
+            if len(array) != len(neighbors):
+                raise ValueError(
+                    "The number of neighbors and exchange values does not coincide."
+                )
+            *array_indices, value_indices = self._get_neighbor_indices(
+                list(neighbors), tol=tol
+            )
+        else:
+            if self._exchange_values.shape[:2] != array.shape[:2]:
+                raise ValueError(
+                    f"The shape of the array is incompatible with '{self.exchange_values.shape}'"
+                )
+            array_indices, value_indices = (
+                [slice(None), slice(None)],
+                (slice(None), slice(None)),
+            )
+
+        if name == "Jiso":
+            self._exchange_values[*array_indices, 3] = array[value_indices]
+        elif name == "Biquad":
+            self._exchange_values[*array_indices, 4:6] = array[value_indices]
+        elif name == "DMI":
+            self._exchange_values[*array_indices, 6:9] = array[value_indices]
+        elif name == "Jani":
+            self._exchange_values[*array_indices, 9:] = array[value_indices].reshape(
+                array.shape[:2] + (9,)
+            )
+        else:
+            raise ValueError(f"Unrecognized exchange array name: '{name}'.")
+
+    @property
+    def Jiso(self):
+        return self._exchange_values[:, :, 3]
+
+    @property
+    def Biquad(self):
+        return self._exchange_values[:, :, 4:6]
+
+    @property
+    def DMI(self):
+        return self._exchange_values[:, :, 6:9]
+
+    @property
+    def Jani(self):
+        matrix_shape = self._exchange_values.shape[:2] + (3, 3)
+        return self._exchange_values[:, :, 9:].reshape(matrix_shape)
+
+    def exchange_tensor(self, anisotropic=True):
+        """
+        Compute the exchange interaction tensor.
+
+        Parameters
+        ----------
+        anisotropic : bool, optional
+            Whether to include anisotropic interactions, default True
+
+        Returns
+        -------
+        numpy.ndarray
+            Exchange interaction tensor
+        """
+        shape = self._exchange_values.shape[:2] + (3, 3)
+        tensor = np.zeros(shape, dtype=float)
+
+        if anisotropic and not self.collinear:
+            # anisotropic exchange tensor
+            tensor += self._exchange_values[:, :, 9:].reshape(shape)
+            # DMI
+            pos_indices = ([1, 2, 0], [2, 0, 1])
+            neg_indices = ([2, 0, 1], [1, 2, 0])
+            tensor[:, :, *pos_indices] += self._exchange_values[:, :, 6:9]
+            tensor[:, :, *neg_indices] -= self._exchange_values[:, :, 6:9]
+        # isotropic exchange
+        diag_indices = ([0, 1, 2], [0, 1, 2])
+        tensor[:, :, *diag_indices] += self._exchange_values[:, :, 3, None]
+
+        return tensor
+
+    def Jq(self, kpoints, anisotropic=True):
+        """
+        Compute the exchange interactions in reciprocal space.
+
+        Parameters
+        ----------
+        kpoints : array_like
+            k-points at which to evaluate the exchange interactions
+        anisotropic : bool, optional
+            Whether to include anisotropic interactions, default True
+
+        Returns
+        -------
+        numpy.ndarray
+            Exchange interactions in reciprocal space
+        """
+        vectors = self._exchange_values[:, :, :3].copy()
+        tensor = self.exchange_tensor(anisotropic=anisotropic)
+
+        if self._Q is not None:
+            phi = 2 * np.pi * vectors.round(3).astype(int) @ self._Q
+            rv = np.einsum("ij,k->ijk", phi, self._n)
+            R = (
+                Rotation.from_rotvec(rv.reshape(-1, 3))
+                .as_matrix()
+                .reshape(vectors.shape[:2] + (3, 3))
+            )
+            np.einsum("nmij,nmjk->nmik", tensor, R, out=tensor)
+
+        exp_summand = np.exp(2j * np.pi * vectors @ kpoints.T)
+        Jexp = exp_summand[:, :, :, None, None] * tensor[:, :, None]
+        Jq = np.sum(Jexp, axis=1)
+
+        pairs = np.array(self._pairs)
+        idx = np.where(pairs[:, 0] == pairs[:, 1])
+        Jq[idx] /= 2
+
+        return Jq
+
+    def Hq(self, kpoints, anisotropic=True, u=uz):
+        """
+        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
+        u : array_like, optional
+            Reference direction for spin quantization axis
+
+        Returns
+        -------
+        numpy.ndarray
+            Magnon Hamiltonian matrix at each k-point
+        """
+        if self.collinear:
+            magmoms = np.zeros((self._index_spin.size, 3))
+            magmoms[:, 2] = self.magmoms[self._index_spin]
+        else:
+            magmoms = self.magmoms[self._index_spin]
+        magmoms /= np.linalg.norm(magmoms, axis=-1)[:, None]
+
+        U, V = get_rotation_arrays(magmoms, u=u)
+
+        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=uz):
+        H = self.Hq(kpoints, anisotropic=anisotropic, u=u)
+        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 = uz,
+    ):
+        pbc = self._pbc if pbc is None else pbc
+
+        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, u=u)
+
+        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_tb2j(
+        cls,
+        pickle_file: str = "TB2J.pickle",
+        pbc: tuple = (True, True, True),
+        anisotropic: bool = False,
+        quadratic: bool = False,
+    ):
+        from pickle import load
+
+        try:
+            with open(pickle_file, "rb") as File:
+                content = load(File)
+        except FileNotFoundError:
+            raise FileNotFoundError(
+                f"No such file or directory: '{pickle_file}'. Please provide a valid .pickle file."
+            )
+        else:
+            correct_content(content)
+
+        magmoms = content["magmoms"] if content["colinear"] else content["spinat"]
+        magnetic_elements = {
+            content["atoms"].numbers[i]
+            for i, j in enumerate(content["index_spin"])
+            if j > -1
+        }
+
+        exchange = cls(
+            atoms=content["atoms"],
+            magmoms=magmoms,
+            pbc=pbc,
+            collinear=content["colinear"],
+            magnetic_elements=magnetic_elements,
+        )
+
+        num_pairs = len(exchange.interacting_pairs)
+        bkeys = branched_keys(content["distance_dict"].keys(), num_pairs)
+
+        vectors = [
+            [content["distance_dict"][key][0] for key in branch] for branch in bkeys
+        ]
+        exchange.set_vectors(vectors, cartesian=True)
+        Jiso = [[content["exchange_Jdict"][key] for key in branch] for branch in bkeys]
+        exchange.set_exchange_array("Jiso", Jiso)
+
+        if not content["colinear"] and anisotropic:
+            Jani = [[content["Jani_dict"][key] for key in branch] for branch in bkeys]
+            exchange.set_exchange_array("Jani", Jani)
+            DMI = [[content["dmi_ddict"][key] for key in branch] for branch in bkeys]
+            exchange.set_exchange_array("DMI", DMI)
+        if quadratic:
+            Biquad = [
+                [content["biquadratic_Jdict"][key] for key in branch]
+                for branch in bkeys
+            ]
+            exchange.set_exchange_array("Biquad", Biquad)
+
+        return exchange
+
+
+def plot_tb2j_magnon_bands(
+    pickle_file: str = "TB2J.pickle",
+    path: str = None,
+    npoints: int = 300,
+    special_points: dict = None,
+    anisotropic: bool = False,
+    quadratic: bool = False,
+    pbc: tuple = (True, True, True),
+    filename: str = None,
+):
+    """
+    Load TB2J data and plot magnon band structure in one step.
+
+    This is a convenience function that combines loading TB2J data and plotting
+    magnon bands. It first loads the magnetic structure and exchange interactions
+    from a TB2J pickle file, then calculates and plots the magnon band structure.
+
+    Parameters
+    ----------
+    pickle_file : str, optional
+        Path to the TB2J pickle file, default "TB2J.pickle"
+    path : str, optional
+        High-symmetry k-point path for band structure plot
+        (e.g., "GXMG" for a square lattice)
+    npoints : int, optional
+        Number of k-points for band structure calculation, default 300
+    special_points : dict, optional
+        Dictionary of special k-points for custom paths
+    anisotropic : bool, optional
+        Whether to include anisotropic interactions, default False
+    quadratic : bool, optional
+        Whether to include biquadratic interactions, default False
+    pbc : tuple, optional
+        Periodic boundary conditions, default (True, True, True)
+    filename : str, optional
+        If provided, save the plot to this file
+
+    Returns
+    -------
+    exchange : ExchangeIO
+        The ExchangeIO instance containing the loaded data and plot
+
+    Example
+    -------
+    >>> # Basic usage with default parameters
+    >>> plot_tb2j_magnon_bands()
+
+    >>> # Custom path and saving to file
+    >>> plot_tb2j_magnon_bands(
+    ...     path="GXMG",
+    ...     anisotropic=True,
+    ...     filename="magnon_bands.png"
+    ... )
+    """
+    # Load the TB2J data
+    exchange = ExchangeIO.load_tb2j(
+        pickle_file=pickle_file, pbc=pbc, anisotropic=anisotropic, quadratic=quadratic
+    )
+
+    # Plot the magnon bands
+    exchange.plot_magnon_bands(
+        path=path, npoints=npoints, special_points=special_points, filename=filename
+    )
+
+    return exchange