atoms.py

from pandas import DataFrame
import pandas as pd
from parsers import term_frac


class EnergyLevel:
    # TODO: isotope shifts?
    def __init__(self, df, df_index, name='term', I=0.0, A_coeff=0.0, B_coeff=0.0, C_coeff=0.0):
        self._initialized = False
        # get parameters from datafile
        my_row = df.iloc[df_index]

        (self.configuration, self.term, self.S, self.L, self.K,
         self.j1, self.j2, self.J, self.parity, self.level, self.lande) = my_row.values
        if name == 'term':
            self.name = self.term + term_frac(self.J)
        elif name == 'full':
            self.name = self.configuration + ' ' + self.term + term_frac(self.J)
        else:
            self.name = name
        self.A_coeff = A_coeff
        self.B_coeff = B_coeff
        self.C_coeff = C_coeff
        self.coupling = self.get_coupling()
        self.I = I
        self.Fs = self.get_Fs()
        self.hf_levels, self.hf_shifts = self.get_hyperfine_data()
        self.z_levels, self.z_shifts = self.get_zeeman_data()

        self._initialized = True

    def __setattr__(self, key, value):
        self.__dict__[key] = value
        if self._initialized:
            if key in ['level', 'lande', 'A_coeff', 'B_coeff', 'C_coeff']:
                self.hf_levels, self.hf_shifts = self.get_hyperfine_data()
            if key in ['level', 'lande']:
                self.z_levels, self.z_shifts = self.get_zeeman_data()

    def get_Fs(self):
        Fs = []
        F = abs(self.I - self.J)
        while F <= abs(self.I + self.J):
            Fs.append(F)
            F += 1
        return Fs

    def get_coupling(self):
        if self.L is not None:
            coupling = 'LS'
        elif self.K is not None:
            coupling = 'JK'
        else:
            coupling = 'jj'
        return coupling

    def get_hyperfine_data(self):
        hf_levels = {}
        hf_shifts = {}
        I, J = self.I, self.J
        if I == 0:
            hf_levels[J] = self.level
            hf_shifts[J] = 0
        else:
            for F in self.Fs:
                #  The code here follows the results on page 14 of Matthew Hoffman's thesis
                IdotJ = 0.5 * (F*(F + 1) - J*(J + 1) - I*(I + 1))

                FM1 = IdotJ

                if J <= 0.5 or I <= 0.5:
                    FE2 = 0
                else:
                    FE2 = (3*IdotJ**2 + 1.5*IdotJ - I*(I+1)*J*(J+1))/(2.0*I*(2.0*I-1.0)*J*(2.0*J-1.0))

                if J <= 1 or I <= 1:
                    FM3 = 0
                else:
                    FM3 = (10*IdotJ**3 + 20*IdotJ**2 + 2*IdotJ*(-3*I*(I+1)*J*(J+1)+I*(I+1)+J*(J+1)+3)
                           - 5*I*(I+1)*J*(J+1))/(I*(I-1)*J*(J-1)*(2*J-1))

                shift = self.A_coeff*FM1 + self.B_coeff*FE2 + self.C_coeff*FM3
                hf_shifts[F] = shift
                hf_levels[F] = shift + self.level
        return hf_levels, hf_shifts

    def get_zeeman_data(self):
        z_levels = {}
        z_shifts = {}
        for F in self.Fs:
            try:
                g_F = self.lande * (F * (F + 1) + self.J * (self.J + 1) - self.I * (self.I + 1)) / (2 * F * (F + 1))
            except ZeroDivisionError:
                g_F = 0
            z_levels[F] = {}
            z_shifts[F] = {}
            m_F = -F
            while m_F <= F:
                z_shifts[F][m_F] = g_F * m_F * 1.4e-6
                z_levels[F][m_F] = self.hf_levels[F] + z_shifts[F][m_F]
                m_F += 1
        return z_levels, z_shifts

    def set_name(self, name):
        """Sets the name of the level. Can be 'term', 'full', or a custom name"""
        if name == 'term':
            self.name = self.term + term_frac(self.J)
        elif name == 'full':
            self.name = self.configuration + ' ' + self.term + term_frac(self.J)
        else:
            self.name = name

    def data_table(self, hf=True, zeeman=True, b_field=0):
        table = DataFrame(columns=['configuration', 'term', 'level', 'I',
                                   'L', 'S', 'J', 'F', 'm_F', 'J_frac', 'F_frac', 'm_F_frac', 'hf', 'z'])
        if not hf:
            line = DataFrame(data={'configuration': [self.configuration], 'term': [self.term], 'level': self.level, 'I': self.I,
                                   'L': [self.L], 'S': [self.S], 'J': [self.J], 'F': [None], 'm_F': [None],
                                   'J_frac': [term_frac(self.J)], 'F_frac': [None], 'm_F_frac': [None],
                                   'hf': [0.0], 'z': [0.0]})
            table = table.append(line, ignore_index=True)
        else:
            for F in self.Fs:
                if not zeeman:
                    hyperfine = self.hf_shifts[F]
                    level = self.hf_levels[F]
                    line = DataFrame(data={'configuration': [self.configuration], 'term': [self.term], 'level': level, 'I': self.I,
                                           'L': [self.L], 'S': [self.S], 'J': [self.J], 'F': [F], 'm_F': [None],
                                           'J_frac': [term_frac(self.J)], 'F_frac': [term_frac(F)], 'm_F_frac': [None],
                                           'hf': [hyperfine], 'z': [0.0]})
                    table = table.append(line, ignore_index=True)
                else:
                    for m_F in self.z_shifts[F].keys():
                        z = self.z_shifts[F][m_F]
                        hyperfine = self.hf_shifts[F]
                        level = self.level + self.hf_shifts[F] + self.z_shifts[F][m_F]*b_field
                        line = DataFrame(
                            data={'configuration': [self.configuration], 'term': [self.term], 'level': level, 'I': self.I,
                                  'L': [self.L], 'S': [self.S], 'J': [self.J], 'F': [F], 'm_F': [m_F],
                                  'J_frac': [term_frac(self.J)], 'F_frac': [term_frac(F)], 'm_F_frac': [term_frac(m_F)],
                                  'hf': [hyperfine], 'z': [z]})
                        table = table.append(line, ignore_index=True)
        return table


class Transition:
    def __init__(self, level_0, F_0, m_F_0, level_1, F_1, m_F_1):
        self.level_0, self.level_1 = level_0, level_1
        self.F_0, self.m_F_0, self.F_1, self.m_F_1 = F_0, m_F_0, F_1, m_F_1
        self.J_0, self.J_1 = level_0.J, level_1.J
        self.L_0, self.L_1 = level_0.L, level_1.L
        self.parity_0, self.parity_1 = level_0.parity, level_1.parity
        self.name = level_0.name + str(F_0) + str(m_F_0) + '->' + level_1.name + str(F_1) + str(m_F_1)

    def data_table(self):

        table_0 = self.level_0.data_table()
        table_1 = self.level_1.data_table()
        line_0 = table_0.loc[(table_0['m_F'] == self.m_F_0) & (table_0['F'] == self.F_0)]
        line_1 = table_1.loc[(table_1['m_F'] == self.m_F_1) & (table_1['F'] == self.F_1)]
        line_0 = line_0.drop(['J_frac', 'F_frac', 'm_F_frac'], axis=1)
        line_1 = line_1.drop(['J_frac', 'F_frac', 'm_F_frac'], axis=1)
        if line_0.empty or line_1.empty:
            raise ValueError("The selected quantum numbers don't yield a transition."
                             " Check that they exist in the specified levels")
        line_0.columns = [str(col) + '_0' for col in line_0.columns]
        line_1.columns = [str(col) + '_1' for col in line_1.columns]
        line_0 = line_0.reset_index(drop=True)
        line_1 = line_1.reset_index(drop=True)

        data_table = pd.concat([line_0, line_1], axis=1, ignore_index=False)

        delta_l = abs(data_table['level_0'] - data_table['level_1'])
        wavelength = 299792.458/delta_l

        data_table['delta_l'] = [delta_l]
        data_table['wavelength'] = [wavelength]
        data_table['name'] = [self.name]
        return data_table

    def get_type(self):
        # TODO: finish this method sometime when I have time to figure out what all the relevant selection rules are
        d_m = self.m_F_1 - self.m_F_0
        d_F = self.F_1 - self.F_0
        d_J = self.J_1 - self.J_0

        transition_type = 'unknown'
        if abs(d_J) <= 1 and not (self.J_0 == 0 and self.J_1 == 0):
            if abs(d_F) <= 1:
                if self.parity_0 != self.parity_1:
                    transition_type = 'E1'

        return 'dm: {}, dF: {}, dJ:{}'.format(d_m, d_F, d_J), transition_type


class Atom:
    def __init__(self, name, levels=(), transitions=()):
        self.name = name
        self.levels = {}
        for level in levels:
            if level.name in self.levels.keys():
                raise Exception('{} is already a level in this atom. '.format(level.name))
            self.levels[level.name] = level

        for lvl in self.levels.values():
            self.I = lvl.I
            print("ahere")
            break

        self.transitions = {}
        for transition in transitions:
            self.transitions[transition.name] = transition
        if len(levels) > 0:
            self.rezero()

    def rezero(self):
        """Takes the minimum hyperfine level in the atom and sets it to zero, shifting all others appropriately"""
        gs = list(self.levels.values())[0]
        gs_level = min(gs.hf_levels.values())
        for state in self.levels.values():
            if state.level < gs_level:
                gs = state
                gs_level = min(gs.hf_levels.values())
        for state in self.levels.values():
            state.level -= gs_level

    def add_level(self, levels):
        for level in levels:
            self.I = level.I
            self.levels[level.name] = level
        self.rezero()

    def remove_level(self, mode='keys', *levels):
        if mode == 'keys':
            for key in levels:
                self.levels.pop(key)
        elif mode == 'levels':
            for level in levels:
                self.levels.pop(level.name)
        else:
            raise ValueError('Unrecognized mode')
        self.rezero()

    def add_transition(self, transitions):
        for transition in transitions:
            self.transitions[transition.name] = transition

    def remove_transition(self, mode='keys', *transitions):
        if mode == 'keys':
            for key in transitions:
                self.transitions.pop(key)
        elif mode == 'transitions':
            for transition in transitions:
                self.transitions.pop(transition.name)
        else:
            raise ValueError('Unrecognized mode')


if __name__ == '__main__':
    from parsers import parse_NIST_levels

    df = parse_NIST_levels('YbII_NIST_levels.csv')

    S12_171 = EnergyLevel(df, 0, I=0.5, A_coeff=12.645e-3)  # 2S1/2
    P12_171 = EnergyLevel(df, 8, I=0.5, A_coeff=2.1079e-3)  # 2P1/2
    cycling_171 = Transition(level_0=S12_171, F_0=0, m_F_0=0, level_1=P12_171, F_1=1, m_F_1=0)
    print(cycling_171.name)