Diatomic-py: A python module for calculating the rotational and hyperfine structure of $^1\Sigma$ molecules

We present a computer program to calculate the quantised rotational and hyperfine energy levels of $^{1}\Sigma $ diatomic molecules in the presence of dc electric, dc magnetic, and off-resonant optical fields. Our program is applicable to the bialkali molecules used in ongoing state-of-the-art experiments with ultracold molecular gases. We include functions for the calculation of space-fixed electric dipole moments, magnetic moments and transition dipole moments.


Introduction
A new generation of experiments is now able to produce ultracold 1 Σ ground-state molecules by associating pairs of alkali atoms [1,2,3,4,5,6,7,8,9,10,11,12,13,14]. Precise knowledge of the rotational and hyperfine structure is critical for nearly all the foreseeable applications of these molecules (see e.g. [15]), and particularly for those that require long-lived quantum coherences between states [16,17,18,19,20,21]. Basic calculation of the rotational structure of diatomic molecules is well understood [22], but the addition of hyperfine structure, together with external fields that prevent conservation of total angular momentum, greatly increases the size of the Hilbert space over which calculations must be performed. Experimental studies of ultracold molecules often involve external dc magnetic and electric fields and often also need an off-resonant optical field for trapping. There is a need in the ultracold-molecule community for accessible and open-source tools to perform these calculations.
In this work we present a flexible Python-based program for calculating the rotational and hyperfine structure of 1 Σ molecules in external electromagnetic fields. Our code contains a module to automate the construction of the Hamiltonian, which is then diagonalized using functions from the numpy stack [23,24]. We also include functions that simplify the calculation of important quantities such as the electric dipole moments for transitions between pairs of states. We give example plots that demonstrate use of the code to calculate Zeeman, dc Stark, and ac Stark maps for the hyperfine states of 1 Σ molecules.

Theoretical background
The model we consider is valid for diatomic molecules in a single vibrational state of a 1 Σ electronic state, as is relevant to most current experiments with ultracold bialkali molecules. For generality, we label the molecule AB, where A and B are the component nuclei. The Hamiltonian of such a molecule is [25] H where H rot describes the rotational structure, H hf describes the hyperfine structure and H ext describes the interaction between the molecule and external fields.
We construct the Hamiltonian in the fully uncoupled basis where the relevant angular momentum quantum numbers are the molecule's rotational angular momentum (N ) and the spin of the two nuclei (i A , i B ). Each of these angular momenta also has a projection onto the z axis of our coordinate system, defined by the direction of the magnetic field, giving six quantum numbers in total: N , M N , i A , m A , i B , m B . As we do not consider effects that could change the nuclear spin (only its projection) where relevant we write our basis states as |N, M N , m A , m B for brevity. We also define the vector operators N , I A and I B that are associated with each of these angular momenta.
In the following subsections, we give the mathematical expressions that describe each of the components of the Hamiltonian 1.

The rotational Hamiltonian, H rot
We describe the rotation of the molecule using a rigid-rotor model; this results in a spectrum of rotational states that have energy E N ≈ B v N (N +1), where B v is the rotational constant of the molecule in vibrational state v. In terms of angular momentum operators, the rotational contribution to the Hamiltonian is [22] This includes a small correction for centrifugal distortion, characterised by the distortion coefficient D v .

The hyperfine Hamiltonian, H hf
The hyperfine component of the Hamiltonian is necessary for molecules with non-zero nuclear spin. It describes interactions between the angular momentum of the two nuclei, and between each nucleus and the rotational angular momentum of the molecule. We can split this component further into terms that describe four different interactions, The first term H quad describes the nuclear electric quadrupole interaction and is written as where eQ j and q j are spherical tensors of rank 2 representing the nuclear quadrupole moment and electric field gradient (at the position of the nucleus) of nucleus j. This component of the Hamiltonian is governed by the molecular coupling constants (eQq) A and (eQq) B that contain both the magnitude of the nuclear electric quadrupole moments and the electric field gradients relevant for nuclei A and B, respectively; note this term is non-zero only when i j > 1/2. The second and third terms describe scalar and tensor interactions between the nuclear spins, This depends on the molecular constants c 3 and c 4 , and a pair of second-rank tensors T 2 that describe the angular dependence of the interactions. In the above the tensor T 2 (C) has components given by T 2 q (C) = C 2 q (θ, φ), where C 2 q is a Racah-normalised 1 spherical harmonic of order 2. The arguments of each component are the polar angle θ of the molecule's internuclear axis (n) from z and the azimuthal angle φ. The Racah-normalised forms of the spherical harmonics are related to the L 2 -normalised versions (Y k q ) through C 2 q (θ, φ) = 4π/5Y 2 q (θ, φ). The tensor T 2 (I A , I B ) is the second-rank tensor product of the two vec- tors I A and I B . This can be written as where I j q are the components of I j . The final term describes spin-rotation interactions that arise due to the magnetic moment of each nucleus interacting with the magnetic field generated by the rotating molecule. This term is given by where c j is the coupling constant for nucleus j.

Interaction between the molecule and external fields, H ext
To interpret current experiments with ultracold molecules, we need to calculate the internal structure in the presence of external electromagnetic fields. We further decompose the component of the Hamiltonian that describes the interactions between the molecule and external fields as where H Z , H dc , and H ac describe the interaction with dc magnetic, dc electric, and non-resonant optical fields, respectively. To describe the effect of a dc magnetic field B, we construct the Hamiltonian Here the first term accounts for the magnetic moment generated by the rotation of the molecule, described by the rotational g-factor g r . The second term accounts for the magnetic moments associated with the nuclear spins, characterised by the nuclear g-factors g j , shielded isotropically by a factor σ j .
For polar molecules, dc electric fields couple strongly to the rotational angular momentum of the molecule. They can be used to orient the molecule in the laboratory frame, resulting in space-fixed dipoles that produce strong interactions over long range. The coupling between a dc electric field E dc and the angular momentum of the molecule is described by the Hamiltonian where µ 0 is the magnitude of the electric dipole moment in the frame of the molecule andn is a unit vector that points along the internuclear axis of the molecule. Finally, interactions between the molecule and off-resonant optical fields are important for molecules confined to optical traps. Here, there is an interaction between the molecule and the oscillating electric field E ac of the trap light. For linearly polarised light the Hamiltonian is where α describes the molecular polarisability tensor, which depends on the wavelength of the light. The magnitude of E ac is related to the laser intensity by I ac = |E ac | 2 /( 0 c). The polarisability of the molecules is generally anisotropic; for a molecule oriented at an angle θ to the laser polarisation, defined by the angle β from the z axis, the polarisability along the axis of polarization is where α (0) and α (2) describe isotropic and anisotropic components of the molecular polarisability and P 2 (x) = (3x 2 −1)/2 is the second-order Legendre polynomial.

Program structure
The Diatomic-py package has two key modules: hamiltonian and calculate. The former contains functions used to construct the Hamiltonian, while the latter allows the efficient calculation of key quantities from the eigenvalues and eigenstates found by diagonalizing the Hamiltonian.

Hamiltonian
The Hamiltonian matrix is a two-dimensional, square array of dimension (2×i A +1)×(2×i B +1)×(N max +1) 2 . To create this object, the user supplies a Python dictionary containing the relevant molecular constants and a value for N max , the highest-energy rotational state to include in the basis set, to the function build hamiltonians. We include current values of the constants for a selection of experimentally relevant bialkali molecules from [16,26,27].
We now briefly describe the process that the build hamiltonians function follows. The first step of the calculation is to populate nine twodimensional numpy.ndarray objects using the generate vecs function, one for each Cartesian component of the three angular momentum vectors. To perform this population we use the standard definition of the ladder operators and the relations between spherical and Cartesian components of a general angular momentum j: Following this initialisation step, we store each component of the various j as elements in a length-3 list. The final initialisation step is to create a single unified state space, which is done via repeated use of numpy.kron and identity matrices of appropriate dimensions. This transforms each of the 3 vectors from their own, independent, state spaces into the combined space The remainder of the functions in hamiltonian implement the expressions from section 2 with each of the terms in (1) calculated by a separate function call and the component of H AB represented by a numpy.ndarray object. These component objects are shown in Fig. 1 for an exemplar molecule where i A = i B = 3/2 for the first 3 rotational states. The advantage of this approach is that, having abstracted the angular momentum to a vector once, it need not be repeated; this allows a speed-up using the fast vector processing of numpy and scipy. Similarly, for the terms in H ext , we initially treat the electric, magnetic or optical fields as unit vectors, such that these terms can be scaled later.
To assemble the total Hamiltonian, H rot and H hf are simply added together to form H. To calculate H ext , we first construct field-independent matrices H Z /|B|, H dc /|E dc |, and H ac /I ac . These are then combined to form H ext , by multiplying them by |B|, |E dc |, and I ac as required. Each of these three terms is then added to H to form the total Hamiltonian H AB . To diagonalise the Hamiltonian, we recommend using numpy.linalg.eigh as it can not only handle simultaneous multi-processing of multiple field magnitudes but also exploits the Hermitian property of the Hamiltonian matrix to speed up calculations. This function is based on the syevd and heevd routines in the LAPACK linear algebra package for Fortran 90. As an exam-ple, to generate and diagonalize the Hamiltonians needed for a Zeeman plot covering a magnetic field range of 1 to 500 G, with a constant electric field E dc = 5 kV cm −1 and off-resonant light intensity I ac = 2.5 kW cm −2 , we run the script: import numpy import diatomic.hamiltonian as hamiltonian from diatomic.constants import Rb87Cs133 shielding "a0":2020*4*pi*eps0*bohr**3, # isotropic pol, 1064nm "a2":1997*4*pi*eps0*bohr**3, # anisotropic pol, 1064nm "Beta":0} # laser polarisation angle wrt z All quantities in the dictionary are defined in SI units. We include sample dictionaries in an additional module constants.py from which we import Rb87Cs133 in this example. To perform the equivalent calculation for other molecules either a different set of constants can be imported from constants.py or a custom dictionary should be defined. Note that diatomic-py cannot calculate the values of α (0) and α (2) and so these must be supplied, by the user, for each wavelength λ. Where available the constants we have supplied are for λ = 1064 nm.

Calculate
Calculate contains functions that deal with the result of the Hamiltonian diagonalisation. This includes a set of three functions label states N MN, label states I MI, label states F MF. These take the array of eigenvectors and evaluate the expectation values of N · N and N z ; I · I and I z ; and F · F and F z , respectively, and can be used to assign quantum numbers to the eigenstates. Here I = I A + I B is the operator for the total nuclear spin, and F = N + I is that for the total angular momentum.
The function transition dipole moment calculates the transition dipole moment (in units of the molecule-frame dipole moment) between one eigenstate |i and a range of others |f . This calculation is performed by first constructing the space-fixed electric dipole operator µ and then calculating i| µ |f by matrix multiplication. We also include the functions magnetic moment and electric moment that calculate the lab-frame magnetic and electric dipole moments (in SI units) for each eigenstate. Each of these functions constructs the appropriate dipole moment operator µ z for either the electric or magnetic dipole moment pointed along the quantisation axis (z) before calculating the expectation value i| µ z |i . Finally, we include a function for alternative ordering of the energy levels. By default the energy levels are returned in order of ascending energy, such that two levels that cross one another exchange indices. To prevent this, and return levels with indices that reflect the character of the states rather than energy, we provide the function sort smooth. This orders energy levels such that, for each state ψ p , the overlap ψ p k |ψ p k+1 is maximal, where k is an index that increments with the independent variable of the calculation. For each k the function calculates the overlap of each eigenstate ψ p k with all others for k + 1 i.e. the matrix product C T k C k+1 = O k where C k is the eigenvector matrix at field k. By finding the index q for which the value of q| O k |p is maximal, we can locate which pair of eigenstates has the largest overlap. If p = q then we infer that two energy levels have crossed and swap the indices such that ψ q k+1 ↔ ψ p k+1 and similarly with the eigenenergies E q k+1 ↔ E q k+1 .

Benchmarking
The accuracy of the calculation increases with the number of basis states included in the calculation. However, this also increases the time for computation. The dc Stark component of the Hamiltonian causes the largest mixing between rotational states, and therefore calculations for molecules in large dc electric fields are the most sensitive to the number of states included in the calculation. For the purpose of these tests we use the molecule 87 Rb 133 Cs; we do not anticipate that there will be much variation in the convergence or time for different bialkali species. We control the number of basis states by changing the variable N max , i.e. the quantum number of the highest-energy rotational state included in the calculation. For each result, we plot the change in energy of the state when the new rotational state is included in the calculation.
For our example, we consider the 87 Rb 133 Cs molecule in a magnetic field B = 181.5 G, a dc electric field E dc = 5 kV cm −1 , and an optical field, with λ = 1064 nm, with laser intensity I ac = 2.5 kW cm −2 , that is linearly polarised orthogonal to the magnetic and electric fields. This is a typical configuration at an energy scale relevant to prior experimental work [28]. We find that for this configuration, the energy of N = 0 changes by less than h × 1 Hz for N max > 7, better than the uncertainty on experimental measurements of the absolute binding energy [29] or typical uncertainty in rotational spectroscopy [26,30].
In Fig. 2(b), we show the time that each calculation took using an Intel i5-8350U CPU @ 1.7 GHz with 8 GB of RAM. We break the run time up into the construction phase, where the Hamiltonian is constructed, and the numerical diagonalization using numpy.linalg.eigh. We see that the Hamiltonian construction takes an order of magnitude longer than the diagonalization for all calculations. However, for a calculation with multiple field magnitudes,  23 Na 87 Rb, (d) 23 Na 40 K. The electric field is applied in addition to a parallel dc magnetic field of 200 G. The N = 0 hyperfine structure is shown in the lower panels, with the high-field hyperfine ground state indicated in bold. The N = 1 structure is shown in the upper panels. The relative transition strength is coded as in Fig. 3. this construction step must be performed only once. We anticipate that the duration of the diagonalization stage would scale linearly with the number of field magnitudes being studied and as the cube of the number of basis states i.e. as (N max + 1) 6 . Note that the largest calculations shown, with N max = 9, take only a few minutes in total to complete, demonstrating the utility of the code without access to high-performance computing facilities.

Examples
In this section we briefly describe example calculations that are relevant to current research into controlling the quantum states of ultracold bialkali molecules.

Zeeman and dc Stark effects
In Fig. 3 we show the Zeeman structure for N = 0, 1 for a selection of bialkali molecules up to 600 G in increments of 10 G. To perform the calculations, we generate a 3D numpy array constructed by layering the 2D Hamiltonian from each magnetic field to be evaluated, and simultaneously diagonalising all Hamiltonians in a single call of numpy.linalg.eigh. For each molecule, we highlight the spin-stretched hyperfine state that becomes the absolute ground state in the limit of large magnetic field, and also the states with N = 1 that are connected to it by an allowed one-photon transition. In Fig. 4, we show similar dc Stark maps calculated using a similar approach for the same molecules in a 200 G magnetic field. For each calculation, we take the field from 0 V cm −1 to 250 V cm −1 in increments of 2.5 V cm −1 .

Transition energies and transition dipole moments
Electric dipole transitions between rotational states may be driven using resonant microwaves. In Fig. 5 we show the hyperfine states of 87 Rb 133 Cs in a 181.5 G magnetic field. We choose the absolute ground state (N = 0, M F = 5) as the initial state from which we want to calculate the transition dipole moments of all allowed transitions. From here, there are allowed onephoton microwave transitions to N = 1 states with M F = 4, 5, 6, depending on the polarisation of the driving field with respect to the quantisation axis, as defined by the magnetic field. For each transition we show the energy of the final state with respect to the initial state along with the transition dipole moment in units of the molecule-frame dipole moment (for 87 Rb 133 Cs, µ 0 = 1.23 D [6]). In each case, the off-resonant laser field has a wavelength λ = 1064 nm and is linearly polarised at an angle β with respect to z. Sub-levels that are not accessible are shown in yellow as a function of laser intensity. The relative transition strengths for microwaves polarised along z are shown as a blue colour map. For the calculations for 87 Rb 133 Cs the magnetic field is fixed to be 181.5 G, appropriate for matching the experiments of Blackmore et al. [28]. In the calculations for 40 K 87 Rb the magnetic field is fixed to 549.5 G, appropriate for comparison with the experiments of Neyenhuis et al. [16]. The dashed black line indicates the intensity used for the calculations shown in Fig. 7.

ac Stark effects
The ac Stark effect is important for optically trapped molecules. Using our code, transition frequencies can be calculated as a function of either laser intensity I ac or polarisation angle β. We give examples of each of these calculations for the molecules 87 Rb 133 Cs and 40 K 87 Rb at experimentally relevant magnetic fields of 181.5 G and 545.9 G respectively, and laser wavelength λ = 1064 nm. In each case we highlight allowed transitions for microwaves polarised along z from an initial state with N = 0; we choose (N = 0, M F = 5) for 87 Rb 133 Cs and (N = 0, m K = −4, m Rb = 1/2) for 40 K 87 Rb. In Fig. 6 we show the relevant transition frequencies as a function of the laser intensity for β = 0 • and 90 • . For the case of 87 Rb 133 Cs at β = 90 • we observe a complex pattern of avoided crossings as the trapping light mixes states with different M F . This behaviour matches that observed in experiments [31,28]. In Fig. 7 we show a similar plot, but varying β with fixed I ac = 2.35 kW cm −2 . To perform this calculation we recalculate the component H  In each case the off-resonant laser field has a wavelength λ = 1064 nm, an intensity of 2.35 kW cm −2 , and is linearly polarised at an angle β with respect to z. The relative transition strengths for microwaves polarised along z are shown as a blue colour map. For the calculations shown in (a) the magnetic field is fixed to be 181.5 G, appropriate for matching the experiments of Blackmore et al. [28]. For the calculations shown in (b) the magnetic field is fixed to 549.5 G, appropriate for comparison with the experiments of Neyenhuis et al. [16]. The dashed black lines indicate the transition frequencies for the strongest transition in N = 1 in the absence of the laser field.
The horizontal dashed line shows the transition frequency for the strongest allowed transition in the absence of the trap light. For 40 K 87 Rb, there is one strong transition that is allowed across all polarisation angles. This transition frequency crosses the free-space value when β = 52 • , in agreement with experimental observations of Neyenhuis et al. [16].

Installation and usage
Installation of the package is performed by using the .whl files from the GitHub repository [32] or directly from the Python Package Index using pip install diatomic.

Conclusions and outlook
We have presented a Python code that automates the construction of a Hamiltonian that describes the rotational and hyperfine structure of 1 Σ molecules. The Hamiltonian includes terms to describe interactions between the molecule and external dc magnetic, dc electric and the off-resonant optical fields necessary for trapping the molecules. This facilitates the straightforward calculation of Zeeman, dc Stark, and ac Stark maps of the hyperfine structure that can be readily compared with measurements from current experiments. Additional functions for the calculation of static magnetic and electric dipole moments of states are also provided.
Useful future additions to the code may include: greater flexibility in the geometry of the applied fields, e.g. non-parallel dc magnetic and dc electric fields; simulation of dressing by near-resonant microwave fields; extension to 2 Σ molecules, relevant to experiments on laser-cooled molecules.

Conflicts of Interest
The authors declare that they have no known competing financial interests or personal relationships that could have appeared to influence the work reported in this paper.