{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Atoms & Transitions\n",
    "\n",
    "In this example we illustrate the use of the ``atomsmltr.atoms`` subpackage."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Work with existing atoms\n",
    "\n",
    "Some atomic species are already implemented and defined ine the ``atomsmltr.atoms.collection`` module. They can be directly imported from the ``atoms`` subpackage :"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from atomsmltr.atoms import Ytterbium, Strontium"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, it is straightforward to initiate those atoms:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "─────────────\n",
      "| Ytterbium |\n",
      "─────────────\n",
      ". Parameters :\n",
      "  ├── mass (kg) : 2.89e-25\n",
      "  └── mass (au) : 173.94\n",
      "\n",
      ". Transition list :\n",
      "  ├── main\n",
      "  └── intercombination\n",
      "\n",
      ". 'main' transition :\n",
      "  ├── λ : 398.91 nm\n",
      "  ├── Γ : 2π × 2.89e+07 Hz\n",
      "  ├── Isat : 59.51 mw/cm²\n",
      "  └── lande factor g : 1.035\n",
      "\n",
      ". 'intercombination' transition :\n",
      "  ├── λ : 555.80 nm\n",
      "  ├── Γ : 2π × 1.82e+05 Hz\n",
      "  ├── Isat : 0.14 mw/cm²\n",
      "  └── lande factor g : 1.493\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "yb = Ytterbium()\n",
    "yb.print_info()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "─────────────\n",
      "| Strontium |\n",
      "─────────────\n",
      ". Parameters :\n",
      "  ├── mass (kg) : 1.46e-25\n",
      "  └── mass (au) : 88.00\n",
      "\n",
      ". Transition list :\n",
      "  ├── main\n",
      "  └── intercombination\n",
      "\n",
      ". 'main' transition :\n",
      "  ├── λ : 460.86 nm\n",
      "  ├── Γ : 2π × 3.20e+07 Hz\n",
      "  ├── Isat : 42.73 mw/cm²\n",
      "  └── lande factor g : 1.0\n",
      "\n",
      ". 'intercombination' transition :\n",
      "  ├── λ : 689.45 nm\n",
      "  ├── Γ : 2π × 7.46e+03 Hz\n",
      "  ├── Isat : 0.00 mw/cm²\n",
      "  └── lande factor g : 1.5\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "sr = Strontium()\n",
    "sr.print_info()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Atoms properties\n",
    "\n",
    "The ``Atom`` object carries some atomic physical properties, as well as a list of transitions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ytterbium mass : 2.89e-25 kg\n",
      "Ytterbium mass : 174 a.u.\n",
      "Transitions : ['main', 'intercombination']\n"
     ]
    }
   ],
   "source": [
    "from atomsmltr.atoms import Ytterbium\n",
    "\n",
    "yb = Ytterbium()\n",
    "\n",
    "# - some physical properties\n",
    "print(f\"Ytterbium mass : {yb.mass :.3g} kg\")\n",
    "print(f\"Ytterbium mass : {yb.mass_au :.3g} a.u.\")\n",
    "\n",
    "# - implemented transitions\n",
    "print(f\"Transitions : {yb.list_transitions()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "the associated transitions can be retrieved using the ``.trans`` property"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "────────\n",
      "| main |\n",
      "────────\n",
      ". Parameters :\n",
      "  ├── λ : 398.91 nm\n",
      "  ├── Γ : 2π × 2.89e+07 Hz\n",
      "  ├── Isat : 59.51 mw/cm²\n",
      "  └── lande factor g : 1.035\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "main = yb.trans[\"main\"]\n",
    "main.print_info()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Atoms methods\n",
    "\n",
    "It is possible to compute the radiation pressure felt by an atom on one of its transitions using the ``.get_radiation_pressure()`` method :"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "radiation pressure = 1.42e-19 N\n"
     ]
    }
   ],
   "source": [
    "from atomsmltr.atoms import Ytterbium\n",
    "import numpy as np\n",
    "\n",
    "# - settings\n",
    "laser_intensity = 1e4  # W/m^2\n",
    "laser_detuning = 2 * np.pi * 1e6  # rad/s\n",
    "mag_field = 1e-4  # T\n",
    "polarization = np.array((0, 0, 1))  # (pi, sigma+, sigma-)\n",
    "\n",
    "# - compute\n",
    "f = Ytterbium().get_radiation_pressure(\n",
    "    transition=\"main\",\n",
    "    intensity=laser_intensity,\n",
    "    mag_field=mag_field,\n",
    "    polarization=polarization,\n",
    "    detuning=laser_detuning,\n",
    ")\n",
    "\n",
    "print(f\"radiation pressure = {f:.3g} N\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Transitions properties & methods\n",
    "\n",
    "The ``Transition`` object carries information on an atomic transition:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "vacuum wavelength : 4.61e-07 m\n",
      "natural linewidth : 2.01e+08 rad/s\n",
      "saturation intensity : 427 W/m^2\n",
      "saturation intensity : 42.7 mW/cm^2\n",
      "Lande g factor : 1.000\n",
      "wavenumber : 1.36e+07 m^-1\n"
     ]
    }
   ],
   "source": [
    "from atomsmltr.atoms import Strontium\n",
    "\n",
    "# - get strontium main transition\n",
    "main = Strontium().trans[\"main\"]\n",
    "\n",
    "# - access some properties\n",
    "print(f\"vacuum wavelength : {main.wavelength :.3g} m\")\n",
    "print(f\"natural linewidth : {main.Gamma :.3g} rad/s\")\n",
    "print(f\"saturation intensity : {main.Isat :.3g} W/m^2\")\n",
    "print(f\"saturation intensity : {main.Isat_mW_per_cm2 :.3g} mW/cm^2\")\n",
    "print(f\"Lande g factor : {main.lande_factor :.3f}\")\n",
    "print(f\"wavenumber : {main.k :.3g} m^-1\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and some useful methods too:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "s=0.25\n",
      "scattering rate = 1.87e+07 s^-1\n"
     ]
    }
   ],
   "source": [
    "# - saturation parameter\n",
    "intensity = 0.5 * main.Isat  # W/m^2\n",
    "detuning = 0.5 * main.Gamma  # rad/s\n",
    "s = main.get_saturation_parameter(intensity, detuning)\n",
    "\n",
    "print(f\"{s=}\")\n",
    "\n",
    "# - scattering rate\n",
    "intensity = 0.5 * main.Isat  # W/m^2\n",
    "detuning = 0.5 * main.Gamma  # rad/s\n",
    "mag_field = 1e-4  # T\n",
    "polarization = np.array((0, 0, 1))  # (pi, sigma+, sigma-)\n",
    "\n",
    "rate = main.get_scattering_rate(intensity, mag_field, polarization, detuning)\n",
    "print(f\"scattering rate = {rate:.3g} s^-1\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define custom atoms and transitions\n",
    "\n",
    "We strongly advise users that would need atoms or transitions that are not implemented to write dedicated classes and include them in the ``atoms.collection`` module.\n",
    "\n",
    "This being said, atoms and transitions can also be declared on the fly:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "─────────\n",
      "| Yb174 |\n",
      "─────────\n",
      ". Parameters :\n",
      "  ├── mass (kg) : 2.89e-25\n",
      "  └── mass (au) : 174.00\n",
      "\n",
      ". Transition list :\n",
      "  └── main\n",
      "\n",
      ". 'main' transition :\n",
      "  ├── λ : 399.00 nm\n",
      "  ├── Γ : 2π × 3.00e+07 Hz\n",
      "  ├── Isat : 61.73 mw/cm²\n",
      "  └── lande factor g : 1.5\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from scipy import constants as csts\n",
    "from atomsmltr.atoms import Atom, J0J1Transition\n",
    "\n",
    "# - create an ytterbium atom from scratch\n",
    "\n",
    "# create the atom\n",
    "yb174 = Atom(mass=174 * csts.m_u, name=\"Yb174\")\n",
    "\n",
    "# crate a transition\n",
    "yb_main = J0J1Transition(\n",
    "    wavelength=399e-9, Gamma=2 * csts.pi * 30e6, lande_factor=1.5, tag=\"main\"\n",
    ")\n",
    "\n",
    "# add main to the atom transition collection\n",
    "yb174.add_transition(yb_main)\n",
    "\n",
    "# check what we have done\n",
    "yb174.print_info()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "atomsmltr-KmviZRuT-py3.12",
   "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}