satlas2.models.hfsModel — satlas2 0.1.5 documentation

Toggle navigation sidebar
Toggle in-page Table of Contents

#

# Source code for satlas2.models.hfsModel

    """
    Implementation of the HFSModel class, currently only supplied with a Voigt profile.

    .. moduleauthor:: Wouter Gins <wouter.gins@kuleuven.be>
    """
    from __future__ import annotations

    from typing import Tuple

    import numpy as np
    import uncertainties as unc
    from numpy.typing import ArrayLike
    from scipy.special import voigt_profile, wofz
    from sympy.physics.wigner import wigner_3j, wigner_6j

    from satlas2.core import Model, Parameter

    __all__ = ['HFS']

    sqrt2 = 2**0.5
    sqrt2log2t2 = 2 * np.sqrt(2 * np.log(2))
    log2 = np.log(2)

    class HFS(Model):
        """Initializes a hyperfine spectrum Model with the given hyperfine parameters.

        Parameters
        ----------
        I : float
            Integer or half-integer value of the nuclear spin
        J : ArrayLike
            A sequence of 2 spins, respectively the J value of the lower state
            and the J value of the higher state
        A : ArrayLike, optional
            A sequence of 2 A values, respectively for the lower and the higher state, by default [0, 0]
        B : ArrayLike, optional
            A sequence of 2 B values, respectively for the lower and the higher state, by default [0, 0]
        C : ArrayLike, optional
            A sequence of 2 C values, respectively for the lower and the higher state, by default [0, 0]
        df : float, optional
            The centroid of the spectrum, by default 0
        fwhmg : float, optional
            The Gaussian FWHM of the Voigt profile, by default 50
        fwhml : float, optional
            The Lorentzian FWHM of the Voigt profile, by default 50
        name : str, optional
            Name of the model, by default 'HFS'
        N : int, optional
            Number of sidepeaks to be generated, by default None
        offset : float, optional
            Offset in units of x for the sidepeak, by default 0
        poisson : float, optional
            The poisson factor for the sidepeaks, by default 0
        scale : float, optional
            The amplitude of the entire spectrum, by default 1.0
        racah : bool, optional
            Use individual amplitudes are setting the Racah intensities, by default True
        prefunc : callable, optional
            Transformation to be applied on the input before evaluation, by default None
            """
        def __init__(self,
                     I: float,
                     J: ArrayLike,
                     A: ArrayLike = [0, 0],
                     B: ArrayLike = [0, 0],
                     C: ArrayLike = [0, 0],
                     df: float = 0,
                     fwhmg: float = 50,
                     fwhml: float = 50,
                     name: str = 'HFS',
                     N: int = None,
                     offset: float = 0,
                     poisson: float = 0,
                     scale: float = 1.0,
                     racah: bool = True,
                     prefunc: callable = None):
            super().__init__(name, prefunc=prefunc)
            J1, J2 = J
            lower_F = np.arange(abs(I - J1), I + J1 + 1, 1)
            upper_F = np.arange(abs(I - J2), I + J2 + 1, 1)

            self.lines = []
            self.intensities = {}
            self.scaling_Al = {}
            self.scaling_Bl = {}
            self.scaling_Cl = {}
            self.scaling_Au = {}
            self.scaling_Bu = {}
            self.scaling_Cu = {}

            for i, F1 in enumerate(lower_F):
                for j, F2 in enumerate(upper_F):
                    if abs(F2 - F1) <= 1 and not F2 == F1 == 0.0:
                        if F1 % 1 == 0:
                            F1_str = '{:.0f}'.format(F1)
                        else:
                            F1_str = '{:.0f}_2'.format(2 * F1)

                        if F2 % 1 == 0:
                            F2_str = '{:.0f}'.format(F2)
                        else:
                            F2_str = '{:.0f}_2'.format(2 * F2)

                        line = '{}to{}'.format(F1_str, F2_str)
                        self.lines.append(line)

                        C1, D1, E1 = self.calcShift(I, J1, F1)
                        C2, D2, E2 = self.calcShift(I, J2, F2)

                        self.scaling_Al[line] = C1
                        self.scaling_Bl[line] = D1
                        self.scaling_Cl[line] = E1
                        self.scaling_Au[line] = C2
                        self.scaling_Bu[line] = D2
                        self.scaling_Cu[line] = E2

                        intens = float((2 * F1 + 1) * (2 * F2 + 1) * \
                             wigner_6j(J2, F2, I, F1, J1, 1.0) ** 2) # DO NOT REMOVE CAST TO FLOAT!!!
                        self.intensities['Amp' + line] = Parameter(value=intens,
                                                                   min=0,
                                                                   vary=not racah)

            norm = max([p.value for p in self.intensities.values()])
            for n, v in self.intensities.items():
                v.value /= norm

            pars = {
                'centroid': Parameter(value=df),
                'Al': Parameter(value=A[0]),
                'Au': Parameter(value=A[1]),
                'Bl': Parameter(value=B[0]),
                'Bu': Parameter(value=B[1]),
                'Cl': Parameter(value=C[0]),
                'Cu': Parameter(value=C[1]),
                'FWHMG': Parameter(value=fwhmg, min=0.01),
                'FWHML': Parameter(value=fwhml, min=0.01),
                'scale': Parameter(value=scale, min=0, vary=racah)
            }
            if N is not None:
                pars['N'] = Parameter(value=N, vary=False)
                pars['Offset'] = Parameter(value=offset)
                pars['Poisson'] = Parameter(value=poisson, min=0, max=1)
                self.f = self.fShifted
            else:
                self.f = self.fUnshifted
            pars = {**pars, **self.intensities}

            self.params = pars

            if I < 1.5 or J1 < 1.5:
                self.params['Cl'].vary = False
            if I < 1.5 or J2 < 1.5:
                self.params['Cu'].vary = False
            if I < 1 or J1 < 1:
                self.params['Bl'].vary = False
            if I < 1 or J2 < 1:
                self.params['Bu'].vary = False
            if I == 0 or J1 == 0:
                self.params['Al'].vary = False
            if I == 0 or J2 == 0:
                self.params['Au'].vary = False

        def fUnshifted(self, x: ArrayLike) -> ArrayLike:
            """Calculate the response for an unshifted spectrum

            Parameters
            ----------
            x : ArrayLike

            Returns
            -------
            ArrayLike
            """
            centroid = self.params['centroid'].value
            Al = self.params['Al'].value
            Au = self.params['Au'].value
            Bl = self.params['Bl'].value
            Bu = self.params['Bu'].value
            Cl = self.params['Cl'].value
            Cu = self.params['Cu'].value
            FWHMG = self.params['FWHMG'].value
            FWHML = self.params['FWHML'].value
            scale = self.params['scale'].value

            result = np.zeros(len(x))
            x = self.transform(x)
            for line in self.lines:
                pos = centroid + Au * self.scaling_Au[line] + Bu * self.scaling_Bu[
                    line] + Cu * self.scaling_Cu[line] - Al * self.scaling_Al[
                        line] - Bl * self.scaling_Bl[line] - Cl * self.scaling_Cl[
                            line]
                result += scale * self.params['Amp' + line].value * self.peak(
                    x - pos, FWHMG, FWHML)

            return result

        def fShifted(self, x: ArrayLike) -> ArrayLike:
            """Calculate the response with :attr:`N` sidepeaks with an offset
            of :attr:`offset`

            Parameters
            ----------
            x : ArrayLike

            Returns
            -------
            ArrayLike
            """
            centroid = self.params['centroid'].value
            Al = self.params['Al'].value
            Au = self.params['Au'].value
            Bl = self.params['Bl'].value
            Bu = self.params['Bu'].value
            Cl = self.params['Cl'].value
            Cu = self.params['Cu'].value
            FWHMG = self.params['FWHMG'].value
            FWHML = self.params['FWHML'].value
            scale = self.params['scale'].value
            N = self.params['N'].value
            offset = self.params['Offset'].value
            poisson = self.params['Poisson'].value

            result = np.zeros(len(x))
            x = self.transform(x)
            for line in self.lines:
                pos = centroid + Au * self.scaling_Au[line] + Bu * self.scaling_Bu[
                    line] + Cu * self.scaling_Cu[line] - Al * self.scaling_Al[
                        line] - Bl * self.scaling_Bl[line] - Cl * self.scaling_Cl[
                            line]
                for i in range(N + 1):
                    result += self.params['Amp' + line].value * self.peak(
                        self.transform(x - i * offset) - pos, FWHMG,
                        FWHML) * (poisson**i) / np.math.factorial(i)
                result *= scale

            return result

        def peak(self, x: ArrayLike, FWHMG: float, FWHML: float) -> ArrayLike:
            """Calculates the Voigt profile given the Gaussian
            and Lorentzian FWHM

            Parameters
            ----------
            x : ArrayLike
                Evaluation points
            FWHMG : float
                Gaussian FWHM
            FWHML : float
                Lorentzian FWHM

            Returns
            -------
            ArrayLike
            """
            sigma, gamma = FWHMG / sqrt2log2t2, FWHML / 2
            return voigt_profile(x, sigma, gamma) / voigt_profile(0, sigma, gamma)

        def calcShift(self, I: float, J: float, F: int) -> ArrayLike:
            """Calculate the coefficients for the energy shift due to the hyperfine
            interaction up to the octupole moment. A general equation is used
            so extending to higher orders is possible.

            Parameters
            ----------
            I : float
                Nuclear spin
            J : float
                Electronic spin
            F : int
                Hyperfine level spin

            Returns
            -------
            ArrayLike
                Individual coefficients, in ascending order
            """
            phase = (-1)**(I + J + F)
            contrib = []
            for k in range(1, 4):
                n = float(wigner_6j(I, J, F, J, I, k))
                d = float(
                    wigner_3j(I, k, I, -I, 0, I) * wigner_3j(J, k, J, -J, 0, J))
                shift = phase * n / d
                if not np.isfinite(shift):
                    contrib.append(0)
                else:
                    if k == 1:
                        shift = shift * (I * J)
                    elif k == 2:
                        shift = shift / 4
                    contrib.append(shift)
            return contrib

        def pos(self) -> ArrayLike:
            """Returns the positions of the peaks in MHz in the hyperfine spectrum

            Returns
            -------
            ArrayLike
            """
            centroid = self.params['centroid'].value
            Al = self.params['Al'].value
            Au = self.params['Au'].value
            Bl = self.params['Bl'].value
            Bu = self.params['Bu'].value
            Cl = self.params['Cl'].value
            Cu = self.params['Cu'].value
            pos = []
            for line in self.lines:
                p = centroid + Au * self.scaling_Au[line] + Bu * self.scaling_Bu[
                    line] + Cu * self.scaling_Cu[line] - Al * self.scaling_Al[
                        line] - Bl * self.scaling_Bl[line] - Cl * self.scaling_Cl[
                            line]
                pos.append(p)
            return pos

        def calculateFWHM(self) -> Tuple[float, float]:
            """Calculate the total FWHM of the profiles, with uncertainty,
            taking the correlations into account.

            Returns
            -------
            Tuple[float, float]
                Tuple of the form (value, uncertainty)
            """
            G, Gu = self.params['FWHMG'].value, self.params['FWHMG'].unc
            L, Lu = self.params['FWHML'].value, self.params['FWHML'].unc
            try:
                correl = self.params['FWHMG'].correl['FWHML']
            except KeyError:
                correl = 0
            G, L = unc.correlated_values_norm([(G, Gu), (L, Lu)],
                                              np.array([[1, correl], [correl, 1]]))
            fwhm = 0.5346 * L + (0.2166 * L * L + G * G)**0.5
            return fwhm.nominal_value, fwhm.std_dev