# -*- coding: utf-8 -*-
from setuptools import setup

packages = \
['centrex_tlf',
 'centrex_tlf.couplings',
 'centrex_tlf.hamiltonian',
 'centrex_tlf.hamiltonian.B_coupled_Omega',
 'centrex_tlf.lindblad',
 'centrex_tlf.states',
 'centrex_tlf.transitions']

package_data = \
{'': ['*']}

install_requires = \
['numpy>=1.24.1,<2.0.0',
 'pandas>=1.5.3,<2.0.0',
 'scipy>=1.10.0,<2.0.0',
 'sympy>=1.11.1,<2.0.0']

setup_kwargs = {
    'name': 'centrex-tlf',
    'version': '0.1.0',
    'description': '',
    'long_description': '# CeNTREX-TlF-Hamiltonian\nCode for generating the CeNTREX TlF States, Hamiltonians, Transitions, Couplings and Lindblad equations.\n\nConsists of two modules:\n* `states`\n* `hamiltonian`\n* `transitions`\n* `couplings`\n* `lindblad`\n\n`states` has code to generate states and the classes that describe the `CoupledBasisState`, `UncoupledBasisState` and `State`; where `State` holds multiple `CoupledBasisStates` or `UncoupledBasisStates` with different amplitudes, i.e. when superpositions arise.\n\n## Dependencies\n* `numpy`\n* `scipy`\n* `sympy`\n* `pandas`\n  \n## Installation\n`python -m pip install .`  \nwhere `.` is the path to the directory. To install directly from `Github` use:  \n`python -m pip install git+https://github.com/ograsdijk/CeNTREX-TlF`\n\n# `states`\n`states` contains the functions and classes to represent the TlF states:  \n`CoupledBasisState` is a class representing a TlF state with coupled quantum numbers, i.e. F, mF, F1, J, I1, I2, Ω, P.  \n`UncoupledBasisState` is a class representing a TlF state with uncoupled quantum numbers, i.e. J, mJ, I1, m1, I2, m2, Ω, P.  \nFinally `State` is a class representing a collection of states, since in most cases the TlF molecules are in a superposition state.\n\n```Python\nfrom centrex_tlf import states\nstates.CoupledBasisState(F=1, mF=0, F1 = 1/2, J = 0, I1 = 1/2, I2 = 1/2, Omega = 0, P = 1)\n```\nor using some of the functions to generate all hyperfine substates in a given J level: \n```Python\nfrom centrex_tlf import states\nQN = states.generate_uncoupled_states_ground(Js = [0,1])\n```\nwhich returns an array containing the UncoupledBasisStates\n```python\narray([|X, J = 0, mJ = 0, I₁ = 1/2, m₁ = -1/2, I₂ = 1/2, m₂ = -1/2, P = +, Ω = 0>,\n       |X, J = 0, mJ = 0, I₁ = 1/2, m₁ = -1/2, I₂ = 1/2, m₂ = 1/2, P = +, Ω = 0>,\n       |X, J = 0, mJ = 0, I₁ = 1/2, m₁ = 1/2, I₂ = 1/2, m₂ = -1/2, P = +, Ω = 0>,\n       |X, J = 0, mJ = 0, I₁ = 1/2, m₁ = 1/2, I₂ = 1/2, m₂ = 1/2, P = +, Ω = 0>,\n       |X, J = 1, mJ = -1, I₁ = 1/2, m₁ = -1/2, I₂ = 1/2, m₂ = -1/2, P = -, Ω = 0>,\n       |X, J = 1, mJ = -1, I₁ = 1/2, m₁ = -1/2, I₂ = 1/2, m₂ = 1/2, P = -, Ω = 0>,\n       |X, J = 1, mJ = -1, I₁ = 1/2, m₁ = 1/2, I₂ = 1/2, m₂ = -1/2, P = -, Ω = 0>,\n       |X, J = 1, mJ = -1, I₁ = 1/2, m₁ = 1/2, I₂ = 1/2, m₂ = 1/2, P = -, Ω = 0>,\n       |X, J = 1, mJ = 0, I₁ = 1/2, m₁ = -1/2, I₂ = 1/2, m₂ = -1/2, P = -, Ω = 0>,\n       |X, J = 1, mJ = 0, I₁ = 1/2, m₁ = -1/2, I₂ = 1/2, m₂ = 1/2, P = -, Ω = 0>,\n       |X, J = 1, mJ = 0, I₁ = 1/2, m₁ = 1/2, I₂ = 1/2, m₂ = -1/2, P = -, Ω = 0>,\n       |X, J = 1, mJ = 0, I₁ = 1/2, m₁ = 1/2, I₂ = 1/2, m₂ = 1/2, P = -, Ω = 0>,\n       |X, J = 1, mJ = 1, I₁ = 1/2, m₁ = -1/2, I₂ = 1/2, m₂ = -1/2, P = -, Ω = 0>,\n       |X, J = 1, mJ = 1, I₁ = 1/2, m₁ = -1/2, I₂ = 1/2, m₂ = 1/2, P = -, Ω = 0>,\n       |X, J = 1, mJ = 1, I₁ = 1/2, m₁ = 1/2, I₂ = 1/2, m₂ = -1/2, P = -, Ω = 0>,\n       |X, J = 1, mJ = 1, I₁ = 1/2, m₁ = 1/2, I₂ = 1/2, m₂ = 1/2, P = -, Ω = 0>],\n      dtype=object)\n```\nState objects, which are superpositions of BasisStates are also generated easily:\n```Python\nsuperposition = 1*QN[0] + 0.1j*QN[1]\n```\nwhich returns\n```Python\n1.00 x |X, J = 0, mJ = 0, I₁ = 1/2, m₁ = -1/2, I₂ = 1/2, m₂ = -1/2, P = +, Ω = 0>\n0.00+0.10j x |X, J = 0, mJ = 0, I₁ = 1/2, m₁ = -1/2, I₂ = 1/2, m₂ = 1/2, P = +, Ω = 0>\n```\nA subset of `State`, `CoupledBasisStates` can be selected with the `QuantumSelector` as follows:\n```Python\nQN = states.generate_coupled_states_ground(Js = [0,1])\nqn_select = states.QuantumSelector(J = 1, mF = 0, electronic = states.ElectronicState.X)\nqn_select.get_indices(QN)\n```\nwhich returns all the indices with `J=1` and `mJ=0`:\n```python\narray([ 4,  6,  9, 13], dtype=int64)\n```\n# `hamiltonian`\n`hamiltonian` contains the functions to generate TlF hamiltonians in the X and B state in either coupled or uncoupled form.  \nGenerating a ground state X hamiltonian can be accomplished easily using some convenience functions:\n```Python\nfrom centrex_tlf import states, hamiltonian\n\n# generate the hyperfine sublevels in J=0 and J=1\nQN = states.generate_uncoupled_states_ground(Js = [0,1])\n\n# generate a dictionary with X hamiltonian terms\nH = hamiltonian.generate_uncoupled_hamiltonian_X(QN)\n\n# create a function outputting the hamiltonian as a function of E and B\nHfunc = hamiltonian.generate_uncoupled_hamiltonian_X_function(H)\n```\nAll functions generating hamiltonians only require a list or array of TlF states. Generating the hamiltonian only for certain hyperfine sublevels is hence also straightforward. The function `calculate_uncoupled_hamiltonian_X` calculates the hamiltonians from scratch, whereas `generate_uncoupled_hamiltonian_X` pulls the non-zero elements from an sqlite database.\n\nTo convert a hamiltonian from one basis to another transformation matrices can be generated or calculated\n(`generate_transform_matrix` pulls non-zero matrix elements from an sqlite database, `calculate_transform_matrix` does the full element wise calculation):\n```Python\nfrom centrex_tlf import states, hamiltonian\n\n# generate the hyperfine sublevels in J=0 and J=1\nQN = states.generate_uncoupled_states_ground(Js = [0,1])\n# generate the coupled hyperfine sublevels in J=0 and J=1\nQNc = states.generate_coupled_states_ground(Js = [0,1])\n\n# generate a dictionary with X hamiltonian terms\nH = hamiltonian.generate_uncoupled_hamiltonian_X(QN)\nHfunc = hamiltonian.generate_uncoupled_hamiltonian_X_function(H)\nH0 = Hfunc(E = [0,0,0], B = [0,0,1e-3])\n\n# generate the transformation matrix\ntransform = hamiltonian.generate_transform_matrix(QN, QNc)\n\n# calculate the transformed matrix\nH0c = transform.conj().T@H0@transform\n```\nThis is mostly used for optical bloch simulations where the coupled states representation is more convenient.\n\n## Stark Shift Example\nTo calculate the energy levels as a function of the electric field the following code can be used, which calculates all energies up to `J=6` but only plots the `|J=2, mJ=0>` hyperfine levels. These are the states focussed by the electrostatic quadrupole lens in the CeNTREX experiment.\n![Quadrupole Lens States](quadrupole_lens_states.png)\n```Python\nimport numpy as np\nimport matplotlib.pyplot as plt\n\nfrom centrex_tlf import states, hamiltonian\n\n# generate states up to J=6\nQN = states.generate_uncoupled_states_ground(Js=np.arange(7))\n\n# generate the X hamiltonian terms\nH = hamiltonian.generate_uncoupled_hamiltonian_X(QN)\n\n# create a function outputting the hamiltonian as a function of E and B\nHfunc = hamiltonian.generate_uncoupled_hamiltonian_X_function(H)\n\n# V/cm\nEz = np.linspace(0, 50e3, 101)\n\n# generate the Hamiltonian for (almost) zero field, add a small field to make states\n# non-degenerate\nHi = Hfunc(E=[0, 0, 1e-3], B=[0, 0, 1e-3])\nE, V = np.linalg.eigh(Hi)\n\n# get the true superposition-states of the system\nQN_states = hamiltonian.matrix_to_states(V, QN)\n\n# original eigenvectors used in tracking states as energies change order\nV_track = V.copy()\n\n# indices of the J=2, mJ=0 states focused by the lens\nindices_J2_mJ0 = [\n    idx\n    for idx, s in enumerate(QN_states)\n    if s.largest.J == 2 and s.largest.mJ == 0\n]\n\nindices_J012 = [\n    idx for idx, s in enumerate(QN_states) if s.largest.J in [0, 1, 2]\n]\n\n# empty array for storing energies\nenergy = np.empty([Ez.size, len(QN)], dtype=np.complex128)\n\n# iterate over the electric field values\nfor idx, Ei in enumerate(Ez):\n    Hi = Hfunc(E=[0, 0, Ei], B=[0, 0, 1e-3])\n    E, V = np.linalg.eigh(Hi)\n\n    # sort indices to keep the state order the same\n    indices = np.argmax(np.abs(V_track.conj().T @ V), axis=1)\n    energy[idx, :] = E[indices]\n    V_track[:, :] = V[:, indices]\n\n# plot the J=2, mJ=0 Stark curves\nfig, ax = plt.subplots(figsize=(12, 8))\nax.plot(\n    Ez,\n    (energy.real[:, indices_J2_mJ0] - energy.real[:, indices_J2_mJ0][0, 0])\n    / (2 * np.pi * 1e9),\n)\nax.set_xlabel("E [V/cm]")\nax.set_ylabel("Energy [GHz]")\nax.set_title("|J=2, mJ=0> Stark Curve")\nax.grid(True)\nplt.show()\n```\n\n# `couplings`\nCode for generating the CeNTREX TlF couplings. \nIncludes code for generating branching ratios, electric dipole coupling elements and coupling fields\n\n## Generating branching ratios\nThe code below generates branching ratios from `|J\'=1, F1\'=1/2, mF=0>` to all states in the `J=1` manifold.\n```Python\nfrom centrex_tlf import states, couplings\n\nexcited_state = states.CoupledBasisState(\n    J=1, F1=1 / 2, F=1, mF=0, I1=1 / 2, I2=1 / 2, Omega=1, P=1\n)\nqn_select = states.QuantumSelector(J=1)\nground_states = [1*s for s in states.generate_coupled_states_X(qn_select)]\nbr = couplings.calculate_br(1 * excited_state, ground_states)\n```\n\n## Generating couplings\nThe code below generates the coupling fields for the `J=1` manifold to the `J\'=1, F1\'=1/2, F\'=1` manifold. The returned value is a dataclass `CouplingFields` containing the following fields:  \n* `ground_main`\n* `excited_main`\n* `main_coupling`: the electric dipole coupling between `ground_main` and `excited_main`\n* `ground_states`: list of all ground states\n* `excited_states`: list of all excited states\n* `fields`: a list of `CouplingField` dataclasses with the following fields:  \n  * `polarization`: polarization vector\n  * `field`: coupling field in the `ground_states` + `excited_states` basis\n\n```Python\nfrom centrex_tlf import states, couplings\n\nqn_select = states.QuantumSelector(J=1)\nground_states = states.generate_coupled_states_X(qn_select)\n\nqn_select = states.QuantumSelector(J=1, F1=1 / 2, F=1, P=1, Ω=1)\nexcited_states = states.generate_coupled_states_B(qn_select)\n\n# the generate_coupling_field_* functions requires lists as inputs, not np.ndarrays\nQN = list(1 * np.append(ground_states, excited_states))\nground_states = [1*s for s in  ground_states]\nexcited_states = [1*s for s in excited_states]\n\nH_rot = np.eye(len(QN), dtype=complex) * np.arange(len(QN))\nV_ref = np.eye(len(QN))\npol_vecs = [np.array([0.0, 0.0, 1.0]), np.array([1.0, 0.0, 0.0])]\nnormalize_pol = True\n\ncoupling = couplings.generate_coupling_field_automatic(\n    ground_states_approx = ground_states, \n    excited_states_approx = excited_states, \n    QN_basis = QN,\n    H_rot = H_rot, \n    QN = QN, \n    V_ref = V_ref, \n    pol_vecs = pol_vecs, \n    normalize_pol = normalize_pol\n)\n```\n',
    'author': 'ograsdijk',
    'author_email': 'o.grasdijk@gmail.com',
    'maintainer': 'None',
    'maintainer_email': 'None',
    'url': 'https://github.com/ograsdijk/CeNTREX-TlF',
    'packages': packages,
    'package_data': package_data,
    'install_requires': install_requires,
    'python_requires': '>=3.8,<3.12',
}


setup(**setup_kwargs)
