diff --git a/src/quantem/spectroscopy/dataset3deels.py b/src/quantem/spectroscopy/dataset3deels.py index 99b9eb3d..17bff3f6 100644 --- a/src/quantem/spectroscopy/dataset3deels.py +++ b/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, ) diff --git a/src/quantem/spectroscopy/dataset3dspectroscopy.py b/src/quantem/spectroscopy/dataset3dspectroscopy.py index e9f33cb0..aee3aaba 100644 --- a/src/quantem/spectroscopy/dataset3dspectroscopy.py +++ b/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 ( diff --git a/src/quantem/spectroscopy/dataset3dxeds.py b/src/quantem/spectroscopy/dataset3dxeds.py index a515a516..6d99c676 100644 --- a/src/quantem/spectroscopy/dataset3dxeds.py +++ b/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 ( diff --git a/src/quantem/spectroscopy/spectroscopy_models.py b/src/quantem/spectroscopy/spectroscopy_models.py index 195b7cd1..90c5cb7c 100644 --- a/src/quantem/spectroscopy/spectroscopy_models.py +++ b/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 diff --git a/src/quantem/spectroscopy/spectroscopy_visualzitions.py b/src/quantem/spectroscopy/spectroscopy_visualizations.py similarity index 100% rename from src/quantem/spectroscopy/spectroscopy_visualzitions.py rename to src/quantem/spectroscopy/spectroscopy_visualizations.py diff --git a/src/quantem/spectroscopy/xeds_fitting.py b/src/quantem/spectroscopy/xeds_fitting.py index 8d78da04..0c912775 100644 --- a/src/quantem/spectroscopy/xeds_fitting.py +++ b/src/quantem/spectroscopy/xeds_fitting.py @@ -11,8 +11,10 @@ from quantem.spectroscopy.spectroscopy_models import ( CubeSpectrumModel, MeanSpectrumModel, + energy_dependent_line_widths, total_variation_in_plane, xeds_data_loss, + build_element_peak_basis ) @@ -328,6 +330,9 @@ def _fit_mean_model_pytorch( default_lr_adam, default_lr_lbfgs, verbose=False, + width_model="constant", + peak_width_b=None, + train_peak_width=True, ): """Fit one mean spectrum with a small non-negative basis model.""" optimizer_name = _normalize_choice(optimizer, "optimizer", {"adam", "lbfgs"}) @@ -367,6 +372,9 @@ def _fit_mean_model_pytorch( peak_width_init=float(peak_width), amplitude_init=amplitude_init, background_init=background_init, + width_model=width_model, + peak_width_b_init=None if peak_width_b is None else float(peak_width_b), + train_peak_width=bool(train_peak_width), ).to(device=energy_axis.device, dtype=energy_axis.dtype) model.element_names = list(basis_data["element_names"]) model.line_labels = list(basis_data["line_labels"]) @@ -404,6 +412,25 @@ def _fit_mean_model_pytorch( element_amplitudes = model.element_amplitudes() background_coeffs = model.background_coeffs() + width_params = model.peak_width_params() + if width_params is not None: + width_values = [ + float(torch.nn.functional.softplus(p).detach().cpu().item()) + for p in width_params + ] + else: + # widths are fixed: report the stored constants ([a, b] for energy_dependent, scalar otherwise) + width_values = model._peak_width_fixed.detach().cpu().reshape(-1).tolist() + + if width_model == "energy_dependent": + width_a, width_b = width_values + # scalar anchor for downstream scalar consumers: the E->0 offset of width(E) + peak_width_out = width_a + else: + (width_a,) = width_values + width_b = None + peak_width_out = width_a + return { "model": model, "loss_history": np.asarray(loss_history), @@ -412,7 +439,10 @@ def _fit_mean_model_pytorch( "final_background_raw": final_background_raw.detach(), "element_amplitudes": element_amplitudes.detach(), "background_coeffs": background_coeffs.detach(), - "peak_width": float(model.peak_width().detach().cpu().item()), + "peak_width": peak_width_out, + "peak_width_a": width_a, + "peak_width_b": width_b, + "width_model": width_model, "element_names": list(basis_data["element_names"]), "line_labels": list(basis_data["line_labels"]), "line_energies": basis_data["line_energies"].detach(), @@ -437,10 +467,20 @@ def fit_spectrum_mean_pytorch( normalize_target=None, show_plot=True, verbose=True, + width_model="constant", + peak_width_b=None, + freeze_peak_width_global=False, ): - """Fit the mean spectrum first. This is the stable starting point.""" + """Fit the mean spectrum first. This is the stable starting point. + + width_model="constant" fits one shared Gaussian FWHM (peak_width); + width_model="energy_dependent" fits width(E) = a + b * sqrt(E) per X-ray line, with a + initialized from peak_width and b from peak_width_b. Set + freeze_peak_width_global=True to keep the widths fixed at their init values. + """ optimizer_name = _normalize_choice(optimizer, "optimizer", {"adam", "lbfgs"}) loss_name = _normalize_choice(loss, "loss", {"mse", "poisson"}) + width_model = _normalize_choice(width_model, "width_model", {"constant", "energy_dependent"}) if normalize_target is None: normalize_target = loss_name == "mse" @@ -465,6 +505,9 @@ def fit_spectrum_mean_pytorch( default_lr_adam=5e-2, default_lr_lbfgs=1.0, verbose=verbose, + width_model=width_model, + peak_width_b=peak_width_b, + train_peak_width=not freeze_peak_width_global, ) if verbose: @@ -477,7 +520,11 @@ def fit_spectrum_mean_pytorch( concentrations = mean_fit["element_amplitudes"].detach().cpu().numpy() element_names = list(mean_fit["element_names"]) edge_indices = np.arange(len(element_names), dtype=int) - peak_widths = np.asarray([mean_fit["peak_width"]], dtype=np.float32) + peak_widths = ( + np.asarray([mean_fit["peak_width_a"], mean_fit["peak_width_b"]], dtype=np.float32) + if mean_fit["width_model"] == "energy_dependent" + else np.asarray([mean_fit["peak_width"]], dtype=np.float32) + ) if show_plot: _plot_fit_summary( @@ -502,6 +549,9 @@ def fit_spectrum_mean_pytorch( "edge_element_indices": edge_indices, "peak_widths": peak_widths, "peak_width": float(mean_fit["peak_width"]), + "peak_width_a": mean_fit["peak_width_a"], + "peak_width_b": mean_fit["peak_width_b"], + "width_model": mean_fit["width_model"], "background_coeffs": mean_fit["background_coeffs"].detach().cpu().numpy(), "line_labels": list(mean_fit["line_labels"]), "energy_axis": energy_axis_np, @@ -531,14 +581,30 @@ def fit_spectrum_pytorch( lr_local=None, device=None, constrain_background=0.1, + width_model="constant", + peak_width_b=None, + freeze_peak_width_global=None, ): - """Fit XEDS with a mean-first then full-cube workflow.""" + """Fit XEDS with a mean-first then full-cube workflow. + + width_model="constant" fits one shared Gaussian FWHM (peak_width); + width_model="energy_dependent" fits width(E) = a + b * sqrt(E) per X-ray line, with a + initialized from peak_width and b from peak_width_b (e.g. a=0.02 keV, + b=0.043 for a typical Fano-limited detector). freeze_peak_width freezes the + cube-stage widths; freeze_peak_width_global also freezes the mean-stage + widths (defaults to freeze_peak_width when not set). + """ effective_optimizer_global = _normalize_choice( optimizer_global, "optimizer_global", {"adam", "lbfgs"} ) effective_optimizer_local = _normalize_choice( optimizer_local, "optimizer_local", {"adam", "lbfgs"} ) + effective_width_model = _normalize_choice(width_model, "width_model", {"constant", "energy_dependent"}) + # freeze_peak_width_global defaults to freeze_peak_width so one flag freezes both stages + effective_freeze_global = ( + bool(freeze_peak_width) if freeze_peak_width_global is None else bool(freeze_peak_width_global) + ) effective_loss_global = ( _normalize_choice(loss_global, "loss_global", {"poisson", "mse"}) if loss_global is not None @@ -575,6 +641,9 @@ def fit_spectrum_pytorch( normalize_target=effective_loss_global == "mse", show_plot=show_plot, verbose=verbose, + width_model=effective_width_model, + peak_width_b=peak_width_b, + freeze_peak_width_global=effective_freeze_global, ) device = _resolve_torch_device(device) @@ -604,6 +673,9 @@ def fit_spectrum_pytorch( default_lr_adam=5e-2, default_lr_lbfgs=1.0, verbose=verbose, + width_model=effective_width_model, + peak_width_b=peak_width_b, + train_peak_width=not effective_freeze_global, ) scan_row, scan_col, _ = spectra.shape @@ -627,10 +699,12 @@ def fit_spectrum_pytorch( line_element_indices=global_fit["line_element_indices"], n_elements=n_elements, background_degree=polynomial_background_degree, - peak_width_init=global_fit["peak_width"], + peak_width_init=global_fit["peak_width_a"], + peak_width_b_init=global_fit["peak_width_b"], amplitude_init=amplitude_init, background_init=background_init, train_peak_width=not freeze_peak_width, + width_model=effective_width_model, ).to(device=energy_axis.device, dtype=energy_axis.dtype) fitter = XEDSFitter(cube_model) @@ -674,7 +748,16 @@ def local_loss(): with torch.no_grad(): predicted, _peak_part, background_part = cube_model() abundance_maps = cube_model.element_amplitudes().detach().cpu().numpy() - peak_widths = np.asarray([cube_model.peak_width().detach().cpu().item()], dtype=np.float32) + cube_width_params = cube_model.peak_width_params() + if cube_width_params is not None: + cube_width_values = [ + float(torch.nn.functional.softplus(p).detach().cpu().item()) + for p in cube_width_params + ] + else: + # widths are fixed: report the stored constants ([a, b] for energy_dependent, scalar otherwise) + cube_width_values = cube_model._peak_width_fixed.detach().cpu().reshape(-1).tolist() + peak_widths = np.asarray(cube_width_values, dtype=np.float32) mean_input_spectrum = spectra[valid_pixel_mask].mean(dim=0).detach().cpu().numpy() mean_fitted_spectrum = predicted[valid_pixel_mask].mean(dim=0).detach().cpu().numpy() mean_background_spectrum = ( @@ -715,6 +798,7 @@ def local_loss(): "abundance_maps": abundance_maps, "element_names": list(global_fit["element_names"]), "peak_widths": peak_widths, + "width_model": effective_width_model, "loss_history": local_loss_history, "global_loss_history": global_loss_history, "valid_pixel_mask": valid_pixel_mask.detach().cpu().numpy(),