xarpes 0.3.4__py3-none-any.whl → 0.5.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/spectral.py DELETED
@@ -1,2175 +0,0 @@
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 all the spectral quantities."""
13
-
14
- import numpy as np
15
- from igor2 import binarywave
16
- from .plotting import get_ax_fig_plt, add_fig_kwargs
17
- from .functions import fit_leastsq, extend_function
18
- from .distributions import FermiDirac, Linear
19
- from .constants import uncr, pref, dtor, kilo, stdv
20
-
21
- class BandMap:
22
- r"""
23
- Class for the band map from the ARPES experiment.
24
-
25
- Parameters
26
- ----------
27
- datafile : str, optional
28
- Path to an IGOR binary wave file.
29
- intensities : ndarray, optional
30
- 2D array of intensities [E, angle].
31
- angles : ndarray, optional
32
- 1D array of emission angles in degrees.
33
- ekin : ndarray, optional
34
- 1D array of kinetic energies in eV.
35
- enel : ndarray, optional
36
- 1D array of electron energies in eV.
37
- energy_resolution : float, optional
38
- Energy-resolution standard deviation [eV].
39
- angle_resolution : float, optional
40
- Angular-resolution standard deviation [deg].
41
- temperature : float, optional
42
- Sample temperature [K].
43
- hnuminphi : float, optional
44
- Photon energy minus work function [eV].
45
- hnuminphi_std : float, optional
46
- Standard deviation on ``hnuminphi`` [eV].
47
- transpose : bool, optional
48
- If True, transpose the input data.
49
- flip_ekin : bool, optional
50
- If True, flip the energy axis.
51
- flip_angles : bool, optional
52
- If True, flip the angle axis.
53
-
54
- Attributes
55
- ----------
56
- intensities : ndarray
57
- 2D intensity map [energy, angle].
58
- angles : ndarray
59
- Emission angles in degrees.
60
- ekin : ndarray
61
- Kinetic-energy axis in eV.
62
- enel : ndarray
63
- Electron-energy axis in eV.
64
- hnuminphi : float or None
65
- Photon energy minus work function.
66
- hnuminphi_std : float or None
67
- Standard deviation on ``hnuminphi``.
68
-
69
- """
70
-
71
- def __init__(self, datafile=None, intensities=None, angles=None,
72
- ekin=None, enel=None, energy_resolution=None,
73
- angle_resolution=None, temperature=None, hnuminphi=None,
74
- hnuminphi_std=None, transpose=False, flip_ekin=False,
75
- flip_angles=False):
76
-
77
- # --- IO / file load -------------------------------------------------
78
- if datafile is not None:
79
- data = binarywave.load(datafile)
80
- self.intensities = data['wave']['wData']
81
-
82
- fnum, anum = data['wave']['wave_header']['nDim'][0:2]
83
- fstp, astp = data['wave']['wave_header']['sfA'][0:2]
84
- fmin, amin = data['wave']['wave_header']['sfB'][0:2]
85
-
86
- if self.intensities.shape != (fnum, anum):
87
- raise ValueError('nDim and shape of wData do not match.')
88
-
89
- if transpose:
90
- self.intensities = self.intensities.T
91
- fnum, anum = anum, fnum
92
- fstp, astp = astp, fstp
93
- fmin, amin = amin, fmin
94
-
95
- if flip_ekin:
96
- self.intensities = self.intensities[::-1, :]
97
-
98
- if flip_angles:
99
- self.intensities = self.intensities[:, ::-1]
100
-
101
- self.angles = np.linspace(amin, amin + (anum - 1) * astp, anum)
102
- file_ekin = np.linspace(fmin, fmin + (fnum - 1) * fstp, fnum)
103
- else:
104
- file_ekin = None
105
-
106
- # --- Required arrays if not using datafile -------------------------
107
- if intensities is not None:
108
- self.intensities = intensities
109
- elif datafile is None:
110
- raise ValueError('Please provide datafile or intensities.')
111
-
112
- if angles is not None:
113
- self.angles = angles
114
- elif datafile is None:
115
- raise ValueError('Please provide datafile or angles.')
116
-
117
- # --- Initialize energy axes (raw slots) ----------------------------
118
- self._ekin = None
119
- self._enel = None
120
-
121
- # Apply user overrides or file ekin
122
- if ekin is not None and enel is not None:
123
- raise ValueError('Provide only one of ekin or enel, not both.')
124
-
125
- if ekin is not None:
126
- self._ekin = ekin
127
- elif enel is not None:
128
- self._enel = enel
129
- elif file_ekin is not None:
130
- self._ekin = file_ekin
131
- else:
132
- raise ValueError('Please provide datafile, ekin, or enel.')
133
-
134
- # Scalars / metadata
135
- self.energy_resolution = energy_resolution
136
- self.angle_resolution = angle_resolution
137
- self.temperature = temperature
138
-
139
- # Work-function combo and its std
140
- self._hnuminphi = None
141
- self._hnuminphi_std = None
142
- self.hnuminphi = hnuminphi
143
- self.hnuminphi_std = hnuminphi_std
144
-
145
- # --- 1) Track which axis is authoritative --------------------------
146
- self._ekin_explicit = ekin is not None or (file_ekin is not None
147
- and enel is None)
148
- self._enel_explicit = enel is not None
149
-
150
- # --- 2) Derive missing axis if possible ----------------------------
151
- if self._ekin is None and self._enel is not None \
152
- and self._hnuminphi is not None:
153
- self._ekin = self._enel + self._hnuminphi
154
- if self._enel is None and self._ekin is not None \
155
- and self._hnuminphi is not None:
156
- self._enel = self._ekin - self._hnuminphi
157
-
158
- # -------------------- Properties: data arrays ---------------------------
159
- @property
160
- def intensities(self):
161
- return self._intensities
162
-
163
- @intensities.setter
164
- def intensities(self, x):
165
- self._intensities = x
166
-
167
- @property
168
- def angles(self):
169
- return self._angles
170
-
171
- @angles.setter
172
- def angles(self, x):
173
- self._angles = x
174
-
175
- # -------------------- 3) Resolution / temperature ----------------------
176
- @property
177
- def angle_resolution(self):
178
- return self._angle_resolution
179
-
180
- @angle_resolution.setter
181
- def angle_resolution(self, x):
182
- self._angle_resolution = x
183
-
184
- @property
185
- def energy_resolution(self):
186
- return self._energy_resolution
187
-
188
- @energy_resolution.setter
189
- def energy_resolution(self, x):
190
- self._energy_resolution = x
191
-
192
- @property
193
- def temperature(self):
194
- return self._temperature
195
-
196
- @temperature.setter
197
- def temperature(self, x):
198
- self._temperature = x
199
-
200
- # -------------------- 4) Sync ekin / enel / hnuminphi ------------------
201
- @property
202
- def ekin(self):
203
- if self._ekin is None and self._enel is not None \
204
- and self._hnuminphi is not None:
205
- return self._enel + self._hnuminphi
206
- return self._ekin
207
-
208
- @ekin.setter
209
- def ekin(self, x):
210
- if getattr(self, "_enel_explicit", False):
211
- raise AttributeError('enel is explicit; set hnuminphi instead.')
212
- self._ekin = x
213
- self._ekin_explicit = True
214
- if not getattr(self, "_enel_explicit", False) \
215
- and self._hnuminphi is not None and x is not None:
216
- self._enel = x - self._hnuminphi
217
-
218
- @property
219
- def enel(self):
220
- if self._enel is None and self._ekin is not None \
221
- and self._hnuminphi is not None:
222
- return self._ekin - self._hnuminphi
223
- return self._enel
224
-
225
- @enel.setter
226
- def enel(self, x):
227
- if getattr(self, "_ekin_explicit", False):
228
- raise AttributeError('ekin is explicit; set hnuminphi instead.')
229
- self._enel = x
230
- self._enel_explicit = True
231
- if not getattr(self, "_ekin_explicit", False) \
232
- and self._hnuminphi is not None and x is not None:
233
- self._ekin = x + self._hnuminphi
234
-
235
- @property
236
- def hnuminphi(self):
237
- r"""Returns the photon energy minus the work function in eV if it has
238
- been set, either during instantiation, with the setter, or by fitting
239
- the Fermi-Dirac distribution to the integrated weight.
240
-
241
- Returns
242
- -------
243
- hnuminphi : float, None
244
- Kinetic energy minus the work function [eV]
245
-
246
- """
247
- return self._hnuminphi
248
-
249
- @hnuminphi.setter
250
- def hnuminphi(self, x):
251
- r"""TBD
252
- """
253
- self._hnuminphi = x
254
- # Re-derive the non-explicit axis if possible
255
- if not getattr(self, "_ekin_explicit", False) \
256
- and self._enel is not None and x is not None:
257
- self._ekin = self._enel + x
258
- if not getattr(self, "_enel_explicit", False) \
259
- and self._ekin is not None and x is not None:
260
- self._enel = self._ekin - x
261
-
262
- @property
263
- def hnuminphi_std(self):
264
- r"""Returns standard deviation of the photon energy minus the work
265
- function in eV.
266
-
267
- Returns
268
- -------
269
- hnuminphi_std : float
270
- Standard deviation of energy minus the work function [eV]
271
-
272
- """
273
- return self._hnuminphi_std
274
-
275
- @hnuminphi_std.setter
276
- def hnuminphi_std(self, x):
277
- r"""Manually sets the standard deviation of photon energy minus the
278
- work function in eV.
279
-
280
- Parameters
281
- ----------
282
- hnuminphi_std : float
283
- Standard deviation of energy minus the work function [eV]
284
-
285
- """
286
- self._hnuminphi_std = x
287
-
288
- def shift_angles(self, shift):
289
- r"""
290
- Shifts the angles by the specified amount in degrees. Used to shift
291
- from the detector angle to the material angle.
292
-
293
- Parameters
294
- ----------
295
- shift : float
296
- Angular shift [degrees]
297
-
298
- """
299
- self.angles = self.angles + shift
300
-
301
- def mdc_set(self, angle_min, angle_max, energy_value=None,
302
- energy_range=None):
303
- r"""Returns a set of MDCs. Documentation is to be further completed.
304
-
305
- Parameters
306
- ----------
307
- angle_min : float
308
- Minimum angle of integration interval [degrees]
309
- angle_max : float
310
- Maximum angle of integration interval [degrees]
311
-
312
- Returns
313
- -------
314
- angle_range : ndarray
315
- Array of size n containing the angular values
316
- energy_range : ndarray
317
- Array of size m containing the energy values
318
- mdcs : ndarray
319
- Array of size n x m containing the MDC intensities
320
-
321
- """
322
-
323
- if (energy_value is None and energy_range is None) or \
324
- (energy_value is not None and energy_range is not None):
325
- raise ValueError('Please provide either energy_value or ' +
326
- 'energy_range.')
327
-
328
- angle_min_index = np.abs(self.angles - angle_min).argmin()
329
- angle_max_index = np.abs(self.angles - angle_max).argmin()
330
- angle_range_out = self.angles[angle_min_index:angle_max_index + 1]
331
-
332
- if energy_value is not None:
333
- energy_index = np.abs(self.enel - energy_value).argmin()
334
- enel_range_out = self.enel[energy_index]
335
- mdcs = self.intensities[energy_index,
336
- angle_min_index:angle_max_index + 1]
337
-
338
- if energy_range:
339
- energy_indices = np.where((self.enel >= np.min(energy_range))
340
- & (self.enel <= np.max(energy_range))) \
341
- [0]
342
- enel_range_out = self.enel[energy_indices]
343
- mdcs = self.intensities[energy_indices,
344
- angle_min_index:angle_max_index + 1]
345
-
346
- return mdcs, angle_range_out, self.angle_resolution, \
347
- enel_range_out, self.hnuminphi
348
-
349
-
350
- @add_fig_kwargs
351
- def plot(self, abscissa='momentum', ordinate='electron_energy',
352
- self_energies=None, ax=None, markersize=None, **kwargs):
353
- r"""
354
- Plot the band map. Optionally attach a collection of self-energies,
355
- e.g. a CreateSelfEnergies instance or any iterable of self-energy
356
- objects. They are stored on `self` for later overlay plotting.
357
-
358
- When self-energies are present and ``abscissa='momentum'``, their
359
- MDC maxima are overlaid with 95 % confidence intervals.
360
- """
361
- import warnings
362
-
363
- valid_abscissa = ('angle', 'momentum')
364
- valid_ordinate = ('kinetic_energy', 'electron_energy')
365
-
366
- if abscissa not in valid_abscissa:
367
- raise ValueError(
368
- f"Invalid abscissa '{abscissa}'. "
369
- f"Valid options: {valid_abscissa}"
370
- )
371
- if ordinate not in valid_ordinate:
372
- raise ValueError(
373
- f"Invalid ordinate '{ordinate}'. "
374
- f"Valid options: {valid_ordinate}"
375
- )
376
-
377
- # Optionally store self-energies on the instance
378
- if self_energies is not None:
379
-
380
- # MDC maxima are defined in momentum space, not angle space
381
- if abscissa == 'angle' and isinstance(
382
- self_energies, (list, tuple, CreateSelfEnergies)
383
- ):
384
- raise ValueError( "MDC maxima cannot be plotted against "
385
- "angles; they are defined in momentum space. Use " \
386
- "abscissa='momentum' when passing a list of self-energies.")
387
-
388
- if not isinstance(self_energies, CreateSelfEnergies):
389
- self_energies = CreateSelfEnergies(self_energies)
390
-
391
- self._self_energies = self_energies
392
- elif not hasattr(self, "_self_energies"):
393
- self._self_energies = None
394
-
395
- ax, fig, plt = get_ax_fig_plt(ax=ax)
396
- # Below, **kwargs is
397
-
398
- Angl, Ekin = np.meshgrid(self.angles, self.ekin)
399
-
400
- if abscissa == 'angle':
401
- ax.set_xlabel('Angle ($\\degree$)')
402
- if ordinate == 'kinetic_energy':
403
- mesh = ax.pcolormesh(Angl, Ekin, self.intensities,
404
- shading='auto', cmap=plt.get_cmap('bone').reversed(),
405
- **kwargs)
406
- ax.set_ylabel('$E_{\\mathrm{kin}}$ (eV)')
407
- elif ordinate == 'electron_energy':
408
- Enel = Ekin - self.hnuminphi
409
- mesh = ax.pcolormesh(Angl, Enel, self.intensities,
410
- shading='auto', cmap=plt.get_cmap('bone').reversed(),
411
- **kwargs)
412
- ax.set_ylabel('$E-\\mu$ (eV)')
413
-
414
- elif abscissa == 'momentum':
415
- ax.set_xlabel(r'$k_{//}$ ($\mathrm{\AA}^{-1}$)')
416
-
417
- with warnings.catch_warnings(record=True) as wlist:
418
- warnings.filterwarnings("always",
419
- message=("The input coordinates to pcolormesh are "
420
- "interpreted as cell centers, but are not "
421
- "monotonically increasing or decreasing."),
422
- category=UserWarning)
423
-
424
- Mome = np.sqrt(Ekin / pref) * np.sin(Angl * dtor)
425
-
426
- if ordinate == 'kinetic_energy':
427
- mesh = ax.pcolormesh(Mome, Ekin, self.intensities,
428
- shading='auto', cmap=plt.get_cmap('bone').reversed(),
429
- **kwargs)
430
- ax.set_ylabel('$E_{\\mathrm{kin}}$ (eV)')
431
-
432
- elif ordinate == 'electron_energy':
433
- Enel = Ekin - self.hnuminphi
434
- mesh = ax.pcolormesh(Mome, Enel, self.intensities,
435
- shading='auto', cmap=plt.get_cmap('bone').reversed(),
436
- **kwargs)
437
- ax.set_ylabel('$E-\\mu$ (eV)')
438
-
439
- if any("cell centers" in str(w.message) for w in wlist):
440
- warnings.warn("Conversion from angle to momenta causes warping "
441
- "of the cell centers. \n Cell edges of the "
442
- "mesh plot may look irregular.", UserWarning,
443
- stacklevel=2)
444
-
445
- if abscissa == 'momentum' and self._self_energies is not None:
446
- for self_energy in self._self_energies:
447
- mdc_maxima = getattr(self_energy, "mdc_maxima", None)
448
-
449
- # If this self-energy doesn't contain maxima, don't plot them
450
- if mdc_maxima is None:
451
- continue
452
-
453
- peak_sigma = getattr(self_energy, "peak_positions_sigma",
454
- None)
455
- xerr = stdv * peak_sigma if peak_sigma is not None else None
456
-
457
- if ordinate == 'kinetic_energy':
458
- y_vals = self_energy.ekin_range
459
- else: # electron energy
460
- y_vals = self_energy.enel_range
461
-
462
- x_vals = mdc_maxima
463
- label = getattr(self_energy, "label", None)
464
-
465
- if xerr is not None:
466
- ax.errorbar(x_vals, y_vals, xerr=xerr, fmt='o',
467
- linestyle='', label=label,
468
- markersize=markersize)
469
- else:
470
- ax.plot(x_vals, y_vals, linestyle='', marker='o',
471
- label=label, markersize=markersize)
472
-
473
- handles, labels = ax.get_legend_handles_labels()
474
- if any(labels):
475
- ax.legend()
476
-
477
- handles, labels = ax.get_legend_handles_labels()
478
- if any(labels):
479
- ax.legend()
480
-
481
- plt.colorbar(mesh, ax=ax, label='counts (-)')
482
- return fig
483
-
484
-
485
- @add_fig_kwargs
486
- def fit_fermi_edge(self, hnuminphi_guess, background_guess=0.0,
487
- integrated_weight_guess=1.0, angle_min=-np.inf,
488
- angle_max=np.inf, ekin_min=-np.inf,
489
- ekin_max=np.inf, ax=None, **kwargs):
490
- r"""Fits the Fermi edge of the band map and plots the result.
491
- Also sets hnuminphi, the kinetic energy minus the work function in eV.
492
- The fitting includes an energy convolution with an abscissa range
493
- expanded by 5 times the energy resolution standard deviation.
494
-
495
- Parameters
496
- ----------
497
- hnuminphi_guess : float
498
- Initial guess for kinetic energy minus the work function [eV]
499
- background_guess : float
500
- Initial guess for background intensity [counts]
501
- integrated_weight_guess : float
502
- Initial guess for integrated spectral intensity [counts]
503
- angle_min : float
504
- Minimum angle of integration interval [degrees]
505
- angle_max : float
506
- Maximum angle of integration interval [degrees]
507
- ekin_min : float
508
- Minimum kinetic energy of integration interval [eV]
509
- ekin_max : float
510
- Maximum kinetic energy of integration interval [eV]
511
- ax : Matplotlib-Axes / NoneType
512
- Axis for plotting the Fermi edge on. Created if not provided by
513
- the user.
514
-
515
- Other parameters
516
- ----------------
517
- **kwargs : dict, optional
518
- Additional arguments passed on to add_fig_kwargs.
519
-
520
- Returns
521
- -------
522
- fig : Matplotlib-Figure
523
- Figure containing the Fermi edge fit
524
-
525
- """
526
- from scipy.ndimage import gaussian_filter
527
-
528
- ax, fig, plt = get_ax_fig_plt(ax=ax)
529
-
530
- min_angle_index = np.argmin(np.abs(self.angles - angle_min))
531
- max_angle_index = np.argmin(np.abs(self.angles - angle_max))
532
-
533
- min_ekin_index = np.argmin(np.abs(self.ekin - ekin_min))
534
- max_ekin_index = np.argmin(np.abs(self.ekin - ekin_max))
535
-
536
- energy_range = self.ekin[min_ekin_index:max_ekin_index]
537
-
538
- integrated_intensity = np.trapz(
539
- self.intensities[min_ekin_index:max_ekin_index,
540
- min_angle_index:max_angle_index], axis=1)
541
-
542
- fdir_initial = FermiDirac(temperature=self.temperature,
543
- hnuminphi=hnuminphi_guess,
544
- background=background_guess,
545
- integrated_weight=integrated_weight_guess,
546
- name='Initial guess')
547
-
548
- parameters = np.array(
549
- [hnuminphi_guess, background_guess, integrated_weight_guess])
550
-
551
- extra_args = (self.temperature,)
552
-
553
- popt, pcov = fit_leastsq(
554
- parameters, energy_range, integrated_intensity, fdir_initial,
555
- self.energy_resolution, None, *extra_args)
556
-
557
- # Update hnuminphi; automatically sets self.enel
558
- self.hnuminphi = popt[0]
559
- self.hnuminphi_std = np.sqrt(np.diag(pcov)[0])
560
-
561
- fdir_final = FermiDirac(temperature=self.temperature,
562
- hnuminphi=self.hnuminphi, background=popt[1],
563
- integrated_weight=popt[2],
564
- name='Fitted result')
565
-
566
- ax.set_xlabel(r'$E_{\mathrm{kin}}$ (-)')
567
- ax.set_ylabel('Counts (-)')
568
- ax.set_xlim([ekin_min, ekin_max])
569
-
570
- ax.plot(energy_range, integrated_intensity, label='Data')
571
-
572
- extend, step, numb = extend_function(energy_range,
573
- self.energy_resolution)
574
-
575
- initial_result = gaussian_filter(fdir_initial.evaluate(extend),
576
- sigma=step)[numb:-numb if numb else None]
577
-
578
- final_result = gaussian_filter(fdir_final.evaluate(extend),
579
- sigma=step)[numb:-numb if numb else None]
580
-
581
- ax.plot(energy_range, initial_result, label=fdir_initial.name)
582
- ax.plot(energy_range, final_result, label=fdir_final.name)
583
-
584
- ax.legend()
585
-
586
- return fig
587
-
588
-
589
- @add_fig_kwargs
590
- def correct_fermi_edge(self, hnuminphi_guess=None, background_guess=0.0,
591
- integrated_weight_guess=1.0, angle_min=-np.inf,
592
- angle_max=np.inf, ekin_min=-np.inf, ekin_max=np.inf,
593
- slope_guess=0, offset_guess=None,
594
- true_angle=0, ax=None, **kwargs):
595
- r"""TBD
596
- hnuminphi_guess should be estimate at true angle
597
-
598
- Parameters
599
- ----------
600
- hnuminphi_guess : float, optional
601
- Initial guess for kinetic energy minus the work function [eV].
602
-
603
- Other parameters
604
- ----------------
605
- **kwargs : dict, optional
606
- Additional arguments passed on to add_fig_kwargs.
607
-
608
- Returns
609
- -------
610
- fig : Matplotlib-Figure
611
- Figure containing the Fermi edge fit
612
-
613
- """
614
- from scipy.ndimage import map_coordinates
615
-
616
- if hnuminphi_guess is None:
617
- raise ValueError('Please provide an initial guess for ' +
618
- 'hnuminphi.')
619
-
620
- # Here some loop where it fits all the Fermi edges
621
- angle_min_index = np.abs(self.angles - angle_min).argmin()
622
- angle_max_index = np.abs(self.angles - angle_max).argmin()
623
-
624
- ekin_min_index = np.abs(self.ekin - ekin_min).argmin()
625
- ekin_max_index = np.abs(self.ekin - ekin_max).argmin()
626
-
627
- Intensities = self.intensities[ekin_min_index:ekin_max_index + 1,
628
- angle_min_index:angle_max_index + 1]
629
- angle_range = self.angles[angle_min_index:angle_max_index + 1]
630
- energy_range = self.ekin[ekin_min_index:ekin_max_index + 1]
631
-
632
- angle_shape = angle_range.shape
633
- nmps = np.zeros(angle_shape)
634
- stds = np.zeros(angle_shape)
635
-
636
- hnuminphi_left = hnuminphi_guess - (true_angle - angle_min) \
637
- * slope_guess
638
-
639
- fdir_initial = FermiDirac(temperature=self.temperature,
640
- hnuminphi=hnuminphi_left,
641
- background=background_guess,
642
- integrated_weight=integrated_weight_guess,
643
- name='Initial guess')
644
-
645
- parameters = np.array(
646
- [hnuminphi_left, background_guess, integrated_weight_guess])
647
-
648
- extra_args = (self.temperature,)
649
-
650
- for indx in range(angle_max_index - angle_min_index + 1):
651
- edge = Intensities[:, indx]
652
-
653
- parameters, pcov = fit_leastsq(
654
- parameters, energy_range, edge, fdir_initial,
655
- self.energy_resolution, None, *extra_args)
656
-
657
- nmps[indx] = parameters[0]
658
- stds[indx] = np.sqrt(np.diag(pcov)[0])
659
-
660
- # Offset at true angle if not set before
661
- if offset_guess is None:
662
- offset_guess = hnuminphi_guess - slope_guess * true_angle
663
-
664
- parameters = np.array([offset_guess, slope_guess])
665
-
666
- lin_fun = Linear(offset_guess, slope_guess, 'Linear')
667
-
668
- popt, pcov = fit_leastsq(parameters, angle_range, nmps, lin_fun, None,
669
- stds)
670
-
671
- linsp = lin_fun(angle_range, popt[0], popt[1])
672
-
673
- # Update hnuminphi; automatically sets self.enel
674
- self.hnuminphi = lin_fun(true_angle, popt[0], popt[1])
675
- self.hnuminphi_std = np.sqrt(true_angle**2 * pcov[1, 1] + pcov[0, 0]
676
- + 2 * true_angle * pcov[0, 1])
677
-
678
- Angl, Ekin = np.meshgrid(self.angles, self.ekin)
679
-
680
- ax, fig, plt = get_ax_fig_plt(ax=ax)
681
-
682
- ax.set_xlabel('Angle ($\degree$)')
683
- ax.set_ylabel('$E_{\mathrm{kin}}$ (eV)')
684
- mesh = ax.pcolormesh(Angl, Ekin, self.intensities,
685
- shading='auto', cmap=plt.get_cmap('bone').reversed(),
686
- zorder=1)
687
-
688
- ax.errorbar(angle_range, nmps, yerr=uncr * stds, zorder=1)
689
- ax.plot(angle_range, linsp, zorder=2)
690
-
691
- cbar = plt.colorbar(mesh, ax=ax, label='counts (-)')
692
-
693
- # Fermi-edge correction
694
- rows, cols = self.intensities.shape
695
- shift_values = popt[1] * self.angles / (self.ekin[0] - self.ekin[1])
696
- row_coords = np.arange(rows).reshape(-1, 1) - shift_values
697
- col_coords = np.arange(cols).reshape(1, -1).repeat(rows, axis=0)
698
- self.intensities = map_coordinates(self.intensities,
699
- [row_coords, col_coords], order=1)
700
-
701
- return fig
702
-
703
-
704
- class MDCs:
705
- r"""
706
- Container for momentum distribution curves (MDCs) and their fits.
707
-
708
- This class stores the MDC intensity maps, angular and energy grids, and
709
- the aggregated fit results produced by :meth:`fit_selection`.
710
-
711
- Parameters
712
- ----------
713
- intensities : ndarray
714
- MDC intensity data. Typically a 2D array with shape
715
- ``(n_energy, n_angle)`` or a 1D array for a single curve.
716
- angles : ndarray
717
- Angular grid corresponding to the MDCs [degrees].
718
- angle_resolution : float
719
- Angular step size or effective angular resolution [degrees].
720
- enel : ndarray or float
721
- Electron binding energies of the MDC slices [eV].
722
- Can be a scalar for a single MDC.
723
- hnuminphi : float
724
- Photon energy minus work function, used to convert ``enel`` to
725
- kinetic energy [eV].
726
-
727
- Attributes
728
- ----------
729
- intensities : ndarray
730
- MDC intensity data (same object as passed to the constructor).
731
- angles : ndarray
732
- Angular grid [degrees].
733
- angle_resolution : float
734
- Angular step size or resolution [degrees].
735
- enel : ndarray or float
736
- Electron binding energies [eV], as given at construction.
737
- ekin : ndarray or float
738
- Kinetic energies [eV], computed as ``enel + hnuminphi``.
739
- hnuminphi : float
740
- Photon energy minus work function [eV].
741
- ekin_range : ndarray
742
- Kinetic-energy values of the slices that were actually fitted.
743
- Set by :meth:`fit_selection`.
744
- individual_properties : dict
745
- Nested mapping of fitted parameters and their uncertainties for each
746
- component and each energy slice. Populated by :meth:`fit_selection`.
747
-
748
- Notes
749
- -----
750
- After calling :meth:`fit_selection`, :attr:`individual_properties` has the
751
- structure::
752
-
753
- {
754
- label: {
755
- class_name: {
756
- 'label': label,
757
- '_class': class_name,
758
- param: [values per energy slice],
759
- param_sigma: [1σ per slice or None],
760
- ...
761
- }
762
- }
763
- }
764
-
765
- where ``param`` is typically one of ``'offset'``, ``'slope'``,
766
- ``'amplitude'``, ``'peak'``, ``'broadening'``, and ``param_sigma`` stores
767
- the corresponding uncertainty for each slice.
768
-
769
- """
770
-
771
- def __init__(self, intensities, angles, angle_resolution, enel, hnuminphi):
772
- # Core input data (read-only)
773
- self._intensities = intensities
774
- self._angles = angles
775
- self._angle_resolution = angle_resolution
776
- self._enel = enel
777
- self._hnuminphi = hnuminphi
778
-
779
- # Derived attributes (populated by fit_selection)
780
- self._ekin_range = None
781
- self._individual_properties = None # combined values + sigmas
782
-
783
- # -------------------- Immutable physics inputs --------------------
784
-
785
- @property
786
- def angles(self):
787
- """Angular axis for the MDCs."""
788
- return self._angles
789
-
790
- @property
791
- def angle_resolution(self):
792
- """Angular step size (float)."""
793
- return self._angle_resolution
794
-
795
- @property
796
- def enel(self):
797
- """Photoelectron binding energies (array-like). Read-only."""
798
- return self._enel
799
-
800
- @enel.setter
801
- def enel(self, _):
802
- raise AttributeError("`enel` is read-only; set it via the constructor.")
803
-
804
- @property
805
- def hnuminphi(self):
806
- """Work-function/photon-energy offset. Read-only."""
807
- return self._hnuminphi
808
-
809
- @hnuminphi.setter
810
- def hnuminphi(self, _):
811
- raise AttributeError("`hnuminphi` is read-only; set it via the constructor.")
812
-
813
- @property
814
- def ekin(self):
815
- """Kinetic energy array: enel + hnuminphi (computed on the fly)."""
816
- return self._enel + self._hnuminphi
817
-
818
- @ekin.setter
819
- def ekin(self, _):
820
- raise AttributeError("`ekin` is derived and read-only.")
821
-
822
- # -------------------- Data arrays --------------------
823
-
824
- @property
825
- def intensities(self):
826
- """2D or 3D intensity map (energy × angle)."""
827
- return self._intensities
828
-
829
- @intensities.setter
830
- def intensities(self, x):
831
- self._intensities = x
832
-
833
- # -------------------- Results populated by fit_selection --------------------
834
-
835
- @property
836
- def ekin_range(self):
837
- """Kinetic-energy slices that were fitted."""
838
- if self._ekin_range is None:
839
- raise AttributeError("`ekin_range` not yet set. Run `.fit_selection()` first.")
840
- return self._ekin_range
841
-
842
- @property
843
- def individual_properties(self):
844
- """
845
- Aggregated fitted parameter values and uncertainties per component.
846
-
847
- Returns
848
- -------
849
- dict
850
- Nested mapping::
851
-
852
- {
853
- label: {
854
- class_name: {
855
- 'label': label,
856
- '_class': class_name,
857
- <param>: [values per slice],
858
- <param>_sigma: [1σ per slice or None],
859
- ...
860
- }
861
- }
862
- }
863
- """
864
- if self._individual_properties is None:
865
- raise AttributeError(
866
- "`individual_properties` not yet set. Run `.fit_selection()` first."
867
- )
868
- return self._individual_properties
869
-
870
- def energy_check(self, energy_value):
871
- r"""
872
- """
873
- if np.isscalar(self.ekin):
874
- if energy_value is not None:
875
- raise ValueError("This dataset contains only one " \
876
- "momentum-distribution curve; do not provide energy_value.")
877
- else:
878
- kinergy = self.ekin
879
- counts = self.intensities
880
- else:
881
- if energy_value is None:
882
- raise ValueError("This dataset contains multiple " \
883
- "momentum-distribution curves. Please provide an energy_value "
884
- "for which to plot the MDCs.")
885
- else:
886
- energy_index = np.abs(self.enel - energy_value).argmin()
887
- kinergy = self.ekin[energy_index]
888
- counts = self.intensities[energy_index, :]
889
-
890
- if not (self.enel.min() <= energy_value <= self.enel.max()):
891
- raise ValueError(
892
- f"Selected energy_value={energy_value:.3f} "
893
- f"is outside the available energy range "
894
- f"({self.enel.min():.3f} – {self.enel.max():.3f}) "
895
- "of the MDC collection."
896
- )
897
-
898
- return counts, kinergy
899
-
900
-
901
- def plot(self, energy_value=None, energy_range=None, ax=None, **kwargs):
902
- """
903
- Interactive or static plot with optional slider and full wrapper
904
- support. Behavior consistent with Jupyter and CLI based on show /
905
- fig_close.
906
- """
907
- import matplotlib.pyplot as plt
908
- from matplotlib.widgets import Slider
909
- import string
910
- import sys
911
- import warnings
912
-
913
- # Wrapper kwargs
914
- title = kwargs.pop("title", None)
915
- savefig = kwargs.pop("savefig", None)
916
- show = kwargs.pop("show", True)
917
- fig_close = kwargs.pop("fig_close", False)
918
- tight_layout = kwargs.pop("tight_layout", False)
919
- ax_grid = kwargs.pop("ax_grid", None)
920
- ax_annotate = kwargs.pop("ax_annotate", False)
921
- size_kwargs = kwargs.pop("size_kwargs", None)
922
-
923
- if energy_value is not None and energy_range is not None:
924
- raise ValueError(
925
- "Provide at most energy_value or energy_range, not both.")
926
-
927
- ax, fig, plt = get_ax_fig_plt(ax=ax)
928
-
929
- angles = self.angles
930
- energies = self.enel
931
-
932
- if np.isscalar(energies):
933
- if energy_value is not None or energy_range is not None:
934
- raise ValueError(
935
- "This dataset contains only one momentum-distribution "
936
- "curve; do not provide energy_value or energy_range."
937
- )
938
-
939
- intensities = self.intensities
940
- ax.scatter(angles, intensities, label="Data")
941
- ax.set_title(f"Energy slice: {energies * kilo:.3f} meV")
942
-
943
- # --- y-only autoscale, preserve x ---
944
- x0, x1 = ax.get_xlim() # keep current x-range
945
- ax.relim(visible_only=True) # recompute data limits
946
- ax.autoscale_view(scalex=False, scaley=True)
947
- ax.set_xlim(x0, x1) # restore x (belt-and-suspenders)
948
-
949
- else:
950
- if (energy_value is not None) and (energy_range is not None):
951
- raise ValueError("Provide either energy_value or energy_range, not both.")
952
-
953
- emin, emax = energies.min(), energies.max()
954
-
955
- # ---- Single-slice path (no slider) ----
956
- if energy_value is not None:
957
- if energy_value < emin or energy_value > emax:
958
- raise ValueError(
959
- f"Requested energy_value {energy_value:.3f} eV is "
960
- f"outside the available energy range "
961
- f"[{emin:.3f}, {emax:.3f}] eV."
962
- )
963
- idx = int(np.abs(energies - energy_value).argmin())
964
- intensities = self.intensities[idx]
965
- ax.scatter(angles, intensities, label="Data")
966
- ax.set_title(f"Energy slice: {energies[idx] * kilo:.3f} meV")
967
-
968
- # --- y-only autoscale, preserve x ---
969
- x0, x1 = ax.get_xlim() # keep current x-range
970
- ax.relim(visible_only=True) # recompute data limits
971
- ax.autoscale_view(scalex=False, scaley=True)
972
- ax.set_xlim(x0, x1) # restore x (belt-and-suspenders)
973
-
974
- # ---- Multi-slice path (slider) ----
975
- else:
976
- if energy_range is not None:
977
- e_min, e_max = energy_range
978
- mask = (energies >= e_min) & (energies <= e_max)
979
- else:
980
- mask = np.ones_like(energies, dtype=bool)
981
-
982
- indices = np.where(mask)[0]
983
- if len(indices) == 0:
984
- raise ValueError("No energies found in the specified selection.")
985
-
986
- intensities = self.intensities[indices]
987
-
988
- fig.subplots_adjust(bottom=0.25)
989
- idx = 0
990
- scatter = ax.scatter(angles, intensities[idx], label="Data")
991
- ax.set_title(f"Energy slice: "
992
- f"{energies[indices[idx]] * kilo:.3f} meV")
993
-
994
- # Suppress single-point slider warning (when len(indices) == 1)
995
- warnings.filterwarnings(
996
- "ignore",
997
- message="Attempting to set identical left == right",
998
- category=UserWarning
999
- )
1000
-
1001
- slider_ax = fig.add_axes([0.2, 0.08, 0.6, 0.04])
1002
- slider = Slider(
1003
- slider_ax, "Index", 0, len(indices) - 1,
1004
- valinit=idx, valstep=1
1005
- )
1006
-
1007
- def update(val):
1008
- i = int(slider.val)
1009
- yi = intensities[i]
1010
-
1011
- scatter.set_offsets(np.c_[angles, yi])
1012
-
1013
- x0, x1 = ax.get_xlim()
1014
-
1015
- yv = np.asarray(yi, dtype=float).ravel()
1016
- mask = np.isfinite(yv)
1017
- if mask.any():
1018
- y_min = float(yv[mask].min())
1019
- y_max = float(yv[mask].max())
1020
- span = y_max - y_min
1021
- frac = plt.rcParams['axes.ymargin']
1022
-
1023
- if span <= 0 or not np.isfinite(span):
1024
- scale = max(abs(y_max), 1.0)
1025
- pad = frac * scale
1026
- else:
1027
- pad = frac * span
1028
-
1029
- ax.set_ylim(y_min - pad, y_max + pad)
1030
-
1031
- # Keep x unchanged
1032
- ax.set_xlim(x0, x1)
1033
-
1034
- # Update title and redraw
1035
- ax.set_title(f"Energy slice: "
1036
- f"{energies[indices[i]] * kilo:.3f} meV")
1037
- fig.canvas.draw_idle()
1038
-
1039
- slider.on_changed(update)
1040
- self._slider = slider
1041
- self._line = scatter
1042
-
1043
- ax.set_xlabel("Angle (°)")
1044
- ax.set_ylabel("Counts (-)")
1045
- ax.legend()
1046
- self._fig = fig
1047
-
1048
- if size_kwargs:
1049
- fig.set_size_inches(size_kwargs.pop("w"),
1050
- size_kwargs.pop("h"), **size_kwargs)
1051
- if title:
1052
- fig.suptitle(title)
1053
- if tight_layout:
1054
- fig.tight_layout()
1055
- if savefig:
1056
- fig.savefig(savefig)
1057
- if ax_grid is not None:
1058
- for axis in fig.axes:
1059
- axis.grid(bool(ax_grid))
1060
- if ax_annotate:
1061
- tags = string.ascii_lowercase
1062
- for i, axis in enumerate(fig.axes):
1063
- axis.annotate(f"({tags[i]})", xy=(0.05, 0.95),
1064
- xycoords="axes fraction")
1065
-
1066
- is_interactive = hasattr(sys, 'ps1') or 'ipykernel' in sys.modules
1067
- is_cli = not is_interactive
1068
-
1069
- if show:
1070
- if is_cli:
1071
- plt.show()
1072
- if fig_close:
1073
- plt.close(fig)
1074
-
1075
- if not show and (fig_close or is_cli):
1076
- return None
1077
- return fig
1078
-
1079
-
1080
- @add_fig_kwargs
1081
- def visualize_guess(self, distributions, energy_value=None,
1082
- matrix_element=None, matrix_args=None,
1083
- ax=None, **kwargs):
1084
- r"""
1085
- """
1086
-
1087
- counts, kinergy = self.energy_check(energy_value)
1088
-
1089
- ax, fig, plt = get_ax_fig_plt(ax=ax)
1090
-
1091
- ax.set_xlabel('Angle ($\\degree$)')
1092
- ax.set_ylabel('Counts (-)')
1093
- ax.set_title(f"Energy slice: "
1094
- f"{(kinergy - self.hnuminphi) * kilo:.3f} meV")
1095
- ax.scatter(self.angles, counts, label='Data')
1096
-
1097
- final_result = self._merge_and_plot(ax=ax,
1098
- distributions=distributions, kinetic_energy=kinergy,
1099
- matrix_element=matrix_element,
1100
- matrix_args=dict(matrix_args) if matrix_args else None,
1101
- plot_individual=True,
1102
- )
1103
-
1104
- residual = counts - final_result
1105
- ax.scatter(self.angles, residual, label='Residual')
1106
- ax.legend()
1107
-
1108
- return fig
1109
-
1110
-
1111
- def fit_selection(self, distributions, energy_value=None, energy_range=None,
1112
- matrix_element=None, matrix_args=None, ax=None, **kwargs):
1113
- r"""
1114
- """
1115
- import matplotlib.pyplot as plt
1116
- from matplotlib.widgets import Slider
1117
- from copy import deepcopy
1118
- import string
1119
- import sys
1120
- import warnings
1121
- from lmfit import Minimizer
1122
- from scipy.ndimage import gaussian_filter
1123
- from .functions import construct_parameters, build_distributions, \
1124
- residual, resolve_param_name
1125
-
1126
- # Wrapper kwargs
1127
- title = kwargs.pop("title", None)
1128
- savefig = kwargs.pop("savefig", None)
1129
- show = kwargs.pop("show", True)
1130
- fig_close = kwargs.pop("fig_close", False)
1131
- tight_layout = kwargs.pop("tight_layout", False)
1132
- ax_grid = kwargs.pop("ax_grid", None)
1133
- ax_annotate = kwargs.pop("ax_annotate", False)
1134
- size_kwargs = kwargs.pop("size_kwargs", None)
1135
-
1136
- ax, fig, plt = get_ax_fig_plt(ax=ax)
1137
-
1138
- energies = self.enel
1139
- new_distributions = deepcopy(distributions)
1140
-
1141
- if energy_value is not None and energy_range is not None:
1142
- raise ValueError(
1143
- "Provide at most energy_value or energy_range, not both.")
1144
-
1145
- if np.isscalar(energies):
1146
- if energy_value is not None or energy_range is not None:
1147
- raise ValueError(
1148
- "This dataset contains only one momentum-distribution "
1149
- "curve; do not provide energy_value or energy_range."
1150
- )
1151
- kinergies = np.atleast_1d(self.ekin)
1152
- intensities = np.atleast_2d(self.intensities)
1153
-
1154
- else:
1155
- if energy_value is not None:
1156
- if (energy_value < energies.min() or energy_value > energies.max()):
1157
- raise ValueError( f"Requested energy_value {energy_value:.3f} eV is "
1158
- f"outside the available energy range "
1159
- f"[{energies.min():.3f}, {energies.max():.3f}] eV." )
1160
- idx = np.abs(energies - energy_value).argmin()
1161
- indices = np.atleast_1d(idx)
1162
- kinergies = self.ekin[indices]
1163
- intensities = self.intensities[indices, :]
1164
-
1165
- elif energy_range is not None:
1166
- e_min, e_max = energy_range
1167
- indices = np.where((energies >= e_min) & (energies <= e_max))[0]
1168
- if len(indices) == 0:
1169
- raise ValueError("No energies found in the specified energy_range.")
1170
- kinergies = self.ekin[indices]
1171
- intensities = self.intensities[indices, :]
1172
-
1173
- else: # Without specifying a range, all MDCs are plotted
1174
- kinergies = self.ekin
1175
- intensities = self.intensities
1176
-
1177
- # Final shape guard
1178
- kinergies = np.atleast_1d(kinergies)
1179
- intensities = np.atleast_2d(intensities)
1180
-
1181
- all_final_results = []
1182
- all_residuals = []
1183
- all_individual_results = [] # List of (n_individuals, n_angles)
1184
-
1185
- aggregated_properties = {}
1186
-
1187
- # map class_name -> parameter names to extract
1188
- param_spec = {
1189
- 'Constant': ('offset',),
1190
- 'Linear': ('offset', 'slope'),
1191
- 'SpectralLinear': ('amplitude', 'peak', 'broadening'),
1192
- 'SpectralQuadratic': ('amplitude', 'peak', 'broadening'),
1193
- }
1194
-
1195
- order = np.argsort(kinergies)[::-1]
1196
- for idx in order:
1197
- kinergy = kinergies[idx]
1198
- intensity = intensities[idx]
1199
- if matrix_element is not None:
1200
- parameters, element_names = construct_parameters(
1201
- new_distributions, matrix_args)
1202
- new_distributions = build_distributions(new_distributions, parameters)
1203
- mini = Minimizer(
1204
- residual, parameters,
1205
- fcn_args=(self.angles, intensity, self.angle_resolution,
1206
- new_distributions, kinergy, self.hnuminphi,
1207
- matrix_element, element_names)
1208
- )
1209
- else:
1210
- parameters = construct_parameters(new_distributions)
1211
- new_distributions = build_distributions(new_distributions, parameters)
1212
- mini = Minimizer(
1213
- residual, parameters,
1214
- fcn_args=(self.angles, intensity, self.angle_resolution,
1215
- new_distributions, kinergy, self.hnuminphi)
1216
- )
1217
-
1218
- outcome = mini.minimize('least_squares')
1219
-
1220
- pcov = outcome.covar
1221
-
1222
- var_names = getattr(outcome, 'var_names', None)
1223
- if not var_names:
1224
- var_names = [n for n, p in outcome.params.items() if p.vary]
1225
- var_idx = {n: i for i, n in enumerate(var_names)}
1226
-
1227
- param_sigma_full = {}
1228
- for name, par in outcome.params.items():
1229
- sigma = None
1230
- if pcov is not None and name in var_idx:
1231
- d = pcov[var_idx[name], var_idx[name]]
1232
- if np.isfinite(d) and d >= 0:
1233
- sigma = float(np.sqrt(d))
1234
- if sigma is None:
1235
- s = getattr(par, 'stderr', None)
1236
- sigma = float(s) if s is not None else None
1237
- param_sigma_full[name] = sigma
1238
-
1239
- # Rebuild the *fitted* distributions from optimized params
1240
- fitted_distributions = build_distributions(new_distributions, outcome.params)
1241
-
1242
- # If using a matrix element, extract slice-specific args from the fit
1243
- if matrix_element is not None:
1244
- new_matrix_args = {key: outcome.params[key].value for key in matrix_args}
1245
- else:
1246
- new_matrix_args = None
1247
-
1248
- # individual curves (smoothed, cropped) and final sum (no plotting here)
1249
- extend, step, numb = extend_function(self.angles, self.angle_resolution)
1250
-
1251
- total_result_ext = np.zeros_like(extend)
1252
- indiv_rows = [] # (n_individuals, n_angles)
1253
- individual_labels = []
1254
-
1255
- for dist in fitted_distributions:
1256
- # evaluate each component on the extended grid
1257
- if getattr(dist, 'class_name', None) == 'SpectralQuadratic':
1258
- if (getattr(dist, 'center_angle', None) is not None) and (
1259
- kinergy is None or self.hnuminphi is None
1260
- ):
1261
- raise ValueError(
1262
- 'Spectral quadratic function is defined in terms '
1263
- 'of a center angle. Please provide a kinetic energy '
1264
- 'and hnuminphi.'
1265
- )
1266
- extended_result = dist.evaluate(extend, kinergy, self.hnuminphi)
1267
- else:
1268
- extended_result = dist.evaluate(extend)
1269
-
1270
- if matrix_element is not None and hasattr(dist, 'index'):
1271
- args = new_matrix_args or {}
1272
- extended_result *= matrix_element(extend, **args)
1273
-
1274
- total_result_ext += extended_result
1275
-
1276
- # smoothed & cropped individual
1277
- individual_curve = gaussian_filter(extended_result, sigma=step)[
1278
- numb:-numb if numb else None
1279
- ]
1280
- indiv_rows.append(np.asarray(individual_curve))
1281
-
1282
- # label
1283
- label = getattr(dist, 'label', str(dist))
1284
- individual_labels.append(label)
1285
-
1286
- # ---- collect parameters for this distribution
1287
- # (Aggregated over slices)
1288
- cls = getattr(dist, 'class_name', None)
1289
- wanted = param_spec.get(cls, ())
1290
-
1291
- # ensure dicts exist
1292
- label_bucket = aggregated_properties.setdefault(label, {})
1293
- class_bucket = label_bucket.setdefault(
1294
- cls, {'label': label, '_class': cls}
1295
- )
1296
-
1297
- # store center_wavevector (scalar) for SpectralQuadratic
1298
- if (
1299
- cls == 'SpectralQuadratic'
1300
- and hasattr(dist, 'center_wavevector')
1301
- ):
1302
- class_bucket.setdefault(
1303
- 'center_wavevector', dist.center_wavevector
1304
- )
1305
-
1306
- # ensure keys for both values and sigmas
1307
- for pname in wanted:
1308
- class_bucket.setdefault(pname, [])
1309
- class_bucket.setdefault(f"{pname}_sigma", [])
1310
-
1311
- # append values and sigmas in the order of slices
1312
- for pname in wanted:
1313
- param_key = resolve_param_name(outcome.params, label, pname)
1314
-
1315
- if param_key is not None and param_key in outcome.params:
1316
- class_bucket[pname].append(outcome.params[param_key].value)
1317
- class_bucket[f"{pname}_sigma"].append(param_sigma_full.get(param_key, None))
1318
- else:
1319
- # Not fitted in this slice → keep the value if present on the dist, sigma=None
1320
- class_bucket[pname].append(getattr(dist, pname, None))
1321
- class_bucket[f"{pname}_sigma"].append(None)
1322
-
1323
- # final (sum) curve, smoothed & cropped
1324
- final_result_i = gaussian_filter(total_result_ext, sigma=step)[
1325
- numb:-numb if numb else None]
1326
- final_result_i = np.asarray(final_result_i)
1327
-
1328
- # Residual for this slice
1329
- residual_i = np.asarray(intensity) - final_result_i
1330
-
1331
- # Store per-slice results
1332
- all_final_results.append(final_result_i)
1333
- all_residuals.append(residual_i)
1334
- all_individual_results.append(np.vstack(indiv_rows))
1335
-
1336
- # --- after the reversed-order loop, restore original (ascending) order ---
1337
- inverse_order = np.argsort(np.argsort(kinergies)[::-1])
1338
-
1339
- # Reorder per-slice arrays/lists computed in the loop
1340
- all_final_results[:] = [all_final_results[i] for i in inverse_order]
1341
- all_residuals[:] = [all_residuals[i] for i in inverse_order]
1342
- all_individual_results[:] = [all_individual_results[i] for i in inverse_order]
1343
-
1344
- # Reorder all per-slice lists in aggregated_properties
1345
- for label_dict in aggregated_properties.values():
1346
- for cls_dict in label_dict.values():
1347
- for key, val in cls_dict.items():
1348
- if isinstance(val, list) and len(val) == len(kinergies):
1349
- cls_dict[key] = [val[i] for i in inverse_order]
1350
-
1351
- self._ekin_range = kinergies
1352
- self._individual_properties = aggregated_properties
1353
-
1354
- if np.isscalar(energies):
1355
- # One slice only: plot MDC, Fit, Residual, and Individuals
1356
- ydata = np.asarray(intensities).squeeze()
1357
- yfit = np.asarray(all_final_results[0]).squeeze()
1358
- yres = np.asarray(all_residuals[0]).squeeze()
1359
- yind = np.asarray(all_individual_results[0])
1360
-
1361
- ax.scatter(self.angles, ydata, label="Data")
1362
- # plot individuals with their labels
1363
- for j, lab in enumerate(individual_labels or []):
1364
- ax.plot(self.angles, yind[j], label=str(lab))
1365
- ax.plot(self.angles, yfit, label="Fit")
1366
- ax.scatter(self.angles, yres, label="Residual")
1367
-
1368
- ax.set_title(f"Energy slice: {energies * kilo:.3f} meV")
1369
- ax.relim() # recompute data limits from all artists
1370
- ax.autoscale_view() # apply autoscaling + axes.ymargin padding
1371
-
1372
- else:
1373
- if energy_value is not None:
1374
- _idx = int(np.abs(energies - energy_value).argmin())
1375
- energies_sel = np.atleast_1d(energies[_idx])
1376
- elif energy_range is not None:
1377
- e_min, e_max = energy_range
1378
- energies_sel = energies[(energies >= e_min)
1379
- & (energies <= e_max)]
1380
- else:
1381
- energies_sel = energies
1382
-
1383
- # Number of slices must match
1384
- n_slices = len(all_final_results)
1385
- assert intensities.shape[0] == n_slices == len(all_residuals) \
1386
- == len(all_individual_results), (f"Mismatch: data \
1387
- {intensities.shape[0]}, fits {len(all_final_results)}, "
1388
- f"residuals {len(all_residuals)}, \
1389
- individuals {len(all_individual_results)}."
1390
- )
1391
- n_individuals = all_individual_results[0].shape[0] \
1392
- if n_slices else 0
1393
-
1394
- fig.subplots_adjust(bottom=0.25)
1395
- idx = 0
1396
-
1397
- # Initial draw (MDC + Individuals + Fit + Residual) at slice 0
1398
- scatter = ax.scatter(self.angles, intensities[idx], label="Data")
1399
-
1400
- individual_lines = []
1401
- if n_individuals:
1402
- for j in range(n_individuals):
1403
- if individual_labels and j < len(individual_labels):
1404
- label = str(individual_labels[j])
1405
- else:
1406
- label = f"Comp {j}"
1407
-
1408
- yvals = all_individual_results[idx][j]
1409
- line, = ax.plot(self.angles, yvals, label=label)
1410
- individual_lines.append(line)
1411
-
1412
- result_line, = ax.plot(self.angles, all_final_results[idx],
1413
- label="Fit")
1414
- resid_scatter = ax.scatter(self.angles, all_residuals[idx],
1415
- label="Residual")
1416
-
1417
- # Title + limits (use only the currently shown slice)
1418
- ax.set_title(f"Energy slice: {energies_sel[idx] * kilo:.3f} meV")
1419
- ax.relim() # recompute data limits from all artists
1420
- ax.autoscale_view() # apply autoscaling + axes.ymargin padding
1421
-
1422
- # Suppress warning when a single MDC is plotted
1423
- warnings.filterwarnings(
1424
- "ignore",
1425
- message="Attempting to set identical left == right",
1426
- category=UserWarning
1427
- )
1428
-
1429
- # Slider over slice index (0..n_slices-1)
1430
- slider_ax = fig.add_axes([0.2, 0.08, 0.6, 0.04])
1431
- slider = Slider(
1432
- slider_ax, "Index", 0, n_slices - 1,
1433
- valinit=idx, valstep=1
1434
- )
1435
-
1436
- def update(val):
1437
- i = int(slider.val)
1438
- # Update MDC points
1439
- scatter.set_offsets(np.c_[self.angles, intensities[i]])
1440
-
1441
- # Update individuals
1442
- if n_individuals:
1443
- Yi = all_individual_results[i] # (n_individuals, n_angles)
1444
- for j, ln in enumerate(individual_lines):
1445
- ln.set_ydata(Yi[j])
1446
-
1447
- # Update fit and residual
1448
- result_line.set_ydata(all_final_results[i])
1449
- resid_scatter.set_offsets(np.c_[self.angles, all_residuals[i]])
1450
-
1451
- ax.relim()
1452
- ax.autoscale_view()
1453
-
1454
- # Update title and redraw
1455
- ax.set_title(f"Energy slice: "
1456
- f"{energies_sel[i] * kilo:.3f} meV")
1457
- fig.canvas.draw_idle()
1458
-
1459
- slider.on_changed(update)
1460
- self._slider = slider
1461
- self._line = scatter
1462
- self._individual_lines = individual_lines
1463
- self._result_line = result_line
1464
- self._resid_scatter = resid_scatter
1465
-
1466
- ax.set_xlabel("Angle (°)")
1467
- ax.set_ylabel("Counts (-)")
1468
- ax.legend()
1469
- self._fig = fig
1470
-
1471
- if size_kwargs:
1472
- fig.set_size_inches(size_kwargs.pop("w"),
1473
- size_kwargs.pop("h"), **size_kwargs)
1474
- if title:
1475
- fig.suptitle(title)
1476
- if tight_layout:
1477
- fig.tight_layout()
1478
- if savefig:
1479
- fig.savefig(savefig)
1480
- if ax_grid is not None:
1481
- for axis in fig.axes:
1482
- axis.grid(bool(ax_grid))
1483
- if ax_annotate:
1484
- tags = string.ascii_lowercase
1485
- for i, axis in enumerate(fig.axes):
1486
- axis.annotate(f"({tags[i]})", xy=(0.05, 0.95),
1487
- xycoords="axes fraction")
1488
-
1489
- is_interactive = hasattr(sys, 'ps1') or 'ipykernel' in sys.modules
1490
- is_cli = not is_interactive
1491
-
1492
- if show:
1493
- if is_cli:
1494
- plt.show()
1495
- if fig_close:
1496
- plt.close(fig)
1497
-
1498
- if not show and (fig_close or is_cli):
1499
- return None
1500
- return fig
1501
-
1502
-
1503
- @add_fig_kwargs
1504
- def fit(self, distributions, energy_value=None, matrix_element=None,
1505
- matrix_args=None, ax=None, **kwargs):
1506
- r"""
1507
- """
1508
- from copy import deepcopy
1509
- from lmfit import Minimizer
1510
- from .functions import construct_parameters, build_distributions, \
1511
- residual
1512
-
1513
- counts, kinergy = self.energy_check(energy_value)
1514
-
1515
- ax, fig, plt = get_ax_fig_plt(ax=ax)
1516
-
1517
- ax.set_xlabel('Angle ($\\degree$)')
1518
- ax.set_ylabel('Counts (-)')
1519
- ax.set_title(f"Energy slice: "
1520
- f"{(kinergy - self.hnuminphi) * kilo:.3f} meV")
1521
-
1522
- ax.scatter(self.angles, counts, label='Data')
1523
-
1524
- new_distributions = deepcopy(distributions)
1525
-
1526
- if matrix_element is not None:
1527
- parameters, element_names = construct_parameters(distributions,
1528
- matrix_args)
1529
- new_distributions = build_distributions(new_distributions, \
1530
- parameters)
1531
- mini = Minimizer(
1532
- residual, parameters,
1533
- fcn_args=(self.angles, counts, self.angle_resolution,
1534
- new_distributions, kinergy, self.hnuminphi,
1535
- matrix_element, element_names))
1536
- else:
1537
- parameters = construct_parameters(distributions)
1538
- new_distributions = build_distributions(new_distributions,
1539
- parameters)
1540
- mini = Minimizer(residual, parameters,
1541
- fcn_args=(self.angles, counts, self.angle_resolution,
1542
- new_distributions, kinergy, self.hnuminphi))
1543
-
1544
- outcome = mini.minimize('least_squares')
1545
- pcov = outcome.covar
1546
-
1547
- # If matrix params were fitted, pass the fitted values to plotting
1548
- if matrix_element is not None:
1549
- new_matrix_args = {key: outcome.params[key].value for key in
1550
- matrix_args}
1551
- else:
1552
- new_matrix_args = None
1553
-
1554
- final_result = self._merge_and_plot(ax=ax,
1555
- distributions=new_distributions, kinetic_energy=kinergy,
1556
- matrix_element=matrix_element, matrix_args=new_matrix_args,
1557
- plot_individual=True)
1558
-
1559
- residual_vals = counts - final_result
1560
- ax.scatter(self.angles, residual_vals, label='Residual')
1561
- ax.legend()
1562
- if matrix_element is not None:
1563
- return fig, new_distributions, pcov, new_matrix_args
1564
- else:
1565
- return fig, new_distributions, pcov
1566
-
1567
-
1568
- def _merge_and_plot(self, ax, distributions, kinetic_energy,
1569
- matrix_element=None, matrix_args=None,
1570
- plot_individual=True):
1571
- r"""
1572
- Evaluate distributions on the extended grid, apply optional matrix
1573
- element, smooth, plot individuals and the summed curve.
1574
-
1575
- Returns
1576
- -------
1577
- final_result : np.ndarray
1578
- Smoothed, cropped total distribution aligned with self.angles.
1579
- """
1580
- from scipy.ndimage import gaussian_filter
1581
-
1582
- # Build extended grid
1583
- extend, step, numb = extend_function(self.angles, self.angle_resolution)
1584
- total_result = np.zeros_like(extend)
1585
-
1586
- for dist in distributions:
1587
- # Special handling for SpectralQuadratic
1588
- if getattr(dist, 'class_name', None) == 'SpectralQuadratic':
1589
- if (getattr(dist, 'center_angle', None) is not None) and (
1590
- kinetic_energy is None or self.hnuminphi is None
1591
- ):
1592
- raise ValueError(
1593
- 'Spectral quadratic function is defined in terms '
1594
- 'of a center angle. Please provide a kinetic energy '
1595
- 'and hnuminphi.'
1596
- )
1597
- extended_result = dist.evaluate(extend, kinetic_energy, \
1598
- self.hnuminphi)
1599
- else:
1600
- extended_result = dist.evaluate(extend)
1601
-
1602
- # Optional matrix element (only for components that advertise an index)
1603
- if matrix_element is not None and hasattr(dist, 'index'):
1604
- args = matrix_args or {}
1605
- extended_result *= matrix_element(extend, **args)
1606
-
1607
- total_result += extended_result
1608
-
1609
- if plot_individual and ax:
1610
- individual = gaussian_filter(extended_result, sigma=step)\
1611
- [numb:-numb if numb else None]
1612
- ax.plot(self.angles, individual, label=getattr(dist, \
1613
- 'label', str(dist)))
1614
-
1615
- # Smoothed, cropped total curve aligned to self.angles
1616
- final_result = gaussian_filter(total_result, sigma=step)[numb:-numb \
1617
- if numb else None]
1618
- if ax:
1619
- ax.plot(self.angles, final_result, label='Distribution sum')
1620
-
1621
- return final_result
1622
-
1623
-
1624
- def expose_parameters(self, select_label, fermi_wavevector=None,
1625
- fermi_velocity=None, bare_mass=None, side=None):
1626
- r"""
1627
- Select and return fitted parameters for a given component label, plus a
1628
- flat export dictionary containing values **and** 1σ uncertainties.
1629
-
1630
- Parameters
1631
- ----------
1632
- select_label : str
1633
- Label to look for among the fitted distributions.
1634
- fermi_wavevector : float, optional
1635
- Optional Fermi wave vector to include.
1636
- fermi_velocity : float, optional
1637
- Optional Fermi velocity to include.
1638
- bare_mass : float, optional
1639
- Optional bare mass to include (used for SpectralQuadratic
1640
- dispersions).
1641
- side : {'left','right'}, optional
1642
- Optional side selector for SpectralQuadratic dispersions.
1643
-
1644
- Returns
1645
- -------
1646
- ekin_range : np.ndarray
1647
- Kinetic-energy grid corresponding to the selected label.
1648
- hnuminphi : float
1649
- Photoelectron work-function offset.
1650
- label : str
1651
- Label of the selected distribution.
1652
- selected_properties : dict or list of dict
1653
- Nested dictionary (or list thereof) containing <param> and
1654
- <param>_sigma arrays. For SpectralQuadratic components, a
1655
- scalar `center_wavevector` is also present.
1656
- exported_parameters : dict
1657
- Flat dictionary of parameters and their uncertainties, plus
1658
- optional Fermi quantities and `side`. For SpectralQuadratic
1659
- components, `center_wavevector` is included and taken directly
1660
- from the fitted distribution.
1661
- """
1662
-
1663
- if self._ekin_range is None:
1664
- raise AttributeError(
1665
- "ekin_range not yet set. Run `.fit_selection()` first."
1666
- )
1667
-
1668
- store = getattr(self, "_individual_properties", None)
1669
- if not store or select_label not in store:
1670
- all_labels = (sorted(store.keys())
1671
- if isinstance(store, dict) else [])
1672
- raise ValueError(
1673
- f"Label '{select_label}' not found in available labels: "
1674
- f"{all_labels}"
1675
- )
1676
-
1677
- # Convert lists → numpy arrays within the selected label’s classes.
1678
- # Keep scalar center_wavevector as a scalar.
1679
- per_class_dicts = []
1680
- for cls, bucket in store[select_label].items():
1681
- dct = {}
1682
- for k, v in bucket.items():
1683
- if k in ("label", "_class"):
1684
- dct[k] = v
1685
- elif k == "center_wavevector":
1686
- # keep scalar as-is, do not wrap in np.asarray
1687
- dct[k] = v
1688
- else:
1689
- dct[k] = np.asarray(v)
1690
- per_class_dicts.append(dct)
1691
-
1692
- selected_properties = (
1693
- per_class_dicts[0] if len(per_class_dicts) == 1 else per_class_dicts
1694
- )
1695
-
1696
- # Flat export dict: simple keys, includes optional extras
1697
- exported_parameters = {
1698
- "fermi_wavevector": fermi_wavevector,
1699
- "fermi_velocity": fermi_velocity,
1700
- "bare_mass": bare_mass,
1701
- "side": side,
1702
- }
1703
-
1704
- # Collect parameters without prefixing by class. This will also include
1705
- # center_wavevector from the fitted SpectralQuadratic class, and since
1706
- # there is no function argument with that name, it cannot be overridden.
1707
- if isinstance(selected_properties, dict):
1708
- for key, val in selected_properties.items():
1709
- if key not in ("label", "_class"):
1710
- exported_parameters[key] = val
1711
- else:
1712
- # If multiple classes, merge sequentially
1713
- # (last overwrites same-name keys).
1714
- for cls_bucket in selected_properties:
1715
- for key, val in cls_bucket.items():
1716
- if key not in ("label", "_class"):
1717
- exported_parameters[key] = val
1718
-
1719
- return (self._ekin_range, self.hnuminphi, select_label,
1720
- selected_properties, exported_parameters)
1721
-
1722
-
1723
- class SelfEnergy:
1724
- r"""Self-energy (ekin-leading; hnuminphi/ekin are read-only)."""
1725
-
1726
- def __init__(self, ekin_range, hnuminphi, label, properties, parameters):
1727
- # core read-only state
1728
- self._ekin_range = ekin_range
1729
- self._hnuminphi = hnuminphi
1730
- self._label = label
1731
-
1732
- # accept either a dict or a single-element list of dicts
1733
- if isinstance(properties, list):
1734
- if len(properties) == 1:
1735
- properties = properties[0]
1736
- else:
1737
- raise ValueError("`properties` must be a dict or a single " \
1738
- "dict in a list.")
1739
-
1740
- # single source of truth for all params (+ their *_sigma)
1741
- self._properties = dict(properties or {})
1742
- self._class = self._properties.get("_class", None)
1743
-
1744
- # ---- enforce supported classes at construction
1745
- if self._class not in ("SpectralLinear", "SpectralQuadratic"):
1746
- raise ValueError(
1747
- f"Unsupported spectral class '{self._class}'. "
1748
- "Only 'SpectralLinear' or 'SpectralQuadratic' are allowed."
1749
- )
1750
-
1751
- # grab user parameters
1752
- self._parameters = dict(parameters or {})
1753
- self._fermi_wavevector = self._parameters.get("fermi_wavevector")
1754
- self._fermi_velocity = self._parameters.get("fermi_velocity")
1755
- self._bare_mass = self._parameters.get("bare_mass")
1756
- self._side = self._parameters.get("side", None)
1757
-
1758
- # ---- class-specific parameter constraints
1759
- if self._class == "SpectralLinear" and (self._bare_mass is not None):
1760
- raise ValueError("`bare_mass` cannot be set for SpectralLinear.")
1761
- if self._class == "SpectralQuadratic" and (self._fermi_velocity is not None):
1762
- raise ValueError("`fermi_velocity` cannot be set for SpectralQuadratic.")
1763
-
1764
- if self._side is not None and self._side not in ("left", "right"):
1765
- raise ValueError("`side` must be 'left' or 'right' if provided.")
1766
- if self._side is not None:
1767
- self._parameters["side"] = self._side
1768
-
1769
- # convenience attributes (read from properties)
1770
- self._amplitude = self._properties.get("amplitude")
1771
- self._amplitude_sigma = self._properties.get("amplitude_sigma")
1772
- self._peak = self._properties.get("peak")
1773
- self._peak_sigma = self._properties.get("peak_sigma")
1774
- self._broadening = self._properties.get("broadening")
1775
- self._broadening_sigma = self._properties.get("broadening_sigma")
1776
- self._center_wavevector = self._properties.get("center_wavevector")
1777
-
1778
- # lazy caches
1779
- self._peak_positions = None
1780
- self._peak_positions_sigma = None
1781
- self._real = None
1782
- self._real_sigma = None
1783
- self._imag = None
1784
- self._imag_sigma = None
1785
-
1786
- def _check_mass_velocity_exclusivity(self):
1787
- """Ensure that fermi_velocity and bare_mass are not both set."""
1788
- if (self._fermi_velocity is not None) and (self._bare_mass is not None):
1789
- raise ValueError(
1790
- "Cannot set both `fermi_velocity` and `bare_mass`: "
1791
- "choose one physical parametrization (SpectralLinear or SpectralQuadratic)."
1792
- )
1793
-
1794
- # ---------------- core read-only axes ----------------
1795
- @property
1796
- def ekin_range(self):
1797
- return self._ekin_range
1798
-
1799
- @property
1800
- def enel_range(self):
1801
- if self._ekin_range is None:
1802
- return None
1803
- hnp = 0.0 if self._hnuminphi is None else self._hnuminphi
1804
- return np.asarray(self._ekin_range) - hnp
1805
-
1806
- @property
1807
- def hnuminphi(self):
1808
- return self._hnuminphi
1809
-
1810
- # ---------------- identifiers ----------------
1811
- @property
1812
- def label(self):
1813
- return self._label
1814
-
1815
- @label.setter
1816
- def label(self, x):
1817
- self._label = x
1818
-
1819
- # ---------------- exported user parameters ----------------
1820
- @property
1821
- def parameters(self):
1822
- """Dictionary with user-supplied parameters (read-only view)."""
1823
- return self._parameters
1824
-
1825
- @property
1826
- def side(self):
1827
- """Optional side selector: 'left' or 'right'."""
1828
- return self._side
1829
-
1830
- @side.setter
1831
- def side(self, x):
1832
- if x is not None and x not in ("left", "right"):
1833
- raise ValueError("`side` must be 'left' or 'right' if provided.")
1834
- self._side = x
1835
- if x is not None:
1836
- self._parameters["side"] = x
1837
- else:
1838
- self._parameters.pop("side", None)
1839
- # affects sign of peak_positions and thus `real`
1840
- self._peak_positions = None
1841
- self._real = None
1842
- self._real_sigma = None
1843
- self._mdc_maxima = None
1844
-
1845
- @property
1846
- def fermi_wavevector(self):
1847
- """Optional k_F; can be set later."""
1848
- return self._fermi_wavevector
1849
-
1850
- @fermi_wavevector.setter
1851
- def fermi_wavevector(self, x):
1852
- self._fermi_wavevector = x
1853
- self._parameters["fermi_wavevector"] = x
1854
- # invalidate dependent cache
1855
- self._real = None
1856
- self._real_sigma = None
1857
-
1858
- @property
1859
- def fermi_velocity(self):
1860
- """Optional v_F; can be set later."""
1861
- return self._fermi_velocity
1862
-
1863
- @fermi_velocity.setter
1864
- def fermi_velocity(self, x):
1865
- if self._class == "SpectralQuadratic":
1866
- raise ValueError("`fermi_velocity` cannot be set for" \
1867
- " SpectralQuadratic.")
1868
- self._fermi_velocity = x
1869
- self._parameters["fermi_velocity"] = x
1870
- # invalidate dependents
1871
- self._imag = None; self._imag_sigma = None
1872
- self._real = None; self._real_sigma = None
1873
-
1874
- @property
1875
- def bare_mass(self):
1876
- """Optional bare mass; used by SpectralQuadratic formulas."""
1877
- return self._bare_mass
1878
-
1879
- @bare_mass.setter
1880
- def bare_mass(self, x):
1881
- if self._class == "SpectralLinear":
1882
- raise ValueError("`bare_mass` cannot be set for SpectralLinear.")
1883
- self._bare_mass = x
1884
- self._parameters["bare_mass"] = x
1885
- # invalidate dependents
1886
- self._imag = None; self._imag_sigma = None
1887
- self._real = None; self._real_sigma = None
1888
-
1889
- # ---------------- optional fit parameters (convenience) ----------------
1890
- @property
1891
- def amplitude(self):
1892
- return self._amplitude
1893
-
1894
- @amplitude.setter
1895
- def amplitude(self, x):
1896
- self._amplitude = x
1897
- self._properties["amplitude"] = x
1898
-
1899
- @property
1900
- def amplitude_sigma(self):
1901
- return self._amplitude_sigma
1902
-
1903
- @amplitude_sigma.setter
1904
- def amplitude_sigma(self, x):
1905
- self._amplitude_sigma = x
1906
- self._properties["amplitude_sigma"] = x
1907
-
1908
- @property
1909
- def peak(self):
1910
- return self._peak
1911
-
1912
- @peak.setter
1913
- def peak(self, x):
1914
- self._peak = x
1915
- self._properties["peak"] = x
1916
- # invalidate dependent cache
1917
- self._peak_positions = None
1918
- self._real = None
1919
- self._mdc_maxima = None
1920
-
1921
- @property
1922
- def peak_sigma(self):
1923
- return self._peak_sigma
1924
-
1925
- @peak_sigma.setter
1926
- def peak_sigma(self, x):
1927
- self._peak_sigma = x
1928
- self._properties["peak_sigma"] = x
1929
- self._peak_positions_sigma = None
1930
- self._real_sigma = None
1931
-
1932
- @property
1933
- def broadening(self):
1934
- return self._broadening
1935
-
1936
- @broadening.setter
1937
- def broadening(self, x):
1938
- self._broadening = x
1939
- self._properties["broadening"] = x
1940
- self._imag = None
1941
-
1942
- @property
1943
- def broadening_sigma(self):
1944
- return self._broadening_sigma
1945
-
1946
- @broadening_sigma.setter
1947
- def broadening_sigma(self, x):
1948
- self._broadening_sigma = x
1949
- self._properties["broadening_sigma"] = x
1950
- self._imag_sigma = None
1951
-
1952
- @property
1953
- def center_wavevector(self):
1954
- """Read-only center wavevector (SpectralQuadratic, if present)."""
1955
- return self._center_wavevector
1956
-
1957
- # ---------------- derived outputs ----------------
1958
- @property
1959
- def peak_positions(self):
1960
- r"""k_parallel = peak * dtor * sqrt(ekin_range / pref) (lazy)."""
1961
- if self._peak_positions is None:
1962
- if self._peak is None or self._ekin_range is None:
1963
- return None
1964
- if self._class == "SpectralQuadratic":
1965
- if self._side is None:
1966
- raise AttributeError(
1967
- "For SpectralQuadratic, set `side` ('left'/'right') "
1968
- "before accessing peak_positions and quantities that "
1969
- "depend on the latter."
1970
- )
1971
- kpar_mag = np.sqrt(self._ekin_range / pref) * \
1972
- np.sin(np.abs(self._peak) * dtor)
1973
- self._peak_positions = (-1.0 if self._side == "left" \
1974
- else 1.0) * kpar_mag
1975
- else:
1976
- self._peak_positions = np.sqrt(self._ekin_range / pref) \
1977
- * np.sin(self._peak * dtor)
1978
- return self._peak_positions
1979
-
1980
- @property
1981
- def peak_positions_sigma(self):
1982
- r"""Std. dev. of k_parallel (lazy)."""
1983
- if self._peak_positions_sigma is None:
1984
- if self._peak_sigma is None or self._ekin_range is None:
1985
- return None
1986
- self._peak_positions_sigma = (np.sqrt(self._ekin_range / pref)
1987
- * np.abs(np.cos(self._peak * dtor))
1988
- * self._peak_sigma * dtor)
1989
- return self._peak_positions_sigma
1990
-
1991
- @property
1992
- def imag(self):
1993
- r"""-Σ'' (lazy)."""
1994
- if self._imag is None:
1995
- if self._broadening is None or self._ekin_range is None:
1996
- return None
1997
- if self._class == "SpectralLinear":
1998
- if self._fermi_velocity is None:
1999
- raise AttributeError("Cannot compute `imag` "
2000
- "(SpectralLinear): set `fermi_velocity` first.")
2001
- self._imag = np.abs(self._fermi_velocity) * np.sqrt(self._ekin_range \
2002
- / pref) * self._broadening
2003
- else:
2004
- if self._bare_mass is None:
2005
- raise AttributeError("Cannot compute `imag` "
2006
- "(SpectralQuadratic): set `bare_mass` first.")
2007
- self._imag = (self._ekin_range * self._broadening) \
2008
- / np.abs(self._bare_mass)
2009
- return self._imag
2010
-
2011
- @property
2012
- def imag_sigma(self):
2013
- r"""Std. dev. of -Σ'' (lazy)."""
2014
- if self._imag_sigma is None:
2015
- if self._broadening_sigma is None or self._ekin_range is None:
2016
- return None
2017
- if self._class == "SpectralLinear":
2018
- if self._fermi_velocity is None:
2019
- raise AttributeError("Cannot compute `imag_sigma` "
2020
- "(SpectralLinear): set `fermi_velocity` first.")
2021
- self._imag_sigma = np.abs(self._fermi_velocity) * \
2022
- np.sqrt(self._ekin_range / pref) * self._broadening_sigma
2023
- else:
2024
- if self._bare_mass is None:
2025
- raise AttributeError("Cannot compute `imag_sigma` "
2026
- "(SpectralQuadratic): set `bare_mass` first.")
2027
- self._imag_sigma = (self._ekin_range * \
2028
- self._broadening_sigma) / np.abs(self._bare_mass)
2029
- return self._imag_sigma
2030
-
2031
- @property
2032
- def real(self):
2033
- r"""Σ' (lazy)."""
2034
- if self._real is None:
2035
- if self._peak is None or self._ekin_range is None:
2036
- return None
2037
- if self._class == "SpectralLinear":
2038
- if self._fermi_velocity is None or self._fermi_wavevector is None:
2039
- raise AttributeError("Cannot compute `real` "
2040
- "(SpectralLinear): set `fermi_velocity` and " \
2041
- "`fermi_wavevector` first.")
2042
- self._real = self.enel_range - self._fermi_velocity * \
2043
- (self.peak_positions - self._fermi_wavevector)
2044
- else:
2045
- if self._bare_mass is None or self._fermi_wavevector is None:
2046
- raise AttributeError("Cannot compute `real` "
2047
- "(SpectralQuadratic): set `bare_mass` and " \
2048
- "`fermi_wavevector` first.")
2049
- self._real = self.enel_range - (pref / \
2050
- self._bare_mass) * (self.peak_positions**2 \
2051
- - self._fermi_wavevector**2)
2052
- return self._real
2053
-
2054
- @property
2055
- def real_sigma(self):
2056
- r"""Std. dev. of Σ' (lazy)."""
2057
- if self._real_sigma is None:
2058
- if self._peak_sigma is None or self._ekin_range is None:
2059
- return None
2060
- if self._class == "SpectralLinear":
2061
- if self._fermi_velocity is None:
2062
- raise AttributeError("Cannot compute `real_sigma` "
2063
- "(SpectralLinear): set `fermi_velocity` first.")
2064
- self._real_sigma = np.abs(self._fermi_velocity) * self.peak_positions_sigma
2065
- else:
2066
- if self._bare_mass is None or self._fermi_wavevector is None:
2067
- raise AttributeError("Cannot compute `real_sigma` "
2068
- "(SpectralQuadratic): set `bare_mass` and " \
2069
- "`fermi_wavevector` first.")
2070
- self._real_sigma = 2 * pref * self.peak_positions_sigma \
2071
- * np.abs(self.peak_positions / self._bare_mass)
2072
- return self._real_sigma
2073
-
2074
- @property
2075
- def mdc_maxima(self):
2076
- """
2077
- MDC maxima (lazy).
2078
-
2079
- SpectralLinear:
2080
- identical to peak_positions
2081
-
2082
- SpectralQuadratic:
2083
- peak_positions + center_wavevector
2084
- """
2085
- if getattr(self, "_mdc_maxima", None) is None:
2086
- if self.peak_positions is None:
2087
- return None
2088
-
2089
- if self._class == "SpectralLinear":
2090
- self._mdc_maxima = self.peak_positions
2091
- elif self._class == "SpectralQuadratic":
2092
- self._mdc_maxima = (
2093
- self.peak_positions + self._center_wavevector
2094
- )
2095
-
2096
- return self._mdc_maxima
2097
-
2098
-
2099
- class CreateSelfEnergies:
2100
- r"""
2101
- Thin container for self-energies with leaf-aware utilities.
2102
- All items are assumed to be leaf self-energy objects with
2103
- a `.label` attribute for identification.
2104
- """
2105
-
2106
- def __init__(self, self_energies):
2107
- self.self_energies = self_energies
2108
-
2109
- # ------ Basic container protocol ------
2110
- def __call__(self):
2111
- return self.self_energies
2112
-
2113
- @property
2114
- def self_energies(self):
2115
- return self._self_energies
2116
-
2117
- @self_energies.setter
2118
- def self_energies(self, x):
2119
- self._self_energies = x
2120
-
2121
- def __iter__(self):
2122
- return iter(self.self_energies)
2123
-
2124
- def __getitem__(self, index):
2125
- return self.self_energies[index]
2126
-
2127
- def __setitem__(self, index, value):
2128
- self.self_energies[index] = value
2129
-
2130
- def __len__(self):
2131
- return len(self.self_energies)
2132
-
2133
- def __deepcopy__(self, memo):
2134
- import copy
2135
- return type(self)(copy.deepcopy(self.self_energies, memo))
2136
-
2137
- # ------ Label-based utilities ------
2138
- def get_by_label(self, label):
2139
- r"""
2140
- Return the self-energy object with the given label.
2141
-
2142
- Parameters
2143
- ----------
2144
- label : str
2145
- Label of the self-energy to retrieve.
2146
-
2147
- Returns
2148
- -------
2149
- obj : SelfEnergy
2150
- The corresponding self-energy instance.
2151
-
2152
- Raises
2153
- ------
2154
- KeyError
2155
- If no self-energy with the given label exists.
2156
- """
2157
- for se in self.self_energies:
2158
- if getattr(se, "label", None) == label:
2159
- return se
2160
- raise KeyError(
2161
- f"No self-energy with label {label!r} found in container."
2162
- )
2163
-
2164
- def labels(self):
2165
- r"""
2166
- Return a list of all labels.
2167
- """
2168
- return [getattr(se, "label", None) for se in self.self_energies]
2169
-
2170
- def as_dict(self):
2171
- r"""
2172
- Return a {label: self_energy} dictionary for convenient access.
2173
- """
2174
- return {se.label: se for se in self.self_energies}
2175
-