refactor: rename visualize module to output module and separate concerns into core, csv_export, and visualization modules

2026-07-07 12:59:11 +03:00
parent 6dc1e3924d
commit 9bf1e1333e
11 changed files with 989 additions and 884 deletions
@@ -0,0 +1,444 @@
+"""
+Visualization module — Plotting functionality for BetaSpectrumAnalyzer.
+"""
+
+from __future__ import annotations
+
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, Dict, Optional
+
+import numpy as np
+import matplotlib.pyplot as plt
+
+from beta_spectrum.logging_utils import get_git_short_hash
+from beta_spectrum.utils import T_to_W
+
+if TYPE_CHECKING:
+    from matplotlib.figure import Figure
+
+    from beta_spectrum.output.core import BetaSpectrumAnalyzer
+
+
+class Visualizer:
+    """Visualization functionality for BetaSpectrumAnalyzer."""
+
+    @staticmethod
+    def plot_analysis(
+        analyzer: BetaSpectrumAnalyzer,
+        save_path: Optional[str] = None,
+        show_components: bool = True,
+    ) -> None:
+        """
+        Create visualization of the spectrum and all correction factors.
+
+        Parameters
+        ----------
+        analyzer : BetaSpectrumAnalyzer
+            The analyzer instance to visualize.
+        save_path : str, optional
+            Path to save the figure.
+        show_components : bool
+            If True, show full debug view.
+            If False, show only the spectrum plot.
+        """
+        total = analyzer.total_spectrum(normalize=True)
+        components = analyzer.components
+        analyzer._timestamp = datetime.now(timezone.utc).strftime(
+            "%Y-%m-%d %H:%M:%S UTC"
+        )
+
+        if show_components:
+            Visualizer._plot_debug_view(analyzer, total, components, save_path)
+        else:
+            Visualizer._plot_spectra_view(analyzer, total, components, save_path)
+
+    @staticmethod
+    def _build_figure_header(analyzer: BetaSpectrumAnalyzer, fig: "Figure") -> None:
+        """Add a compact header with nuclide info, e_step, commit and timestamp outside axes."""
+        parent = analyzer._element_symbol(analyzer.config.Z_parent)
+        daughter = analyzer._element_symbol(analyzer.config.Z_daughter)
+        e_step_keV = analyzer.config.e_step_MeV * 1000
+        header = (
+            f"{parent}{analyzer.config.A_number} -> {daughter}{analyzer.config.A_number}  "
+            f"|  E₀={analyzer.config.endpoint_MeV * 1000:.1f} keV  "
+            f"|  ΔE={e_step_keV:.1f} keV"
+        )
+        fig.text(
+            0.01,
+            0.98,
+            header,
+            transform=fig.transFigure,
+            fontsize=9,
+            verticalalignment="top",
+            horizontalalignment="left",
+            fontfamily="monospace",
+        )
+
+        # Add neutrino mass info to header if non-zero
+        if analyzer.config.neutrino_mass_MeV > 0:
+            nu_mass_text = f"|  m_ν={analyzer.config.neutrino_mass_MeV:.6f} MeV"
+            fig.text(
+                0.01,
+                0.96,
+                nu_mass_text,
+                transform=fig.transFigure,
+                fontsize=8,
+                verticalalignment="top",
+                horizontalalignment="left",
+                fontfamily="monospace",
+                color="red",
+            )
+
+        id_text = f"{get_git_short_hash(6)}  |  {analyzer._timestamp}"
+        fig.text(
+            0.01,
+            0.01,
+            id_text,
+            transform=fig.transFigure,
+            fontsize=8,
+            verticalalignment="bottom",
+            horizontalalignment="left",
+            fontfamily="monospace",
+            color="gray",
+        )
+
+    @staticmethod
+    def _build_figure_title(analyzer: BetaSpectrumAnalyzer, fig: "Figure") -> str:
+        """Build a unified figure title."""
+        parent = analyzer._element_symbol(analyzer.config.Z_parent)
+        daughter = analyzer._element_symbol(analyzer.config.Z_daughter)
+        n = analyzer.config.A_number
+        title = f"Beta-decay: {n}{parent} -> {n}{daughter}"
+        if analyzer._n_branches > 1:
+            # Count active branches (not excluded by neutrino mass or intensity cutoff)
+            active_branches = 0
+            for i, branch in enumerate(analyzer._unified_branches):
+                branch_endpoint_keV = branch.endpoint_MeV * 1000
+                neutrino_mass_keV = analyzer.config.neutrino_mass_MeV * 1000
+                is_excluded_by_mass = branch_endpoint_keV <= neutrino_mass_keV
+                is_excluded_by_intensity = (
+                    branch.intensity < analyzer.config.intensity_cutoff
+                )
+                if not is_excluded_by_mass and not is_excluded_by_intensity:
+                    active_branches += 1
+
+            if active_branches < analyzer._n_branches:
+                title += f" ({active_branches} of {analyzer._n_branches} branches)"
+            else:
+                title += f" ({analyzer._n_branches} branches)"
+        return title
+
+    @staticmethod
+    def _format_branch_label(analyzer: BetaSpectrumAnalyzer, i: int) -> str:
+        """Abbreviated branch label for legend."""
+        branch = analyzer._unified_branches[i]
+        e0 = branch.endpoint_MeV * 1000
+        intensity_pct = branch.intensity * 100
+
+        # Check if branch is excluded
+        branch_endpoint_keV = branch.endpoint_MeV * 1000
+        neutrino_mass_keV = analyzer.config.neutrino_mass_MeV * 1000
+        is_excluded_by_mass = branch_endpoint_keV <= neutrino_mass_keV
+        is_excluded_by_intensity = branch.intensity < analyzer.config.intensity_cutoff
+
+        if is_excluded_by_mass:
+            return f"Br. {i + 1}: {e0:.1f} keV ({intensity_pct:.1f}%) - excluded (endpoint < m_nu)"
+        elif is_excluded_by_intensity:
+            return f"Br. {i + 1}: {e0:.1f} keV ({intensity_pct:.1f}%) - excluded (intensity < cutoff)"
+        else:
+            return f"Br. {i + 1}: {e0:.1f} keV ({intensity_pct:.1f}%)"
+
+    @staticmethod
+    def _plot_spectra_view(
+        analyzer: BetaSpectrumAnalyzer,
+        total: np.ndarray,
+        components: Dict[str, np.ndarray],
+        save_path: Optional[str],
+    ) -> None:
+        """Plot total spectrum with individual branch decomposition."""
+        fig, ax = plt.subplots(figsize=(12, 7))
+
+        branch_colors = [
+            "#1f77b4",
+            "#ff7f0e",
+            "#2ca02c",
+            "#d62728",
+            "#9467bd",
+            "#8c564b",
+            "#e377c2",
+            "#7f7f7f",
+            "#bcbd22",
+            "#17becf",
+        ]
+
+        e_step_keV = analyzer.config.e_step_MeV * 1000
+
+        # Compute branch contributions (universal * intensity * branch_spectrum)
+        branch_contribs = []
+        # Universal components are shared (no per-branch key prefix)
+        universal_names = [
+            "Fermi",
+            "Screening",
+            "FiniteSizeL0",
+            "ChargeDistributionU",
+            "Exchange",
+        ]
+        for i in range(analyzer._n_branches):
+            universal_product = np.ones_like(analyzer.W)
+            for name in universal_names:
+                key = analyzer._get_component_key(i, name)
+                if key in components:
+                    universal_product *= components[key]
+                elif name in components:  # fallback: universal components
+                    universal_product *= components[name]
+            bs = analyzer._get_branch_spectrum(i)
+            contrib = universal_product * analyzer._unified_branches[i].intensity * bs
+            branch_contribs.append(contrib)
+
+        total_raw = sum(branch_contribs)
+
+        ax.plot(
+            analyzer.energies_MeV,
+            total_raw,
+            "k-",
+            lw=2.5,
+            label="Total spectrum",
+            zorder=10,
+        )
+
+        # Only plot active branches (not excluded by neutrino mass or intensity cutoff)
+        active_branch_indices = []
+        for i in range(analyzer._n_branches):
+            branch_endpoint_keV = analyzer._unified_branches[i].endpoint_MeV * 1000
+            neutrino_mass_keV = analyzer.config.neutrino_mass_MeV * 1000
+            is_excluded_by_mass = branch_endpoint_keV <= neutrino_mass_keV
+            is_excluded_by_intensity = (
+                analyzer._unified_branches[i].intensity
+                < analyzer.config.intensity_cutoff
+            )
+            if not is_excluded_by_mass and not is_excluded_by_intensity:
+                active_branch_indices.append(i)
+
+        for i in active_branch_indices:
+            contrib = branch_contribs[i]
+            color = branch_colors[i % len(branch_colors)]
+            ax.plot(
+                analyzer.energies_MeV,
+                contrib,
+                color=color,
+                lw=1.5,
+                alpha=0.7,
+                label=Visualizer._format_branch_label(analyzer, i),
+            )
+
+        ax.set_xlabel(r"Electron kinetic energy $E$ [MeV]", fontsize=12)
+        ax.set_ylabel(f"Normalized Counts / {e_step_keV:.1f} keV", fontsize=12)
+        fig.suptitle(
+            Visualizer._build_figure_title(analyzer, fig),
+            fontsize=14,
+            fontweight="bold",
+        )
+        ax.set_yscale("log")
+        ax.grid(True, alpha=0.3)
+        ax.legend(
+            loc="best",
+            fontsize=9,
+            ncol=min(2, len(active_branch_indices) + 1),
+        )
+
+        Visualizer._build_figure_header(analyzer, fig)
+
+        plt.tight_layout()
+
+        if save_path:
+            fig.savefig(save_path, dpi=150, bbox_inches="tight")
+            print(f"Figure saved to {save_path}")
+
+    @staticmethod
+    def _plot_debug_view(
+        analyzer: BetaSpectrumAnalyzer,
+        total: np.ndarray,
+        components: Dict[str, np.ndarray],
+        save_path: Optional[str],
+    ) -> None:
+        """Debug view: vertical layout with 4 panels."""
+        n_branches = analyzer._n_branches
+        e_step_keV = analyzer.config.e_step_MeV * 1000
+
+        fig, axes = plt.subplots(4, 1, figsize=(12, 16), sharex=True)
+        fig.suptitle(
+            Visualizer._build_figure_title(analyzer, fig),
+            fontsize=14,
+            fontweight="bold",
+        )
+
+        branch_colors = [
+            "#1f77b4",
+            "#ff7f0e",
+            "#2ca02c",
+            "#d62728",
+            "#9467bd",
+            "#8c564b",
+            "#e377c2",
+            "#7f7f7f",
+            "#bcbd22",
+            "#17becf",
+        ]
+
+        # Panel 1: Total spectrum + per-branch spectra
+        branch_contribs = []
+        universal_names = [
+            "Fermi",
+            "Screening",
+            "FiniteSizeL0",
+            "ChargeDistributionU",
+            "Exchange",
+        ]
+        for i in range(n_branches):
+            universal_product = np.ones_like(analyzer.W)
+            for name in universal_names:
+                key = analyzer._get_component_key(i, name)
+                if key in components:
+                    universal_product *= components[key]
+                elif name in components:
+                    universal_product *= components[name]
+            bs = analyzer._get_branch_spectrum(i)
+            contrib = universal_product * analyzer._unified_branches[i].intensity * bs
+            branch_contribs.append(contrib)
+
+        total_raw = sum(branch_contribs)
+
+        ax = axes[0]
+        ax.plot(
+            analyzer.energies_MeV,
+            total_raw,
+            "k-",
+            lw=2,
+            label="Total spectrum",
+            zorder=10,
+        )
+
+        # Only plot active branches (not excluded by neutrino mass or intensity cutoff)
+        active_branch_indices = []
+        for i in range(n_branches):
+            branch_endpoint_keV = analyzer._unified_branches[i].endpoint_MeV * 1000
+            neutrino_mass_keV = analyzer.config.neutrino_mass_MeV * 1000
+            is_excluded_by_mass = branch_endpoint_keV <= neutrino_mass_keV
+            is_excluded_by_intensity = (
+                analyzer._unified_branches[i].intensity
+                < analyzer.config.intensity_cutoff
+            )
+            if not is_excluded_by_mass and not is_excluded_by_intensity:
+                active_branch_indices.append(i)
+
+        for i in active_branch_indices:
+            contrib = branch_contribs[i]
+            color = branch_colors[i % len(branch_colors)]
+            ax.plot(
+                analyzer.energies_MeV,
+                contrib,
+                color=color,
+                lw=1.5,
+                alpha=0.7,
+                label=Visualizer._format_branch_label(analyzer, i),
+            )
+        ax.set_yscale("log")
+        ax.set_ylabel(f"Normalized Counts / {e_step_keV:.1f} keV", fontsize=10)
+        ax.set_title("Spectrum", fontsize=11)
+        ax.grid(True, alpha=0.3)
+        ax.legend(fontsize=8, loc="best", ncol=min(2, len(active_branch_indices) + 1))
+
+        # Panel 2: Fermi function + all phase spaces
+        ax = axes[1]
+        if "Fermi" in components:
+            fermi_vals = components["Fermi"]
+            fermi_scaled = fermi_vals / np.max(fermi_vals)
+            ax.plot(
+                analyzer.energies_MeV,
+                fermi_scaled,
+                "r-",
+                lw=1.5,
+                label="Fermi",
+                alpha=0.8,
+            )
+        for i in active_branch_indices:
+            color = branch_colors[i % len(branch_colors)]
+            branch = analyzer._unified_branches[i]
+            key = analyzer._get_component_key(i, "PhaseSpace")
+            if key in components:
+                ps_vals = components[key]
+                W0_branch = T_to_W(branch.endpoint_MeV)
+                mask = analyzer.W <= W0_branch
+                ps_masked = np.where(mask, ps_vals, 0.0)
+                ps_max = np.max(ps_masked)
+                if ps_max > 0:
+                    ps_scaled = ps_masked / ps_max
+                else:
+                    ps_scaled = ps_masked
+                ax.plot(
+                    analyzer.energies_MeV,
+                    ps_scaled,
+                    color=color,
+                    lw=1.5,
+                    alpha=0.7,
+                    label=f"PS Br. {i + 1}",
+                )
+        ax.set_ylabel("Normalized Factor", fontsize=10)
+        ax.set_title("Fermi & Phase Space", fontsize=11)
+        ax.grid(True, alpha=0.3)
+        ax.legend(fontsize=8, loc="best")
+
+        # Panel 3: Universal W0-independent corrections
+        ax = axes[2]
+        universal_components = [
+            "Screening",
+            "FiniteSizeL0",
+            "ChargeDistributionU",
+            "Exchange",
+        ]
+        colors = ["green", "orange", "purple", "brown"]
+        for name, color in zip(universal_components, colors):
+            if name in components:
+                values = components[name]
+                values_scaled = values / np.max(values)
+                ax.plot(
+                    analyzer.energies_MeV,
+                    values_scaled,
+                    color=color,
+                    lw=1.5,
+                    label=name,
+                    alpha=0.8,
+                )
+        ax.set_ylabel("Normalized Factor", fontsize=10)
+        ax.set_title("Universal Corrections", fontsize=11)
+        ax.grid(True, alpha=0.3)
+        ax.legend(fontsize=8, loc="best")
+
+        # Panel 4: W0-dependent corrections (Radiative for each branch)
+        ax = axes[3]
+        for i in active_branch_indices:
+            color = branch_colors[i % len(branch_colors)]
+            key = analyzer._get_component_key(i, "Radiative")
+            if key in components:
+                rad_vals = components[key]
+                ax.plot(
+                    analyzer.energies_MeV,
+                    rad_vals,
+                    color=color,
+                    lw=1.5,
+                    alpha=0.7,
+                    label=f"Br. {i + 1}",
+                )
+        ax.set_ylabel("Radiative Correction", fontsize=10)
+        ax.set_xlabel(r"Electron kinetic energy $E$ [MeV]", fontsize=10)
+        ax.set_title("Radiative Corrections", fontsize=11)
+        ax.grid(True, alpha=0.3)
+        ax.legend(fontsize=8, loc="best", ncol=min(2, len(active_branch_indices) + 1))
+
+        Visualizer._build_figure_header(analyzer, fig)
+
+        plt.tight_layout()
+
+        if save_path:
+            fig.savefig(save_path, dpi=150, bbox_inches="tight")
+            print(f"Figure saved to {save_path}")