{
 "cells": [
  {
   "cell_type": "raw",
   "id": "0",
   "metadata": {
    "editable": true,
    "raw_mimetype": "text/restructuredtext",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    ".. _notebooks-cavities-reference-phase:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1",
   "metadata": {},
   "source": [
    "# Cavities reference phase"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2",
   "metadata": {},
   "source": [
    "## Definitions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3",
   "metadata": {},
   "source": [
    "There are three ways to define the reference phase of a cavity:\n",
    "- by its relative entry phase $\\phi_{0,\\,\\mathrm{rel}}$,\n",
    "- its absolute entry phase $\\phi_{0,\\,\\mathrm{abs}}$,\n",
    "- its synchronous phase $\\phi_s$."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4",
   "metadata": {},
   "source": [
    "A cavity reference phase controls its behavior when the beam phase differs from the baseline, which happens downstream of a failure if the beam absolute phase was not recovered.\n",
    "- `\"phi_0_abs\"`: the cavity is not rephased.\n",
    "- `\"phi_0_rel\"`: $\\phi_{0,\\,\\mathrm{rel}}$ is preserved, so that the beam always \"sees\" the same RF phase.\n",
    "- `\"phi_s\"`: $\\phi_s$ is preserved, so that the cavity acceptance stays the same."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5",
   "metadata": {},
   "source": [
    "The relation between relative and entry phases is:\n",
    "\\begin{equation}\n",
    "E_0\\cos{\\phi_{0,\\,\\mathrm{abs}}} = E_0\\cos{(\\phi_{0,\\,\\mathrm{rel}} + \\phi_\\mathrm{in})}\n",
    "\\end{equation}\n",
    "where $\\phi_\\mathrm{in}$ is the phase at which the synchronous particle enters the cavity, $E_0$ is the amplitude of the electric field."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6",
   "metadata": {},
   "source": [
    "There is no analytic relation between $phi_s$ and the RF phases $\\phi_{0,\\,\\mathrm{rel}}$, $\\phi_{0,\\,\\mathrm{abs}}$. Hereby, when the reference phase of a cavity is $phi_s$, calculations are performed to find the actual RF phase allowing to retrieve $phi_s$. Those calculations can take some time."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7",
   "metadata": {},
   "source": [
    "The cavities reference phases are set by the `reference_phase_policy` entry in the `beam_calculator` section of the `TOML` file:\n",
    "- `\"phi_0_rel\"`, `\"phi_0_abs\"` or `\"phi_s\"`: use this phase as reference\n",
    "- `\"as_in_original_dat\"`: keep the reference stated in the `DAT` file, following the `SET_SYNC_PHASE` and the `FLAG_PHI_ABS` value of the `FIELD_MAP` commands.\n",
    "  - Different cavities can have different reference phase!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8",
   "metadata": {},
   "source": [
    "In the rest of this notebook, we illustrate these differences and their impact on the beam dynamics."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9",
   "metadata": {},
   "source": [
    "## Preparation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10",
   "metadata": {},
   "source": [
    "### Load libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "from pathlib import Path\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from lightwin.beam_calculation.beam_calculator import BeamCalculator\n",
    "from lightwin.beam_calculation.factory import BeamCalculatorsFactory\n",
    "from lightwin.beam_calculation.simulation_output.simulation_output import SimulationOutput\n",
    "from lightwin.config.config_manager import process_config\n",
    "from lightwin.constants import example_config, example_results\n",
    "from lightwin.core.accelerator.accelerator import Accelerator\n",
    "from lightwin.failures.fault_scenario import FaultScenario, FaultScenarioFactory\n",
    "from lightwin.util.log_manager import set_up_logging\n",
    "from lightwin.ui.workflow_setup import set_up_accelerators\n",
    "from lightwin.visualization import plot\n",
    "\n",
    "plt.rcParams[\"figure.figsize\"] = (15, 5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12",
   "metadata": {},
   "source": [
    "### Set the different solvers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "CONFIG_KEYS = {\n",
    "    'files': 'files',\n",
    "    'beam_calculator': 'generic_envelope1d',\n",
    "    'beam': 'beam',\n",
    "    'plots': 'plots_minimal',\n",
    "    'wtf': 'generic_wtf',\n",
    "    'design_space': 'generic_design_space',\n",
    "}\n",
    "override = {\n",
    "    'plots': {'energy': True, \"add_objectives\": False, \"twiss\": False},\n",
    "    'wtf': {'objective_preset': 'EnergyPhaseMismatch'},\n",
    "}\n",
    "config_common = process_config(example_config, CONFIG_KEYS, warn_mismatch=True, override=override)\n",
    "\n",
    "beam_calculators_factory = BeamCalculatorsFactory(**config_common)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14",
   "metadata": {},
   "source": [
    "#### Linac with absolute reference phase"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "override['beam_calculator'] = {'reference_phase_policy': \"phi_0_abs\"}\n",
    "config_abs = process_config(example_config, CONFIG_KEYS, warn_mismatch=True, override=override)\n",
    "\n",
    "# We will use this for plots representing cavity settings, phase\n",
    "plots_phase = config_abs['plots']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16",
   "metadata": {},
   "source": [
    "#### Linac with relative reference phase"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17",
   "metadata": {},
   "outputs": [],
   "source": [
    "override['beam_calculator'] = {'reference_phase_policy': \"phi_0_rel\"}\n",
    "config_rel = process_config(example_config, CONFIG_KEYS, warn_mismatch=True, override=override)\n",
    "\n",
    "# We will use this for plots representing cavity settings, phase, energy\n",
    "plots_complete = config_rel['plots']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18",
   "metadata": {},
   "source": [
    "### Linac with synchronous reference phase"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19",
   "metadata": {},
   "outputs": [],
   "source": [
    "override['beam_calculator'] = {'reference_phase_policy': \"phi_s\"}\n",
    "config_sync = process_config(example_config, CONFIG_KEYS, warn_mismatch=True, override=override)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20",
   "metadata": {},
   "source": [
    "### Create all BeamCalculators at once"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21",
   "metadata": {},
   "outputs": [],
   "source": [
    "policies = (\"abs\", \"rel\", \"sync\")\n",
    "\n",
    "solvers = beam_calculators_factory.run_all(\n",
    "    (\n",
    "        config_abs[\"beam_calculator\"],\n",
    "        config_rel[\"beam_calculator\"],\n",
    "        config_sync[\"beam_calculator\"],\n",
    "    )\n",
    ")\n",
    "\n",
    "solvers = {\n",
    "    policy: solver for policy, solver in zip(policies, solvers, strict=True)\n",
    "}\n",
    "\n",
    "# solver_abs.id += r' (absolute $\\phi_0$)'\n",
    "# solver_rel.id += r' (relative $\\phi_0$)'\n",
    "# solver_sync.id += r' ($\\phi_s$)'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22",
   "metadata": {},
   "source": [
    "### Set Accelerator objects"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23",
   "metadata": {},
   "source": [
    "Reduce logging level to keep only the warnings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24",
   "metadata": {},
   "outputs": [],
   "source": [
    "logging.getLogger().setLevel(logging.WARNING)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25",
   "metadata": {},
   "source": [
    "Note that LightWin warns us about the different reference phases."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26",
   "metadata": {},
   "outputs": [],
   "source": [
    "accelerators = {\n",
    "    policy: set_up_accelerators(config_common, solver)\n",
    "    for policy, solver in solvers.items()\n",
    "}"
   ]
  },
  {
   "cell_type": "raw",
   "id": "27",
   "metadata": {
    "editable": true,
    "raw_mimetype": "text/restructuredtext",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "Every value in `accelerators` links a :class:`.FaultScenario` index with corresponding :class:`.Accelerator`\\(s):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28",
   "metadata": {},
   "outputs": [],
   "source": [
    "accelerators[\"abs\"]"
   ]
  },
  {
   "cell_type": "raw",
   "id": "29",
   "metadata": {
    "editable": true,
    "raw_mimetype": "text/restructuredtext",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "To ease manipulation in this notebook, we also group `Reference` and `Broken`/`Fixed` :class:`.Accelerator`\\s:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30",
   "metadata": {},
   "outputs": [],
   "source": [
    "reference_accelerators = {\n",
    "    policy: accelerator[0][0]\n",
    "    for policy, accelerator in accelerators.items()\n",
    "}\n",
    "broken_accelerators = {\n",
    "    policy: accelerator[1][0]\n",
    "    for policy, accelerator in accelerators.items()\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31",
   "metadata": {},
   "source": [
    "## Propagate the beam"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32",
   "metadata": {},
   "source": [
    "### Nominal linac"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33",
   "metadata": {},
   "outputs": [],
   "source": [
    "for solver, ref in zip(solvers.values(), reference_accelerators.values(), strict=True):\n",
    "    simulation_output = solver.compute(ref)"
   ]
  },
  {
   "cell_type": "raw",
   "id": "34",
   "metadata": {
    "editable": true,
    "raw_mimetype": "text/restructuredtext",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    ".. note::\n",
    "   In Figure below, the reference accelerator calculated with `solver_abs` is showed twice (`Reference 0_Envelope1D`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35",
   "metadata": {},
   "outputs": [],
   "source": [
    "refs_for_plot = {\n",
    "    0: [reference_accelerators[\"abs\"]], # <- will be plotted twice, that's because plot faactory was not designed for\n",
    "    1: [reference_accelerators[\"rel\"]], #    this kind of plots where the 3 Accelerator are a reference\n",
    "    2: [reference_accelerators[\"sync\"]]\n",
    "}\n",
    "figs = plot.factory(refs_for_plot, plots_phase, save_fig=False, clean_fig=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36",
   "metadata": {},
   "source": [
    "In the nominal linac, there is no difference between the different reference phases."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-info\">\n",
    "The reference $\\phi_0$ that is used is the one defined in the `DAT`.\n",
    "As a matter of a fact, we need to propagate the beam a first time to know the entry phase of the synchronous particle in every cavity.\n",
    "This quantity is necessary to link $\\phi_{0,\\,\\mathrm{abs}}$ with $\\phi_{0,\\,\\mathrm{rel}}$.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38",
   "metadata": {},
   "source": [
    "### Linac with one broken cavity"
   ]
  },
  {
   "cell_type": "raw",
   "id": "39",
   "metadata": {
    "editable": true,
    "raw_mimetype": "text/restructuredtext",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "We create the one :class:`.FaultScenario` per solver; internally, the reference phase of cavities downstream the failure will follow the :class:`.BeamCalculator` setting `\"reference_phase_policy\"`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40",
   "metadata": {},
   "outputs": [],
   "source": [
    "design_space = config_common[\"design_space\"]\n",
    "fault_scenario_factories = {\n",
    "    policy: FaultScenarioFactory(accelerators[policy], solvers[policy], design_space)\n",
    "    for policy in policies\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41",
   "metadata": {},
   "outputs": [],
   "source": [
    "fault_scenarios = {\n",
    "    policy: factory.create(**config_common[\"wtf\"])\n",
    "    for policy, factory in fault_scenario_factories.items()\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42",
   "metadata": {},
   "outputs": [],
   "source": [
    "for policy in policies:\n",
    "    solver = solvers[policy]\n",
    "    broken = broken_accelerators[policy]\n",
    "    simulation_output = solver.compute(broken)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43",
   "metadata": {},
   "source": [
    "### Absolute reference phase"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44",
   "metadata": {},
   "source": [
    "Every cavity downstream the failure will keep its absolute phase. In other words, cavities are not rephased at all."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45",
   "metadata": {},
   "outputs": [],
   "source": [
    "figs = plot.factory(accelerators[\"abs\"], plots_complete, save_fig=False, clean_fig=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46",
   "metadata": {},
   "source": [
    "Cavities in their nominal tuning are green, and the failed cavity is represented in red.\n",
    "Here the beam energy and phase remain very close to the baseline. However, the positive synchronous phases show that high losses would be to expect."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47",
   "metadata": {},
   "source": [
    "### Relative reference phase"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48",
   "metadata": {},
   "source": [
    "Every cavity downstream the failure will keep its relative phase. In other words, if the beams is 20 degrees late at the exit of the failed cavity, all downstream cavities are rephased of 20 degrees."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49",
   "metadata": {},
   "outputs": [],
   "source": [
    "figs = plot.factory(accelerators[\"rel\"], plots_complete, save_fig=False, clean_fig=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50",
   "metadata": {},
   "source": [
    "In this example, the error on the beam energy is progressively damped. In LightWin, rephased cavities are represented with olive color."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51",
   "metadata": {},
   "source": [
    "### Synchronous reference phase"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52",
   "metadata": {},
   "source": [
    "Every cavity downstream the failure will keep its synchronous phase. In other words, all cavities downstream the failure are rephased to maintain the nominal synchronous phases."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53",
   "metadata": {},
   "outputs": [],
   "source": [
    "figs = plot.factory(accelerators[\"sync\"], plots_complete, save_fig=False, clean_fig=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54",
   "metadata": {},
   "source": [
    "Notice that the synchronous phases are unaffected by failed cavity."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55",
   "metadata": {},
   "source": [
    "## Fixing the fault"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56",
   "metadata": {},
   "outputs": [],
   "source": [
    "for fault_scenario in fault_scenarios.values():\n",
    "    for scenario in fault_scenario:\n",
    "        scenario.fix_all()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57",
   "metadata": {},
   "source": [
    "### Absolute phase linac"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58",
   "metadata": {},
   "source": [
    "For the absolute phase linac, we try to retrieve the energy and the phase at the end of the compensation zone, as well as to minimize the mismatch factor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59",
   "metadata": {},
   "outputs": [],
   "source": [
    "figs = plot.factory(accelerators[\"abs\"], plots_complete, save_fig=False, clean_fig=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60",
   "metadata": {},
   "source": [
    "### Relative phase linac"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61",
   "metadata": {},
   "source": [
    "When downstream cavities are rephased, it is no longer mandatory to retrieve nominal beam phase at the exit of the compensation zone. Here, we try to retrieve the kinetic energy and to minimize the mismatch factor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62",
   "metadata": {},
   "outputs": [],
   "source": [
    "figs = plot.factory(accelerators[\"rel\"], plots_complete, save_fig=False, clean_fig=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63",
   "metadata": {},
   "source": [
    "### Synchronous phase linac"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64",
   "metadata": {},
   "source": [
    "Results will be very similar to the relative phase linac."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65",
   "metadata": {},
   "outputs": [],
   "source": [
    "figs = plot.factory(accelerators[\"sync\"], plots_complete, save_fig=False, clean_fig=False)"
   ]
  }
 ],
 "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
}