xarpes 0.6.0__tar.gz → 0.6.1__tar.gz

{xarpes-0.6.0 → xarpes-0.6.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: xarpes
-Version: 0.6.0
+Version: 0.6.1
 Summary: Extraction from angle resolved photoemission spectra
 Author: xARPES Developers
 Requires-Python: >=3.7.0

{xarpes-0.6.0 → xarpes-0.6.1}/doc/notebooks/verification.ipynb

@@ -15,7 +15,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 4,
    "id": "cae844d7",
    "metadata": {},
    "outputs": [
@@ -44,7 +44,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 5,
    "id": "043a5793",
    "metadata": {},
    "outputs": [],
@@ -57,7 +57,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": null,
    "id": "f701b78c",
    "metadata": {},
    "outputs": [],

{xarpes-0.6.0 → xarpes-0.6.1}/xarpes/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.6.0"
+__version__ = "0.6.1"
 
 from importlib import import_module
 

{xarpes-0.6.0 → xarpes-0.6.1}/xarpes/bandmap.py

@@ -14,7 +14,7 @@
 import numpy as np
 from igor2 import binarywave
 from .plotting import get_ax_fig_plt, add_fig_kwargs
-from .functions import fit_leastsq, extend_function
+from .functions import fit_least_squares, extend_function
 from .distributions import FermiDirac, Linear
 from .constants import PREF
 
@@ -746,9 +746,10 @@ class BandMap:
 
         extra_args = (self.temperature,)
 
-        popt, pcov = fit_leastsq(
-        parameters, energy_range, integrated_intensity, fdir_initial,
-        self.energy_resolution, None, *extra_args)
+        popt, pcov, success = fit_least_squares(
+            p0=parameters, xdata=energy_range, ydata=integrated_intensity,
+            function=fdir_initial, resolution=self.energy_resolution,
+            yerr=None, bounds=None, extra_args=extra_args)
 
         # Update hnuminPhi; automatically sets self.enel
         self.hnuminPhi = popt[0]
@@ -846,9 +847,10 @@ class BandMap:
         for indx in range(angle_max_index - angle_min_index + 1):
             edge = Intensities[:, indx]
             
-            parameters, pcov = fit_leastsq(
-            parameters, energy_range, edge, fdir_initial,
-            self.energy_resolution, None, *extra_args)
+            parameters, pcov, success = fit_least_squares(
+                p0=parameters, xdata=energy_range, ydata=edge,
+                function=fdir_initial, resolution=self.energy_resolution,
+                yerr=None, bounds=None, extra_args=extra_args)
 
             nmps[indx] = parameters[0]
             stds[indx] = np.sqrt(np.diag(pcov)[0])
@@ -861,8 +863,9 @@ class BandMap:
         
         lin_fun = Linear(offset_guess, slope_guess, 'Linear')
                     
-        popt, pcov = fit_leastsq(parameters, angle_range, nmps, lin_fun, None,
-                                 stds)
+        popt, pcov, success = fit_least_squares(p0=parameters, xdata=angle_range, 
+                        ydata=nmps, function=lin_fun, resolution=None,
+                                 yerr=stds, bounds=None)
 
         linsp = lin_fun(angle_range, popt[0], popt[1])
 

{xarpes-0.6.0 → xarpes-0.6.1}/xarpes/functions.py

@@ -148,8 +148,8 @@ def extend_function(abscissa_range, abscissa_resolution):
     return extend, step, numb
 
 
-def error_function(p, xdata, ydata, function, resolution, yerr, extra_args):
-    r"""The error function used inside the fit_leastsq function.
+def error_function(p, xdata, ydata, function, resolution, yerr, *extra_args):
+    r"""The error function used inside the fit_least_squares function.
 
     Parameters
     ----------
@@ -187,57 +187,54 @@ def error_function(p, xdata, ydata, function, resolution, yerr, extra_args):
     return residual
 
 
-def fit_leastsq(p0, xdata, ydata, function, resolution=None,
-                yerr=None, *extra_args):
-    r"""Wrapper around scipy.optimize.leastsq.
+def fit_least_squares(p0, xdata, ydata, function, resolution=None, yerr=None,
+                      bounds=None, extra_args=None):
+    r"""Least-squares fit using `scipy.optimize.least_squares`.
 
-    Parameters
-    ----------
-    p0 : ndarray
-        Initial guess for parameters to be optimized.
-    xdata : ndarray
-        Abscissa values the function is evaluated on.
-    ydata : ndarray
-        Measured values to compare to.
-    function : callable
-        Function or class with __call__ method to evaluate.
-    resolution : float or None, optional
-        Convolution resolution (sigma), if applicable.
-    yerr : ndarray or None, optional
-        Standard deviations of ydata. Defaults to ones if None.
-    extra_args : tuple
-        Additional arguments passed to the function.
+    Default behavior is Levenberg–Marquardt (`method="lm"`) when unbounded.
+    If `bounds` is provided, switches to trust-region reflective (`"trf"`).
 
-    Returns
-    -------
-    pfit_leastsq : ndarray
-        Optimized parameters.
-    pcov : ndarray or float
-        Scaled covariance matrix of the optimized parameters.
-        If the covariance could not be estimated, returns np.inf.
+    Returns (pfit, pcov, success) in the same style as the old `leastsq`
+    wrapper, with an additional boolean `success` from SciPy.
     """
-    from scipy.optimize import leastsq
+    from scipy.optimize import least_squares
 
     if yerr is None:
         yerr = np.ones_like(ydata)
 
-    pfit, pcov, infodict, errmsg, success = leastsq(
-        error_function,
-        p0,
-        args=(xdata, ydata, function, resolution, yerr, extra_args),
-        full_output=1
-    )
+    if extra_args is None:
+        extra_args = ()
+
+    def _residuals(p):
+        return error_function(
+            p, xdata, ydata, function, resolution, yerr, *extra_args
+        )
+
+    if bounds is None:
+        res = least_squares(_residuals, p0, method="lm")
+    else:
+        res = least_squares(_residuals, p0, method="trf", bounds=bounds)
+
+    pfit = res.x
+    success = bool(getattr(res, "success", False))
+
+    m = len(ydata)
+    n = pfit.size
 
-    if (len(ydata) > len(p0)) and pcov is not None:
-        s_sq = (
-            error_function(pfit, xdata, ydata, function, resolution,
-                           yerr, extra_args) ** 2
-        ).sum() / (len(ydata) - len(p0))
-        pcov *= s_sq
+    if (m > n) and (res.jac is not None) and res.jac.size:
+        resid = res.fun
+        s_sq = (resid ** 2).sum() / (m - n)
+
+        try:
+            jtj = res.jac.T @ res.jac
+            pcov = np.linalg.inv(jtj) * s_sq
+        except np.linalg.LinAlgError:
+            pcov = np.inf
     else:
         pcov = np.inf
 
-    return pfit, pcov
+    return pfit, pcov, success
+
 
 
 def download_examples():

{xarpes-0.6.0 → xarpes-0.6.1}/xarpes/selfenergies.py

@@ -34,7 +34,7 @@ class SelfEnergy:
             else:
                 raise ValueError(
                     "`properties` must be a dict or a single dict in a list."
-             )
+                )
 
         # single source of truth for all params (+ their *_sigma)
         self._properties = dict(properties or {})
@@ -84,12 +84,19 @@ class SelfEnergy:
         self._imag = None
         self._imag_sigma = None
 
+        # lazy caches for α²F(ω) extraction results
+        self._a2f_spectrum = None
+        self._a2f_model = None
+        self._a2f_omega_range = None
+        self._a2f_alpha_select = None
+        self._a2f_cost = None
+
     def _check_mass_velocity_exclusivity(self):
         """Ensure that fermi_velocity and bare_mass are not both set."""
         if (self._fermi_velocity is not None) and (self._bare_mass is not None):
             raise ValueError(
-            "Cannot set both `fermi_velocity` and  `bare_mass`: choose one "
-            "physical parametrization (SpectralLinear or SpectralQuadratic)."
+                "Cannot set both `fermi_velocity` and  `bare_mass`: choose one "
+                "physical parametrization (SpectralLinear or SpectralQuadratic)."
             )
 
     # ---------------- core read-only axes ----------------
@@ -289,8 +296,10 @@ class SelfEnergy:
                 self._peak_positions = ((-1.0 if self._side == "left"
                                         else 1.0) * kpar_mag)
             else:
-                self._peak_positions = (np.sqrt(self._ekin_range / PREF)
-                * np.sin(np.deg2rad(self._peak)))
+                self._peak_positions = (
+                    np.sqrt(self._ekin_range / PREF)
+                    * np.sin(np.deg2rad(self._peak))
+                )
         return self._peak_positions
     
 
@@ -346,6 +355,31 @@ class SelfEnergy:
                 return None
             self._real_sigma = self._compute_real_sigma()
         return self._real_sigma
+    
+    @property
+    def a2f_spectrum(self):
+        """Cached α²F(ω) spectrum from last extraction (or None)."""
+        return self._a2f_spectrum
+
+    @property
+    def a2f_model(self):
+        """Cached MEM model spectrum from last extraction (or None)."""
+        return self._a2f_model
+
+    @property
+    def a2f_omega_range(self):
+        """Cached ω grid for the last extraction (or None)."""
+        return self._a2f_omega_range
+
+    @property
+    def a2f_alpha_select(self):
+        """Cached selected alpha from last extraction (or None)."""
+        return self._a2f_alpha_select
+
+    @property
+    def a2f_cost(self):
+        """Cached cost from last bayesian_loop (or None)."""
+        return self._a2f_cost
 
 
     def _compute_imag(self, fermi_velocity=None, bare_mass=None):
@@ -553,15 +587,37 @@ class SelfEnergy:
         imag_label = rf"$-\Sigma_{{\mathrm{{{safe_label}}}}}''(E)$"
 
         return real_label, imag_label
+    
+    def _a2f_legend_labels(self):
+        """Return (a2f_label, model_label) for legend with safe subscripts."""
+        se_label = getattr(self, "label", None)
+
+        if se_label is None:
+            return r"$\alpha^2F(\omega)$", r"$m(\omega)$"
+
+        safe_label = str(se_label).replace("_", r"\_")
+        if safe_label == "":
+            return r"$\alpha^2F(\omega)$", r"$m(\omega)$"
+
+        a2f_label = rf"$\alpha^2F_{{\mathrm{{{safe_label}}}}}(\omega)$"
+        model_label = rf"$m_{{\mathrm{{{safe_label}}}}}(\omega)$"
+        return a2f_label, model_label
 
     @add_fig_kwargs
-    def plot_real(self, ax=None, **kwargs):
+    def plot_real(self, ax=None, scale="eV", resolution_range="absent", **kwargs):
         r"""Plot the real part Σ' of the self-energy as a function of E-μ.
 
         Parameters
         ----------
         ax : Matplotlib-Axes or None
             Axis to plot on. Created if not provided by the user.
+        scale : {"eV", "meV"}
+            Units for both axes. If "meV", x and y (and yerr) are multiplied by
+            `KILO`.
+        resolution_range : {"absent", "applied"}
+            If "applied", removes points with |E-μ| <= energy_resolution (around
+            the chemical potential). The energy resolution is taken from
+            ``self.energy_resolution`` (in eV) and scaled consistently with `scale`.
         **kwargs :
             Additional keyword arguments passed to ``ax.errorbar``.
 
@@ -574,10 +630,31 @@ class SelfEnergy:
 
         ax, fig, plt = get_ax_fig_plt(ax=ax)
 
+        if scale not in ("eV", "meV"):
+            raise ValueError("scale must be either 'eV' or 'meV'.")
+        if resolution_range not in ("absent", "applied"):
+            raise ValueError("resolution_range must be 'absent' or 'applied'.")
+
+        factor = KILO if scale == "meV" else 1.0
+
         x = self.enel_range
         y = self.real
         y_sigma = self.real_sigma
 
+        if x is not None:
+            x = factor * np.asarray(x, dtype=float)
+        if y is not None:
+            y = factor * np.asarray(y, dtype=float)
+
+        if resolution_range == "applied" and x is not None and y is not None:
+            res = self.energy_resolution
+            if res is not None:
+                keep = np.abs(x) > (factor * float(res))
+                x = x[keep]
+                y = y[keep]
+                if y_sigma is not None:
+                    y_sigma = np.asarray(y_sigma, dtype=float)[keep]
+
         real_label, _ = self._se_legend_labels()
         kwargs.setdefault("label", real_label)
 
@@ -587,23 +664,33 @@ class SelfEnergy:
                     "Warning: some Σ'(E) uncertainty values are missing. "
                     "Error bars omitted at those energies."
                 )
-            kwargs.setdefault("yerr", xprs.sigma_confidence * y_sigma)
+            kwargs.setdefault("yerr", xprs.sigma_confidence * factor * y_sigma)
 
         ax.errorbar(x, y, **kwargs)
-        ax.set_xlabel(r"$E-\mu$ (eV)")
-        ax.set_ylabel(r"$\Sigma'(E)$ (eV)")
+
+        x_unit = "meV" if scale == "meV" else "eV"
+        ax.set_xlabel(rf"$E-\mu$ ({x_unit})")
+        ax.set_ylabel(rf"$\Sigma'(E)$ ({x_unit})")
         ax.legend()
 
         return fig
 
+
     @add_fig_kwargs
-    def plot_imag(self, ax=None, **kwargs):
+    def plot_imag(self, ax=None, scale="eV", resolution_range="absent", **kwargs):
         r"""Plot the imaginary part -Σ'' of the self-energy vs. E-μ.
 
         Parameters
         ----------
         ax : Matplotlib-Axes or None
             Axis to plot on. Created if not provided by the user.
+        scale : {"eV", "meV"}
+            Units for both axes. If "meV", x and y (and yerr) are multiplied by
+            `KILO`.
+        resolution_range : {"absent", "applied"}
+            If "applied", removes points with |E-μ| <= energy_resolution (around
+            the chemical potential). The energy resolution is taken from
+            ``self.energy_resolution`` (in eV) and scaled consistently with `scale`.
         **kwargs :
             Additional keyword arguments passed to ``ax.errorbar``.
 
@@ -616,10 +703,31 @@ class SelfEnergy:
 
         ax, fig, plt = get_ax_fig_plt(ax=ax)
 
+        if scale not in ("eV", "meV"):
+            raise ValueError("scale must be either 'eV' or 'meV'.")
+        if resolution_range not in ("absent", "applied"):
+            raise ValueError("resolution_range must be 'absent' or 'applied'.")
+
+        factor = KILO if scale == "meV" else 1.0
+
         x = self.enel_range
         y = self.imag
         y_sigma = self.imag_sigma
 
+        if x is not None:
+            x = factor * np.asarray(x, dtype=float)
+        if y is not None:
+            y = factor * np.asarray(y, dtype=float)
+
+        if resolution_range == "applied" and x is not None and y is not None:
+            res = self.energy_resolution
+            if res is not None:
+                keep = np.abs(x) > (factor * float(res))
+                x = x[keep]
+                y = y[keep]
+                if y_sigma is not None:
+                    y_sigma = np.asarray(y_sigma, dtype=float)[keep]
+
         _, imag_label = self._se_legend_labels()
         kwargs.setdefault("label", imag_label)
 
@@ -629,31 +737,76 @@ class SelfEnergy:
                     "Warning: some -Σ''(E) uncertainty values are missing. "
                     "Error bars omitted at those energies."
                 )
-            kwargs.setdefault("yerr", xprs.sigma_confidence * y_sigma)
+            kwargs.setdefault("yerr", xprs.sigma_confidence * factor * y_sigma)
 
         ax.errorbar(x, y, **kwargs)
-        ax.set_xlabel(r"$E-\mu$ (eV)")
-        ax.set_ylabel(r"$-\Sigma''(E)$ (eV)")
+
+        x_unit = "meV" if scale == "meV" else "eV"
+        ax.set_xlabel(rf"$E-\mu$ ({x_unit})")
+        ax.set_ylabel(rf"$-\Sigma''(E)$ ({x_unit})")
         ax.legend()
 
         return fig
 
+
     @add_fig_kwargs
-    def plot_both(self, ax=None, **kwargs):
-        r"""Plot Σ'(E) and -Σ''(E) vs. E-μ on the same axis."""
+    def plot_both(self, ax=None, scale="eV", resolution_range="absent", **kwargs):
+        r"""Plot Σ'(E) and -Σ''(E) vs. E-μ on the same axis.
+
+        Parameters
+        ----------
+        ax : Matplotlib-Axes or None
+            Axis to plot on. Created if not provided by the user.
+        scale : {"eV", "meV"}
+            Units for both axes. If "meV", x, y, and yerr are multiplied by
+            `KILO`.
+        resolution_range : {"absent", "applied"}
+            If "applied", removes points with |E-μ| <= energy_resolution (around
+            the chemical potential). The energy resolution is taken from
+            ``self.energy_resolution`` (in eV) and scaled consistently with `scale`.
+        **kwargs :
+            Additional keyword arguments passed to ``ax.errorbar``.
+        """
         from . import settings_parameters as xprs
 
         ax, fig, plt = get_ax_fig_plt(ax=ax)
 
+        if scale not in ("eV", "meV"):
+            raise ValueError("scale must be either 'eV' or 'meV'.")
+        if resolution_range not in ("absent", "applied"):
+            raise ValueError("resolution_range must be 'absent' or 'applied'.")
+
+        factor = KILO if scale == "meV" else 1.0
+
         x = self.enel_range
         real = self.real
         imag = self.imag
         real_sigma = self.real_sigma
         imag_sigma = self.imag_sigma
 
+        if x is not None:
+            x = factor * np.asarray(x, dtype=float)
+        if real is not None:
+            real = factor * np.asarray(real, dtype=float)
+        if imag is not None:
+            imag = factor * np.asarray(imag, dtype=float)
+
+        if resolution_range == "applied" and x is not None:
+            res = self.energy_resolution
+            if res is not None:
+                keep = np.abs(x) > (factor * float(res))
+                x = x[keep]
+                if real is not None:
+                    real = real[keep]
+                if imag is not None:
+                    imag = imag[keep]
+                if real_sigma is not None:
+                    real_sigma = np.asarray(real_sigma, dtype=float)[keep]
+                if imag_sigma is not None:
+                    imag_sigma = np.asarray(imag_sigma, dtype=float)[keep]
+
         real_label, imag_label = self._se_legend_labels()
 
-        # --- plot Σ'
         kw_real = dict(kwargs)
         if real_sigma is not None:
             if np.isnan(real_sigma).any():
@@ -661,11 +814,10 @@ class SelfEnergy:
                     "Warning: some Σ'(E) uncertainty values are missing. "
                     "Error bars omitted at those energies."
                 )
-            kw_real.setdefault("yerr", xprs.sigma_confidence * real_sigma)
+            kw_real.setdefault("yerr", xprs.sigma_confidence * factor * real_sigma)
         kw_real.setdefault("label", real_label)
         ax.errorbar(x, real, **kw_real)
 
-        # --- plot -Σ''
         kw_imag = dict(kwargs)
         if imag_sigma is not None:
             if np.isnan(imag_sigma).any():
@@ -673,27 +825,187 @@ class SelfEnergy:
                     "Warning: some -Σ''(E) uncertainty values are missing. "
                     "Error bars omitted at those energies."
                 )
-            kw_imag.setdefault("yerr", xprs.sigma_confidence * imag_sigma)
+            kw_imag.setdefault("yerr", xprs.sigma_confidence * factor * imag_sigma)
         kw_imag.setdefault("label", imag_label)
         ax.errorbar(x, imag, **kw_imag)
 
-        ax.set_xlabel(r"$E-\mu$ (eV)")
-        ax.set_ylabel(r"$\Sigma'(E),\ -\Sigma''(E)$ (eV)")
+        x_unit = "meV" if scale == "meV" else "eV"
+        ax.set_xlabel(rf"$E-\mu$ ({x_unit})")
+        ax.set_ylabel(rf"$\Sigma'(E),\ -\Sigma''(E)$ ({x_unit})")
+        ax.legend()
+
+        return fig
+
+    @add_fig_kwargs
+    def plot_a2f(self, ax=None, abscissa="forward", **kwargs):
+        ax, fig, plt = get_ax_fig_plt(ax=ax)
+
+        xlim_in = ax.get_xlim()
+        ylim_in = ax.get_ylim()
+
+        if abscissa not in ("forward", "reversed"):
+            raise ValueError("abscissa must be either 'forward' or 'reversed'.")
+
+        omega = self.a2f_omega_range
+        spectrum = self.a2f_spectrum
+
+        if omega is None or spectrum is None:
+            raise AttributeError(
+                "No cached α²F(ω) spectrum found. Run `extract_a2f()` or "
+                "`bayesian_loop()` first."
+            )
+
+        if abscissa == "reversed":
+            omega = -omega[::-1]
+            spectrum = spectrum[::-1]
+
+        a2f_label, _ = self._a2f_legend_labels()
+        kwargs.setdefault("label", a2f_label)
+        ax.plot(omega, spectrum, **kwargs)
+
+        ax.set_xlabel(r"$\omega$ (meV)")
+        ax.set_ylabel(r"$\alpha^2F(\omega)$ (-)")
+
+        self._apply_spectra_axis_defaults(ax, omega, abscissa, xlim_in, ylim_in)
+
+        ax.legend()
+        return fig
+
+    @add_fig_kwargs
+    def plot_model(self, ax=None, abscissa="forward", **kwargs):
+        ax, fig, plt = get_ax_fig_plt(ax=ax)
+
+        xlim_in = ax.get_xlim()
+        ylim_in = ax.get_ylim()
+
+        if abscissa not in ("forward", "reversed"):
+            raise ValueError("abscissa must be either 'forward' or 'reversed'.")
+
+        omega = self.a2f_omega_range
+        model = self.a2f_model
+
+        if omega is None or model is None:
+            raise AttributeError(
+                "No cached model spectrum found. Run `extract_a2f()` or "
+                "`bayesian_loop()` first."
+            )
+
+        if abscissa == "reversed":
+            omega = -omega[::-1]
+            model = model[::-1]
+
+        _, model_label = self._a2f_legend_labels()
+        kwargs.setdefault("label", model_label)
+        ax.plot(omega, model, **kwargs)
+
+        ax.set_xlabel(r"$\omega$ (meV)")
+        ax.set_ylabel(r"$m(\omega)$ (-)")
+
+        self._apply_spectra_axis_defaults(ax, omega, abscissa, xlim_in, ylim_in)
+
+        ax.legend()
+        return fig
+    
+    @add_fig_kwargs
+    def plot_spectra(self, ax=None, abscissa="forward", **kwargs):
+        ax, fig, plt = get_ax_fig_plt(ax=ax)
+
+        xlim_in = ax.get_xlim()
+        ylim_in = ax.get_ylim()
+
+        if abscissa not in ("forward", "reversed"):
+            raise ValueError("abscissa must be either 'forward' or 'reversed'.")
+
+        omega = self.a2f_omega_range
+        spectrum = self.a2f_spectrum
+        model = self.a2f_model
+
+        if omega is None or spectrum is None or model is None:
+            raise AttributeError(
+                "No cached spectra found. Run `extract_a2f()` or `bayesian_loop()` "
+                "first."
+            )
+
+        if abscissa == "reversed":
+            omega = -omega[::-1]
+            spectrum = spectrum[::-1]
+            model = model[::-1]
+
+        kw_a2f = dict(kwargs)
+        kw_model = dict(kwargs)
+        a2f_label, model_label = self._a2f_legend_labels()
+        kw_a2f.setdefault("label", a2f_label)
+        kw_model.setdefault("label", model_label)
+
+        ax.plot(omega, spectrum, **kw_a2f)
+        ax.plot(omega, model, **kw_model)
+
+        self._apply_spectra_axis_defaults(ax, omega, abscissa, xlim_in, ylim_in)
+
+        ax.set_xlabel(r"$\omega$ (meV)")
+        ax.set_ylabel(r"$\alpha^2F_n(\omega),~m_n(\omega)~(-)$")
         ax.legend()
 
         return fig
 
+    @staticmethod
+    def _apply_spectra_axis_defaults(ax, omega, abscissa, xlim_in, ylim_in):
+        """Apply default spectra x-range and y-min, without stomping overrides.
+
+        Defaults are applied only if the incoming axis limits were Matplotlib's
+        defaults (0, 1), i.e. the caller did not pre-set them.
+        """
+        if abscissa not in ("forward", "reversed"):
+            raise ValueError("abscissa must be either 'forward' or 'reversed'.")
+
+        omega_max = float(np.max(np.abs(omega)))
+
+        # --- X defaults (only if user did not pre-set xlim)
+        x0, x1 = xlim_in
+        x_is_default = np.isclose(x0, 0.0) and np.isclose(x1, 1.0)
+        if x_is_default:
+            if abscissa == "forward":
+                ax.set_xlim(0.0, omega_max)
+            else:
+                ax.set_xlim(-omega_max, 0.0)
+
+        # --- Y default: set only the bottom to 0 (only if user did not pre-set)
+        y0, y1 = ylim_in
+        y_is_default = np.isclose(y0, 0.0) and np.isclose(y1, 1.0)
+        if y_is_default:
+            ax.set_ylim(bottom=0.0)
+
 
+    @add_fig_kwargs
     def extract_a2f(self, *, omega_min, omega_max, omega_num, omega_I, omega_M,
-                    mem=None, **mem_kwargs):
+                    mem=None, ax=None, **mem_kwargs):
         r"""
         Extract Eliashberg function α²F(ω) from the self-energy. While working
         with band maps and MDCs is more intuitive in eV, the self-energy
         extraction is performed in eV.
 
+        Parameters
+        ----------
+        ax : Matplotlib-Axes or None
+            Axis to plot on. Created if not provided by the user. (Not used yet;
+            reserved for future plotting.)
+
+        Returns
+        -------
+        spectrum : ndarray
+            Extracted α²F(ω).
+        model : ndarray
+            MEM model spectrum.
+        omega_range : ndarray
+            ω grid used for the extraction.
+        alpha_select : float
+            Selected alpha returned by the chi2kink procedure.
         """
         from . import settings_parameters as xprs
 
+        # Reserve the plot API now; not used yet, but this matches xARPES style.
+        ax, fig, plt = get_ax_fig_plt(ax=ax)
+
         mem_cfg = self._merge_defaults(xprs.mem_defaults, mem, mem_kwargs)
 
         method = mem_cfg["method"]
@@ -724,22 +1036,24 @@ class SelfEnergy:
             f_chi_squared = 2.5 if parts == "both" else 2.0
         else:
             f_chi_squared = float(f_chi_squared)
+        if d_guess <= 0.0:
+            raise ValueError(
+                "chi2kink requires d_guess > 0 to fix the logistic sign "
+                "ambiguity."
+            )
 
         h_n = mem_cfg.get("h_n", None)
         if h_n is None:
             raise ValueError(
-                "`h_n` must be provided explicitly (h_n=... or mem={'h_n': ...}). "
-                "No default is assumed."
+                "`optimisation_parameters` must include 'h_n' for cost evaluation."
             )
 
         from . import (create_model_function, create_kernel_function,
                        singular_value_decomposition, MEM_core)
 
         omega_range = np.linspace(omega_min, omega_max, omega_num)
+        model = create_model_function(omega_range, omega_I, omega_M, omega_S, h_n)
 
-        model = create_model_function(omega_range, omega_I, omega_M, omega_S,
-                                       h_n)
-        
         delta_omega = (omega_max - omega_min) / (omega_num - 1)
         model_in = model * delta_omega
 
@@ -763,7 +1077,6 @@ class SelfEnergy:
 
         energies_eV_masked = energies_eV[mE]
         energies = energies_eV_masked * KILO
-
         k_BT = K_B * self.temperature * KILO
 
         kernel = create_kernel_function(energies, omega_range, k_BT)
@@ -771,16 +1084,13 @@ class SelfEnergy:
         if lambda_el:
             if W is None:
                 if self._class == "SpectralQuadratic":
-                    W = (
-                        PREF * self._fermi_wavevector**2 / self._bare_mass
-                        ) * KILO
+                    W = (PREF * self._fermi_wavevector**2 / self._bare_mass) * KILO
                 else:
                     raise ValueError(
-                    "lambda_el was provided, but W is None. For a linearised "
-                    "band (SpectralLinear), you must also provide W in meV: "
-                    "the electron–electron interaction  scale."
+                        "lambda_el was provided, but W is None. For a linearised "
+                        "band (SpectralLinear), you must also provide W in meV: "
+                        "the electron–electron interaction  scale."
                     )
-            
 
             energies_el = energies_eV_masked * KILO
             real_el, imag_el = self._el_el_self_energy(
@@ -798,14 +1108,12 @@ class SelfEnergy:
             dvec = np.concatenate((real, imag))
             wvec = np.concatenate((real_sigma**(-2), imag_sigma**(-2)))
             H = np.concatenate((np.real(kernel), -np.imag(kernel)))
-
         elif parts == "real":
             real = self.real[mE] * KILO - real_el
             real_sigma = self.real_sigma[mE] * KILO
             dvec = real
             wvec = real_sigma**(-2)
             H = np.real(kernel)
-
         else:  # parts == "imag"
             imag = self.imag[mE] * KILO - impurity_magnitude - imag_el
             imag_sigma = self.imag_sigma[mE] * KILO
@@ -816,14 +1124,55 @@ class SelfEnergy:
         V_Sigma, U, uvec = singular_value_decomposition(H, sigma_svd)
 
         if method == "chi2kink":
-            spectrum_in, _ = self._chi2kink_a2f(
-                    dvec, model_in, uvec, mu, wvec, V_Sigma, U, alpha_min,
-                    alpha_max, alpha_num, a_guess, b_guess, c_guess, d_guess,
-                    f_chi_squared, t_criterion, iter_max, MEM_core
-                    )
+            (spectrum_in, alpha_select, fit_curve, guess_curve,
+            chi2kink_result) = self._chi2kink_a2f(
+                dvec, model_in, uvec, mu, wvec, V_Sigma, U, alpha_min,
+                alpha_max, alpha_num, a_guess, b_guess, c_guess, d_guess,
+                f_chi_squared, t_criterion, iter_max, MEM_core
+            )
+        else:
+            raise NotImplementedError(
+                f"extract_a2f does not support method='{method}'."
+            )
+
+        # --- Plot on ax: always raw chi2 + guess; add fit only if success ---
+        alpha_range = chi2kink_result["alpha_range"]
+        alpha0 = float(alpha_range[0])
+        x_plot = np.log10(alpha_range / alpha0)
+        y_chi2 = chi2kink_result["log_chi_squared"]
+
+        ax.set_xlabel(r"log$_{10}(\alpha)$ (-)")
+        ax.set_ylabel(r"log$_{10}(\chi^2)$ (-)")
+
+        ax.plot(x_plot, y_chi2, label="data")
+        ax.plot(x_plot, guess_curve, label="guess")
+
+        if chi2kink_result["success"]:
+            ax.plot(x_plot, fit_curve, label="fit")
+            ax.axvline(
+                np.log10(alpha_select / alpha0),
+                linestyle="--",
+                label=r"$\alpha_{\rm sel}$",
+            )
+        ax.legend()
+
+        # Abort extraction if fit failed (you asked to terminate in that case)
+        if not chi2kink_result["success"]:
+            raise RuntimeError(
+                "chi2kink logistic fit failed; aborting extract_a2f after "
+                "plotting chi2 and guess."
+            )
 
+        # From here on, we know spectrum_in and alpha_select exist
         spectrum = spectrum_in * omega_num / omega_max
-        return spectrum, model
+
+        self._a2f_spectrum = spectrum
+        self._a2f_model = model
+        self._a2f_omega_range = omega_range
+        self._a2f_alpha_select = alpha_select
+        self._a2f_cost = None
+
+        return fig, spectrum, model, omega_range, alpha_select
     
 
     def bayesian_loop(self, *, omega_min, omega_max, omega_num, omega_I, 
@@ -1193,12 +1542,15 @@ class SelfEnergy:
             print(" | ".join(msg))
 
             return cost_f
-
+        
         class TerminationCallback:
-            def __init__(self, tole, converge_iters, min_steps_for_regression):
+            def __init__(self, tole, converge_iters,
+                         min_steps_for_regression):
                 self.tole = None if tole is None else float(tole)
                 self.converge_iters = int(converge_iters)
-                self.min_steps_for_regression = int(min_steps_for_regression)
+                self.min_steps_for_regression = int(
+                    min_steps_for_regression
+                )
                 self.iter_count = 0
                 self.call_count = 0
 
@@ -1212,36 +1564,41 @@ class SelfEnergy:
                 if current is None:
                     return
 
-                best_cost = float(best["cost"])
-                if np.isfinite(best_cost) and abs(current - best_cost) < self.tole:
-                    self.iter_count += 1
-                else:
-                    self.iter_count = 0
+                best_cost = float(best_global["cost"])
+                if np.isfinite(best_cost):
+                    if abs(current - best_cost) < self.tole:
+                        self.iter_count += 1
+                    else:
+                        self.iter_count = 0
 
                 if self.iter_count >= self.converge_iters:
                     raise ConvergenceException(
-                        f"Converged: |cost-best| < {self.tole:g} for "
+                        "Converged: |cost-best| < "
+                        f"{self.tole:g} for "
                         f"{self.converge_iters} iterations."
                     )
-                
+
                 if self.call_count < self.min_steps_for_regression:
                     return
 
-                current = float(current)
-                init = initial_cost["cost"]
-                if init is None:
+                init_cost = initial_cost["cost"]
+                if init_cost is None:
                     return
 
-                best_cost = float(best["cost"])
+                current = float(current)
+                init_cost = float(init_cost)
+
                 if not np.isfinite(best_cost):
                     return
 
-                if abs(current - init) * relative_best < abs(current - best_cost):
+                if (
+                    abs(current - init_cost) * relative_best
+                    < abs(current - best_cost)
+                ):
                     raise RegressionException(
                         "Regression toward initial guess detected."
                     )
 
-
         callback = TerminationCallback(
             tole=tole,
             converge_iters=converge_iters,
@@ -1345,17 +1702,26 @@ class SelfEnergy:
         print("Optimised parameters:")
         print(args)
 
-        return cost, spectrum, model, alpha_select, params
+        # store inside class methods
+        self._a2f_spectrum = spectrum
+        self._a2f_model = model
+        self._a2f_omega_range = omega_range
+        self._a2f_alpha_select = alpha_select
+        self._a2f_cost = cost
 
+        return spectrum, model, omega_range, alpha_select, cost, params
 
     @staticmethod
     def _merge_defaults(defaults, override_dict=None, override_kwargs=None):
         """Merge defaults with dict + kwargs overrides (kwargs win)."""
         cfg = dict(defaults)
-        if override_dict:
+
+        if override_dict is not None:
             cfg.update(dict(override_dict))
-        if override_kwargs:
+
+        if override_kwargs is not None:
             cfg.update({k: v for k, v in override_kwargs.items() if v is not None})
+
         return cfg
 
     def _prepare_bare(self, fermi_velocity, fermi_wavevector, bare_mass):
@@ -1490,11 +1856,9 @@ class SelfEnergy:
         else:
             f_chi_squared = float(f_chi_squared)
 
-        h_n = mem_cfg.get("h_n", None)
-        if h_n is None:
+        if d_guess <= 0.0:
             raise ValueError(
-                "`h_n` must be provided explicitly (h_n=... or mem={'h_n': ...}). "
-                "No default is assumed."
+                "chi2kink requires d_guess > 0 to fix the logistic sign ambiguity."
             )
 
         if parts not in {"both", "real", "imag"}:
@@ -1622,7 +1986,8 @@ class SelfEnergy:
             dvec = imag_m
             wvec = imag_sig_m**(-2)
 
-        spectrum_in, alpha_select = self._chi2kink_a2f(
+        (spectrum_in, alpha_select, fit_curve, guess_curve,
+            chi2kink_result) = self._chi2kink_a2f(
             dvec, model_in, uvec, mu, wvec, V_Sigma, U, alpha_min, alpha_max,
             alpha_num, a_guess, b_guess, c_guess, d_guess, f_chi_squared,
             t_criterion, iter_max, MEM_core,
@@ -1657,26 +2022,44 @@ class SelfEnergy:
 
 
     @staticmethod
-    def _chi2kink_a2f(dvec, model_in, uvec, mu, wvec, V_Sigma, U,
-                            alpha_min, alpha_max, alpha_num, a_guess, b_guess,
-                            c_guess, d_guess, f_chi_squared, t_criterion, 
-                            iter_max, MEM_core):
-        r"""Compute MEM spectrum using the chi2kink alpha-selection procedure.
+    def _chi2kink_a2f(dvec, model_in, uvec, mu, wvec, V_Sigma, U, alpha_min,
+                    alpha_max, alpha_num, a_guess, b_guess, c_guess, d_guess,
+                    f_chi_squared, t_criterion, iter_max, MEM_core, *,
+                    plot=None):
+        r"""
+        Compute MEM spectrum using the chi2-kink alpha-selection procedure.
 
-        Returns
-        -------
-        spectrum_in : ndarray
-            Selected spectrum from MEM_core evaluated at the chi2kink alpha.
+        Notes
+        -----
+        This routine contains extensive logic to detect failure modes of the
+        chi2-kink logistic fit, including (but not limited to):
+
+        - non-finite or non-positive chi² values,
+        - lack of meaningful parameter updates relative to the initial guess,
+        - absence of improvement in the residual sum of squares,
+        - numerical instabilities or overflows in the logistic model,
+        - invalid or non-finite alpha selection.
+
+        Despite these safeguards, it is **not possible to guarantee** that all
+        failure modes are detected in a nonlinear least-squares problem.
+        Consequently, a reported success should be interpreted as a *necessary*
+        but *not sufficient* condition for physical or numerical reliability.
+
+        Callers **must** inspect the returned ``success`` flag (contained in
+        ``chi2kink_result``) before using the fitted curve, selected alpha, or
+        MEM spectrum. When ``success`` is False, the returned quantities are
+        limited to those required for diagnostic plotting only.
         """
-        from . import (fit_leastsq, chi2kink_logistic)
+        from . import fit_least_squares, chi2kink_logistic
 
-        alpha_range = np.logspace(alpha_min, alpha_max, alpha_num)
+        alpha_range = np.logspace(alpha_min, alpha_max, int(alpha_num))
         chi_squared = np.empty_like(alpha_range, dtype=float)
 
         for i, alpha in enumerate(alpha_range):
-            spectrum_in, uvec = MEM_core(dvec, model_in, uvec, mu, alpha, 
-                wvec, V_Sigma, U, t_criterion, iter_max)
-            
+            spectrum_in, uvec = MEM_core(
+                dvec, model_in, uvec, mu, alpha, wvec, V_Sigma, U,
+                t_criterion, iter_max
+            )
             T = V_Sigma @ (U.T @ spectrum_in)
             chi_squared[i] = wvec @ ((T - dvec) ** 2)
 
@@ -1689,18 +2072,79 @@ class SelfEnergy:
         log_chi_squared = np.log10(chi_squared)
 
         p0 = np.array([a_guess, b_guess, c_guess, d_guess], dtype=float)
-        pfit, pcov = fit_leastsq(
+        pfit, pcov, lsq_success = fit_least_squares(
             p0, log_alpha, log_chi_squared, chi2kink_logistic
         )
 
-        cout = pfit[2]
-        dout = pfit[3]
-        alpha_select = 10 ** (cout - f_chi_squared / dout)
+        with np.errstate(over="ignore", invalid="ignore", divide="ignore"):
+            guess_curve = chi2kink_logistic(log_alpha, *p0)
+
+        # Start from the necessary requirement: least_squares must say success
+        success = bool(lsq_success)
 
-        spectrum_in, uvec = MEM_core(dvec, model_in, uvec, mu, alpha_select, 
-                        wvec, V_Sigma, U, t_criterion, iter_max)
+        # If the guess itself blows up, we can't trust anything
+        if not np.all(np.isfinite(guess_curve)):
+            success = False
+
+        fit_curve_tmp = None
+        if success:
+            pfit = np.asarray(pfit, dtype=float)
+
+            if np.allclose(pfit, p0, rtol=1e-12, atol=0.0):
+                success = False
+            else:
+                with np.errstate(over="ignore", invalid="ignore",
+                                 divide="ignore"):
+                    fit_curve_tmp = chi2kink_logistic(log_alpha, *pfit)
+
+                if not np.all(np.isfinite(fit_curve_tmp)):
+                    success = False
+                else:
+                    r0 = guess_curve - log_chi_squared
+                    r1 = fit_curve_tmp - log_chi_squared
+                    sse0 = float(r0 @ r0)
+                    sse1 = float(r1 @ r1)
+
+                    tol = 1e-12 * max(1.0, sse0)
+                    if (not np.isfinite(sse1)) or (sse1 >= sse0 - tol):
+                        success = False
+
+        alpha_select = None
+        fit_curve = None
+        spectrum_out = None
+
+        if success:
+            fit_curve = fit_curve_tmp
+
+            cout = float(pfit[2])
+            dout = float(pfit[3])
+            exp10 = cout - float(f_chi_squared) / dout
+
+            if (not np.isfinite(exp10)) or (exp10 < -308.0) or (exp10 > 308.0):
+                success = False
+                fit_curve = None
+            else:
+                with np.errstate(over="raise", invalid="raise"):
+                    alpha_select = float(np.power(10.0, exp10))
+
+                spectrum_out, uvec = MEM_core(
+                    dvec, model_in, uvec, mu, alpha_select, wvec, V_Sigma, U,
+                    t_criterion, iter_max
+                )
+
+        chi2kink_result = {
+            "alpha_range": alpha_range,
+            "chi_squared": chi_squared,
+            "log_alpha": log_alpha,
+            "log_chi_squared": log_chi_squared,
+            "p0": p0,
+            "pfit": pfit,
+            "pcov": pcov,
+            "success": bool(success),
+            "alpha_select": alpha_select,
+        }
 
-        return spectrum_in, alpha_select
+        return spectrum_out, alpha_select, fit_curve, guess_curve, chi2kink_result
 
 
     @staticmethod