In [1]:
from AtomNumbers import State, TermEnergy, Rb87
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
import itertools
import scipy.constants as consts
from sympy import *
from sympy.physics.wigner import clebsch_gordan


No database file found at data/RMEdbs/Rb87.sqlite3
No database file found at data/RMEdbs/Rb87.sqlite3

In [2]:
def getRelevantCouplings(couplings, forsters, r, ccut, ecut = 1):
    # get the center molecular state (energy is 0 by design)
    cindex = np.abs(forsters).argsort()[0]
    r3 = r**3
    # make the hamiltonian at r
    H=(np.divide(couplings, r3) + np.diag(forsters))
    # do the thing
    e, v = np.linalg.eig(H)
    v = np.asarray(v)
    #print(e)
    # save the data that has coupling greater than ccut
    data = []
    for i in range(len(v)):
        if (abs(v[cindex][i]) >= ccut) and (abs(e[i]) <= ecut):
            data.append([ r, np.real(e[i]), (np.real(v[cindex][i]))**2, np.abs(v[:,i]).argsort()[-5:][::-1] ])
    
    return data #[ r(um), E(r) (GHz), overlap with center state, largest overlaps from basis states ]

In [3]:
# return a list of states that ED couple to the input states
def next_ED_order(m_states):
  newStates=set() # set class maintains uniqueness automatically
  for ms in m_states:
    m_total = ms[0][0].mj + ms[0][1].mj # total angualr momentum is conserved
    for new_ms in itertools.product(coupled_angmom_states(ms[0][0]), coupled_angmom_states(ms[0][1])):
      if (new_ms[0][-1] + new_ms[1][-1]) == m_total:
        newStates.add(new_ms)
  return newStates

In [4]:
# return a set of angular momentum states that could ED couple to the input state
def coupled_angmom_states(state):
  all_states = []
  ll = state.l
  for l in np.arange(abs(ll-1),ll+1.1,2):
    for j in np.arange(abs(l-0.5),l+0.6):
      if abs(j-state.j) <= 1:
        for mj in np.linspace(-j,j,int(2*j+1)):
          if abs(mj-state.mj) <= 1:
            all_states.append((l,j,mj))
  return all_states

In [5]:
coupled_angmom_states(State(97,2,2.5,2.5))


Out[5]:
[(1.0, 1.5, 1.5),
 (3.0, 2.5, 1.5),
 (3.0, 2.5, 2.5),
 (3.0, 3.5, 1.5),
 (3.0, 3.5, 2.5),
 (3.0, 3.5, 3.5)]

In [6]:
ang_mom_list=list(next_ED_order([[(State(97,2,2.5,2.5),State(97,2,2.5,2.5)),5555]]))

In [7]:
ang_mom_list


Out[7]:
[((3.0, 3.5, 2.5), (3.0, 3.5, 2.5)),
 ((3.0, 2.5, 1.5), (3.0, 3.5, 3.5)),
 ((3.0, 2.5, 2.5), (3.0, 3.5, 2.5)),
 ((3.0, 3.5, 3.5), (3.0, 3.5, 1.5)),
 ((3.0, 3.5, 1.5), (3.0, 3.5, 3.5)),
 ((3.0, 2.5, 2.5), (3.0, 2.5, 2.5)),
 ((3.0, 3.5, 3.5), (3.0, 2.5, 1.5)),
 ((3.0, 3.5, 2.5), (3.0, 2.5, 2.5)),
 ((3.0, 3.5, 3.5), (1.0, 1.5, 1.5)),
 ((1.0, 1.5, 1.5), (3.0, 3.5, 3.5))]

In [19]:
# returns list of states passing the n-level, energy-level, and C3 coupling strength cut at r=R0 (C3/R0**3)
def filter_molecular_states(atom, state1, state2, mol_angmon, F_cut, R0, n1_range, n2_range):
  # save states as tuples first, then convert into State classes
  # since I doubt that the set stuff will work correctly with classes
  newStates = set()
  E0 = abs(consts.c*(TermEnergy(atom, state1)[0] + TermEnergy(atom, state2)[0])) # Hz
  for ang_states in mol_angmon:
    #print ang_states
    p = ang_states[0][-1] - state1.mj
    dl1 = state1.l - ang_states[0][0]
    dl2 = state2.l - ang_states[1][0]
    if (abs(p)<=1) and (abs(dl1)==1) and (abs(dl2)==1):   
      # rme database lookups
      l1b, j1b, mj1b = ang_states[0]
      l2b, j2b, mj2b = ang_states[1]
      # precalculate angular factors
      ang_factor = atom.AME(state1, State(0,l1b,j1b,mj1b), p)
      ang_factor *= atom.AME(state2, State(0,l2b,j2b,mj2b), -p)
      ang_factor *= clebsch_gordan(1,1,2,p,-p,0)
      # calculate threshold to add molecular state to the basis
      if ang_factor != 0:
        cut = abs(N(F_cut*R0**3/(ang_factor*atom.GHz_um3_factor*sqrt(6.0))))
        #print("cut: {}".format(cut))
        # now check all permutations of the n-levels in the range
        s1bs = [ [ (x[0],)+ang_states[0], x[1]] for x in atom.RMEs(state1,n1_range,l1b,j1b) ]
        s2bs = [ [ (x[0],)+ang_states[1], x[1]] for x in atom.RMEs(state2,n2_range,l2b,j2b) ]
        for ms in itertools.product(s1bs,s2bs):
          state1b = State(ms[0][0][0],ms[0][0][1],ms[0][0][2],ms[0][0][3])
          state2b = State(ms[1][0][0],ms[1][0][1],ms[1][0][2],ms[1][0][3])
          # check the energy first since its easier
          E = abs(consts.c*(TermEnergy(atom, state1b)[0] + TermEnergy(atom, state2b)[0])) # Hz
               
          if abs(ms[0][1]*ms[1][1]/((E-E0)*1e-9)) > cut:
            #print("E (GHz): {}".format((E-E0)*1e-9))
            newStates.add((ms[0][0],ms[1][0]))
  return newStates

In [20]:
if __name__ == "__main__":
  sI=State(97,2,2.5,2.5)
  print(filter_molecular_states(Rb87, sI, sI, ang_mom_list, 1, 3, (80,120),(80,120)))


((3.0, 3.5, 2.5), (3.0, 3.5, 2.5))
cut: 113075929.837376
((3.0, 2.5, 1.5), (3.0, 3.5, 3.5))
cut: 468178828.628314
((3.0, 2.5, 2.5), (3.0, 3.5, 2.5))
cut: 276978330.292323
((3.0, 3.5, 3.5), (3.0, 3.5, 1.5))
cut: 296102290.147023
((3.0, 3.5, 1.5), (3.0, 3.5, 3.5))
cut: 296102290.147023
((3.0, 2.5, 2.5), (3.0, 2.5, 2.5))
cut: 678455579.024255
((3.0, 3.5, 3.5), (3.0, 2.5, 1.5))
cut: 468178828.628314
((3.0, 3.5, 2.5), (3.0, 2.5, 2.5))
cut: 276978330.292323
((3.0, 3.5, 3.5), (1.0, 1.5, 1.5))
cut: 66882689.8040448
E (GHz): -0.148413459858
E (GHz): -0.236731656209
((1.0, 1.5, 1.5), (3.0, 3.5, 3.5))
cut: 66882689.8040448
E (GHz): -0.236731656209
E (GHz): -0.148413459858
set([((98, 1.0, 1.5, 1.5), (96, 3.0, 3.5, 3.5)), ((96, 3.0, 3.5, 3.5), (98, 1.0, 1.5, 1.5)), ((99, 1.0, 1.5, 1.5), (95, 3.0, 3.5, 3.5)), ((95, 3.0, 3.5, 3.5), (99, 1.0, 1.5, 1.5))])

In [ ]: