madman.analysis.spectrum

  1r"""Spectrum base classes."""
  2
  3from collections.abc import Mapping, Sequence
  4from numbers import Integral, Real
  5from sys import modules
  6from typing import Any, Self, Type
  7
  8import numpy as np
  9import seaborn as sns
 10from matplotlib import pyplot as plt
 11from scipy.ndimage import gaussian_filter1d
 12
 13from madman.utilities import gen_regular_grid
 14
 15
 16sns.set_theme(
 17    context="talk",
 18    style="white",
 19    rc={"figure.titlesize": "medium", "axes.formatter.useoffset": False},
 20)
 21
 22
 23class EnergySpectrumMeta(type):
 24    r"""Metaclas of EnergySpectrum."""
 25
 26    @property
 27    def rsv(cls) -> Type:
 28        r"""Associated resolved class."""
 29        module = modules[cls.__module__]
 30        name = "Resolved" + cls.__name__
 31        return getattr(module, name)
 32
 33    @property
 34    def band_rsv(cls) -> Type:
 35        r"""Associated band resolved class."""
 36        module = modules[cls.__module__]
 37        name = "BandResolved" + cls.__name__
 38        return getattr(module, name)
 39
 40    @property
 41    def cell_rsv(cls) -> Type:
 42        r"""Associated cell resolved class."""
 43        module = modules[cls.__module__]
 44        name = "CellResolved" + cls.__name__
 45        return getattr(module, name)
 46
 47    @property
 48    def std_axes(cls) -> plt.Axes:
 49        r"""Standard axes to plot energy spectrum."""
 50        _, ax = plt.subplots(tight_layout=True)
 51        ax.set_xlabel(r"Energy / $\mathrm{eV}$")
 52        ax.set_ylabel(r"Values / $\mathrm{n.a.}$")
 53        return ax
 54
 55
 56class EnergySpectrum(metaclass=EnergySpectrumMeta):
 57    r"""Energy spectrum."""
 58
 59    def __init__(
 60        self,
 61        energies: Sequence[Real, ...],
 62        values: Sequence[Real, ...],
 63        *,
 64        d_energy: Real | None = None,
 65    ) -> None:
 66        r"""Initialize EnergySpectrum object.
 67
 68        Args:
 69            energies: Spectrum energies [$\mathrm{eV}$].
 70            values: Spectrum values.
 71            d_energy: Target energy increment [$\mathrm{eV}$].
 72
 73        Raises:
 74            ValueError: If there are duplicate energies.
 75        """
 76        energies = np.asarray(energies, dtype=float).flatten()
 77        values = np.asarray(values, dtype=float).flatten()
 78
 79        sort_i = np.argsort(energies)
 80        energies = energies[sort_i]
 81        values = values[sort_i]
 82
 83        if not np.array_equal(energies, np.unique(energies)):
 84            raise ValueError("Duplicate energies!")
 85
 86        energy_a = np.amin(energies)
 87        energy_b = np.amax(energies)
 88        if d_energy is None:
 89            d_energy = np.diff(energies).min()
 90
 91        self._energies = gen_regular_grid(energy_a, energy_b, d_energy)
 92        self._values = np.interp(self._energies, energies, values)
 93        self._integr = np.trapz(self.values, x=self.energies)
 94
 95    def __add__(self, other: Self) -> Self:
 96        r"""Add energy spectrum.
 97
 98        Args:
 99            other: Energy spectrum to add.
100
101        Returns:
102            Sum of `self` and `other`.
103
104        Raises:
105            TypeError: If `self` and `other` have incompatible types.
106            ValueError: If `self` and `other` have incompatible energies.
107        """
108        if type(self) != type(other):
109            raise TypeError("Incompatible types!")
110        if not np.array_equal(self.energies, other.energies):
111            raise ValueError("Incompatible energies!")
112        return type(self)(self.energies, self.values + other.values)
113
114    @property
115    def energies(self) -> np.ndarray[float]:
116        r"""Spectrum energies [$\mathrm{eV}$]."""
117        return self._energies
118
119    @property
120    def values(self) -> np.ndarray[float]:
121        r"""Spectrum values."""
122        return self._values
123
124    @property
125    def integr(self) -> float:
126        r"""Spectrum integrated value."""
127        return self._integr
128
129    def ret_zero(self) -> Self:
130        r"""Return zero energy spectrum."""
131        return type(self).zero(self.energies)
132
133    def ret_smeared(self, *, sigma: Real) -> Self:
134        r"""Return smeared energy spectrum.
135
136        Args:
137            sigma: Standard deviation for Gaussian kernel.
138
139        Returns:
140            Smeared energy spectrum.
141        """
142        values = gaussian_filter1d(self.values, sigma)
143        return type(self)(self.energies, values)
144
145    def plot(
146        self,
147        *,
148        ax: plt.Axes | None = None,
149        color: str | None = None,
150        label: str | None = None,
151    ) -> plt.Figure:
152        r"""Plot energy spectrum.
153
154        Args:
155            ax: Plot axes; if None, new figure and axes are created.
156            color: Energy spectrum color.
157            label: Energy spectrum label.
158
159        Returns:
160            Plot figure.
161        """
162        if ax is None:
163            ax = type(self).std_axes
164        ax.plot(self.energies, self.values, color=color, label=label)
165        return ax.get_figure()
166
167    @classmethod
168    def zero(
169        cls, energies: Sequence[Real, ...], *, d_energy: Real | None = None
170    ) -> np.ndarray[float]:
171        r"""Zero energy spectrum.
172
173        Args:
174            energies: Spectrum energies [$\mathrm{eV}$].
175            d_energy: Target energy increment [$\mathrm{eV}$].
176
177        Returns:
178            Zero energy spectrum.
179        """
180        values = np.zeros_like(energies)
181        return cls(energies, values, d_energy=d_energy)
182
183    @classmethod
184    def piecewise_constant(
185        cls,
186        energy_a: Real,
187        energy_b: Real,
188        d_energy: Real,
189        *,
190        amplitudes: Sequence[Real, ...],
191        low_bounds: Sequence[Real, ...],
192        upp_bounds: Sequence[Real, ...],
193    ) -> Self:
194        r"""Piecewise constant energy spectrum.
195
196        Args:
197            energy_a: Energy minimum [$\mathrm{eV}$].
198            energy_b: Energy maximum [$\mathrm{eV}$].
199            d_energy: Target energy increment [$\mathrm{eV}$].
200            amplitudes: Piece amplitudes.
201            low_bounds: Piece lower bounds [$\mathrm{eV}$].
202            upp_bounds: Piece upper bounds [$\mathrm{eV}$].
203
204        Returns:
205            Piecewise constant energy spectrum.
206        """
207        energies = gen_regular_grid(energy_a, energy_b, d_energy)
208
209        def rect(x, a, b):
210            return np.heaviside(x - a, 0.5) - np.heaviside(x - b, 0.5)
211
212        values = 0.0
213        for amp, lb, ub in zip(amplitudes, low_bounds, upp_bounds):
214            values += amp * rect(energies, lb, ub)
215        return cls(energies, values)
216
217
218class ResolvedEnergySpectrumMeta(type):
219    r"""Metaclass of ResolvedEnergySpectrum."""
220
221    @property
222    def total(cls) -> Type:
223        r"""Total energy spectrum class."""
224        module = modules[cls.__module__]
225        name = cls.__name__.split("Resolved")[1]
226        return getattr(module, name)
227
228
229class ResolvedEnergySpectrum(metaclass=ResolvedEnergySpectrumMeta):
230    r"""Resolved energy spectrum."""
231
232    def __init__(
233        self,
234        energies: Sequence[Real, ...],
235        rsv_values: Mapping[Any, Sequence[Real, ...]],
236        *,
237        d_energy: Real | None = None,
238    ) -> None:
239        r"""Initialize ResolvedEnergySpectrum object.
240
241        Args:
242            energies: Spectrum energies [$\mathrm{eV}$].
243            rsv_values: Mapping of feature into spectrum values.
244            d_energy: Target energy increment [$\mathrm{eV}$].
245
246        Raises:
247            ValueError: If band index is non-integer.
248            ValueError: If cell role is not in {'v', 'i', 'c'}.
249        """
250        self._rsv_values = {}
251        self._rsv_integr = {}
252        self._total = type(self).total.zero(energies, d_energy=d_energy)
253
254        for feat, values in rsv_values.items():
255            if not isinstance(self, PairResolvedEnergySpectrum):
256
257                if "Band" in type(self).__name__:
258                    if not isinstance(feat, Integral):
259                        raise ValueError("Non-integer band index!")
260                    feat = int(feat)
261
262                if "Cell" in type(self).__name__:
263                    if not feat in {"v", "i", "c"}:
264                        raise ValueError("Cell role not in {'v', 'i', 'c'}")
265
266            spectrum = type(self).total(energies, values, d_energy=d_energy)
267            self._rsv_values[feat] = spectrum.values
268            self._rsv_integr[feat] = np.trapz(
269                spectrum.values, x=spectrum.energies
270            )
271            self._total += spectrum
272
273            if (
274                not isinstance(self, PairResolvedEnergySpectrum)
275                and "Cell" in type(self).__name__
276            ):
277                rsv_values_sorted = {}
278                for feat in ["v", "i", "c"]:
279                    if feat in self._rsv_values:
280                        rsv_values_sorted.update({feat: self._rsv_values[feat]})
281                self._rsv_values = rsv_values_sorted
282
283        self._energies = spectrum.energies
284
285        self._rsv_sp_weight = {}
286        for feat, values in rsv_values.items():
287            sp_weight = self._rsv_integr[feat] / self._total.integr
288            self._rsv_sp_weight[feat] = sp_weight
289
290        self._rsv_sp_select = {}
291        for feat, values in rsv_values.items():
292            sp_select = 1.0 - np.trapz(
293                values * (self._total.values - values), x=self._energies
294            ) / np.sqrt(
295                np.trapz(values**2, x=self._energies)
296                * np.trapz((self._total.values - values) ** 2, x=self._energies)
297            )
298            self._rsv_sp_select[feat] = sp_select
299
300    def __add__(self, other: Self) -> Self:
301        r"""Add resolved energy spectrum.
302
303        Args:
304            other: Resolved energy spectrum to add.
305
306        Returns:
307            Sum of `self` and `other`.
308
309        Raises:
310            TypeError: If `self` and `other` have incompatible types.
311            ValueError: If `self` and `other` have incompatible energies.
312            ValueError: If `self` and `other` have incompatible features.
313        """
314        if type(self) != type(other):
315            raise TypeError("Incompatible types!")
316        if not np.array_equal(self.energies, other.energies):
317            raise ValueError("Incompatible energies!")
318        if self.rsv_values.keys() != other.rsv_values.keys():
319            raise ValueError("Incompatible features!")
320        rsv_values = {
321            feat: values + other.rsv_values[feat]
322            for feat, values in self.rsv_values.items()
323        }
324        return type(self)(self.energies, rsv_values)
325
326    @property
327    def energies(self) -> np.ndarray[float]:
328        r"""Spectrum energies [$\mathrm{eV}$]."""
329        return self._energies
330
331    @property
332    def rsv_values(self) -> dict[Any, np.ndarray[float]]:
333        r"""Resolved spectrum values."""
334        return self._rsv_values
335
336    @property
337    def rsv_integr(self) -> dict[Any, float]:
338        r"""Resolved spectrum integrated values."""
339        return self._rsv_integr
340
341    @property
342    def rsv_sp_weight(self) -> dict[Any, float]:
343        r"""Resolved spectral weight.
344
345        Note:
346            The spectral weight of feature $x$ is defined as:
347
348            $$
349            \frac{\int \mathrm{d}(\hbar \, \omega) \, \Sigma_{x}(\hbar \, \omega)}{\int \mathrm{d}(\hbar \, \omega) \, \Sigma_{\text{tot}}(\hbar \, \omega)}
350            $$
351
352            where $\hbar \, \omega$ is the photon energy, $\Sigma_{x}$ is the
353            partial spectrum og feature $x$, and $\Sigma_{\text{tot}}$ is the
354            total spectrum.
355        """
356        return self._rsv_sp_weight
357
358    @property
359    def rsv_sp_select(self) -> dict[Any, float]:
360        r"""Resolved spectral selectivity.
361
362        Note:
363            The spectral selectivity of feature $x$ is defined as:
364
365            $$
366            1 - \frac{\int \mathrm{d}(\hbar \, \omega) \, \Sigma_{x}(\hbar \, \omega) \, \left[ \Sigma_{\text{tot}}(\hbar \, \omega) - \Sigma_{x}(\hbar \, \omega) \right]}{\sqrt{\int \mathrm{d}(\hbar \, \omega) \, \Sigma_{x}^{2}(\hbar \, \omega) \, \int \mathrm{d}(\hbar \, \omega) \, \left[ \Sigma_{\text{tot}}(\hbar \, \omega) - \Sigma_{x}(\hbar \, \omega)\right]^{2}}}
367            $$
368
369            where $\hbar \, \omega$ is the photon energy, $\Sigma_{x}$ is the
370            partial spectrum og feature $x$, and $\Sigma_{\text{tot}}$ is the total
371            spectrum.
372        """
373        return self._rsv_sp_select
374
375    @property
376    def total(self) -> EnergySpectrum:
377        r"""Total energy spectrum."""
378        return self._total
379
380    def ret_zero(self) -> Self:
381        r"""Return zero resolved energy spectrum."""
382        return type(self).zero(self.energies, self.rsv_values)
383
384    def ret_smeared(self, *, sigma: Real) -> Self:
385        r"""Return smeared resolved energy spectrum.
386
387        Args:
388            sigma: Standard deviation for Gaussian kernel.
389
390        Returns:
391            Smeared resolved energy spectrum.
392        """
393        rsv_values = {
394            feat: gaussian_filter1d(values, sigma)
395            for feat, values in self.rsv_values.items()
396        }
397        return type(self)(self.energies, rsv_values)
398
399    def plot(
400        self,
401        *,
402        ax: plt.Axes | None = None,
403        stack: bool = False,
404    ) -> plt.Figure:
405        r"""Plot resolved energy spectrum.
406
407        Args:
408            ax: Plot axes; if None, new figure and axes are created.
409            stack: If True, a stackplot is generated.
410
411        Returns:
412            Plot figure.
413
414        Raises:
415            TypeError: If `stack` is not boolean.
416        """
417        if ax is None:
418            ax = type(self).total.std_axes
419
420        if stack is True:
421            ax.stackplot(
422                self.energies,
423                self.rsv_values.values(),
424                labels=self.rsv_values.keys(),
425            )
426        elif stack is False:
427            for feat, values in self.rsv_values.items():
428                ax.plot(self.energies, values, label=feat)
429        else:
430            raise TypeError("`stack` must be boolean!")
431
432        ax.plot(
433            self.energies,
434            self.total.values,
435            linestyle="--",
436            color="k",
437            label="total",
438        )
439        ax.legend(loc="upper left", bbox_to_anchor=[1.0, 1.0], frameon=False)
440        return ax.get_figure()
441
442    @classmethod
443    def zero(
444        cls,
445        energies: Sequence[Real, ...],
446        features: Sequence[Any, ...],
447        *,
448        d_energy: Real | None = None,
449    ) -> Self:
450        r"""Zero resolved energy spectrum.
451
452        Args:
453            energies: Resolved spectrum energies [$\mathrm{eV}$].
454            features: Resolved spectrum features.
455            d_energy: Target energy increment [$\mathrm{eV}$].
456
457        Returns:
458            Zero resolved energy spectrum.
459        """
460        rsv_values = {feat: np.zeros_like(energies) for feat in features}
461        return cls(energies, rsv_values, d_energy=d_energy)
462
463
464class BandResolvedEnergySpectrum(ResolvedEnergySpectrum):
465    r"""Band resolved energy spectrum.
466
467    Note:
468        'Band' is hereby intended as a range of electron energies with no zero
469        electron density of states values such that any larger range contains
470        zero electron density of states values.
471        To distinguish, a band of a dispersion relation (e.g. the electronic
472        band structure of a crystalline solid) is hereby referred as 'branch'.
473    """
474
475
476class CellResolvedEnergySpectrum(ResolvedEnergySpectrum):
477    r"""Cell resolved energy spectrum.
478
479    Note:
480        The energy spectrum is resolved according to the components role in a
481        solar cells.
482        The following convention is adopted:
483
484        - 'v': valence bands;
485        - 'i': intermediate bands;
486        - 'c': conduction bands.
487
488        See
489        [Shockley & Queisser (1961)](https://dx.doi.org/10.1063%2F1.1736034)
490        and
491        [Levy & Honsberg (2008)](http://dx.doi.org/10.1103%2FPhysRevB.78.165122)
492        for details.
493    """
494
495
496class PairResolvedEnergySpectrum(ResolvedEnergySpectrum):
497    r"""Pair resolved energy spectrum."""
498
499    def __init__(
500        self,
501        energies: Sequence[Real, ...],
502        rsv_values: Mapping[Sequence[Any, Any], Sequence[Real, ...]],
503        *,
504        d_energy: Real | None = None,
505    ) -> None:
506        r"""Initialize ResolvedEnergySpectrum object.
507
508        Args:
509            energies: Spectrum energies [$\mathrm{eV}$].
510            rsv_values: Mapping of feature pair into spectrum values.
511            d_energy: Target energy increment [$\mathrm{eV}$].
512
513        Raises:
514            TypeError: If a feature pair is not a sequence.
515            ValueError: If a feature pair is not a pair.
516            ValueError: If a band index pair is not a pair of integers.
517            ValueError: If a cell role pair is not in {'ii', 'vi', 'ic', 'vc'}.
518        """
519        rsv_values_new = {}
520        for pair, values in rsv_values.items():
521            if not isinstance(pair, Sequence):
522                raise TypeError("Not a sequence!")
523            if not len(pair) == 2:
524                raise ValueError("Not a pair!")
525
526            if "Band" in type(self).__name__:
527                i, j = pair
528                if not isinstance(i, Integral) or not isinstance(j, Integral):
529                    raise ValueError("Non-integer band pair index!")
530                pair = tuple([int(i), int(j)])
531
532            elif "Cell" in type(self).__name__:
533                if pair not in {"ii", "vi", "ic", "vc"}:
534                    raise ValueError(
535                        "Cell role pair not in {'ii', 'vi', 'ic', 'vc'}!"
536                    )
537
538            else:
539                pair = tuple(pair)
540
541            spectrum = type(self).total(energies, values, d_energy=d_energy)
542            rsv_values_new[pair] = spectrum.values
543
544        if "Cell" in type(self).__name__:
545            rsv_values_new_sorted = {}
546            for pair in ["ii", "vi", "ic", "vc"]:
547                if pair in rsv_values_new:
548                    rsv_values_new_sorted.update({pair: rsv_values_new[pair]})
549            rsv_values_new = rsv_values_new_sorted
550
551        super().__init__(energies, rsv_values_new, d_energy=d_energy)
552
553    def plot(
554        self,
555        *,
556        ax: plt.Axes | None = None,
557        stack: bool = False,
558    ) -> plt.Figure:
559        r"""Plot pair resolved energy spectrum.
560
561        Args:
562            ax: Plot axes; if None, new figure and axes are created.
563            stack: If True, a stackplot is generated.
564
565        Returns:
566            Plot figure.
567
568        Raises:
569            TypeError: If `stack` is not a boolean.
570        """
571        if ax is None:
572            ax = type(self).total.std_axes
573
574        if stack is True:
575            ax.stackplot(
576                self.energies,
577                self.rsv_values.values(),
578                labels=[f"{pair[0]} → {pair[1]}" for pair in self.rsv_values],
579            )
580        elif stack is False:
581            for pair, values in self.rsv_values.items():
582                ax.plot(self.energies, values, label=f"{pair[0]} → {pair[1]}")
583        else:
584            raise TypeError("`stack` must be boolean!")
585
586        ax.plot(
587            self.energies,
588            self.total.values,
589            linestyle="--",
590            color="k",
591            label="total",
592        )
593        ax.legend(loc="upper left", bbox_to_anchor=[1.0, 1.0], frameon=False)
594        return ax.get_figure()
595
596
597class BandPairResolvedEnergySpectrum(PairResolvedEnergySpectrum):
598    r"""Band pair resolved energy spectrum.
599
600    Note:
601        See `madman.analysis.spectrum.BandResolvedEnergySpectrum` for details.
602    """
603
604
605class CellPairResolvedEnergySpectrum(PairResolvedEnergySpectrum):
606    r"""Cell pair resolved energy spectrum.
607
608    Note:
609        See `madman.analysis.spectrum.CellResolvedEnergySpectrum` for details.
610    """