TB2J 0.9.9.9__py3-none-any.whl → 0.9.9.10__py3-none-any.whl

TB2J/magnon/magnon3.py

@@ -1,11 +1,58 @@
-from dataclasses import dataclass
+import json
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import List, Optional, Tuple, Union
 
 import numpy as np
+import tomli
+import tomli_w
 from scipy.spatial.transform import Rotation
 
-from ..io_exchange import SpinIO
-from ..mathutils import Hermitize, get_rotation_arrays
-from .plot import BandsPlot
+from TB2J.io_exchange import SpinIO
+from TB2J.magnon.magnon_math import get_rotation_arrays
+from TB2J.mathutils.auto_kpath import auto_kpath
+from TB2J.magnon.magnon_band import MagnonBand
+
+
+@dataclass
+class MagnonParameters:
+    """Parameters for magnon band structure calculations"""
+
+    path: str = "TB2J_results"
+    kpath: str = None
+    npoints: int = 300
+    filename: str = "magnon_bands.png"
+    Jiso: bool = True
+    Jani: bool = False
+    DMI: bool = False
+    Q: Optional[List[float]] = None
+    uz_file: Optional[str] = None
+    n: Optional[List[float]] = None
+
+    @classmethod
+    def from_toml(cls, filename: str) -> "MagnonParameters":
+        """Load parameters from a TOML file"""
+        with open(filename, "rb") as f:
+            data = tomli.load(f)
+        return cls(**data)
+
+    def to_toml(self, filename: str):
+        """Save parameters to a TOML file"""
+        # Convert to dict and remove None values
+        data = {k: v for k, v in asdict(self).items() if v is not None}
+        with open(filename, "wb") as f:
+            tomli_w.dump(data, f)
+
+    def __post_init__(self):
+        """Validate parameters after initialization"""
+        if self.Q is not None and len(self.Q) != 3:
+            raise ValueError("Q must be a list of 3 numbers")
+        if self.n is not None and len(self.n) != 3:
+            raise ValueError("n must be a list of 3 numbers")
+
+        # Convert path to absolute path if uz_file is relative to it
+        if self.uz_file and not Path(self.uz_file).is_absolute():
+            self.uz_file = str(Path(self.path) / self.uz_file)
 
 
 @dataclass
@@ -19,9 +66,11 @@ class Magnon:
     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)
+    cell: np.ndarray
+    _Q: np.ndarray
+    _uz: np.ndarray
+    _n: np.ndarray
+    pbc: tuple = (True, True, True)
 
     def set_reference(self, Q, uz, n):
         """
@@ -102,7 +151,7 @@ class Magnon:
                 # 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)
+                JRprime[iR] = np.einsum(" ijxy, yb -> ijxb", JR[iR], Rmat)
 
         nkpt = kpoints.shape[0]
         Jq = np.zeros((nkpt, self.nspin, self.nspin, 3, 3), dtype=complex)
@@ -113,15 +162,13 @@ class Magnon:
                 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
+        Jq_copy = Jq.copy()
+        Jq.swapaxes(-1, -2)  # swap xyz
+        Jq.swapaxes(-3, -4)  # swap ij
+        Jq = (Jq.conj() + Jq_copy) / 2.0
         return Jq
 
-    def Hq(self, kpoints, anisotropic=True):
+    def Hq(self, kpoints):
         """
         Compute the magnon Hamiltonian in reciprocal space.
 
@@ -142,20 +189,25 @@ class Magnon:
 
         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))
+        J0 = -self.Jq(np.zeros((1, 3)))[0]
+        # J0 = -Hermitize(J0)[:, :, 0]
+        # Jq = -Hermitize(self.Jq(kpoints, anisotropic=anisotropic))
+
+        Jq = -self.Jq(kpoints)
+        print(f"J0 shape: {J0.shape}")
 
         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)
+        B = np.einsum("ix,kijxy,jy->kij", U, Jq, U)
+        A1 = np.einsum("ix,kijxy,jy->kij", U, Jq, U.conj())
+        A2 = np.einsum("ix,kijxy,jy->kij", U.conj(), Jq, U)
 
-        return np.block([[A1 - C, B], [B.swapaxes(-1, -2).conj(), A2 - C]])
+        H = np.block([[A1 - C, B], [B.swapaxes(-1, -2).conj(), A2 - C]])
+        print(f"H shape: {H.shape}")
+        return H
 
-    def _magnon_energies(self, kpoints, anisotropic=True, u=None):
+    def _magnon_energies(self, kpoints, u=None):
         """Calculate magnon energies"""
-        H = self.Hq(kpoints, anisotropic=anisotropic)
+        H = self.Hq(kpoints)
         n = H.shape[-1] // 2
         I = np.eye(n)
 
@@ -176,8 +228,9 @@ class Magnon:
 
         g = np.block([[1 * I, 0 * I], [0 * I, -1 * I]])
         KH = K.swapaxes(-1, -2).conj()
-
+        # Why only n:?
         return np.linalg.eigvalsh(KH @ g @ K)[:, n:] + min_eig
+        # return np.linalg.eigvalsh(KH @ g @ K)[:, :] + min_eig
 
     def get_magnon_bands(
         self,
@@ -192,33 +245,82 @@ class Magnon:
         anisotropic: bool = True,
         u: np.array = None,
     ):
-        """Get magnon band structure"""
-        pbc = self._pbc if pbc is None else pbc
+        """Get magnon band structure.
+
+        Parameters
+        ----------
+        kpoints : np.array, optional
+            Explicit k-points to calculate bands at. If empty, generates k-points from path.
+        path : str, optional
+            String specifying the k-path. If None, generates automatically using auto_kpath.
+        npoints : int, optional
+            Number of k-points along the path. Default is 300.
+        special_points : dict, optional
+            Dictionary of special points coordinates.
+        tol : float, optional
+            Tolerance for k-point comparisons. Default is 2e-4.
+        pbc : tuple, optional
+            Periodic boundary conditions. Default is None.
+        cartesian : bool, optional
+            Whether k-points are in cartesian coordinates. Default is False.
+        labels : list, optional
+            List of k-point labels. Default is None.
+        anisotropic : bool, optional
+            Whether to include anisotropic interactions. Default is True.
+        u : np.array, optional
+            Quantization axis. Default is None.
+
+        Returns
+        -------
+        tuple
+            - labels : list of (index, name) tuples for special k-points
+            - bands : array of band energies
+            - xlist : list of arrays with x-coordinates for plotting (if using auto_kpath)
+        """
+        pbc = self.pbc if pbc is None else pbc
+        pbc = [True, True, True]
         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]
-            ]
+            if path is None:
+                # Use auto_kpath to generate path automatically
+                xlist, kptlist, Xs, knames, spk = auto_kpath(self.cell, None, npoints=npoints)
+                kpoints = np.concatenate(kptlist)
+                # Create labels from special points
+                labels = []
+                current_pos = 0
+                for i, (x, k) in enumerate(zip(xlist, kptlist)):
+                    for name in knames:
+                        matches = np.where((k == spk[name]).all(axis=1))[0]
+                        if matches.size > 0:
+                            labels.append((matches[0] + current_pos, name))
+                    current_pos += len(k)
+            else:
+                bandpath = 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
+            kpoints = np.linalg.solve(self.cell.T, kpoints.T).T
 
-        bands = self._magnon_energies(kpoints, anisotropic=anisotropic)
+        bands = self._magnon_energies(kpoints)
+        print(f"bands shape: {bands.shape}")
 
-        return labels, bands
+        if path is None and kpoints.size == 0:  # Fixed condition
+            # When using auto_kpath, return xlist for segmented plotting
+            return labels, bands, xlist
+        else:
+            return labels, bands, None
 
     def plot_magnon_bands(self, **kwargs):
         """
@@ -227,12 +329,37 @@ class Magnon:
         Parameters
         ----------
         **kwargs
-            Additional keyword arguments passed to get_magnon_bands and plotting functions
+            Additional keyword arguments passed to get_magnon_bands and plotting functions.
+            Supported plotting options:
+            - filename : str, optional
+                Output filename for saving the plot
+            - ax : matplotlib.axes.Axes, optional
+                Axes for plotting. If None, creates new figure
+            - show : bool, optional
+                Whether to show the plot on screen
         """
         filename = kwargs.pop("filename", None)
-        kpath, bands = self.get_magnon_bands(**kwargs)
-        bands_plot = BandsPlot(bands, kpath)
-        bands_plot.plot(filename=filename)
+        kpath_labels, bands, xlist = self.get_magnon_bands(**kwargs)
+        
+        # Get k-points and special points
+        if 'path' in kwargs and kwargs['path'] is None:
+            _, kptlist, _, _, spk = auto_kpath(self.cell, None, npoints=kwargs.get('npoints', 300))
+            kpoints = np.concatenate(kptlist)
+        else:
+            bandpath = self.cell.bandpath(path=kwargs.get('path', 'GXMG'), npoints=kwargs.get('npoints', 300))
+            kpoints = bandpath.kpts
+            spk = bandpath.special_points.copy()
+            spk[r"$\Gamma$"] = spk.pop("G", np.zeros(3))
+
+        bands_plot = MagnonBand(
+            energies=bands * 1000,  # Convert to meV
+            kpoints=kpoints,
+            kpath_labels=kpath_labels,
+            special_points=spk,
+            xcoords=xlist,
+        )
+        
+        return bands_plot.plot(filename=filename, **kwargs)
 
     @classmethod
     def load_from_io(cls, exc: SpinIO, **kwargs):
@@ -251,11 +378,23 @@ class Magnon:
         Magnon
             Initialized Magnon instance
         """
+        # Get magnetic moments for magnetic atoms
+        magmoms = exc.get_magnetic_moments()
+        nspin = len(magmoms)  # Number of magnetic atoms
+
+        cell = exc.atoms.get_cell()
+        pbc = exc.atoms.get_pbc()
+
         return cls(
             nspin=exc.nspin,
-            magmom=exc.magmoms,
+            magmom=magmoms,
             Rlist=exc.Rlist,
-            JR=exc.get_full_Jtensor_for_Rlist(order="ij33", **kwargs),
+            JR=exc.get_full_Jtensor_for_Rlist(order="ij33", asr=False, **kwargs),
+            cell=cell,
+            _Q=np.zeros(3),  # Default propagation vector
+            _uz=np.array([[0.0, 0.0, 1.0]]),  # Default quantization axis
+            _n=np.array([0.0, 0.0, 1.0]),  # Default rotation axis
+            pbc=pbc,
         )
 
     @classmethod
@@ -294,7 +433,7 @@ def test_magnon(path="TB2J_results"):
 
     # Load magnon calculator from TB2J results
     print(f"Loading exchange parameters from {path}...")
-    magnon = Magnon.from_TB2J_results(path=path, iso_only=True)
+    magnon = Magnon.from_TB2J_results(path=path, Jiso=True, Jani=False, DMI=False)
 
     # Define high-symmetry points for a cube
     kpoints = np.array(
@@ -325,10 +464,361 @@ def test_magnon(path="TB2J_results"):
     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
+        # print(f"Energies: {energies[i] * 1000:.3f} meV")  # Convert to meV
+        print(f"Energies: {energies[i] * 1000} meV")  # Convert to meV
+
+    print("\nPlotting magnon bands...")
+    magnon.plot_magnon_bands(
+        # kpoints=kpoints,
+        # labels=klabels,
+        path="GHPGPH,PN",
+        filename="magnon_bands.png",
+    )
 
     return magnon, Jq, energies
 
 
+def create_plot_script(filename: str):
+    """Create a Python script for plotting the saved band structure data.
+
+    Parameters
+    ----------
+    filename : str
+        Base filename (without extension) to use for the plot script
+    """
+    script_name = f"plot_{filename}.py"
+    script = '''#!/usr/bin/env python3
+"""Simple script to plot magnon band structure from saved data."""
+
+from TB2J.magnon.magnon_band import MagnonBand
+import matplotlib.pyplot as plt
+
+def plot_magnon_bands(input_file, output_file=None, ax=None, color='blue', show=True):
+    """Load and plot magnon band structure.
+    
+    Parameters
+    ----------
+    input_file : str
+        JSON file containing band structure data
+    output_file : str, optional
+        Output file for saving the plot
+    ax : matplotlib.axes.Axes, optional
+        Axes for plotting. If None, creates new figure
+    color : str, optional
+        Color of the band lines (default: blue)
+    show : bool, optional
+        Whether to show the plot on screen (default: True)
+    
+    Returns
+    -------
+    matplotlib.axes.Axes
+        The plotting axes
+    """
+    # Load band structure data
+    bands = MagnonBand.load(input_file)
+    
+    # Create plot
+    ax = bands.plot(
+        ax=ax,
+        filename=output_file,
+        color=color,
+        show=show
+    )
+    return ax
+
+if __name__ == "__main__":
+    # Usage example
+    # Example usage
+    import matplotlib.pyplot as plt
+    
+    # Create a figure and axis (optional)
+    fig, ax = plt.subplots(figsize=(6, 4))
+    
+    # Plot bands with custom color on given axis
+    plot_magnon_bands(
+        input_file="magnon_bands.json",
+        output_file="magnon_bands.png",
+        ax=ax,
+        color='red',
+        show=True
+    )
+'''
+
+    with open(script_name, 'w') as f:
+        f.write(script)
+    
+    import os
+    os.chmod(script_name, 0o755)  # Make executable
+
+def save_bands_data(
+    kpoints: np.ndarray,
+    energies: np.ndarray,
+    kpath_labels: List[Tuple[int, str]],
+    special_points: dict,
+    xcoords: Optional[Union[np.ndarray, List[np.ndarray]]],
+    filename: str,
+):
+    """Save magnon band structure data to a JSON file using MagnonBand class.
+
+    Parameters
+    ----------
+    kpoints : array_like
+        Array of k-points coordinates
+    energies : array_like
+        Array of band energies (in meV)
+    kpath_labels : list of (int, str)
+        List of tuples containing k-point indices and their labels
+    special_points : dict
+        Dictionary of special points and their coordinates
+    xcoords : array_like or list of arrays
+        x-coordinates for plotting (can be segmented)
+    filename : str
+        Output filename
+    """
+    from TB2J.magnon.magnon_band import MagnonBand  # Using same import as above
+
+    bands = MagnonBand(
+        energies=energies,
+        kpoints=kpoints,
+        kpath_labels=kpath_labels,
+        special_points=special_points,
+        xcoords=xcoords,
+    )
+    bands.save(filename)
+
+    # Create plotting script
+    base_name = filename.rsplit('.', 1)[0]
+    create_plot_script(base_name)
+
+    print(f"Band structure data saved to {filename}")
+    print(f"Created plotting script: plot_{base_name}.py")
+    print("Usage: ")
+    print(f"See plot_{base_name}.py for example usage")
+
+    return bands
+
+
+def plot_magnon_bands_from_TB2J(
+    params: MagnonParameters,
+):
+    """
+    Load TB2J results and plot magnon band structure along a specified k-path.
+
+    Parameters
+    ----------
+    path : str, optional
+        Path to TB2J results directory, default is "TB2J_results"
+    kpath : str, optional
+        String specifying the k-path, e.g. "GXMR" for Gamma-X-M-R path
+        Default is "GXMR"
+    npoints : int, optional
+        Number of k-points along the path, default is 300
+    filename : str, optional
+        Output file name for the plot, default is "magnon_bands.png"
+    Jiso : bool, optional
+        Include isotropic exchange interactions, default is True
+    Jani : bool, optional
+        Include anisotropic exchange interactions, default is False
+    DMI : bool, optional
+        Include Dzyaloshinskii-Moriya interactions, default is False
+    Q : array-like, optional
+        Propagation vector [Qx, Qy, Qz], default is [0, 0, 0]
+    uz_file : str, optional
+        Path to file containing quantization axes for each spin (natom×3 array)
+        If not provided, default [0, 0, 1] will be used for all spins
+    n : array-like, optional
+        Normal vector for rotation [nx, ny, nz], default is [0, 0, 1]
+
+    Returns
+    -------
+    magnon : Magnon
+        The Magnon instance used for calculations
+    """
+    # Load magnon calculator from TB2J results
+    print(f"Loading exchange parameters from {params.path}...")
+    magnon = Magnon.from_TB2J_results(
+        path=params.path, Jiso=params.Jiso, Jani=params.Jani, DMI=params.DMI
+    )
+
+    # Set reference vectors if provided
+    if any(x is not None for x in [params.Q, params.uz_file, params.n]):
+        Q = [0, 0, 0] if params.Q is None else params.Q
+        n = [0, 0, 1] if params.n is None else params.n
+
+        # Handle quantization axes
+        if params.uz_file is not None:
+            uz = np.loadtxt(params.uz_file)
+            if uz.shape[1] != 3:
+                raise ValueError(
+                    f"Quantization axes file should contain a natom×3 array. Got shape {uz.shape}"
+                )
+            if uz.shape[0] != magnon.nspin:
+                raise ValueError(
+                    f"Number of spins in uz file ({uz.shape[0]}) does not match the system ({magnon.nspin})"
+                )
+        else:
+            # Default: [0, 0, 1] for all spins
+            uz = np.array([[0.0, 0.0, 1.0] for _ in range(magnon.nspin)])
+
+        magnon.set_reference(Q, uz, n)
+
+    # Get band structure data
+    print(f"\nCalculating bands along path {params.kpath}...")
+    kpath_labels, bands, xlist = magnon.get_magnon_bands(
+        path=params.kpath,
+        npoints=params.npoints,
+    )
+
+    # Convert energies to meV
+    bands_meV = bands * 1000
+
+    # Save band structure data and create plot
+    data_file = params.filename.rsplit(".", 1)[0] + ".json"
+    print(f"\nSaving band structure data to {data_file}")
+
+    # Get k-points and special points
+    if params.kpath is None:
+        _, kptlist, _, _, spk = auto_kpath(magnon.cell, None, npoints=params.npoints)
+        kpoints = np.concatenate(kptlist)
+    else:
+        bandpath = magnon.cell.bandpath(path=params.kpath, npoints=params.npoints)
+        kpoints = bandpath.kpts
+        spk = bandpath.special_points
+        spk[r"$\Gamma$"] = spk.pop("G", np.zeros(3))
+
+    magnon_bands = save_bands_data(
+        kpoints=kpoints,
+        energies=bands_meV,
+        kpath_labels=kpath_labels,
+        special_points=spk,
+        xcoords=xlist,
+        filename=data_file,
+    )
+
+    # Plot band structure
+    print(f"Plotting bands to {params.filename}")
+    magnon_bands.plot(filename=params.filename)
+
+    return magnon
+
+
+def main():
+    """Command line interface for plotting magnon bands from TB2J results."""
+    import argparse
+
+    parser = argparse.ArgumentParser(
+        description="Plot magnon band structure from TB2J results"
+    )
+
+    # Add a mutually exclusive group for config file vs. command line arguments
+    group = parser.add_mutually_exclusive_group()
+    group.add_argument(
+        "--config",
+        type=str,
+        help="Path to TOML configuration file",
+    )
+    group.add_argument(
+        "--save-config",
+        type=str,
+        help="Save default configuration to specified TOML file",
+    )
+
+    # Command line arguments (used if no config file is provided)
+    parser.add_argument(
+        "--path",
+        default="TB2J_results",
+        help="Path to TB2J results directory (default: TB2J_results)",
+    )
+    parser.add_argument(
+        "--kpath",
+        default=None,
+        help="k-path specification (default: auto-detected from type of cell)",
+    )
+    parser.add_argument(
+        "--npoints",
+        type=int,
+        default=300,
+        help="Number of k-points along the path (default: 300)",
+    )
+    parser.add_argument(
+        "--output",
+        default="magnon_bands.png",
+        help="Output file name (default: magnon_bands.png)",
+    )
+    parser.add_argument(
+        "--Jiso",
+        action="store_true",
+        default=True,
+        help="Include isotropic exchange interactions (default: True)",
+    )
+    parser.add_argument(
+        "--no-Jiso",
+        action="store_false",
+        dest="Jiso",
+        help="Exclude isotropic exchange interactions",
+    )
+    parser.add_argument(
+        "--Jani",
+        action="store_true",
+        default=False,
+        help="Include anisotropic exchange interactions (default: False)",
+    )
+    parser.add_argument(
+        "--DMI",
+        action="store_true",
+        default=False,
+        help="Include Dzyaloshinskii-Moriya interactions (default: False)",
+    )
+    parser.add_argument(
+        "--Q",
+        nargs=3,
+        type=float,
+        metavar=("Qx", "Qy", "Qz"),
+        help="Propagation vector [Qx, Qy, Qz] (default: [0, 0, 0])",
+    )
+    parser.add_argument(
+        "--uz-file",
+        type=str,
+        help="Path to file containing quantization axes for each spin (natom×3 array)",
+    )
+    parser.add_argument(
+        "--n",
+        nargs=3,
+        type=float,
+        metavar=("nx", "ny", "nz"),
+        help="Normal vector for rotation [nx, ny, nz] (default: [0, 0, 1])",
+    )
+
+    args = parser.parse_args()
+
+    # Handle configuration file options
+    if args.save_config:
+        # Create default parameters and save to file
+        params = MagnonParameters()
+        params.to_toml(args.save_config)
+        print(f"Saved default configuration to {args.save_config}")
+        return
+
+    if args.config:
+        # Load parameters from config file
+        params = MagnonParameters.from_toml(args.config)
+    else:
+        # Create parameters from command line arguments
+        params = MagnonParameters(
+            path=args.path,
+            kpath=args.kpath,
+            npoints=args.npoints,
+            filename=args.output,
+            Jiso=args.Jiso,
+            Jani=args.Jani,
+            DMI=args.DMI,
+            Q=args.Q if args.Q is not None else None,
+            uz_file=args.uz_file,
+            n=args.n if args.n is not None else None,
+        )
+
+    plot_magnon_bands_from_TB2J(params)
+
+
 if __name__ == "__main__":
-    test_magnon()
+    main()