# -*- coding: utf-8 -*-
"""
qspec.models._splitter
Created on 14.03.2022
@author: Patrick Mueller
Splitter classes for lineshape models.
"""
import os
from string import ascii_uppercase
import numpy as np
from qspec.tools import merge_intervals
from qspec.physics import get_f
from qspec.algebra import wigner_6j, a, b, c
from qspec.models._base import Model, Summed
from qspec.models._spectrum import LorentzQI
__all__ = ['gen_splitter_model', 'get_all_f', 'hf_coeff', 'hf_trans', 'hf_shift', 'hf_int', 'Splitter',
'SplitterSummed', 'Hyperfine', 'HyperfineQI', 'HyperfineMixed']
[docs]def gen_splitter_model(qi: bool = False, hf_mixing: bool = False):
if qi and hf_mixing:
pass
elif qi:
return HyperfineQI
elif hf_mixing:
return HyperfineMixed
else:
return Hyperfine
raise ValueError('Specified splitter model not available.')
[docs]def get_all_f(i, j):
i, j = np.array(i, float).flatten(), np.array(j, float).flatten()
return sorted(set(f + abs(_i - _j) for _i in i for _j in j for f in range(int(_i + _j - abs(_i - _j) + 1))))
[docs]def hf_coeff(i, j, f):
""" Return the tuple of hyperfine coefficients for A and B-factor for a given quantum state """
# First and third order are taken from https://journals.aps.org/pra/abstract/10.1103/PhysRevA.103.032826.
# Second order from https://link.springer.com/referencework/10.1007/978-0-387-26308-3.
if i < 0.5 or j < 0.5:
return tuple()
# Magnetic dipole, the corresponding hyperfine constant is A = mu / (IJ) * <T_1>.
k = f * (f + 1) - i * (i + 1) - j * (j + 1)
co_a = 0.5 * k
if i < 1 or j < 1:
return co_a,
# Electric quadrupole, the corresponding hyperfine constant is B = 2eQ * <T_2>.
co_b = (0.75 * k * (k + 1) - j * (j + 1) * i * (i + 1)) / (2 * i * (2 * i - 1) * j * (2 * j - 1))
if i < 1.5 or j < 1.5:
return co_a, co_b
# Magnetic octupole, the corresponding hyperfine constant is C = -Omega * <T_3>.
co_c = k ** 3 + 4 * k ** 2 + 0.8 * k * (-3 * i * (i + 1) * j * (j + 1) + i * (i + 1) + j * (j + 1) + 3) \
- 4 * i * (i + 1) * j * (j + 1)
co_c /= i * (i - 1) * (2 * i - 1) * j * (j - 1) * (2 * j - 1)
co_c *= 1.25
if i < 2 or j < 2:
return co_a, co_b, co_c
return co_a, co_b, co_c # Highest implemented order.
[docs]def hf_trans(i, j_l, j_u):
"""
Calculate all allowed hyperfine transitions and their hyperfine coefficients.
Returns (f_l, f_u, coAl, coBl, coAu, coBu)
"""
return [[(f_l, f_u), hf_coeff(i, j_l, f_l), hf_coeff(i, j_u, f_u)]
for f_l in get_f(i, j_l) for f_u in get_f(i, j_u)
if abs(f_u - f_l) < 1.1 and not f_u == f_l == 0]
[docs]def hf_shift(hyper_l, hyper_u, coeff_l, coeff_u):
"""
:param hyper_l: The hyperfine structure constants of the lower state (Al, Bl, Cl, ...).
:param hyper_u: The hyperfine structure constants of the upper state (Au, Bu, Cu, ...).
:param coeff_l: The coefficients of the lower state to be multiplied by the constants (coAl, coBl, coCl, ...).
:param coeff_u: The coefficients of the lower state to be multiplied by the constants (coAu, coBu, coCu, ...).
:returns: The hyperfine structure shift of an optical transition.
"""
return sum(const * coeff for const, coeff in zip(hyper_u, coeff_u)) \
- sum(const * coeff for const, coeff in zip(hyper_l, coeff_l))
[docs]def hf_int(i, j_l, j_u, transitions):
""" Calculate relative line intensities """
return [np.around((2 * f_u + 1) * (2 * f_l + 1) * (wigner_6j(j_l, f_l, i, f_u, j_u, 1, as_sympy=False) ** 2),
decimals=9) for (f_l, f_u), *r in transitions]
[docs]class Splitter(Model):
def __init__(self, model, i, j_l, j_u, name):
super().__init__(model=model)
self.type = 'Splitter'
self.i = i
self.j_l = j_l
self.j_u = j_u
self.name = name
self.racah_indices = []
self.racah_intensities = []
[docs] def racah(self):
for i, intensity in zip(self.racah_indices, self.racah_intensities):
self.vals[i] = intensity
[docs]class SplitterSummed(Summed):
def __init__(self, splitter_models):
if any(not isinstance(model, Splitter) for model in splitter_models):
raise TypeError('All models passed to \'SplitterSummed\' must have type \'Splitter\'.')
super().__init__(splitter_models, labels=['({})'.format(model.name) for model in splitter_models]
if len(splitter_models) > 1 else None)
[docs] def racah(self):
i0 = 0
for model in self.models:
for i, intensity in zip(model.racah_indices, model.racah_intensities):
self.set_val(i0 + i, intensity, force=True)
i0 += model.size
self.set_vals(self.vals, force=True)
[docs]class Hyperfine(Splitter):
def __init__(self, model, i, j_l, j_u, name):
super().__init__(model, i, j_l, j_u, name)
self.type = 'Hyperfine'
self.transitions = hf_trans(self.i, self.j_l, self.j_u)
self.racah_intensities = hf_int(self.i, self.j_l, self.j_u, self.transitions)
self.n_l = len(self.transitions[0][1])
self.n_u = len(self.transitions[0][2])
for i in range(self.n_l):
self._add_arg('{}l'.format(ascii_uppercase[i]), 0., False, False)
for i in range(self.n_u):
self._add_arg('{}u'.format(ascii_uppercase[i]), 0., False, False)
for i in range(min([self.n_l, self.n_u])):
fix = '{} / {}'.format('{}u'.format(ascii_uppercase[i]), '{}l'.format(ascii_uppercase[i]))
self._add_arg('{}_ratio'.format(ascii_uppercase[i]), 0., fix, False)
for i, (t, intensity) in enumerate(zip(self.transitions, self.racah_intensities)):
self.racah_indices.append(self._index)
self._add_arg('int({}, {})'.format(t[0][0], t[0][1]), intensity, i == 0, False)
[docs] def evaluate(self, x, *args, **kwargs):
const_l = tuple(args[self.model.size + i] for i in range(self.n_l))
const_u = tuple(args[self.model.size + self.n_l + i] for i in range(self.n_u))
return np.sum([args[i] * self.model.evaluate(x - hf_shift(const_l, const_u, t[1], t[2]), *args, **kwargs)
for i, t in zip(self.racah_indices, self.transitions)], axis=0)
[docs] def min(self):
const_l = tuple(self.vals[self.model.size + i] for i in range(self.n_l))
const_u = tuple(self.vals[self.model.size + self.n_l + i] for i in range(self.n_u))
return self.model.min() + min(hf_shift(const_l, const_u, t[1], t[2]) for t in self.transitions)
[docs] def max(self):
const_l = tuple(self.vals[self.model.size + i] for i in range(self.n_l))
const_u = tuple(self.vals[self.model.size + self.n_l + i] for i in range(self.n_u))
return self.model.max() + max(hf_shift(const_l, const_u, t[1], t[2]) for t in self.transitions)
[docs] def intervals(self):
const_l = tuple(self.vals[self.model.size + i] for i in range(self.n_l))
const_u = tuple(self.vals[self.model.size + self.n_l + i] for i in range(self.n_u))
shifts = [hf_shift(const_l, const_u, t[1], t[2]) for t in self.transitions]
return merge_intervals([[self.model.min() + shift, self.model.max() + shift] for shift in shifts])
[docs] def racah(self):
for i, intensity in zip(self.racah_indices, self.racah_intensities):
self.vals[i] = intensity
def load_qi(filepath):
if not os.path.isfile(filepath):
return None
with open(os.path.join(filepath), 'r') as file:
ret = eval(file.read())
return ret['a'], ret['b'], ret['c']
[docs]class HyperfineQI(Splitter):
def __init__(self, model, i, j_l, j_u, name, qi_path=None):
super().__init__(LorentzQI(), i, j_l, j_u, name)
self.type = 'HyperfineQI'
self.qi_path = qi_path
self.file = f'qi_{name}.txt'
self.transitions = hf_trans(self.i, self.j_l, self.j_u)
self.racah_intensities = [0., ]
save = False
if self.qi_path is None or not os.path.isfile(os.path.join(self.qi_path, self.file)):
print(f'Calculating QI A-matrix of isotope {name} ... ')
self.a_qi = [a(self.i, self.j_l, f_l, self.j_u, f_u, as_sympy=False)
for f_l in get_f(self.i, self.j_l) for f_u in get_f(self.i, self.j_u)
if abs(f_u - f_l) < 1.1 and not f_u == f_l == 0]
print(f'Calculating QI B-matrix of isotope {name} ... ')
self.b_qi = [b(self.i, self.j_l, f_l, self.j_u, f_u, as_sympy=False)
for f_l in get_f(self.i, self.j_l) for f_u in get_f(self.i, self.j_u)
if abs(f_u - f_l) < 1.1 and not f_u == f_l == 0]
print(f'Calculating QI C-matrix of isotope {name} ... ')
self.c_qi = [[c(self.i, self.j_l, f_l, self.j_u, f1_u, f2_u, as_sympy=False)
for i2, f2_u in enumerate(get_f(self.i, self.j_u)) if i2 > i1
if abs(f1_u - f_l) < 1.1 and abs(f2_u - f_l) < 1.1
and not f1_u == f_l == 0 and not f2_u == f_l == 0]
for f_l in get_f(self.i, self.j_l) for i1, f1_u in enumerate(get_f(self.i, self.j_u))]
if qi_path is not None:
save = True
else:
self.a_qi, self.b_qi, self.c_qi = load_qi(os.path.join(self.qi_path, self.file))
if save:
qi_dict = {'a': self.a_qi, 'b': self.b_qi, 'c': self.c_qi}
with open(os.path.join(self.qi_path, self.file), 'w') as file:
file.write(str(qi_dict))
self.transitions_qi = [[[(f_l, f2_u), hf_coeff(self.i, self.j_l, f_l), hf_coeff(self.i, self.j_u, f2_u)]
for i2, f2_u in enumerate(get_f(self.i, self.j_u)) if i2 > i1
if abs(f1_u - f_l) < 1.1 and abs(f2_u - f_l) < 1.1
and not f1_u == f_l == 0 and not f2_u == f_l == 0]
for f_l in get_f(self.i, self.j_l) for i1, f1_u in enumerate(get_f(self.i, self.j_u))]
self.n_l = len(self.transitions[0][1])
self.n_u = len(self.transitions[0][2])
for i in range(self.n_l):
self._add_arg('{}l'.format(ascii_uppercase[i]), 0., False, False)
for i in range(self.n_u):
self._add_arg('{}u'.format(ascii_uppercase[i]), 0., False, False)
for i in range(min([self.n_l, self.n_u])):
fix = '{} / {}'.format('{}u'.format(ascii_uppercase[i]), '{}l'.format(ascii_uppercase[i]))
self._add_arg('{}_ratio'.format(ascii_uppercase[i]), 0., fix, False)
self.racah_indices.append(self._index)
self._add_arg('geo', 0., False, False)
[docs] def evaluate(self, x, *args, **kwargs):
const_l = tuple(args[self.model.size + i] for i in range(self.n_l))
const_u = tuple(args[self.model.size + self.n_l + i] for i in range(self.n_u))
return np.sum([(_a + _b * args[self.racah_indices[0]])
* self.model.evaluate(x - hf_shift(const_l, const_u, t[1], t[2]), *args, **kwargs)
+ np.sum([_c * args[self.racah_indices[0]]
* self.model.evaluate_qi(x - hf_shift(const_l, const_u, t[1], t[2]),
hf_shift(const_l, const_u, _t[1], _t[2])
- hf_shift(const_l, const_u, t[1], t[2]), *args, **kwargs)
for _t, _c in zip(t_list, c_list)], axis=0)
for t, _a, _b, t_list, c_list in zip(
self.transitions, self.a_qi, self.b_qi, self.transitions_qi, self.c_qi)], axis=0) / np.max(self.a_qi)
[docs] def min(self):
const_l = tuple(self.vals[self.model.size + i] for i in range(self.n_l))
const_u = tuple(self.vals[self.model.size + self.n_l + i] for i in range(self.n_u))
return self.model.min() + min(hf_shift(const_l, const_u, t[1], t[2]) for t in self.transitions)
[docs] def max(self):
const_l = tuple(self.vals[self.model.size + i] for i in range(self.n_l))
const_u = tuple(self.vals[self.model.size + self.n_l + i] for i in range(self.n_u))
return self.model.max() + max(hf_shift(const_l, const_u, t[1], t[2]) for t in self.transitions)
[docs] def intervals(self):
const_l = tuple(self.vals[self.model.size + i] for i in range(self.n_l))
const_u = tuple(self.vals[self.model.size + self.n_l + i] for i in range(self.n_u))
shifts = [hf_shift(const_l, const_u, t[1], t[2]) for t in self.transitions]
return merge_intervals([[self.model.min() + shift, self.model.max() + shift] for shift in shifts])
[docs] def racah(self):
self.vals[self.racah_indices[0]] = 0.
[docs]class HyperfineMixed(Splitter):
"""
Hyperfine-mixing model based on https://doi.org/10.1103/PhysRevA.55.2728 [1].
"""
def __init__(self, model, i, j_l, j_u, name, config):
super().__init__(model, i, j_l, j_u, name)
self.type = 'HyperfineMixed'
self.config = config
self.states = ['l', 'u']
to_mhz = 13074.70
self.enabled = {s: self.config['enabled_{}'.format(s)] for s in self.states}
self.J = {s: np.array(self.config['J{}'.format(s)]) for s in self.states}
self.F = {s: get_all_f(self.i, self.J[s]) for s in self.states}
self.T = {s: np.array(self.config['T{}'.format(s)], dtype=float) * to_mhz for s in self.states}
self.M = 0. if self.i == 0 else np.sqrt((2 * self.i + 1) * (self.i + 1) / self.i) * self.config['mu']
self.fs = {s: np.array(self.config['f{}'.format(s)]).flatten() for s in self.states}
self.mask_J = {s: [np.array([i for i, j in enumerate(self.config['J{}'.format(s)])
if abs(self.i - j) - 0.1 < f < self.i + j + 0.1], dtype=int) for f in self.F[s]]
for s in self.states}
self.W = {s: [np.array([[(-1) ** (self.i + self.J[s][i] + f)
* wigner_6j(self.i, self.J[s][i], f, self.J[s][j], self.i, 1)
for j in self.mask_J[s][k]] for i in self.mask_J[s][k]], dtype=float)
* self.T[s][np.ix_(self.mask_J[s][k], self.mask_J[s][k])] * self.M
for k, f in enumerate(self.F[s])] for s in self.states} # Eq. (2.5) in [1] without diagonal.
self.transitions = [[(_m0, _m1), (self.J['l'][_m0], self.J['u'][_m1]), (f0, f1),
hf_coeff(self.i, self.J['l'][_m0], f0),
hf_coeff(self.i, self.J['u'][_m1], f1)]
for i, (m0, f0) in enumerate(zip(self.mask_J['l'], self.F['l']))
for j, (m1, f1) in enumerate(zip(self.mask_J['u'], self.F['u']))
for _m0 in m0 for _m1 in m1
if abs(f1 - f0) < 1.1 and abs(self.J['u'][_m1] - self.J['l'][_m0]) < 1.1]
self.wt_map = {s: np.array([[(1 - 2 * int(s == 'l')) * int(t[0][int(s == 'u')] == _m
and t[2][int(s == 'u')] == f)
for m, f in zip(self.mask_J[s], self.F[s]) for _m in m]
for t in self.transitions], dtype=float) for s in self.states}
self.racah_intensities = [a(self.i, t[1][0], t[2][0], t[1][1], t[2][1]) for t in self.transitions]
self.order = {s: [int(min((j // 0.5, self.i // 0.5))) for j in self.J[s]] for s in self.states}
self.hf_args_map = [[] for _ in self.transitions]
for s in self.states:
if self.enabled[s] and self.i > 0:
for i, j in enumerate(self.J[s]):
self._add_arg('FS_{}{}({})'.format(s, i, j), self.fs[s][i], i == 0, False)
else:
for i, j in enumerate(self.J[s]):
for k in range(self.order[s][i]):
for t, m in zip(self.transitions, self.hf_args_map):
if t[0][int(s == 'u')] == i:
m.append(self._index)
self._add_arg('{}_{}{}({})'.format(ascii_uppercase[k], s, i, j), 0., False, False)
for i, (t, intensity) in enumerate(zip(self.transitions, self.racah_intensities)):
self.racah_indices.append(self._index)
self._add_arg('int{}([{}, {}] -> [{}, {}])'.format(i, t[1][0], t[2][0], t[1][1], t[2][1]),
intensity, i == 0, False)
[docs] def x0(self, *args):
x0 = np.zeros(len(self.transitions))
for s in self.states:
if self.enabled[s]:
_x0 = [np.linalg.eigh( # Eigenvalues and eigenvectors of Eq. (2.5) in [1].
np.diag([args[self.p['FS_{}{}({})'.format(s, _m, self.J[s][_m])]] for _m in m]) + w)
for m, w in zip(self.mask_J[s], self.W[s])]
# Eigenvalues are returned in ascending order! Sort based on eigenvectors.
# Invert order
# a = [a, b, c]
# order = [2, 0, 1]
# desired a = [b, c, a]
# required order = [1, 2, 0]
orders = [np.argmax(np.abs(w[1]), axis=0) for w in _x0]
orders = [np.array([j for i in range(order.size) for j, o in enumerate(order) if i == o])
for w, order in zip(_x0, orders)]
_x0 = np.concatenate(tuple(w[0][order] for w, order in zip(_x0, orders)))
_x0 = self.wt_map[s] @ _x0
else:
_x0 = np.array([sum(args[i] * coeff for i, coeff in zip(m, t[3 + int(s == 'u')]))
for m, t in zip(self.hf_args_map, self.transitions)])
_x0 *= (1 - 2 * int(s == 'l'))
x0 += _x0
# print(f'[{", ".join([str(_x0) for _x0 in x0])}]')
return x0
[docs] def evaluate(self, x, *args, **kwargs):
return np.sum([args[i] * self.model.evaluate(x - _x0, *args, **kwargs)
for i, _x0 in zip(self.racah_indices, self.x0(*args))], axis=0)
[docs] def min(self):
return self.model.min() + np.min(self.x0(*self.vals))
[docs] def max(self):
return self.model.max() + np.max(self.x0(*self.vals))
[docs] def intervals(self):
return merge_intervals([[self.model.min() + _x0, self.model.max() + _x0] for _x0 in self.x0(*self.vals)])
[docs] def racah(self):
for i, intensity in zip(self.racah_indices, self.racah_intensities):
self.vals[i] = intensity