Source code for mag4.magnon.plot

"""Magnon band-structure plotting (matplotlib, optional)."""

from __future__ import annotations

import logging
from typing import Dict

import numpy as np

log = logging.getLogger(__name__)


[docs] def plot_magnon_bands(magnon_bands: np.ndarray, kpath: Dict, S: float, n_sub: int, name: str) -> str: """Plot the magnon band structure and save to ``<name>_magnon_bands.png``. Returns the output path, or ``""`` if matplotlib is unavailable. """ try: import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt except ImportError: log.error("matplotlib is required. pip install matplotlib") return "" distances = kpath["distances"] tick_pos = kpath["tick_positions"] tick_lbl = kpath["tick_labels"] N_q, N_bands = magnon_bands.shape fig_w = max(7, min(14, len(tick_lbl) * 0.7)) fig, ax = plt.subplots(figsize=(fig_w, 5)) band_color = "steelblue" break_indices = kpath.get("break_indices", set()) sorted_breaks = sorted(break_indices) plot_segs = [] prev = 0 for b in sorted_breaks: if b > prev: plot_segs.append((prev, b + 1)) prev = b plot_segs.append((prev, N_q)) plot_segs = [(s, e) for s, e in plot_segs if e - s > 1] log.info(" Plot segments: %d", len(plot_segs)) for k in range(N_bands): label = f"{N_bands} magnon branches" if k == 0 else None first_seg = True for start, end in plot_segs: ax.plot(distances[start:end], magnon_bands[start:end, k], color=band_color, linewidth=1.2, alpha=0.85, label=label if first_seg else None) first_seg = False ax.axhline(0, color="black", linewidth=0.6, linestyle="--", alpha=0.5) for xpos in tick_pos: ax.axvline(xpos, color="black", linewidth=0.5, alpha=0.4) ax.set_xlim(distances[0], distances[-1]) ax.set_ylabel("Magnon energy (meV)", fontsize=11) ax.set_xlabel("") fig.canvas.draw() ax_width_inch = fig.get_figwidth() * ax.get_position().width x_range = distances[-1] - distances[0] min_spacing = (1.5 * 7.0 / fig.dpi / ax_width_inch) * x_range display_pos, display_lbl = [], [] for i, (pos, lbl) in enumerate(zip(tick_pos, tick_lbl)): is_last = (i == len(tick_pos) - 1) if display_pos and (pos - display_pos[-1]) < min_spacing and not is_last: prev_lbl = display_lbl[-1] if lbl not in prev_lbl.split("|"): display_lbl[-1] = prev_lbl + "|" + lbl display_pos[-1] = (display_pos[-1] + pos) / 2.0 else: display_pos.append(pos) display_lbl.append(lbl) ax.set_xticks(tick_pos) ax.set_xticklabels([""] * len(tick_pos)) for pos, lbl in zip(display_pos, display_lbl): ax.text(pos, -0.07, lbl, ha="center", va="top", fontsize=10, transform=ax.get_xaxis_transform()) path_str = " – ".join(tick_lbl) band_word = "bands" if N_bands > 1 else "band" ax.set_title( f"Magnon band structure — LSWT (S = {S}, {N_bands} {band_word})\n" f"Path: {path_str}", fontsize=10, linespacing=1.6) ax.legend(fontsize=9, loc="upper right") plt.tight_layout() out_path = f"{name}_magnon_bands.png" fig.savefig(out_path, dpi=150, bbox_inches="tight") plt.close(fig) log.info("Magnon band structure saved to %s", out_path) return out_path