{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Configuration\n",
    "\n",
    "Here we illustrate how to create experimental configurations, using the ``atomsmltr.simulation.configurator`` subpackage"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introduction to configurations\n",
    "\n",
    "Configuration objects are at the heart of this package. They are the cement that allows to combine all the environment objects (lasers, magnetic fields, zones) and feed them into a simulation.\n",
    "\n",
    "Here is how to create a configuration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from atomsmltr.simulation import Configuration\n",
    "from atomsmltr.atoms import Ytterbium\n",
    "\n",
    "# init a configuration for ytterbium atoms\n",
    "config = Configuration(atom=Ytterbium())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "objects can be added to the configuration using the ``add_objects()`` method or the ``+=`` operator. The Configuration object will automatically sort them according to their type. \n",
    "\n",
    "**ATTENTION** objects are stored in internal dictionnaries, using their ``tag`` property as a key. Therefore, it is important to **define tags for the objects** and to make sure that objects do not have the same tags (otherwise, an error will be raised)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from atomsmltr.environment import Limits, MagneticOffset, MagneticGradient, GaussianLaserBeam\n",
    "\n",
    "# let's define two lasers\n",
    "laser1 = GaussianLaserBeam(tag=\"laser1\")\n",
    "laser2 = GaussianLaserBeam(tag=\"laser2\")\n",
    "\n",
    "# two magnetic fields\n",
    "mag_offset = MagneticOffset((0,1,0), tag=\"offset\")\n",
    "mag_gradient = MagneticGradient((0,0,0), 0.5, (0,1,0), (0,0,1), tag=\"gradient\")\n",
    "# one limit\n",
    "xlim = Limits(0, 10, axis=0, target=\"position\")\n",
    "\n",
    "# and add them\n",
    "config.add_objects([laser1, mag_offset])   # we can provide a list\n",
    "config.add_objects(xlim)  # a single object\n",
    "config += mag_gradient, laser2  # or use the operator :)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "now we can print the information"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# ------------ CONFIG INFO > START ------------ #\n",
      "────────────────────────\n",
      "| General informations |\n",
      "────────────────────────\n",
      ". atom :\n",
      "  └── name : Ytterbium\n",
      "\n",
      ". lasers :\n",
      "  ├── laser1\n",
      "  └── laser2\n",
      "\n",
      ". magnetic fields :\n",
      "  ├── offset\n",
      "  └── gradient\n",
      "\n",
      ". forces :\n",
      "  └── empty\n",
      "\n",
      ". zones :\n",
      "  └── vazewa\n",
      "\n",
      "\n",
      "────────────────────\n",
      "| atom | 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",
      "────────────────────────\n",
      "| laser | tag='laser1' |\n",
      "────────────────────────\n",
      ". Parameters :\n",
      "  ├── type : Gaussian beam\n",
      "  ├── tag : laser1\n",
      "  ├── waist (m) : 0.001\n",
      "  ├── power (W) : 0.001\n",
      "  ├── waist position (m) : [0 0 0]\n",
      "  ├── direction type : vector\n",
      "  ├── direction : [0 0 1]\n",
      "  ├── unit vector : [0. 0. 1.]\n",
      "  ├── unit vector phi : π × 0.0\n",
      "  ├── unit vector theta : π × 0.0\n",
      "  └── Rayleigh length : 7.9 m\n",
      "\n",
      ". Polarization settings :\n",
      "  └── type : Vertical\n",
      "\n",
      ". Polarization vector :\n",
      "  ├── coords : (1.00, 0.00, 0.00)\n",
      "  ├── polar angle u : 0.50 pi\n",
      "  └── azimt angle v : 0.00 pi\n",
      "\n",
      ". Projections (amplitudes) :\n",
      "  ├── vertical : 1.00+0.00j\n",
      "  ├── horizontal : 0.00+0.00j\n",
      "  ├── circular left : 0.71+0.00j\n",
      "  └── circular right : 0.71+0.00j\n",
      "\n",
      ". Projections (squared norm) :\n",
      "  ├── vertical : 1.00\n",
      "  ├── horizontal : 0.00\n",
      "  ├── circular left : 0.50\n",
      "  └── circular right : 0.50\n",
      "\n",
      "\n",
      "────────────────────────\n",
      "| laser | tag='laser2' |\n",
      "────────────────────────\n",
      ". Parameters :\n",
      "  ├── type : Gaussian beam\n",
      "  ├── tag : laser2\n",
      "  ├── waist (m) : 0.001\n",
      "  ├── power (W) : 0.001\n",
      "  ├── waist position (m) : [0 0 0]\n",
      "  ├── direction type : vector\n",
      "  ├── direction : [0 0 1]\n",
      "  ├── unit vector : [0. 0. 1.]\n",
      "  ├── unit vector phi : π × 0.0\n",
      "  ├── unit vector theta : π × 0.0\n",
      "  └── Rayleigh length : 7.9 m\n",
      "\n",
      ". Polarization settings :\n",
      "  └── type : Vertical\n",
      "\n",
      ". Polarization vector :\n",
      "  ├── coords : (1.00, 0.00, 0.00)\n",
      "  ├── polar angle u : 0.50 pi\n",
      "  └── azimt angle v : 0.00 pi\n",
      "\n",
      ". Projections (amplitudes) :\n",
      "  ├── vertical : 1.00+0.00j\n",
      "  ├── horizontal : 0.00+0.00j\n",
      "  ├── circular left : 0.71+0.00j\n",
      "  └── circular right : 0.71+0.00j\n",
      "\n",
      ". Projections (squared norm) :\n",
      "  ├── vertical : 1.00\n",
      "  ├── horizontal : 0.00\n",
      "  ├── circular left : 0.50\n",
      "  └── circular right : 0.50\n",
      "\n",
      "\n",
      "─────────────────────────────────\n",
      "| magnetic field | tag='offset' |\n",
      "─────────────────────────────────\n",
      ". Parameters :\n",
      "  ├── type : constant field\n",
      "  ├── tag : offset\n",
      "  ├── field_value (T) : [0 1 0]\n",
      "  └── norm (T) : 1\n",
      "\n",
      "\n",
      "───────────────────────────────────\n",
      "| magnetic field | tag='gradient' |\n",
      "───────────────────────────────────\n",
      ". Parameters :\n",
      "  ├── type : perfect gradient\n",
      "  ├── tag : gradient\n",
      "  ├── slope (T/m) : 0.5\n",
      "  ├── gradient direction : [0. 1. 0.]\n",
      "  ├── field direction : [0. 0. 1.]\n",
      "  ├── origin (m) : [0 0 0]\n",
      "  └── offset (T) :   0\n",
      "\n",
      "\n",
      "───────────────────────\n",
      "| zone | tag='vazewa' |\n",
      "───────────────────────\n",
      ". Parameters :\n",
      "  ├── type : 1D limits\n",
      "  ├── tag : vazewa\n",
      "  ├── in_tag\n",
      "  ├── out_tag : vazewa\n",
      "  ├── target : position\n",
      "  ├── action : stop\n",
      "  ├── min : 0.0\n",
      "  ├── max : 10.0\n",
      "  ├── axis : 0\n",
      "  └── inverted : False\n",
      "\n",
      "\n",
      "────────────────────────\n",
      "| Atom-light couplings |\n",
      "────────────────────────\n",
      ". transition > 'main' :\n",
      "  └── empty\n",
      "\n",
      ". transition > 'intercombination' :\n",
      "  └── empty\n",
      "\n",
      "\n",
      "# ------------ CONFIG INFO > STOP  ------------ #\n"
     ]
    }
   ],
   "source": [
    "config.print_info()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Manage objects in a collection\n",
    "\n",
    "Currently, our convention is that **copies** of objects are stored in the configuration object. This way, if an object is modified after it has been added in the configuration, this modification won't affect the configuration.\n",
    "\n",
    "It might be that we change this behaviour at some point... this is to be discussed.\n",
    "\n",
    "The configuration object comes with a lot of useful methods to manage the objects. Here are some examples ; for more information, please have a look at the **API documentation**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from atomsmltr.simulation import Configuration\n",
    "from atomsmltr.atoms import Ytterbium\n",
    "from atomsmltr.environment import Limits, MagneticOffset, MagneticGradient, GaussianLaserBeam\n",
    "\n",
    "# - INIT\n",
    "# let's define some objects\n",
    "laser1 = GaussianLaserBeam(tag=\"laser1\")\n",
    "laser2 = GaussianLaserBeam(tag=\"laser2\")\n",
    "mag_offset = MagneticOffset((0,1,0), tag=\"offset\")\n",
    "\n",
    "# init a configuration for ytterbium atoms\n",
    "config = Configuration(atom=Ytterbium())\n",
    "\n",
    "# add objects\n",
    "config += laser1, laser2, mag_offset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['laser1', 'laser2']\n",
      "['laser2']\n"
     ]
    }
   ],
   "source": [
    "# - REMOVE AN OBJECT\n",
    "print(config.list_lasers())\n",
    "config.rm_laser(\"laser1\")\n",
    "print(config.list_lasers())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[]\n",
      "['offset']\n",
      "[]\n"
     ]
    }
   ],
   "source": [
    "# - REMOVE ALL OBJECTS\n",
    "# all lasers\n",
    "config.rm_all_lasers()\n",
    "print(config.list_lasers())\n",
    "print(config.list_magnetic_fields())\n",
    "\n",
    "# everything\n",
    "config.rm_all_objects()\n",
    "print(config.list_magnetic_fields())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3.99e-07\n",
      "3.99e-07\n",
      "5.56e-07\n"
     ]
    }
   ],
   "source": [
    "# - UPDATE AN OBJECT\n",
    "# reset\n",
    "config.rm_all_objects()\n",
    "\n",
    "# add a laser > a copy is stored\n",
    "laser1.wavelength = 399e-9\n",
    "config += laser1\n",
    "\n",
    "# check wavelength of laser stored\n",
    "print(config.get_laser_copy(\"laser1\").wavelength)\n",
    "\n",
    "# update the laser > the stored version is not affected\n",
    "laser1.wavelength = 556e-9\n",
    "print(config.get_laser_copy(\"laser1\").wavelength)\n",
    "\n",
    "# to change the config, laser has to be updated\n",
    "config.update_objects(laser1)\n",
    "print(config.get_laser_copy(\"laser1\").wavelength)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting up atom-light couplings\n",
    "\n",
    "In order to give the user a full control on atom-light interactions, we decided to make the configuration object ignorant of any atom-light interaction coupling. Hence, it is not sufficient to add an atom and a laser with compatible wavelengths for the two to interact.\n",
    "\n",
    "Atom-light coupling are given by hand using the ``add_atomlight_coupling()`` method, associating a laser and a transition using their *tags*."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "────────────────────────\n",
      "| Atom-light couplings |\n",
      "────────────────────────\n",
      ". transition > 'main' :\n",
      "  └── laser '399' : detuning=-3.63e+08 (-2.00Γ)\n",
      "\n",
      ". transition > 'intercombination' :\n",
      "  └── empty\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from atomsmltr.simulation import Configuration\n",
    "from atomsmltr.atoms import Ytterbium\n",
    "from atomsmltr.environment import PlaneWaveLaserBeam\n",
    "\n",
    "# init a configuration for ytterbium atoms\n",
    "config = Configuration(atom=Ytterbium())\n",
    "\n",
    "# get ytterbium main transition information\n",
    "main = config.atom.trans[\"main\"]\n",
    "\n",
    "# create a plane wave laser\n",
    "laser = PlaneWaveLaserBeam()\n",
    "laser.wavelength = main.wavelength\n",
    "laser.set_power_from_I(main.Isat)  # set power to have I=Isat\n",
    "laser.tag = \"399\"\n",
    "\n",
    "# add laser to config\n",
    "config += laser\n",
    "\n",
    "# create a coupling\n",
    "config.add_atomlight_coupling(\n",
    "    laser=\"399\",\n",
    "    transition=\"main\",\n",
    "    detuning=-2 * main.Gamma,\n",
    ")\n",
    "\n",
    "# show info\n",
    "config.print_atomlight_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
}