{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Walkthrough 3: Optical barriers #"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introduction ##"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next level of control of quantum matter is made possible by altering the potential energy as a function of position by applying far-detuned light fields to the atom ensemble during the manipulation phase.  In the current hardware implementation, the \"painted\" light introduced to the atom ensemble is blue-detuned from resonance, and thus represents a *repulsive* potential energy to the atoms via the AC Stark shift.  The spatially-dependent potential energy depends on the local intensity of the light field, thus dynamic control of the local intensity alters the potential energy landscape in which the atoms are trapped.  In an abstract sense, through Oqtant users have control of this light field by creating and manipulating one or more `Barrier` objects, which are superimposed on the background magnetic potential that traps the atoms.\n",
    "\n",
    "In this walkthrough, we will explore application of quasi-1D optical barriers on Oqtant's cigar-shaped atomic ensembles.  Users have control over the dynamic position (in microns), width (in microns), energetic height (in kHz), and shape of the barrier objects along the long dimension of the ensemble.  The effective 1D potential is achieved by quickly rastering/scanning the applied light along one short axis, while the light propagates along the other short axis.  Application of these barriers allows for the ensemble to be split, recombined, swept to one side, or any number of other desired manipulations.  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Imports and user authentication ###"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from oqtant.oqtant_client import get_oqtant_client\n",
    "from oqtant.util.auth import notebook_login\n",
    "from oqtant.schemas.quantum_matter import QuantumMatterFactory as qmf\n",
    "\n",
    "oqtant_account = notebook_login()\n",
    "oqtant_account"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client = get_oqtant_client(oqtant_account.access_token)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Barrier objects  ##"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Barriers are included in QuantumMatter and associated OqtantJob objects by passing a list of `Barrier` objects during instantiation.  We begin by exploring how these (effectively 1D) objects are described and instantiated.   "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Creating a Barrier object ###"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us begin by creating a Gaussian-shaped barrier centered at a position of 10 (microns) with an energetic height of 30 (kHz) and a width of 3 (microns).  Here, a position of 0 corresponds to the geometric center of the atom ensemble along the long dimension.  Let us plan on having a lifetime of 10 ms for the manipulation phase of our QuantumMatter object, with our barrier being applied starting at t = 2.5 ms and lasting until t = 7.5 ms, i.e. a barrier lifetime of 5 ms.  Recall that the manipulation phase starts, by definition, at t = 0.  Just as for QuantumMatter objects themselves, we create Barrier objects using our QuantumMatterFactory (qmf):  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "barrier = qmf.create_barrier(\n",
    "    position=10,\n",
    "    height=30,\n",
    "    width=3,\n",
    "    birth=2.5,  # barrier initially applied here\n",
    "    lifetime=5,  # before being destroyed after this amount of time\n",
    "    shape=\"GAUSSIAN\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualizing a Barrier object at fixed times ###"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What does our barrier look like?  We can use the `Barrier.show_potential()` method to find out.  This method accepts a list of times for which to plot the position-dependent potential energy represented by the barrier:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "barrier.show_potential(times=[0, 5])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that at t = 0, our barrier hasn't been \"born\" yet, and the potential energy is zero everywhere.  However,  at t = 5 our barrier is now present (and will remain present until the end of its lifetime at t = 7.5).  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualizing a Barrier's dynamics ###"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also view a barrier's dynamics over the time of the manipulation phase.    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "barrier.show_dynamics()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see the constructed values (of position, width, and height) for times between our barrier's \"birth\" (at t = 2.5) and its \"death\" (at t = 7.5).  Outside of its lifetime, this barrier will not exist (no light will be projected onto the atoms, at least from this particular barrier).  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Scripting a Barrier's behavior ###"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So far, our barrier is *static*, i.e. its position, width, and height are all constants during its lifetime.  If we want to add some dynamics to the barrier, we can use the `Barrier.evolve()` method, which allows us to generically script a barrier's behavior in time.    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "barrier = qmf.create_barrier(\n",
    "    position=-10,\n",
    "    height=50,\n",
    "    width=3,\n",
    "    birth=1,  # start with these parameters at t = 1\n",
    "    lifetime=2,  # and hold with those conditions for 2 ms\n",
    "    shape=\"GAUSSIAN\",\n",
    ")\n",
    "\n",
    "# add dynamics, automatically increasing the barrier lifetime\n",
    "barrier.evolve(duration=2.5, position=25, height=25, width=5)\n",
    "\n",
    "# and more dynamics after that\n",
    "barrier.evolve(duration=5, position=-25, height=5, width=1)\n",
    "\n",
    "# show the barrier at the following times, along with the overall dynamics\n",
    "barrier.show_dynamics()\n",
    "barrier.show_potential(times=[2, 5, 9])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that we now have a highly dynamic barrier whose position, height, and width change over the course of its lifetime.  Any parameters that are to remain unchanged during a scripted evolution step can be omitted from the call to `Barrier.evolve()`.    \n",
    "\n",
    "NOTE: a barrier object's shape is not dynamically controllable.  However, such behavior is supported by so-called `Landscape` objects, introduced in the following walkthrough.\n",
    "\n",
    "NOTE: A barrier object's *lifetime* will automatically be extended by the *duration* parameter passed to the `Barrier.evolve()` method.  \n",
    "\n",
    "If you wish to immediately start evolving the barrier after creation, you can omit the *lifetime* parameter passed to the instantiation and then ensure to evolve the barrier as desired:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# start with these conditions\n",
    "barrier = qmf.create_barrier(position=0, height=25, width=7, shape=\"GAUSSIAN\", birth=1)\n",
    "\n",
    "# and immediately start evolving\n",
    "barrier.evolve(duration=2.5, position=25, height=25, width=5)\n",
    "\n",
    "# show me\n",
    "barrier.show_dynamics()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Available Barrier shapes ###"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Supported Barrier shape options include GAUSSIAN, SQUARE, and LORENTZIAN.  In all three cases, the width parameter means something slightly different, but for consistency generally refers to the appropriate *half width* of the functional form.  The below table summarizes the potential energy ($U(x)$) as a function of position ($x$), barrier height $h$ , position $x_0$, and width $w$, for the currently supported barrier shapes:\n",
    "\n",
    "| specified shape | $U(x)$ functional form | width ($w$) interpretation               |\n",
    "|-----------------|---------------------------------------------|---------------------|\n",
    "| GAUSSIAN        | $ h e^{-\\frac{(x - x_0)^2}{2 w^2}} $        | gaussian width      |\n",
    "| LORENTZIAN      | $ \\frac{h}{1 + (x - x_0)^{2} / w^{2}} $     | half-width half-max |\n",
    "| SQUARE          | $h$ if $abs(x-x_0) \\leq w$, else $0$        | half width            |\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Barriers of different shape can be created as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "shapes = [\"GAUSSIAN\", \"LORENTZIAN\", \"SQUARE\"]\n",
    "for shape in shapes:\n",
    "    barrier = qmf.create_barrier(\n",
    "        position=0, height=50, width=3, shape=shape, lifetime=10\n",
    "    )\n",
    "    barrier.show_potential([5])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Adding Barriers to QuantumMatter objects ##"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "All that we need to do to add one or more *Barrier* objects to a *QuantumMatter* object is pass a list of them as a `barriers` input parameter during instantiation:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Instantiation ###"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define the first barrier, with initial hold + evolution\n",
    "barrier1 = qmf.create_barrier(\n",
    "    position=10, height=30, width=3, birth=2.5, lifetime=5, shape=\"GAUSSIAN\"\n",
    ")\n",
    "barrier1.evolve(duration=2, position=20, height=20, width=5)\n",
    "\n",
    "# and a second one, also with initial hold + evolution\n",
    "barrier2 = qmf.create_barrier(\n",
    "    position=-10, height=30, width=3, birth=2.5, lifetime=7.5, shape=\"SQUARE\"\n",
    ")\n",
    "barrier2.evolve(duration=4, position=-25, height=5, width=5)\n",
    "\n",
    "# and construct the desired QuantumMatter object\n",
    "barrier_matter = qmf.create_quantum_matter(\n",
    "    temperature=100,\n",
    "    lifetime=15,\n",
    "    time_of_flight=10,\n",
    "    barriers=[barrier1, barrier2],\n",
    "    name=\"Now with barriers!\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualizing multi-barrier dynamics of the QuantumMatter object ###"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The dynamics of all barrier objects in a QuantumMatter object can be visualized simultaneously using the `QuantumMatter.show_barrier_dynamics()` method:  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "barrier_matter.show_barrier_dynamics()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualizing the total potential energy ###"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It is also useful to visualize the overall *total* spatial energy profile at specific times in the manipulation phase, including contributions from both the (magnetic) trapping potential as well as the (optical) barrier contributions.  This can be achieved using the `QuantumMatter.show_potential()` method.  Just like the similar method available to Barrier objects themselves, this method accepts a list of times for which to plot the potential energy:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "barrier_matter.show_potential(times=[0, 5, 14])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we can see, at t = 0 the potential energy profile is identical to the purely-magnetic trapping potential seen in the last walkthrough as no barriers have yet been created.  At t = 5, we have two barriers of different shapes centered at +/- 10 microns, superimposed onto the background magnetic trapping potential.  At t = 14 ms, barrier 1 has completed its lifetime and is no longer present, leaving only barrier 2, which by that time has shifted over to a position of -25, increased its width, and reduced its height.  \n",
    "\n",
    "WARNING: The total *optical* potential energy at a single position is limited to be 100 kHz in the current hardware implementation.  Tall barriers that are close enough to spatially overlap can encounter this limit.  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## In-trap imaging ##"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In addition to time-of-flight imaging, where the atom ensemble is released from the trapping potential some time before the image of the atoms is taken, Oqtant also supports the option to take the image of the atoms *in-trap*.  This option allows for direct imaging of atom ensembles that are swept, split, or shifted, while or shortly after barriers are being applied.  These manipulations can act as the basis for studying quantum dynamics, atomtronics, and developing quantum sensors.  In-trap imaging allows images of the atoms to be taken during these manipulations.\n",
    "\n",
    "NOTE: While in-trap imaging is useful for revealing the spatial distribution of atoms within the ensemble at the time they are imaged, it is less useful for extracting quantitative information such as temperature, condensed fraction, etc.  Hence, these calculations are not done when this imaging option is specified. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a simple static barrier\n",
    "barrier = qmf.create_barrier(\n",
    "    position=-10, height=30, width=3, shape=\"GAUSSIAN\", birth=2.5, lifetime=7.5\n",
    ")\n",
    "\n",
    "# and QuantumMatter with a lifetime so that the barrier exists right up to the end\n",
    "barrier_matter = qmf.create_quantum_matter(\n",
    "    temperature=100,\n",
    "    lifetime=barrier.death,\n",
    "    barriers=[barrier],\n",
    "    image=\"IN_TRAP\",  # take an in-trap image, time_of_flight input parameter is not relevant in this case\n",
    "    note=\"in-trap imaging with barrier\",\n",
    ")\n",
    "\n",
    "job_id = client.submit(barrier_matter, track=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The output data for jobs that use the in-trap imaging option is somewhat different than those that use the default time of flight imaging option:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "my_barrier_job = client.get_job(job_id)\n",
    "print(my_barrier_job.output_fields)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As discussed above, no temperature or ensemble population calculations are done in this case.  Instead, the output is limited to just the OD image, which can be viewed as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "my_barrier_job.plot_it()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Advanced topics and discussion ##"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Hardware implementation limits and effects ###"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Projected barrier resolution ####"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "In the hardware implementation, the optical projection system (that projects the barrier light) has finite resolution.  In its current form, the smallest diameter feature that can be projected, in terms of the 1/e^2 diameter, is around 2.5 microns.  Barrier programs with very thin barriers can reveal these hardware design choices, resulting in, among other effects, barrier shapes losing fidelity as the barrier width becomes narrower.  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Spatial resolution of barrier positions and widths ####"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Currently, the effective one-dimensional barriers are calculated and projected on an *spatial* grid with a spacing of 1 micron.  This means, for instance, that if a smoothly increasing barrier width is specified that it will actually increase in discrete steps -- both due to the temporal update time of recalculating barrier parameters every 100 microseconds (dynamics limitations, discussed above) and due to this spatial calculation grid.  The largest jumps in, e.g. width, will occur for very narrow barriers as their width transitions across integer values (depending on the barrier center) -- at the next timestep the barrier width will suddenly have a new position/frequency contribution and it will increase in a discrete manner.  A similar phenomenon occurs as a barrier is scanned in position.   "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Dynamics limitations ####"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Another hardware/implementation limit can be observed if very fast dynamics are programmed into barrier objects.  The optical projection system has finite bandwidth and, as a result, dynamic barrier widths/positions/heights are only updated every 100 microseconds.  Between these periodic updates, the instantaneous barrier parameters, determined by the dynamics of the projected light, are static.  The corresponding dynamics frequency of 10 kHz, and the short-axis projected light raster frequency of 20 kHz, are both much faster than the natural frequency of atoms in the trap along the long axis, which is around 40 Hz.  Thus, the potential energy applied by these light fields is effectively averaged for atoms in the ensemble.  \n",
    "\n",
    "We can show this \"update rate\" limitation in the visualization tools for barrier dynamics etc. by passing the optional *corrected* boolean flag, which defaults to *False*, to our *show()* methods.  This can be applied at the level of a Barrier object or an entire program:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define the first barrier, with initial hold + evolution\n",
    "barrier1 = qmf.create_barrier(\n",
    "    position=10, height=30, width=3, birth=2.5, lifetime=5, shape=\"GAUSSIAN\"\n",
    ")\n",
    "barrier1.evolve(duration=2, position=20, height=20, width=5)\n",
    "\n",
    "# and a second one, also with initial hold + evolution\n",
    "barrier2 = qmf.create_barrier(\n",
    "    position=-10, height=30, width=3, birth=2.5, lifetime=7.5, shape=\"SQUARE\"\n",
    ")\n",
    "barrier2.evolve(duration=4, position=-25, height=5, width=5)\n",
    "\n",
    "# and construct the desired QuantumMatter object\n",
    "barrier_matter = qmf.create_quantum_matter(\n",
    "    temperature=100,\n",
    "    lifetime=15,\n",
    "    time_of_flight=10,\n",
    "    barriers=[barrier1, barrier2],\n",
    "    name=\"Now with barriers!\",\n",
    ")\n",
    "\n",
    "barrier_matter.show_barrier_dynamics(corrected=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can observe the discrete steps of barrier parameters at the update period of 100 microseconds.  These updates are more obvious when one or more of the parameters are changing rapidly.  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Peeking under the hood ###"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the background of Oqtant, barrier programs have nearly identical structure as quantum programs.  The only difference is that the *optical_barriers* field has been populated.  As before, the underlying data structure can be useful for advanced program construction or manipulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(barrier_matter.model_dump())"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
