Addition of optional energy-dependent peak width model to XEDS pytorc…

src/quantem/spectroscopy/dataset3deels.py

@@ -8,25 +8,25 @@
 
 from quantem.core.visualization import show_2d
 from quantem.spectroscopy.dataset3dspectroscopy import Dataset3dspectroscopy
-from quantem.spectroscopy.spectroscopy_visualzitions import (
+from quantem.spectroscopy.spectroscopy_visualizations import (
     interpret_thickness_quality as _visualize_thickness_quality,
 )
-from quantem.spectroscopy.spectroscopy_visualzitions import (
+from quantem.spectroscopy.spectroscopy_visualizations import (
     plot_absolute_thickness as _visualize_absolute_thickness,
 )
-from quantem.spectroscopy.spectroscopy_visualzitions import (
+from quantem.spectroscopy.spectroscopy_visualizations import (
     plot_absolute_zlp_shift as _visualize_absolute_zlp_shift,
 )
-from quantem.spectroscopy.spectroscopy_visualzitions import (
+from quantem.spectroscopy.spectroscopy_visualizations import (
     plot_dual_eels_picker as _visualize_dual_eels_picker,
 )
-from quantem.spectroscopy.spectroscopy_visualzitions import (
+from quantem.spectroscopy.spectroscopy_visualizations import (
     plot_quantem_diagnostic as _visualize_quantem_diagnostic,
 )
-from quantem.spectroscopy.spectroscopy_visualzitions import (
+from quantem.spectroscopy.spectroscopy_visualizations import (
     plot_zlp_drift_diagnostics as _visualize_zlp_drift_diagnostics,
 )
-from quantem.spectroscopy.spectroscopy_visualzitions import (
+from quantem.spectroscopy.spectroscopy_visualizations import (
     visualize_thickness_windows as _visualize_thickness_windows,
 )
 

src/quantem/spectroscopy/dataset3dspectroscopy.py

@@ -9,19 +9,19 @@
 
 from quantem.core.datastructures.dataset3d import Dataset3d
 from quantem.core.utils.filter import median_filter
-from quantem.spectroscopy.spectroscopy_visualzitions import (
+from quantem.spectroscopy.spectroscopy_visualizations import (
     _plot_background_subtraction as _visualize_background_subtraction,
 )
-from quantem.spectroscopy.spectroscopy_visualzitions import (
+from quantem.spectroscopy.spectroscopy_visualizations import (
     _plot_pca_results as _visualize_pca_results,
 )
-from quantem.spectroscopy.spectroscopy_visualzitions import (
+from quantem.spectroscopy.spectroscopy_visualizations import (
     plot_attached_spectrum as _visualize_attached_spectrum,
 )
-from quantem.spectroscopy.spectroscopy_visualzitions import (
+from quantem.spectroscopy.spectroscopy_visualizations import (
     show_energy_window_map as _visualize_energy_window_map,
 )
-from quantem.spectroscopy.spectroscopy_visualzitions import (
+from quantem.spectroscopy.spectroscopy_visualizations import (
     show_mean_spectrum as _visualize_mean_spectrum,
 )
 from quantem.spectroscopy.utils import (

src/quantem/spectroscopy/dataset3dxeds.py

@@ -11,7 +11,7 @@
 from quantem.spectroscopy.dataset3dxeds_fitting import (
     peak_autoid as _peak_autoid,
 )
-from quantem.spectroscopy.spectroscopy_visualzitions import (
+from quantem.spectroscopy.spectroscopy_visualizations import (
     show_spectrum_images as _visualize_spectrum_images,
 )
 from quantem.spectroscopy.xeds_fitting import (

src/quantem/spectroscopy/spectroscopy_models.py

@@ -48,6 +48,15 @@ def build_polynomial_background_basis(
     return torch.stack([x**power for power in range(degree + 1)], dim=0)
 
 
+def energy_dependent_line_widths(
+    line_energies: torch.Tensor,
+    width_a: torch.Tensor,
+    width_b: torch.Tensor,
+) -> torch.Tensor:
+    """Energy-dependent Gaussian FWHM: width(E) = a + b * sqrt(E)."""
+    return width_a + width_b * torch.sqrt(torch.clamp(line_energies, min=0.0))
+
+
 def build_element_peak_basis(
     energy_axis: torch.Tensor,
     line_energies: torch.Tensor,
@@ -63,7 +72,12 @@ def build_element_peak_basis(
         raise ValueError("n_elements must be >= 1")
 
     width = torch.clamp(peak_width, min=1e-4)
-    sigma = width / 2.355
+    if width.ndim == 0:
+        sigma = width / 2.355
+    else:
+        if width.numel() != line_energies.numel():
+            raise ValueError("peak_width must have one entry per X-ray line")
+        sigma = (width / 2.355).reshape(-1, 1)
 
     centers = line_energies.unsqueeze(1)
     channels = energy_axis.unsqueeze(0)
@@ -94,7 +108,11 @@ def total_variation_in_plane(maps: torch.Tensor) -> torch.Tensor:
 
 
 class MeanSpectrumModel(nn.Module, PPLR):
-    """Non-negative mean-spectrum model with one shared peak width."""
+    """Non-negative mean-spectrum model with constant or energy-dependent peak widths.
+
+    ``width_model="constant"`` fits one shared FWHM; ``width_model="energy_dependent"`` fits
+    width(E) = a + b * sqrt(E) with per-line FWHM from two non-negative params.
+    """
 
     def __init__(
         self,
@@ -107,8 +125,15 @@ def __init__(
         peak_width_init: float,
         amplitude_init: torch.Tensor,
         background_init: torch.Tensor,
+        width_model: str = "constant",
+        peak_width_b_init: float | None = None,
+        train_peak_width: bool = True,
     ):
         super().__init__()
+        if width_model not in ("constant", "energy_dependent"):
+            raise ValueError("width_model must be 'constant' or 'energy_dependent'")
+        self.width_model = width_model
+        self.train_peak_width = bool(train_peak_width)
         self.n_elements = int(n_elements)
         self.register_buffer("energy_axis", energy_axis)
         self.register_buffer("line_energies", line_energies)
@@ -124,11 +149,42 @@ def __init__(
             dtype=energy_axis.dtype,
             device=energy_axis.device,
         )
-        self.peak_width_raw = nn.Parameter(inverse_softplus(peak_width_tensor))
+        if self.width_model == "energy_dependent":
+            width_b = (
+                peak_width_tensor
+                if peak_width_b_init is None
+                else torch.as_tensor(peak_width_b_init, dtype=energy_axis.dtype, device=energy_axis.device)
+            )
+            if self.train_peak_width:
+                self.peak_width_a_raw = nn.Parameter(inverse_softplus(peak_width_tensor))
+                self.peak_width_b_raw = nn.Parameter(inverse_softplus(width_b))
+                self.register_buffer("_peak_width_fixed", torch.empty(0, device=energy_axis.device))
+            else:
+                self.peak_width_a_raw = None
+                self.peak_width_b_raw = None
+                self.register_buffer(
+                    "_peak_width_fixed",
+                    torch.stack([peak_width_tensor.reshape(()), width_b.reshape(())]),
+                )
+        elif self.train_peak_width:
+            self.peak_width_raw = nn.Parameter(inverse_softplus(peak_width_tensor))
+            self.register_buffer("_peak_width_fixed", torch.empty(0, device=energy_axis.device))
+        else:
+            self.peak_width_raw = None
+            self.register_buffer("_peak_width_fixed", peak_width_tensor.reshape(()))
+
         self.element_amplitudes_raw = nn.Parameter(inverse_softplus(amplitude_init))
         self.background_coeffs_raw = nn.Parameter(inverse_softplus(background_init))
 
     def peak_width(self) -> torch.Tensor:
+        if self.width_model == "energy_dependent":
+            if self.peak_width_a_raw is None:
+                width_a, width_b = self._peak_width_fixed[0], self._peak_width_fixed[1]
+            else:
+                width_a, width_b = F.softplus(self.peak_width_a_raw), F.softplus(self.peak_width_b_raw)
+            return energy_dependent_line_widths(self.line_energies, width_a, width_b)
+        if self.peak_width_raw is None:
+            return self._peak_width_fixed
         return F.softplus(self.peak_width_raw)
 
     def element_amplitudes(self) -> torch.Tensor:
@@ -156,20 +212,37 @@ def background_spectrum(self) -> torch.Tensor:
     def forward(self) -> torch.Tensor:
         return self.peak_spectrum() + self.background_spectrum()
 
+    def peak_width_params(self) -> list[nn.Parameter] | None:
+        if self.width_model == "energy_dependent":
+            if self.peak_width_a_raw is None:
+                return None
+            return [self.peak_width_a_raw, self.peak_width_b_raw]
+        return [self.peak_width_raw] if self.peak_width_raw is not None else None
+
     def get_params(self) -> dict[str, list[nn.Parameter]]:
-        return {
+        groups = {
             "amplitudes": [self.element_amplitudes_raw],
             "background": [self.background_coeffs_raw],
-            "peak_width": [self.peak_width_raw],
         }
+        width_params = self.peak_width_params()
+        if width_params is not None:
+            groups["peak_width"] = width_params
+        return groups
 
     @property
     def param_keys(self) -> list[str]:
-        return ["amplitudes", "background", "peak_width"]
+        keys = ["amplitudes", "background"]
+        if self.peak_width_params() is not None:
+            keys.append("peak_width")
+        return keys
 
 
 class CubeSpectrumModel(nn.Module, PPLR):
-    """Non-negative full-cube XEDS model that reuses the mean-fit basis."""
+    """Non-negative full-cube XEDS model that reuses the mean-fit basis.
+
+    ``width_model="constant"`` fits one shared FWHM; ``width_model="energy_dependent"`` fits
+    width(E) = a + b * sqrt(E) with per-line FWHM from two non-negative params.
+    """
 
     def __init__(
         self,
@@ -183,8 +256,13 @@ def __init__(
         amplitude_init: torch.Tensor,
         background_init: torch.Tensor,
         train_peak_width: bool,
+        width_model: str = "constant",
+        peak_width_b_init: float | None = None,
     ):
         super().__init__()
+        if width_model not in ("constant", "energy_dependent"):
+            raise ValueError("width_model must be 'constant' or 'energy_dependent'")
+        self.width_model = width_model
         self.n_elements = int(n_elements)
 
         self.register_buffer("energy_axis", energy_axis)
@@ -201,7 +279,24 @@ def __init__(
             dtype=energy_axis.dtype,
             device=energy_axis.device,
         )
-        if train_peak_width:
+        if self.width_model == "energy_dependent":
+            width_b = (
+                peak_width_tensor
+                if peak_width_b_init is None
+                else torch.as_tensor(peak_width_b_init, dtype=energy_axis.dtype, device=energy_axis.device)
+            )
+            if train_peak_width:
+                self.peak_width_a_raw = nn.Parameter(inverse_softplus(peak_width_tensor))
+                self.peak_width_b_raw = nn.Parameter(inverse_softplus(width_b))
+                self.register_buffer("_peak_width_fixed", torch.empty(0, device=energy_axis.device))
+            else:
+                self.peak_width_a_raw = None
+                self.peak_width_b_raw = None
+                self.register_buffer(
+                    "_peak_width_fixed",
+                    torch.stack([peak_width_tensor.reshape(()), width_b.reshape(())]),
+                )
+        elif train_peak_width:
             self.peak_width_raw = nn.Parameter(inverse_softplus(peak_width_tensor))
             self.register_buffer("_peak_width_fixed", torch.empty(0, device=energy_axis.device))
         else:
@@ -212,6 +307,12 @@ def __init__(
         self.background_coeffs_raw = nn.Parameter(inverse_softplus(background_init))
 
     def peak_width(self) -> torch.Tensor:
+        if self.width_model == "energy_dependent":
+            if self.peak_width_a_raw is None:
+                width_a, width_b = self._peak_width_fixed[0], self._peak_width_fixed[1]
+            else:
+                width_a, width_b = F.softplus(self.peak_width_a_raw), F.softplus(self.peak_width_b_raw)
+            return energy_dependent_line_widths(self.line_energies, width_a, width_b)
         if self.peak_width_raw is None:
             return self._peak_width_fixed
         return F.softplus(self.peak_width_raw)
@@ -243,18 +344,26 @@ def forward(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         background_part = self.background_spectra()
         return peak_part + background_part, peak_part, background_part
 
+    def peak_width_params(self) -> list[nn.Parameter] | None:
+        if self.width_model == "energy_dependent":
+            if self.peak_width_a_raw is None:
+                return None
+            return [self.peak_width_a_raw, self.peak_width_b_raw]
+        return [self.peak_width_raw] if self.peak_width_raw is not None else None
+
     def get_params(self) -> dict[str, list[nn.Parameter]]:
         groups = {
             "amplitudes": [self.element_amplitudes_raw],
             "background": [self.background_coeffs_raw],
         }
-        if self.peak_width_raw is not None:
-            groups["peak_width"] = [self.peak_width_raw]
+        width_params = self.peak_width_params()
+        if width_params is not None:
+            groups["peak_width"] = width_params
         return groups
 
     @property
     def param_keys(self) -> list[str]:
         keys = ["amplitudes", "background"]
-        if self.peak_width_raw is not None:
+        if self.peak_width_params() is not None:
             keys.append("peak_width")
         return keys