{
 "cells": [
  {
   "cell_type": "raw",
   "id": "0",
   "metadata": {
    "editable": true,
    "raw_mimetype": "text/restructuredtext",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    ".. _notebooks-lightwin-example:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1",
   "metadata": {},
   "source": [
    "# LightWin compensation example"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2",
   "metadata": {},
   "source": [
    "This notebook showcases how to compensate for a single failure."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4",
   "metadata": {},
   "source": [
    "Note that the first import is only used to import packaged data file.\n",
    "See also: https://learn.scientific-python.org/development/patterns/data-files/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from importlib import resources\n",
    "from pprint import pprint\n",
    "\n",
    "# Actual mandatory imports for LightWin\n",
    "from lightwin.config.config_manager import process_config\n",
    "from lightwin.ui.workflow_setup import run_simulation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "## Presentation of the linac"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "The structure of the linac is stored in a `DAT` file.\n",
    "It follows the same specifications as TraceWin structure file, see [associated help page](../usage.compatibility_with_tracewin.rst) for more information.\n",
    "The example can be found in `data/ads/ads.dat` (or `src/lightwin/data/ads/ads.dat` if you built LightWin from source).\n",
    "Here are the first lines of the file:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "dat_res = resources.files(\"lightwin.data.ads\") / \"ads.dat\"\n",
    "dat_content = dat_res.read_text(encoding=\"utf-8\")\n",
    "print(dat_content[:892])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "This is a high-power, ADS-like proton acceleractor characterized by a very high longitudinal acceptance."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "## Configuration of LightWin"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "In order to ensure reproducibility and portability of results, all the configuration is set in a `TOML` configuration.\n",
    "The full specifications are listed in [Configuration](../configuration/configuration.rst)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "toml_filepath = resources.files(\"lightwin.data.ads\") / \"lightwin.toml\"  # Or provide Path to TOML file"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13",
   "metadata": {},
   "source": [
    "Below, we associate LightWin's configuration objects with sections in the `TOML` file (the names between brackets)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "toml_keys = {\n",
    "    \"files\": \"files\",\n",
    "    \"beam\": \"beam\",\n",
    "    \"beam_calculator\": \"generic_envelope1d\",\n",
    "    \"wtf\": \"generic_wtf\",\n",
    "    \"design_space\": \"fit_phi_s_design_space\",\n",
    "    \"plots\": \"plots_minimal\",\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15",
   "metadata": {},
   "source": [
    "This operation will check validity of the configuration, resolve paths, and  convert the `TOML` appropriate entries to a dictionary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Adapt some configuration entries for this notebook\n",
    "override = {\n",
    "    \"beam_calculator\": {\n",
    "        \"flag_cython\": True,  # Override flag_cython=False in [generic_envelope1d]\n",
    "    },\n",
    "    \"wtf\": {\n",
    "        \"k\": 5,               # Override k=3 in [generic_wtf]\n",
    "    },\n",
    "    \"plots\": {\n",
    "        \"kwargs\":\n",
    "        {\n",
    "            \"lw\": 5,          # Increase size of lines in plots\n",
    "        },\n",
    "    },\n",
    "}\n",
    "config = process_config(toml_filepath, toml_keys, override=override)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "### `files` ([doc](../configuration/files.rst))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "Defines the files to use, as well as where results should be saved."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "pprint(config[\"files\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "### `beam` ([doc](../configuration/beam.rst))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "Define the beam properties."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "pprint(config[\"beam\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "### `beam_calculator`([doc](../configuration/beam_calculator.rst))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "Define how the propagation of the beam in the linac will be calculated.\n",
    "Note that you can define a `beam_calculator_post` in a similar fashion; it will be used to re-compute the propagation of the beam once the compensation settings are found."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "pprint(config[\"beam_calculator\"])"
   ]
  },
  {
   "cell_type": "raw",
   "id": "26",
   "metadata": {
    "editable": true,
    "raw_mimetype": "text/restructuredtext",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "Some comments on the settings we chose:\n",
    "\n",
    "* :class:`.Envelope1D` is usually a good choice, as transverse dynamics are not very important for failure compensation.\n",
    "* ``n_steps_per_cell=40`` with ``method=\"RK4\"`` are a good balance between precision and rapidity.\n",
    "* ``export_phase=\"phi_0_abs\"`` and ``reference_phase_policy=\"phi_0_abs\"`` to avoid implicit rephasing of the linac. See also: `dedicated notebook`_.\n",
    "* ``flag_cython=True`` to speed up the calculations. The speed-up factor is 2 to 4 on my machine. See also :ref:`documentation<Cython-setup>`.\n",
    "\n",
    ".. _dedicated notebook: cavities_reference_phase.ipynb"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "### `plots` ([doc](../configuration/plots.rst))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28",
   "metadata": {
    "editable": true,
    "raw_mimetype": "",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "Define the plots to create at the end of the simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "pprint(config[\"plots\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "### `wtf` ([doc](../configuration/wtf.rst))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "Stands for \"what to fit\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "pprint(config[\"wtf\"])"
   ]
  },
  {
   "cell_type": "raw",
   "id": "33",
   "metadata": {
    "editable": true,
    "raw_mimetype": "text/restructuredtext",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "\n",
    "* ``failed=[[\"FM11\"]]`` and ``id_nature=\"name\"`` define one fault scenario, with one failed cavity which name is \"FM11\".\n",
    "\n",
    "   - We try to follow TraceWin naming conventions.\n",
    "\n",
    "* ``strategy=\"k out of n\"`` with ``k=5`` to compensate every failed cavity with the five closest cavities. The compensation strategies are defined in the :mod:`.strategy` module.\n",
    "\n",
    "   - We overrode ``k=3`` from original ``lightwin.toml`` file.\n",
    "     The optimization process will take more time, but three compensating cavities are not enough to retrieve a perfect beam.\n",
    "\n",
    "* ``objective_preset=\"EnergyPhaseMismatch\"`` defines three objectives to minimize: energy and phase difference with nominal linac at the end of the last compensating cavity and longitudinal mismatch factor. The objectives are defined in the :mod:`.objective.factory` module.\n",
    "* ``optimisation_algorithm=\"downhill_simplex\"`` to use :class:`.DownhillSimplex` optimization algorithm.\n",
    "\n",
    "   - Note that you can provide keyword arguments to personalize the execution of the algorithm.\n",
    "\n",
    "   - Downhill simplex is very fast and effective for this kind of problem, with few variables.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "### `design_space` ([doc](../configuration/design_space.rst))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "Defines the design space, *i.e.* the optimization variables and their bounds."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "pprint(config[\"design_space\"])"
   ]
  },
  {
   "cell_type": "raw",
   "id": "37",
   "metadata": {
    "editable": true,
    "raw_mimetype": "text/restructuredtext",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "With ``design_space_preset=\"sync_phase_amplitude\"``, the variables will be the synchronous phase and the field amplitude of every cavity.\n",
    "This generally works well with optimization algorithm that do not support constraints such as :class:`.DownhillSimplex`, because it allows one to bound the synchronous phase even without constraints.\n",
    "See also :mod:`.design_space.factory` for the various design space presets.\n",
    "\n",
    ".. note::\n",
    "   You can provide a file to define different amplitude/synchronous phase bounds for every cavity.\n",
    "   This is particularly useful for real-life problems, where every cavity has it's own history and behavior."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "## Actual usage"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39",
   "metadata": {},
   "source": [
    "The following function will break and fix the linac.\n",
    "Check out the source code of [the workflow_setup module](../../lightwin/lightwin.ui.workflow_setup.rst) if you want to understand and personnalize this workflow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "config[\"plots\"][\"kwargs\"] = {\"lw\": 10}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "fault_scenarios = run_simulation(config)"
   ]
  },
  {
   "cell_type": "raw",
   "id": "42",
   "metadata": {
    "editable": true,
    "raw_mimetype": "text/restructuredtext",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    ".. todo::\n",
    "   Better display of Figures in ``HTML`` documentation."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "## Notes on objects hierarchy"
   ]
  },
  {
   "cell_type": "raw",
   "id": "44",
   "metadata": {
    "editable": true,
    "raw_mimetype": "text/restructuredtext",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "The unique :class:`.FaultScenario` we studied can be accessed *via*:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "fault_scenario = fault_scenarios[0]"
   ]
  },
  {
   "cell_type": "raw",
   "id": "46",
   "metadata": {
    "editable": true,
    "raw_mimetype": "text/restructuredtext",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "It holds two :class:`.Accelerator` objects, each containing a :class:`.ListOfElements`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "reference_accelerator = fault_scenario.ref_acc\n",
    "reference_elts = reference_accelerator.elts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "fixed_accelerator = fault_scenario.fix_acc\n",
    "fixed_elts = fixed_accelerator.elts"
   ]
  },
  {
   "cell_type": "raw",
   "id": "49",
   "metadata": {
    "editable": true,
    "raw_mimetype": "text/restructuredtext",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "The results are stored in a :class:`.SimulationOutput` object.\n",
    "These latter objects are stored in a dictionary, which keys are the name of the :class:`.BeamCalculator` that were used.\n",
    "In this case, we should only have ``\"CyEnvelope1D_0\"``.\n",
    "A second simulation with the :class:`.TraceWin` solver would add a second :class:`.SimulationOutput`; corresponding key would be named ``\"TraceWin_1\"``.\n",
    "The number at the end of the name represents the order in which the :class:`.BeamCalculator` were defined."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "fixed_simulation_outputs = fixed_accelerator.simulation_outputs\n",
    "pprint(fixed_simulation_outputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "solver_name = list(fixed_simulation_outputs.keys())[0]\n",
    "fixed_simulation_output = fixed_simulation_outputs[solver_name]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "## Accessing data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "The easiest way to access the calculated data is to use the [get method](https://lightwin.readthedocs.io/en/latest/manual/get_method.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "fixed_simulation_output.get(\"w_kin\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "fixed_simulation_output.get(\"phi_s\", to_deg=True, elt=\"FM53\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "fixed_simulation_output.get(\"beta\", phase_space_name=\"phiw\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "fixed_simulation_output.get(\"envelope_energy_zdelta\", elt=\"FM48\", pos=\"out\", to_numpy=False)"
   ]
  },
  {
   "cell_type": "raw",
   "id": "58",
   "metadata": {
    "editable": true,
    "raw_mimetype": "text/restructuredtext",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    ".. note::\n",
    "   In most case, you will want to use :meth:`.SimulationOutput.get` method over :meth:`.Accelerator.get` or :meth:`.ListOfElements.get`, as stated in the message below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "fixed_accelerator.get(\"beta\", phase_space_name=\"phiw\")"
   ]
  }
 ],
 "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.14.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}