xarpes 0.2.4__py3-none-any.whl → 0.6.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
xarpes/mdcs.py ADDED
@@ -0,0 +1,1078 @@
1
+ # Copyright (C) 2025 xARPES Developers
2
+ # This program is free software under the terms of the GNU GPLv3 license.
3
+
4
+ # get_ax_fig_plt and add_fig_kwargs originate from pymatgen/util/plotting.py.
5
+ # Copyright (C) 2011-2024 Shyue Ping Ong and the pymatgen Development Team
6
+ # Pymatgen is released under the MIT License.
7
+
8
+ # See also abipy/tools/plotting.py.
9
+ # Copyright (C) 2021 Matteo Giantomassi and the AbiPy Group
10
+ # AbiPy is free software under the terms of the GNU GPLv2 license.
11
+
12
+ """File containing the MDCs class."""
13
+
14
+ import numpy as np
15
+ from .plotting import get_ax_fig_plt, add_fig_kwargs
16
+ from .functions import extend_function
17
+ from .constants import KILO
18
+
19
+ class MDCs:
20
+ r"""
21
+ Container for momentum distribution curves (MDCs) and their fits.
22
+
23
+ This class stores the MDC intensity maps, angular and energy grids, and
24
+ the aggregated fit results produced by :meth:`fit_selection`.
25
+
26
+ Parameters
27
+ ----------
28
+ intensities : ndarray
29
+ MDC intensity data. Typically a 2D array with shape
30
+ ``(n_energy, n_angle)`` or a 1D array for a single curve.
31
+ angles : ndarray
32
+ Angular grid corresponding to the MDCs [degrees].
33
+ angle_resolution : float
34
+ Angular step size or effective angular resolution [degrees].
35
+ energy_resolution : float
36
+ Energy resolution associated with the MDCs [eV].
37
+ temperature: float
38
+ Temperature associated with the band map [K].
39
+ enel : ndarray or float
40
+ Electron binding energies of the MDC slices [eV].
41
+ Can be a scalar for a single MDC.
42
+ hnuminPhi : float
43
+ Photon energy minus work function, used to convert ``enel`` to
44
+ kinetic energy [eV].
45
+
46
+ Attributes
47
+ ----------
48
+ intensities : ndarray
49
+ MDC intensity data (same object as passed to the constructor).
50
+ angles : ndarray
51
+ Angular grid [degrees].
52
+ angle_resolution : float
53
+ Angular step size or resolution [degrees].
54
+ enel : ndarray or float
55
+ Electron binding energies [eV], as given at construction.
56
+ ekin : ndarray or float
57
+ Kinetic energies [eV], computed as ``enel + hnuminPhi``.
58
+ hnuminPhi : float
59
+ Photon energy minus work function [eV].
60
+ ekin_range : ndarray
61
+ Kinetic-energy values of the slices that were actually fitted.
62
+ Set by :meth:`fit_selection`.
63
+ individual_properties : dict
64
+ Nested mapping of fitted parameters and their uncertainties for each
65
+ component and each energy slice. Populated by :meth:`fit_selection`.
66
+
67
+ Notes
68
+ -----
69
+ After calling :meth:`fit_selection`, :attr:`individual_properties` has the
70
+ structure::
71
+
72
+ {
73
+ label: {
74
+ class_name: {
75
+ 'label': label,
76
+ '_class': class_name,
77
+ param: [values per energy slice],
78
+ param_sigma: [1σ per slice or None],
79
+ ...
80
+ }
81
+ }
82
+ }
83
+
84
+ where ``param`` is typically one of ``'offset'``, ``'slope'``,
85
+ ``'amplitude'``, ``'peak'``, ``'broadening'``, and ``param_sigma`` stores
86
+ the corresponding uncertainty for each slice.
87
+
88
+ """
89
+
90
+ def __init__(self, intensities, angles, angle_resolution,
91
+ energy_resolution, temperature, enel, hnuminPhi):
92
+ self._intensities = intensities
93
+ self._angles = angles
94
+ self._angle_resolution = angle_resolution
95
+ self._energy_resolution = energy_resolution
96
+ self._temperature = temperature
97
+ self._enel = enel
98
+ self._hnuminPhi = hnuminPhi
99
+
100
+ # Derived attributes (populated by fit_selection)
101
+ self._ekin_range = None
102
+ self._individual_properties = None # combined values + sigmas
103
+
104
+ # -------------------- Immutable physics inputs --------------------
105
+
106
+ @property
107
+ def angles(self):
108
+ """Angular axis for the MDCs."""
109
+ return self._angles
110
+
111
+ @property
112
+ def angle_resolution(self):
113
+ """Angular resolution (float)."""
114
+ return self._angle_resolution
115
+
116
+ @angle_resolution.setter
117
+ def angle_resolution(self, _):
118
+ """Setter for the angle resolution. This raises an attribute error
119
+ as the angle resolution needs to be derived from the band map."""
120
+ raise AttributeError("`angle_resolution` is read-only; set it via the "
121
+ "constructor.")
122
+
123
+ @property
124
+ def energy_resolution(self):
125
+ """Energy resolution (float)."""
126
+ return self._energy_resolution
127
+
128
+ @energy_resolution.setter
129
+ def energy_resolution(self, _):
130
+ """Setter for the energy resolution. This raises an attribute error
131
+ as the energy resolution needs to be derived from the band map."""
132
+ raise AttributeError("`energy_resolution` is read-only; set it via the "
133
+ "constructor.")
134
+
135
+ @property
136
+ def temperature(self):
137
+ """Temperature (float)."""
138
+ return self._temperature
139
+
140
+ @temperature.setter
141
+ def temperature(self, _):
142
+ """Setter for the temperature. This raises an attribute error as the
143
+ temperature needs to be derived from the band map."""
144
+ raise AttributeError("`temperature` is read-only; set it via the "
145
+ "constructor.")
146
+
147
+ @property
148
+ def enel(self):
149
+ """Photoelectron binding energies (array-like). Read-only."""
150
+ return self._enel
151
+
152
+ @enel.setter
153
+ def enel(self, _):
154
+ raise AttributeError("`enel` is read-only; set it via the " \
155
+ "constructor.")
156
+
157
+ @property
158
+ def hnuminPhi(self):
159
+ """Work-function/photon-energy offset. Read-only."""
160
+ return self._hnuminPhi
161
+
162
+ @hnuminPhi.setter
163
+ def hnuminPhi(self, _):
164
+ raise AttributeError("`hnuminPhi` is read-only; set it via the constructor.")
165
+
166
+ @property
167
+ def ekin(self):
168
+ """Kinetic energy array: enel + hnuminPhi (computed on the fly)."""
169
+ return self._enel + self._hnuminPhi
170
+
171
+ @ekin.setter
172
+ def ekin(self, _):
173
+ raise AttributeError("`ekin` is derived and read-only.")
174
+
175
+ # -------------------- Data arrays --------------------
176
+
177
+ @property
178
+ def intensities(self):
179
+ """2D or 3D intensity map (energy × angle)."""
180
+ return self._intensities
181
+
182
+ @intensities.setter
183
+ def intensities(self, x):
184
+ self._intensities = x
185
+
186
+ # -------------------- Results populated by fit_selection --------------------
187
+
188
+ @property
189
+ def ekin_range(self):
190
+ """Kinetic-energy slices that were fitted."""
191
+ if self._ekin_range is None:
192
+ raise AttributeError("`ekin_range` not yet set. Run `.fit_selection()` first.")
193
+ return self._ekin_range
194
+
195
+ @property
196
+ def individual_properties(self):
197
+ """
198
+ Aggregated fitted parameter values and uncertainties per component.
199
+
200
+ Returns
201
+ -------
202
+ dict
203
+ Nested mapping::
204
+
205
+ {
206
+ label: {
207
+ class_name: {
208
+ 'label': label,
209
+ '_class': class_name,
210
+ <param>: [values per slice],
211
+ <param>_sigma: [1σ per slice or None],
212
+ ...
213
+ }
214
+ }
215
+ }
216
+ """
217
+ if self._individual_properties is None:
218
+ raise AttributeError(
219
+ "`individual_properties` not yet set. Run `.fit_selection()` first."
220
+ )
221
+ return self._individual_properties
222
+
223
+ def energy_check(self, energy_value):
224
+ r"""
225
+ """
226
+ if np.isscalar(self.ekin):
227
+ if energy_value is not None:
228
+ raise ValueError("This dataset contains only one " \
229
+ "momentum-distribution curve; do not provide energy_value.")
230
+ else:
231
+ kinergy = self.ekin
232
+ counts = self.intensities
233
+ else:
234
+ if energy_value is None:
235
+ raise ValueError("This dataset contains multiple " \
236
+ "momentum-distribution curves. Please provide an energy_value "
237
+ "for which to plot the MDCs.")
238
+ else:
239
+ energy_index = np.abs(self.enel - energy_value).argmin()
240
+ kinergy = self.ekin[energy_index]
241
+ counts = self.intensities[energy_index, :]
242
+
243
+ if not (self.enel.min() <= energy_value <= self.enel.max()):
244
+ raise ValueError(
245
+ f"Selected energy_value={energy_value:.3f} "
246
+ f"is outside the available energy range "
247
+ f"({self.enel.min():.3f} – {self.enel.max():.3f}) "
248
+ "of the MDC collection."
249
+ )
250
+
251
+ return counts, kinergy
252
+
253
+
254
+ def plot(self, energy_value=None, energy_range=None, ax=None, **kwargs):
255
+ """
256
+ Interactive or static plot with optional slider and full wrapper
257
+ support. Behavior consistent with Jupyter and CLI based on show /
258
+ fig_close.
259
+ """
260
+ import matplotlib.pyplot as plt
261
+ from matplotlib.widgets import Slider
262
+ import string
263
+ import sys
264
+ import warnings
265
+
266
+ # Wrapper kwargs
267
+ title = kwargs.pop("title", None)
268
+ savefig = kwargs.pop("savefig", None)
269
+ show = kwargs.pop("show", True)
270
+ fig_close = kwargs.pop("fig_close", False)
271
+ tight_layout = kwargs.pop("tight_layout", False)
272
+ ax_grid = kwargs.pop("ax_grid", None)
273
+ ax_annotate = kwargs.pop("ax_annotate", False)
274
+ size_kwargs = kwargs.pop("size_kwargs", None)
275
+
276
+ if energy_value is not None and energy_range is not None:
277
+ raise ValueError(
278
+ "Provide at most energy_value or energy_range, not both.")
279
+
280
+ ax, fig, plt = get_ax_fig_plt(ax=ax)
281
+
282
+ angles = self.angles
283
+ energies = self.enel
284
+
285
+ if np.isscalar(energies):
286
+ if energy_value is not None or energy_range is not None:
287
+ raise ValueError(
288
+ "This dataset contains only one momentum-distribution "
289
+ "curve; do not provide energy_value or energy_range."
290
+ )
291
+
292
+ intensities = self.intensities
293
+ ax.scatter(angles, intensities, label="Data")
294
+ ax.set_title(f"Energy slice: {energies * KILO:.3f} meV")
295
+
296
+ # --- y-only autoscale, preserve x ---
297
+ x0, x1 = ax.get_xlim() # keep current x-range
298
+ ax.relim(visible_only=True) # recompute data limits
299
+ ax.autoscale_view(scalex=False, scaley=True)
300
+ ax.set_xlim(x0, x1) # restore x (belt-and-suspenders)
301
+
302
+ else:
303
+ if (energy_value is not None) and (energy_range is not None):
304
+ raise ValueError("Provide either energy_value or energy_range, not both.")
305
+
306
+ emin, emax = energies.min(), energies.max()
307
+
308
+ # ---- Single-slice path (no slider) ----
309
+ if energy_value is not None:
310
+ if energy_value < emin or energy_value > emax:
311
+ raise ValueError(
312
+ f"Requested energy_value {energy_value:.3f} eV is "
313
+ f"outside the available energy range "
314
+ f"[{emin:.3f}, {emax:.3f}] eV."
315
+ )
316
+ idx = int(np.abs(energies - energy_value).argmin())
317
+ intensities = self.intensities[idx]
318
+ ax.scatter(angles, intensities, label="Data")
319
+ ax.set_title(f"Energy slice: {energies[idx] * KILO:.3f} meV")
320
+
321
+ # --- y-only autoscale, preserve x ---
322
+ x0, x1 = ax.get_xlim() # keep current x-range
323
+ ax.relim(visible_only=True) # recompute data limits
324
+ ax.autoscale_view(scalex=False, scaley=True)
325
+ ax.set_xlim(x0, x1) # restore x (belt-and-suspenders)
326
+
327
+ # ---- Multi-slice path (slider) ----
328
+ else:
329
+ if energy_range is not None:
330
+ e_min, e_max = energy_range
331
+ mask = (energies >= e_min) & (energies <= e_max)
332
+ else:
333
+ mask = np.ones_like(energies, dtype=bool)
334
+
335
+ indices = np.where(mask)[0]
336
+ if len(indices) == 0:
337
+ raise ValueError("No energies found in the specified selection.")
338
+
339
+ intensities = self.intensities[indices]
340
+
341
+ fig.subplots_adjust(bottom=0.25)
342
+ idx = 0
343
+ scatter = ax.scatter(angles, intensities[idx], label="Data")
344
+ ax.set_title(f"Energy slice: "
345
+ f"{energies[indices[idx]] * KILO:.3f} meV")
346
+
347
+ # Suppress single-point slider warning (when len(indices) == 1)
348
+ warnings.filterwarnings(
349
+ "ignore",
350
+ message="Attempting to set identical left == right",
351
+ category=UserWarning
352
+ )
353
+
354
+ slider_ax = fig.add_axes([0.2, 0.08, 0.6, 0.04])
355
+ slider = Slider(
356
+ slider_ax, "Index", 0, len(indices) - 1,
357
+ valinit=idx, valstep=1
358
+ )
359
+
360
+ def update(val):
361
+ i = int(slider.val)
362
+ yi = intensities[i]
363
+
364
+ scatter.set_offsets(np.c_[angles, yi])
365
+
366
+ x0, x1 = ax.get_xlim()
367
+
368
+ yv = np.asarray(yi, dtype=float).ravel()
369
+ mask = np.isfinite(yv)
370
+ if mask.any():
371
+ y_min = float(yv[mask].min())
372
+ y_max = float(yv[mask].max())
373
+ span = y_max - y_min
374
+ frac = plt.rcParams['axes.ymargin']
375
+
376
+ if span <= 0 or not np.isfinite(span):
377
+ scale = max(abs(y_max), 1.0)
378
+ pad = frac * scale
379
+ else:
380
+ pad = frac * span
381
+
382
+ ax.set_ylim(y_min - pad, y_max + pad)
383
+
384
+ # Keep x unchanged
385
+ ax.set_xlim(x0, x1)
386
+
387
+ # Update title and redraw
388
+ ax.set_title(f"Energy slice: "
389
+ f"{energies[indices[i]] * KILO:.3f} meV")
390
+ fig.canvas.draw_idle()
391
+
392
+ slider.on_changed(update)
393
+ self._slider = slider
394
+ self._line = scatter
395
+
396
+ ax.set_xlabel("Angle (°)")
397
+ ax.set_ylabel("Counts (-)")
398
+ ax.legend()
399
+ self._fig = fig
400
+
401
+ if size_kwargs:
402
+ fig.set_size_inches(size_kwargs.pop("w"),
403
+ size_kwargs.pop("h"), **size_kwargs)
404
+ if title:
405
+ fig.suptitle(title)
406
+ if tight_layout:
407
+ fig.tight_layout()
408
+ if savefig:
409
+ fig.savefig(savefig)
410
+ if ax_grid is not None:
411
+ for axis in fig.axes:
412
+ axis.grid(bool(ax_grid))
413
+ if ax_annotate:
414
+ tags = string.ascii_lowercase
415
+ for i, axis in enumerate(fig.axes):
416
+ axis.annotate(f"({tags[i]})", xy=(0.05, 0.95),
417
+ xycoords="axes fraction")
418
+
419
+ is_interactive = hasattr(sys, 'ps1') or 'ipykernel' in sys.modules
420
+ is_cli = not is_interactive
421
+
422
+ if show:
423
+ if is_cli:
424
+ plt.show()
425
+ if fig_close:
426
+ plt.close(fig)
427
+
428
+ if not show and (fig_close or is_cli):
429
+ return None
430
+ return fig
431
+
432
+
433
+ @add_fig_kwargs
434
+ def visualize_guess(self, distributions, energy_value=None,
435
+ matrix_element=None, matrix_args=None,
436
+ ax=None, **kwargs):
437
+ r"""
438
+ """
439
+
440
+ counts, kinergy = self.energy_check(energy_value)
441
+
442
+ ax, fig, plt = get_ax_fig_plt(ax=ax)
443
+
444
+ ax.set_xlabel('Angle ($\\degree$)')
445
+ ax.set_ylabel('Counts (-)')
446
+ ax.set_title(f"Energy slice: "
447
+ f"{(kinergy - self.hnuminPhi) * KILO:.3f} meV")
448
+ ax.scatter(self.angles, counts, label='Data')
449
+
450
+ final_result = self._merge_and_plot(ax=ax,
451
+ distributions=distributions, kinetic_energy=kinergy,
452
+ matrix_element=matrix_element,
453
+ matrix_args=dict(matrix_args) if matrix_args else None,
454
+ plot_individual=True,
455
+ )
456
+
457
+ residual = counts - final_result
458
+ ax.scatter(self.angles, residual, label='Residual')
459
+ ax.legend()
460
+
461
+ return fig
462
+
463
+
464
+ def fit_selection(self, distributions, energy_value=None, energy_range=None,
465
+ matrix_element=None, matrix_args=None, ax=None, **kwargs):
466
+ r"""
467
+ """
468
+ import matplotlib.pyplot as plt
469
+ from matplotlib.widgets import Slider
470
+ from copy import deepcopy
471
+ import string
472
+ import sys
473
+ import warnings
474
+ from lmfit import Minimizer
475
+ from scipy.ndimage import gaussian_filter
476
+ from .functions import construct_parameters, build_distributions, \
477
+ residual, resolve_param_name
478
+
479
+ # Wrapper kwargs
480
+ title = kwargs.pop("title", None)
481
+ savefig = kwargs.pop("savefig", None)
482
+ show = kwargs.pop("show", True)
483
+ fig_close = kwargs.pop("fig_close", False)
484
+ tight_layout = kwargs.pop("tight_layout", False)
485
+ ax_grid = kwargs.pop("ax_grid", None)
486
+ ax_annotate = kwargs.pop("ax_annotate", False)
487
+ size_kwargs = kwargs.pop("size_kwargs", None)
488
+
489
+ ax, fig, plt = get_ax_fig_plt(ax=ax)
490
+
491
+ energies = self.enel
492
+ new_distributions = deepcopy(distributions)
493
+
494
+ if energy_value is not None and energy_range is not None:
495
+ raise ValueError(
496
+ "Provide at most energy_value or energy_range, not both.")
497
+
498
+ if np.isscalar(energies):
499
+ if energy_value is not None or energy_range is not None:
500
+ raise ValueError(
501
+ "This dataset contains only one momentum-distribution "
502
+ "curve; do not provide energy_value or energy_range."
503
+ )
504
+ kinergies = np.atleast_1d(self.ekin)
505
+ intensities = np.atleast_2d(self.intensities)
506
+
507
+ else:
508
+ if energy_value is not None:
509
+ if (energy_value < energies.min() or energy_value > energies.max()):
510
+ raise ValueError( f"Requested energy_value {energy_value:.3f} eV is "
511
+ f"outside the available energy range "
512
+ f"[{energies.min():.3f}, {energies.max():.3f}] eV." )
513
+ idx = np.abs(energies - energy_value).argmin()
514
+ indices = np.atleast_1d(idx)
515
+ kinergies = self.ekin[indices]
516
+ intensities = self.intensities[indices, :]
517
+
518
+ elif energy_range is not None:
519
+ e_min, e_max = energy_range
520
+ indices = np.where((energies >= e_min) & (energies <= e_max))[0]
521
+ if len(indices) == 0:
522
+ raise ValueError("No energies found in the specified energy_range.")
523
+ kinergies = self.ekin[indices]
524
+ intensities = self.intensities[indices, :]
525
+
526
+ else: # Without specifying a range, all MDCs are plotted
527
+ kinergies = self.ekin
528
+ intensities = self.intensities
529
+
530
+ # Final shape guard
531
+ kinergies = np.atleast_1d(kinergies)
532
+ intensities = np.atleast_2d(intensities)
533
+
534
+ all_final_results = []
535
+ all_residuals = []
536
+ all_individual_results = [] # List of (n_individuals, n_angles)
537
+
538
+ aggregated_properties = {}
539
+
540
+ # map class_name -> parameter names to extract
541
+ param_spec = {
542
+ 'Constant': ('offset',),
543
+ 'Linear': ('offset', 'slope'),
544
+ 'SpectralLinear': ('amplitude', 'peak', 'broadening'),
545
+ 'SpectralQuadratic': ('amplitude', 'peak', 'broadening'),
546
+ }
547
+
548
+ order = np.argsort(kinergies)[::-1]
549
+ for idx in order:
550
+ kinergy = kinergies[idx]
551
+ intensity = intensities[idx]
552
+ if matrix_element is not None:
553
+ parameters, element_names = construct_parameters(
554
+ new_distributions, matrix_args)
555
+ new_distributions = build_distributions(new_distributions, parameters)
556
+ mini = Minimizer(
557
+ residual, parameters,
558
+ fcn_args=(self.angles, intensity, self.angle_resolution,
559
+ new_distributions, kinergy, self.hnuminPhi,
560
+ matrix_element, element_names)
561
+ )
562
+ else:
563
+ parameters = construct_parameters(new_distributions)
564
+ new_distributions = build_distributions(new_distributions, parameters)
565
+ mini = Minimizer(
566
+ residual, parameters,
567
+ fcn_args=(self.angles, intensity, self.angle_resolution,
568
+ new_distributions, kinergy, self.hnuminPhi)
569
+ )
570
+
571
+ outcome = mini.minimize('least_squares')
572
+
573
+ pcov = outcome.covar
574
+
575
+ var_names = getattr(outcome, 'var_names', None)
576
+ if not var_names:
577
+ var_names = [n for n, p in outcome.params.items() if p.vary]
578
+ var_idx = {n: i for i, n in enumerate(var_names)}
579
+
580
+ param_sigma_full = {}
581
+ for name, par in outcome.params.items():
582
+ sigma = None
583
+ if pcov is not None and name in var_idx:
584
+ d = pcov[var_idx[name], var_idx[name]]
585
+ if np.isfinite(d) and d >= 0:
586
+ sigma = float(np.sqrt(d))
587
+ if sigma is None:
588
+ s = getattr(par, 'stderr', None)
589
+ sigma = float(s) if s is not None else None
590
+ param_sigma_full[name] = sigma
591
+
592
+ # Rebuild the *fitted* distributions from optimized params
593
+ fitted_distributions = build_distributions(new_distributions, outcome.params)
594
+
595
+ # If using a matrix element, extract slice-specific args from the fit
596
+ if matrix_element is not None:
597
+ new_matrix_args = {key: outcome.params[key].value for key in matrix_args}
598
+ else:
599
+ new_matrix_args = None
600
+
601
+ # individual curves (smoothed, cropped) and final sum (no plotting here)
602
+ extend, step, numb = extend_function(self.angles, self.angle_resolution)
603
+
604
+ total_result_ext = np.zeros_like(extend)
605
+ indiv_rows = [] # (n_individuals, n_angles)
606
+ individual_labels = []
607
+
608
+ for dist in fitted_distributions:
609
+ # evaluate each component on the extended grid
610
+ if getattr(dist, 'class_name', None) == 'SpectralQuadratic':
611
+ if (getattr(dist, 'center_angle', None) is not None) and (
612
+ kinergy is None or self.hnuminPhi is None
613
+ ):
614
+ raise ValueError(
615
+ 'Spectral quadratic function is defined in terms '
616
+ 'of a center angle. Please provide a kinetic energy '
617
+ 'and hnuminPhi.'
618
+ )
619
+ extended_result = dist.evaluate(extend, kinergy, self.hnuminPhi)
620
+ else:
621
+ extended_result = dist.evaluate(extend)
622
+
623
+ if matrix_element is not None and hasattr(dist, 'index'):
624
+ args = new_matrix_args or {}
625
+ extended_result *= matrix_element(extend, **args)
626
+
627
+ total_result_ext += extended_result
628
+
629
+ # smoothed & cropped individual
630
+ individual_curve = gaussian_filter(extended_result, sigma=step)[
631
+ numb:-numb if numb else None
632
+ ]
633
+ indiv_rows.append(np.asarray(individual_curve))
634
+
635
+ # label
636
+ label = getattr(dist, 'label', str(dist))
637
+ individual_labels.append(label)
638
+
639
+ # ---- collect parameters for this distribution
640
+ # (Aggregated over slices)
641
+ cls = getattr(dist, 'class_name', None)
642
+ wanted = param_spec.get(cls, ())
643
+
644
+ # ensure dicts exist
645
+ label_bucket = aggregated_properties.setdefault(label, {})
646
+ class_bucket = label_bucket.setdefault(
647
+ cls, {'label': label, '_class': cls}
648
+ )
649
+
650
+ # store center_wavevector (scalar) for SpectralQuadratic
651
+ if (
652
+ cls == 'SpectralQuadratic'
653
+ and hasattr(dist, 'center_wavevector')
654
+ ):
655
+ class_bucket.setdefault(
656
+ 'center_wavevector', dist.center_wavevector
657
+ )
658
+
659
+ # ensure keys for both values and sigmas
660
+ for pname in wanted:
661
+ class_bucket.setdefault(pname, [])
662
+ class_bucket.setdefault(f"{pname}_sigma", [])
663
+
664
+ # append values and sigmas in the order of slices
665
+ for pname in wanted:
666
+ param_key = resolve_param_name(outcome.params, label, pname)
667
+
668
+ if param_key is not None and param_key in outcome.params:
669
+ class_bucket[pname].append(outcome.params[param_key].value)
670
+ class_bucket[f"{pname}_sigma"].append(param_sigma_full.get(param_key, None))
671
+ else:
672
+ # Not fitted in this slice → keep the value if present on the dist, sigma=None
673
+ class_bucket[pname].append(getattr(dist, pname, None))
674
+ class_bucket[f"{pname}_sigma"].append(None)
675
+
676
+ # final (sum) curve, smoothed & cropped
677
+ final_result_i = gaussian_filter(total_result_ext, sigma=step)[
678
+ numb:-numb if numb else None]
679
+ final_result_i = np.asarray(final_result_i)
680
+
681
+ # Residual for this slice
682
+ residual_i = np.asarray(intensity) - final_result_i
683
+
684
+ # Store per-slice results
685
+ all_final_results.append(final_result_i)
686
+ all_residuals.append(residual_i)
687
+ all_individual_results.append(np.vstack(indiv_rows))
688
+
689
+ # --- after the reversed-order loop, restore original (ascending) order ---
690
+ inverse_order = np.argsort(np.argsort(kinergies)[::-1])
691
+
692
+ # Reorder per-slice arrays/lists computed in the loop
693
+ all_final_results[:] = [all_final_results[i] for i in inverse_order]
694
+ all_residuals[:] = [all_residuals[i] for i in inverse_order]
695
+ all_individual_results[:] = [all_individual_results[i] for i in inverse_order]
696
+
697
+ # Reorder all per-slice lists in aggregated_properties
698
+ for label_dict in aggregated_properties.values():
699
+ for cls_dict in label_dict.values():
700
+ for key, val in cls_dict.items():
701
+ if isinstance(val, list) and len(val) == len(kinergies):
702
+ cls_dict[key] = [val[i] for i in inverse_order]
703
+
704
+ self._ekin_range = kinergies
705
+ self._individual_properties = aggregated_properties
706
+
707
+ if np.isscalar(energies):
708
+ # One slice only: plot MDC, Fit, Residual, and Individuals
709
+ ydata = np.asarray(intensities).squeeze()
710
+ yfit = np.asarray(all_final_results[0]).squeeze()
711
+ yres = np.asarray(all_residuals[0]).squeeze()
712
+ yind = np.asarray(all_individual_results[0])
713
+
714
+ ax.scatter(self.angles, ydata, label="Data")
715
+ # plot individuals with their labels
716
+ for j, lab in enumerate(individual_labels or []):
717
+ ax.plot(self.angles, yind[j], label=str(lab))
718
+ ax.plot(self.angles, yfit, label="Fit")
719
+ ax.scatter(self.angles, yres, label="Residual")
720
+
721
+ ax.set_title(f"Energy slice: {energies * KILO:.3f} meV")
722
+ ax.relim() # recompute data limits from all artists
723
+ ax.autoscale_view() # apply autoscaling + axes.ymargin padding
724
+
725
+ else:
726
+ if energy_value is not None:
727
+ _idx = int(np.abs(energies - energy_value).argmin())
728
+ energies_sel = np.atleast_1d(energies[_idx])
729
+ elif energy_range is not None:
730
+ e_min, e_max = energy_range
731
+ energies_sel = energies[(energies >= e_min)
732
+ & (energies <= e_max)]
733
+ else:
734
+ energies_sel = energies
735
+
736
+ # Number of slices must match
737
+ n_slices = len(all_final_results)
738
+ assert intensities.shape[0] == n_slices == len(all_residuals) \
739
+ == len(all_individual_results), (f"Mismatch: data \
740
+ {intensities.shape[0]}, fits {len(all_final_results)}, "
741
+ f"residuals {len(all_residuals)}, \
742
+ individuals {len(all_individual_results)}."
743
+ )
744
+ n_individuals = all_individual_results[0].shape[0] \
745
+ if n_slices else 0
746
+
747
+ fig.subplots_adjust(bottom=0.25)
748
+ idx = 0
749
+
750
+ # Initial draw (MDC + Individuals + Fit + Residual) at slice 0
751
+ scatter = ax.scatter(self.angles, intensities[idx], label="Data")
752
+
753
+ individual_lines = []
754
+ if n_individuals:
755
+ for j in range(n_individuals):
756
+ if individual_labels and j < len(individual_labels):
757
+ label = str(individual_labels[j])
758
+ else:
759
+ label = f"Comp {j}"
760
+
761
+ yvals = all_individual_results[idx][j]
762
+ line, = ax.plot(self.angles, yvals, label=label)
763
+ individual_lines.append(line)
764
+
765
+ result_line, = ax.plot(self.angles, all_final_results[idx],
766
+ label="Fit")
767
+ resid_scatter = ax.scatter(self.angles, all_residuals[idx],
768
+ label="Residual")
769
+
770
+ # Title + limits (use only the currently shown slice)
771
+ ax.set_title(f"Energy slice: {energies_sel[idx] * KILO:.3f} meV")
772
+ ax.relim() # recompute data limits from all artists
773
+ ax.autoscale_view() # apply autoscaling + axes.ymargin padding
774
+
775
+ # Suppress warning when a single MDC is plotted
776
+ warnings.filterwarnings(
777
+ "ignore",
778
+ message="Attempting to set identical left == right",
779
+ category=UserWarning
780
+ )
781
+
782
+ # Slider over slice index (0..n_slices-1)
783
+ slider_ax = fig.add_axes([0.2, 0.08, 0.6, 0.04])
784
+ slider = Slider(
785
+ slider_ax, "Index", 0, n_slices - 1,
786
+ valinit=idx, valstep=1
787
+ )
788
+
789
+ def update(val):
790
+ i = int(slider.val)
791
+ # Update MDC points
792
+ scatter.set_offsets(np.c_[self.angles, intensities[i]])
793
+
794
+ # Update individuals
795
+ if n_individuals:
796
+ Yi = all_individual_results[i] # (n_individuals, n_angles)
797
+ for j, ln in enumerate(individual_lines):
798
+ ln.set_ydata(Yi[j])
799
+
800
+ # Update fit and residual
801
+ result_line.set_ydata(all_final_results[i])
802
+ resid_scatter.set_offsets(np.c_[self.angles, all_residuals[i]])
803
+
804
+ ax.relim()
805
+ ax.autoscale_view()
806
+
807
+ # Update title and redraw
808
+ ax.set_title(f"Energy slice: "
809
+ f"{energies_sel[i] * KILO:.3f} meV")
810
+ fig.canvas.draw_idle()
811
+
812
+ slider.on_changed(update)
813
+ self._slider = slider
814
+ self._line = scatter
815
+ self._individual_lines = individual_lines
816
+ self._result_line = result_line
817
+ self._resid_scatter = resid_scatter
818
+
819
+ ax.set_xlabel("Angle (°)")
820
+ ax.set_ylabel("Counts (-)")
821
+ ax.legend()
822
+ self._fig = fig
823
+
824
+ if size_kwargs:
825
+ fig.set_size_inches(size_kwargs.pop("w"),
826
+ size_kwargs.pop("h"), **size_kwargs)
827
+ if title:
828
+ fig.suptitle(title)
829
+ if tight_layout:
830
+ fig.tight_layout()
831
+ if savefig:
832
+ fig.savefig(savefig)
833
+ if ax_grid is not None:
834
+ for axis in fig.axes:
835
+ axis.grid(bool(ax_grid))
836
+ if ax_annotate:
837
+ tags = string.ascii_lowercase
838
+ for i, axis in enumerate(fig.axes):
839
+ axis.annotate(f"({tags[i]})", xy=(0.05, 0.95),
840
+ xycoords="axes fraction")
841
+
842
+ is_interactive = hasattr(sys, 'ps1') or 'ipykernel' in sys.modules
843
+ is_cli = not is_interactive
844
+
845
+ if show:
846
+ if is_cli:
847
+ plt.show()
848
+ if fig_close:
849
+ plt.close(fig)
850
+
851
+ if not show and (fig_close or is_cli):
852
+ return None
853
+ return fig
854
+
855
+
856
+ @add_fig_kwargs
857
+ def fit(self, distributions, energy_value=None, matrix_element=None,
858
+ matrix_args=None, ax=None, **kwargs):
859
+ r"""
860
+ """
861
+ from copy import deepcopy
862
+ from lmfit import Minimizer
863
+ from .functions import construct_parameters, build_distributions, \
864
+ residual
865
+
866
+ counts, kinergy = self.energy_check(energy_value)
867
+
868
+ ax, fig, plt = get_ax_fig_plt(ax=ax)
869
+
870
+ ax.set_xlabel('Angle ($\\degree$)')
871
+ ax.set_ylabel('Counts (-)')
872
+ ax.set_title(f"Energy slice: "
873
+ f"{(kinergy - self.hnuminPhi) * KILO:.3f} meV")
874
+
875
+ ax.scatter(self.angles, counts, label='Data')
876
+
877
+ new_distributions = deepcopy(distributions)
878
+
879
+ if matrix_element is not None:
880
+ parameters, element_names = construct_parameters(distributions,
881
+ matrix_args)
882
+ new_distributions = build_distributions(new_distributions, \
883
+ parameters)
884
+ mini = Minimizer(
885
+ residual, parameters,
886
+ fcn_args=(self.angles, counts, self.angle_resolution,
887
+ new_distributions, kinergy, self.hnuminPhi,
888
+ matrix_element, element_names))
889
+ else:
890
+ parameters = construct_parameters(distributions)
891
+ new_distributions = build_distributions(new_distributions,
892
+ parameters)
893
+ mini = Minimizer(residual, parameters,
894
+ fcn_args=(self.angles, counts, self.angle_resolution,
895
+ new_distributions, kinergy, self.hnuminPhi))
896
+
897
+ outcome = mini.minimize('least_squares')
898
+ pcov = outcome.covar
899
+
900
+ # If matrix params were fitted, pass the fitted values to plotting
901
+ if matrix_element is not None:
902
+ new_matrix_args = {key: outcome.params[key].value for key in
903
+ matrix_args}
904
+ else:
905
+ new_matrix_args = None
906
+
907
+ final_result = self._merge_and_plot(ax=ax,
908
+ distributions=new_distributions, kinetic_energy=kinergy,
909
+ matrix_element=matrix_element, matrix_args=new_matrix_args,
910
+ plot_individual=True)
911
+
912
+ residual_vals = counts - final_result
913
+ ax.scatter(self.angles, residual_vals, label='Residual')
914
+ ax.legend()
915
+ if matrix_element is not None:
916
+ return fig, new_distributions, pcov, new_matrix_args
917
+ else:
918
+ return fig, new_distributions, pcov
919
+
920
+
921
+ def _merge_and_plot(self, ax, distributions, kinetic_energy,
922
+ matrix_element=None, matrix_args=None,
923
+ plot_individual=True):
924
+ r"""
925
+ Evaluate distributions on the extended grid, apply optional matrix
926
+ element, smooth, plot individuals and the summed curve.
927
+
928
+ Returns
929
+ -------
930
+ final_result : np.ndarray
931
+ Smoothed, cropped total distribution aligned with self.angles.
932
+ """
933
+ from scipy.ndimage import gaussian_filter
934
+
935
+ # Build extended grid
936
+ extend, step, numb = extend_function(self.angles, self.angle_resolution)
937
+ total_result = np.zeros_like(extend)
938
+
939
+ for dist in distributions:
940
+ # Special handling for SpectralQuadratic
941
+ if getattr(dist, 'class_name', None) == 'SpectralQuadratic':
942
+ if (getattr(dist, 'center_angle', None) is not None) and (
943
+ kinetic_energy is None or self.hnuminPhi is None
944
+ ):
945
+ raise ValueError(
946
+ 'Spectral quadratic function is defined in terms '
947
+ 'of a center angle. Please provide a kinetic energy '
948
+ 'and hnuminPhi.'
949
+ )
950
+ extended_result = dist.evaluate(extend, kinetic_energy, \
951
+ self.hnuminPhi)
952
+ else:
953
+ extended_result = dist.evaluate(extend)
954
+
955
+ # Optional matrix element (only for components that advertise an index)
956
+ if matrix_element is not None and hasattr(dist, 'index'):
957
+ args = matrix_args or {}
958
+ extended_result *= matrix_element(extend, **args)
959
+
960
+ total_result += extended_result
961
+
962
+ if plot_individual and ax:
963
+ individual = gaussian_filter(extended_result, sigma=step)\
964
+ [numb:-numb if numb else None]
965
+ ax.plot(self.angles, individual, label=getattr(dist, \
966
+ 'label', str(dist)))
967
+
968
+ # Smoothed, cropped total curve aligned to self.angles
969
+ final_result = gaussian_filter(total_result, sigma=step)[numb:-numb \
970
+ if numb else None]
971
+ if ax:
972
+ ax.plot(self.angles, final_result, label='Distribution sum')
973
+
974
+ return final_result
975
+
976
+
977
+ def expose_parameters(self, select_label, fermi_wavevector=None,
978
+ fermi_velocity=None, bare_mass=None, side=None):
979
+ r"""
980
+ Select and return fitted parameters for a given component label, plus a
981
+ flat export dictionary containing values **and** 1σ uncertainties.
982
+
983
+ Parameters
984
+ ----------
985
+ select_label : str
986
+ Label to look for among the fitted distributions.
987
+ fermi_wavevector : float, optional
988
+ Optional Fermi wave vector to include.
989
+ fermi_velocity : float, optional
990
+ Optional Fermi velocity to include.
991
+ bare_mass : float, optional
992
+ Optional bare mass to include (used for SpectralQuadratic
993
+ dispersions).
994
+ side : {'left','right'}, optional
995
+ Optional side selector for SpectralQuadratic dispersions.
996
+
997
+ Returns
998
+ -------
999
+ ekin_range : np.ndarray
1000
+ Kinetic-energy grid corresponding to the selected label.
1001
+ hnuminPhi : float
1002
+ Photoelectron work-function offset.
1003
+ energy_resolution : float
1004
+ Energy resolution associated with the extracted self-energy data.
1005
+ temperature : float
1006
+ Temperature [K] associated with the extracted self-energy data.
1007
+ label : str
1008
+ Label of the selected distribution.
1009
+ selected_properties : dict or list of dict
1010
+ Nested dictionary (or list thereof) containing <param> and
1011
+ <param>_sigma arrays. For SpectralQuadratic components, a
1012
+ scalar `center_wavevector` is also present.
1013
+ exported_parameters : dict
1014
+ Flat dictionary of parameters and their uncertainties, plus
1015
+ optional Fermi quantities and `side`. For SpectralQuadratic
1016
+ components, `center_wavevector` is included and taken directly
1017
+ from the fitted distribution.
1018
+ """
1019
+
1020
+ if self._ekin_range is None:
1021
+ raise AttributeError(
1022
+ "ekin_range not yet set. Run `.fit_selection()` first."
1023
+ )
1024
+
1025
+ store = getattr(self, "_individual_properties", None)
1026
+ if not store or select_label not in store:
1027
+ all_labels = (sorted(store.keys())
1028
+ if isinstance(store, dict) else [])
1029
+ raise ValueError(
1030
+ f"Label '{select_label}' not found in available labels: "
1031
+ f"{all_labels}"
1032
+ )
1033
+
1034
+ # Convert lists → numpy arrays within the selected label’s classes.
1035
+ # Keep scalar center_wavevector as a scalar.
1036
+ per_class_dicts = []
1037
+ for cls, bucket in store[select_label].items():
1038
+ dct = {}
1039
+ for k, v in bucket.items():
1040
+ if k in ("label", "_class"):
1041
+ dct[k] = v
1042
+ elif k == "center_wavevector":
1043
+ # keep scalar as-is, do not wrap in np.asarray
1044
+ dct[k] = v
1045
+ else:
1046
+ dct[k] = np.asarray(v)
1047
+ per_class_dicts.append(dct)
1048
+
1049
+ selected_properties = (
1050
+ per_class_dicts[0] if len(per_class_dicts) == 1 else per_class_dicts
1051
+ )
1052
+
1053
+ # Flat export dict: simple keys, includes optional extras
1054
+ exported_parameters = {
1055
+ "fermi_wavevector": fermi_wavevector,
1056
+ "fermi_velocity": fermi_velocity,
1057
+ "bare_mass": bare_mass,
1058
+ "side": side,
1059
+ }
1060
+
1061
+ # Collect parameters without prefixing by class. This will also include
1062
+ # center_wavevector from the fitted SpectralQuadratic class, and since
1063
+ # there is no function argument with that name, it cannot be overridden.
1064
+ if isinstance(selected_properties, dict):
1065
+ for key, val in selected_properties.items():
1066
+ if key not in ("label", "_class"):
1067
+ exported_parameters[key] = val
1068
+ else:
1069
+ # If multiple classes, merge sequentially
1070
+ # (last overwrites same-name keys).
1071
+ for cls_bucket in selected_properties:
1072
+ for key, val in cls_bucket.items():
1073
+ if key not in ("label", "_class"):
1074
+ exported_parameters[key] = val
1075
+
1076
+ return (self._ekin_range, self.hnuminPhi, self.energy_resolution,
1077
+ self.temperature, select_label, selected_properties,
1078
+ exported_parameters)