docs/usage/structure.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Structure and Angular Momentum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Adding Structure\n",
    "\n",
    "So far we've looked at simple 2 and 3 level systems, but to accurately model a physical system we may need to consider complex structures. For example, the hyperfine structure of alkali metals used commonly in optical experiments. It is possible to define structure in MaxwellBloch and provide strength factors for the field couplings and decays between pairs of sublevels.\n",
    "\n",
    "We'll take the example of a Rubidium 87 atom addressed with a weak field resonant with the $5\\mathrm{S}_{1/2} F=1 \\rightarrow 5\\mathrm{P}_{1/2} F=1$ transition. We can neglect the other upper level in the fine structure manifold, $5\\mathrm{P}_{1/2} F={2}$, as it is far from resonance. However, the decay to $5\\mathrm{S}_{1/2} F={2}$ should be considered.\n",
    "\n",
    "If we start with the case of an isotropic field (with equal components in all three polarizations), we can neglect the substructure and just add the hyperfine levels, as the coupling how the population is distributed among the sublevels. For now, the relative strengths of the transitions are given in the following table\n",
    "\n",
    "|     | F'=1           | F'=2           | \n",
    "| --- | ---            | ---            |\n",
    "| F=1 | $\\tfrac{1}{6}$ | $\\tfrac{5}{6}$ |\n",
    "| F=2 | $\\tfrac{1}{2}$ | $\\tfrac{1}{2}$ |\n",
    "\n",
    "and below (see **Note on Strength Factors**) we'll see how these are computed.\n",
    "\n",
    "We will index the three hyperfine levels we require as in the following table.\n",
    "\n",
    "| idx | Hyperfine level      |\n",
    "| --- | -----------------         |\n",
    "| 0   | $5\\mathrm{S}_{1/2} F~=~1$ |\n",
    "| 1   | $5\\mathrm{S}_{1/2} F~=~2$ |\n",
    "| 2   | $5\\mathrm{P}_{1/2} F'~=~1$|\n",
    "\n",
    "How do we include those transition strengths in the system? Both the fields and decays have optional `factors` mapped to them for this purpose. For the decays, we have transitions $\\left|0\\right> \\rightarrow \\left|2\\right>$ and $\\left|1\\right> \\rightarrow \\left|2\\right>$. As per the table above, $S_{11} = \\tfrac{1}{6}$ and $S_{21} = \\tfrac{1}{2}$. Our factors for the isotropic field are then $\\sqrt{\\tfrac{1}{3}\\cdot\\tfrac{1}{6}}$ and $\\sqrt{\\tfrac{1}{3}\\cdot\\tfrac{1}{2}}$ respectively: the square roots of these strength factors multiplied by a third (as any given polarization of the field only interacts with one of the three components of the dipole moment.\n",
    "\n",
    "[steck87]: https://steck.us/alkalidata/rubidium87numbers.1.6.pdf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(np.sqrt(1/6/3))\n",
    "print(np.sqrt(1/2/3))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll put those `channels` and their respective `factors` into the MBSolve object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mb_solve_json = \"\"\"\n",
    "{\n",
    "  \"atom\": {\n",
    "    \"decays\": [\n",
    "      {\n",
    "        \"rate\": 1.0,\n",
    "        \"channels\": [[0,2], [1,2]],\n",
    "        \"factors\": [0.23570226039551587, 0.408248290463863]\n",
    "      }\n",
    "    ],\n",
    "    \"fields\": [\n",
    "      {\n",
    "        \"coupled_levels\": [[0, 2]],\n",
    "        \"factors\": [0.23570226039551584],\n",
    "        \"rabi_freq\": 1e-3,\n",
    "        \"rabi_freq_t_args\": {\n",
    "          \"ampl\": 1.0,\n",
    "          \"centre\": 0.0,\n",
    "          \"fwhm\": 1.0\n",
    "        },\n",
    "        \"rabi_freq_t_func\": \"gaussian\"\n",
    "      }\n",
    "    ],\n",
    "    \"num_states\": 3,\n",
    "    \"initial_state\": [0.5, 0.5, 0]\n",
    "  },\n",
    "  \"t_min\": -2.0,\n",
    "  \"t_max\": 10.0,\n",
    "  \"t_steps\": 100,\n",
    "  \"z_min\": -0.2,\n",
    "  \"z_max\": 1.2,\n",
    "  \"z_steps\": 100,\n",
    "  \"z_steps_inner\": 1,\n",
    "  \"interaction_strengths\": [\n",
    "    1.0e2\n",
    "  ],\n",
    "  \"savefile\": \"structure-1\"\n",
    "}\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from maxwellbloch import mb_solve\n",
    "mbs = mb_solve.MBSolve().from_json_str(mb_solve_json)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Omegas_zt, states_zt = mbs.mbsolve(recalc=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "import seaborn as sns\n",
    "\n",
    "sns.set_style('dark')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(1, figsize=(16, 6))\n",
    "ax = fig.add_subplot(111)\n",
    "# cmap_range = np.linspace(0.0, 1.0e-3, 11)\n",
    "cf = ax.contourf(mbs.tlist, mbs.zlist, \n",
    "                 np.abs(mbs.Omegas_zt[0]/(2*np.pi)),\n",
    "#                  cmap_range, \n",
    "                 cmap=plt.cm.Blues\n",
    "                )\n",
    "ax.set_title('Rabi Frequency ($\\Gamma / 2\\pi $)')\n",
    "ax.set_xlabel('Time ($1/\\Gamma$)')\n",
    "ax.set_ylabel('Distance ($L$)')\n",
    "for y in [0.0, 1.0]:\n",
    "    ax.axhline(y, c='grey', lw=1.0, ls='dotted')\n",
    "plt.colorbar(cf);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Hyperfine Structure for Single (Outer) Electron Atoms\n",
    "\n",
    "In experiments with single outer electron atoms like Rubidium involving polarised (non-isotropic) light, we may need to consider the effect of hyperfine pumping mechanisms. Our state basis must then include the hyperfine structure and magnetic substructure of these fine structure manifolds, consisting of $2F+1$ degenerate sublevels $m_F=\\{F,−F+1, . . . ,F−1,F\\}$ within each hyperfine level.\n",
    "\n",
    "The coupling of each pair of hyperfine sublevels is factored by the _reduced dipole matrix element_ for that pair. This is not the place to go into how we calculate those reduced dipole matrix elements, see [Ogden 2016, chapter 5][thesis] for details.\n",
    "\n",
    "The good news is that MaxwellBloch has helper functions in the `hyperfine` module to\n",
    "\n",
    "- build the hyperfine structure for us, and\n",
    "- calculate the reduced dipole matrix element factors.\n",
    "\n",
    "[thesis]: https://github.com/tpogden/phd-thesis"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First we can use the `hyperfine` module to create the hyperfine structure ($F$ levels). As before we will model a weak field resonant with the $5\\mathrm{S}_{1/2} F=1 \\rightarrow 5\\mathrm{P}_{1/2} F=1$ transition. Again we can neglect the other upper level in the fine structure manifold, $5\\mathrm{P}_{1/2} F={2}$, as it is far from resonance, but the decay to $5\\mathrm{S}_{1/2} F={2}$ will be considered."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from maxwellbloch import hyperfine\n",
    "# Init the two lower F levels\n",
    "Rb87_5s12_F1 = hyperfine.LevelF(I=1.5, J=0.5, F=1, energy=0.0, mF_energies=[])\n",
    "Rb87_5s12_F2 = hyperfine.LevelF(I=1.5, J=0.5, F=2, energy=0.0, mF_energies=[])\n",
    "# Init the upper F level\n",
    "Rb87_5p12_F1 = hyperfine.LevelF(I=1.5, J=0.5, F=1, energy=0.0, mF_energies=[])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see looking at the lower $F=1$ level that the hyperfine structure $m_F=\\{-1, 0 ,1\\}$ is built automatically based on the angular moment numbers $I$, $J$ and $F$. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(Rb87_5s12_F1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If there are levels to be addressed that are off resonance, we would need to add $F$ level energies (e.g. if we included the $F'=2$ upper level). Note that these only need to be relative to the resonance freq. If the mF levels are split (for example with a magnetic field), the individual mf level energies can be added relative to the $F$ level energy.\n",
    "\n",
    "We add the $F$ levels to the atom structure object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Init the atom structure object and add the F levels \n",
    "atom1e = hyperfine.Atom1e()\n",
    "atom1e.add_F_level(Rb87_5s12_F1)\n",
    "atom1e.add_F_level(Rb87_5s12_F2)\n",
    "atom1e.add_F_level(Rb87_5p12_F1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now have everything we need to build our MBSolve object. First, the total number of states:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "NUM_STATES = atom1e.get_num_mF_levels()\n",
    "print(NUM_STATES)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Energies and Detuning"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then the energies of these states, if they have been specified (no need in this case but here for completeness):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ENERGIES = atom1e.get_energies()\n",
    "print(ENERGIES)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As for the solvers without structure, if desired we can specify a detuning:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Tune to be on resonance with the F1 -> F1 transition\n",
    "DETUNING = 0\n",
    "print(DETUNING)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now the levels to be coupled by the field are each F=1 to F'=1 combination pair. atom1e has a helper function for that \n",
    "(based on the order we added them above. `F_level_idxs_a` is for the lower level, `F_level_idxs_b` is for the upper):"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Field Channels and Factors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "FIELD_CHANNELS = atom1e.get_coupled_levels(F_level_idxs_a=[0], F_level_idxs_b=[2])\n",
    "print(FIELD_CHANNELS)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And now the important part, we need to calculate the [Clebsch-Gordan coefficients](https://en.wikipedia.org/wiki/Clebsch%E2%80%93Gordan_coefficients) that couple each $M_F$ substructure pair. The `atom1e` class has a method to calculate these for us. We need to specify the polarisation of the light using $q=\\{-1, 0 ,1\\}$. In this case we'll choose $q=1$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "q = 1 # Field polarisation\n",
    "FIELD_FACTORS = atom1e.get_clebsch_hf_factors(F_level_idxs_a=[0], F_level_idxs_b=[2], q=q)\n",
    "print(FIELD_FACTORS)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Decay Channels and Factors"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Similarly to the levels coupled by the field, we need to specify the pairs of levels that will be channels for spontaneous decay. This time the lower levels are all in both F=1 and F=2, so there are more combinations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DECAY_CHANNELS = atom1e.get_coupled_levels(F_level_idxs_a=[0,1], F_level_idxs_b=[2])\n",
    "print(DECAY_CHANNELS)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The spontaneous decay factors matching those channels (in order) can be obtained with the `get_decay_factors` method. This is equivalent to summing the `get_clebsch_hf_factors_hf` for all polarisations as decay photons are of all polarisations. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DECAY_FACTORS = atom1e.get_decay_factors(F_level_idxs_a=[0, 1], F_level_idxs_b=[2])\n",
    "print(DECAY_FACTORS)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initial State"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the two and three level models we didn't specify an initial state for the system, which defaults to all population in the first (i.e. ground state). Here we need to correctly distribute the population in the ground state $m_F$ levels, which should of course sum to one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "INITIAL_STATE = (\n",
    "    [0.5/3.0]*3 + # s12_F1b\n",
    "    [0.5/5.0]*5 + # s12_F2\n",
    "    [0.0]*3) # p12_F1    \n",
    "print(INITIAL_STATE)\n",
    "print(sum(INITIAL_STATE))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Putting it all into MBSolve"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now have everything we need to build the `MBSolve` object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mb_solve_json = \"\"\"\n",
    "{{\n",
    "  \"atom\": {{\n",
    "    \"decays\": [\n",
    "      {{\n",
    "        \"channels\": {decay_channels},\n",
    "        \"rate\": 1.0,\n",
    "        \"factors\": {decay_factors}\n",
    "      }}\n",
    "    ],\n",
    "    \"energies\": {energies},\n",
    "    \"fields\": [\n",
    "      {{\n",
    "        \"coupled_levels\": {field_channels},\n",
    "        \"factors\": {field_factors},\n",
    "        \"detuning\": {detuning},\n",
    "        \"detuning_positive\": true,\n",
    "        \"label\": \"probe\",\n",
    "        \"rabi_freq\": 1e-3,\n",
    "        \"rabi_freq_t_args\": {{\n",
    "          \"ampl\": 1.0,\n",
    "          \"centre\": 0.0,\n",
    "          \"fwhm\": 1.0\n",
    "        }},\n",
    "        \"rabi_freq_t_func\": \"gaussian\"\n",
    "      }}\n",
    "    ],\n",
    "    \"num_states\": {num_states},\n",
    "    \"initial_state\": {initial_state}\n",
    "  }},\n",
    "  \"t_min\": -2.0,\n",
    "  \"t_max\": 10.0,\n",
    "  \"t_steps\": 100,\n",
    "  \"z_min\": -0.2,\n",
    "  \"z_max\": 1.2,\n",
    "  \"z_steps\": 100,\n",
    "  \"z_steps_inner\": 1,\n",
    "  \"interaction_strengths\": [\n",
    "    1.0e2\n",
    "  ],\n",
    "  \"savefile\": \"structure-2\"\n",
    "}}\n",
    "\"\"\".format(num_states=NUM_STATES, energies=ENERGIES, initial_state=INITIAL_STATE, detuning=DETUNING, \n",
    "    field_channels=FIELD_CHANNELS, field_factors=list(FIELD_FACTORS), decay_channels=DECAY_CHANNELS, \n",
    "           decay_factors=list(DECAY_FACTORS))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from maxwellbloch import mb_solve\n",
    "mbs = mb_solve.MBSolve().from_json_str(mb_solve_json)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Omegas_zt, states_zt = mbs.mbsolve(recalc=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Output\n",
    "\n",
    "We can output the field exactly as we would do for the two-level system."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "import seaborn as sns\n",
    "\n",
    "sns.set_style('dark')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(1, figsize=(16, 6))\n",
    "ax = fig.add_subplot(111)\n",
    "# cmap_range = np.linspace(0.0, 1.0e-3, 11)\n",
    "cf = ax.contourf(mbs.tlist, mbs.zlist, \n",
    "                 np.abs(mbs.Omegas_zt[0]/(2*np.pi)),\n",
    "#                  cmap_range, \n",
    "                 cmap=plt.cm.Blues\n",
    "                )\n",
    "ax.set_title('Rabi Frequency ($\\Gamma / 2\\pi $)')\n",
    "ax.set_xlabel('Time ($1/\\Gamma$)')\n",
    "ax.set_ylabel('Distance ($L$)')\n",
    "for y in [0.0, 1.0]:\n",
    "    ax.axhline(y, c='grey', lw=1.0, ls='dotted')\n",
    "plt.colorbar(cf);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plotting Populations and Coherences"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we want to get the populations in all the levels coupled by a field, we have a helper function `mbs.populations_field`. For example, if we want to look at the total population of all the upper sublevels coupled by the field:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(1, figsize=(16, 6))\n",
    "ax = fig.add_subplot(111)\n",
    "# cmap_range = np.linspace(0.0, 1.0e-3, 11)\n",
    "cf = ax.contourf(mbs.tlist, mbs.zlist, \n",
    "                 np.abs(mbs.populations_field(field_idx=0, upper=False)),\n",
    "#                  cmap_range,\n",
    "                 cmap=plt.cm.Reds\n",
    "                )\n",
    "ax.set_title('Rabi Frequency ($\\Gamma / 2\\pi $)')\n",
    "ax.set_xlabel('Time ($1/\\Gamma$)')\n",
    "ax.set_ylabel('Distance ($L$)')\n",
    "for y in [0.0, 1.0]:\n",
    "    ax.axhline(y, c='grey', lw=1.0, ls='dotted')\n",
    "plt.colorbar(cf);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And for the sum of all the complex coherences on a field transition, we similarly have `mbs.coherences_field`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(1, figsize=(16, 6))\n",
    "ax = fig.add_subplot(111)\n",
    "# cmap_range = np.linspace(0.0, 1.0e-3, 11)\n",
    "cf = ax.contourf(mbs.tlist, mbs.zlist, \n",
    "                 np.imag(mbs.coherences_field(field_idx=0)),\n",
    "#                  cmap_range, \n",
    "                 cmap=plt.cm.Greens\n",
    "                )\n",
    "ax.set_title('Rabi Frequency ($\\Gamma / 2\\pi $)')\n",
    "ax.set_xlabel('Time ($1/\\Gamma$)')\n",
    "ax.set_ylabel('Distance ($L$)')\n",
    "for y in [0.0, 1.0]:\n",
    "    ax.axhline(y, c='grey', lw=1.0, ls='dotted')\n",
    "plt.colorbar(cf);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Note on Stength Factors\n",
    "\n",
    "The strength factor gives the relative strength hyperfine transition. This is the sum $S_{FF'}$ of the matrix elements from a single ground-state sublevel to the levels in a particular $F'$ energy level, and the sum is independent of the ground state sublevel chosen. The sum of $S_{FF'}$ over upper $F'$ levels is $1$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s_11 = atom1e.get_strength_factor(F_level_idx_lower=0, F_level_idx_upper=2)\n",
    "print(s_11)\n",
    "s_12 = atom1e.get_strength_factor(F_level_idx_lower=1, F_level_idx_upper=2)\n",
    "print(s_12)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So this is where the $S_{11} = \\tfrac{1}{6}$ and $S_{21} = \\tfrac{1}{2}$ factors came from."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}