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)


[docs]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
[docs] 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
[docs] 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
[docs] 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)
[docs] 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
[docs] 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
[docs] 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