"""Define a base class for beam propagation computing tools.
The base class :class:`BeamCalculator` allows to compute the propagation of
the beam in a :class:`.ListOfElements`, possibly with a specific
:class:`.SetOfCavitySettings` (optimisation process). It should return a
:class:`.SimulationOutput`.
.. todo::
Precise that BeamParametersFactory and TransferMatrixFactory are mandatory.
"""
import datetime
import logging
import time
from abc import ABC, abstractmethod
from itertools import count
from pathlib import Path
from typing import Any
from lightwin.beam_calculation.parameters.factory import (
ElementBeamCalculatorParametersFactory,
)
from lightwin.beam_calculation.simulation_output.factory import (
SimulationOutputFactory,
)
from lightwin.beam_calculation.simulation_output.simulation_output import (
SimulationOutput,
)
from lightwin.core.accelerator.accelerator import Accelerator
from lightwin.core.elements.field_maps.cavity_settings_factory import (
CavitySettingsFactory,
)
from lightwin.core.list_of_elements.factory import ListOfElementsFactory
from lightwin.core.list_of_elements.list_of_elements import ListOfElements
from lightwin.failures.set_of_cavity_settings import SetOfCavitySettings
from lightwin.util.typing import (
EXPORT_PHASES_T,
REFERENCE_PHASES_T,
BeamKwargs,
)
[docs]
class BeamCalculator(ABC):
"""Store a beam dynamics solver and its results."""
_ids = count(0)
[docs]
def __init__(
self,
flag_phi_abs: bool,
out_folder: Path | str,
default_field_map_folder: Path | str,
beam_kwargs: BeamKwargs,
flag_cython: bool = False,
export_phase: EXPORT_PHASES_T = "as_in_settings",
**kwargs,
) -> None:
r"""Set ``id``, some generic parameters such as results folders.
Parameters
----------
flag_phi_abs :
If the entry phase of the cavities :math:`\phi_0` are absolute or
relative. See the examples for an illustration of what it implies.
out_folder : pathlib.Path | str
Name of the folder where results should be stored, for each
:class:`.Accelerator` under study. This is the name of a folder,
not a full path.
default_field_map_folder : pathlib.Path | str
Where to look for field map files by default.
flag_cython :
If the beam calculator involves loading cython field maps. The
default is False.
beam_kwargs :
The config dictionary holding all the initial beam properties.
export_phase :
"as_in_settings", "as_in_original_dat"], optional
The type of phase you want to export for your ``FIELD_MAP``. The
default is ``"as_in_settings"``, which should be the same phases
as in the original DAT file.
"""
self.flag_phi_abs = flag_phi_abs
self.flag_cython = flag_cython
self.id: str = f"{self.__class__.__name__}_{next(self._ids)}"
self._export_phase: EXPORT_PHASES_T = export_phase
if isinstance(out_folder, str):
out_folder = Path(out_folder)
self.out_folder = out_folder
if isinstance(default_field_map_folder, str):
default_field_map_folder = Path(default_field_map_folder)
self.default_field_map_folder = (
default_field_map_folder.resolve().absolute()
)
self._beam_kwargs = beam_kwargs
self.simulation_output_factory: SimulationOutputFactory
self.list_of_elements_factory: ListOfElementsFactory
self.beam_calc_parameters_factory: (
ElementBeamCalculatorParametersFactory
)
self._set_up_common_factories()
self._set_up_specific_factories()
[docs]
def _set_up_common_factories(self) -> None:
"""Create the factories declared in :meth:`__init__`.
.. todo::
``default_field_map_folder`` has a wrong default value. Should take
path to the ``DAT`` file, that is not known at this point. Maybe
handle this directly in the :class:`.InstructionsFactory` or
whatever.
"""
self.list_of_elements_factory = ListOfElementsFactory(
self.is_a_3d_simulation,
self.is_a_multiparticle_simulation,
beam_kwargs=self._beam_kwargs,
default_field_map_folder=self.default_field_map_folder,
load_field_maps=True, # useless with TraceWin
field_maps_in_3d=False, # not implemented anyway
# Different loading of field maps if Cython
load_cython_field_maps=self.flag_cython,
elements_to_dump=(),
)
[docs]
@abstractmethod
def _set_up_specific_factories(self) -> None:
"""Set up the factories specific to the :class:`.BeamCalculator`."""
[docs]
def run(
self,
elts: ListOfElements,
update_reference_phase: bool = False,
**kwargs,
) -> SimulationOutput:
"""Perform a simulation with default settings.
.. todo::
``update_reference_phase`` is currently unused, because it is not
useful once the propagation has been calculated. So... should I
keep it? Maybe it can be useful in post_optimisation_run_with_this,
or in scripts to convert the phase between the different
references, or when I want to save the .dat?
Parameters
----------
elts :
List of elements in which the beam must be propagated.
update_reference_phase :
To change the reference phase of cavities when it is different from
the one asked in the ``TOML``. To use after the first calculation,
if ``BeamCalculator.flag_phi_abs`` does not correspond to
``CavitySettings.reference``. The default is False.
kwargs
Other keyword arguments passed to :meth:`run_with_this`. As for
now, only used by :class:`.TraceWin`.
Returns
-------
simulation_output : SimulationOutput
Holds energy, phase, transfer matrices (among others) packed into a
single object.
"""
simulation_output = self.run_with_this(
None, elts, use_a_copy_for_nominal_settings=False, **kwargs
)
if update_reference_phase:
elts.force_reference_phases_to(self.reference_phase)
return simulation_output
[docs]
@abstractmethod
def run_with_this(
self,
set_of_cavity_settings: SetOfCavitySettings | None,
elts: ListOfElements,
use_a_copy_for_nominal_settings: bool = True,
) -> SimulationOutput:
"""Perform a simulation with new cavity settings.
Calling it with ``set_of_cavity_settings = None`` shall be the same as
calling the plain ``run`` method.
Parameters
----------
set_of_cavity_settings :
Holds the norms and phases of the compensating cavities.
elts :
List of elements in which the beam should be propagated.
use_a_copy_for_nominal_settings :
To copy the nominal :class:`.CavitySettings` and avoid altering
their nominal counterpart. Set it to True during optimisation, to
False when you want to keep the current settings. The default is
True.
Returns
-------
simulation_output : SimulationOutput
Holds energy, phase, transfer matrices (among others) packed into a
single object.
"""
[docs]
@abstractmethod
def post_optimisation_run_with_this(
self,
optimized_cavity_settings: SetOfCavitySettings,
full_elts: ListOfElements,
**kwargs,
) -> SimulationOutput:
"""Run a simulation a simulation after optimisation is over.
With :class:`.Envelope1D`, it just calls the classic
:meth:`run_with_this`. But with :class:`.TraceWin`, we need to update
the ``optimized_cavity_settings`` as running an optimisation run on a
fraction of the linac is pretty different from running a simulation on
the whole linac.
"""
[docs]
@abstractmethod
def init_solver_parameters(self, accelerator: Accelerator) -> None:
"""Init some :class:`BeamCalculator` solver parameters."""
[docs]
def _generate_simulation_output(self, *args, **kwargs) -> SimulationOutput:
"""Transform the output of ``run`` to a :class:`.SimulationOutput`."""
return self.simulation_output_factory.run(*args, **kwargs)
@property
def reference_phase(self) -> REFERENCE_PHASES_T:
"""Give the reference phase.
.. todo::
Handle reference synchronous phase.
"""
if self.flag_phi_abs:
return "phi_0_abs"
return "phi_0_rel"
@property
@abstractmethod
def is_a_multiparticle_simulation(self) -> bool:
"""Tell if the simulation is a multiparticle simulation."""
pass
@property
@abstractmethod
def is_a_3d_simulation(self) -> bool:
"""Tell if the simulation is in 3D."""
pass
[docs]
def compute(
self,
accelerator: Accelerator,
keep_settings: bool = True,
recompute_reference: bool = True,
output_time: bool = True,
ref_simulation_output: SimulationOutput | None = None,
) -> SimulationOutput:
"""Wrap full process to compute propagation of beam in accelerator.
Parameters
----------
accelerator :
Accelerator under study.
keep_settings :
If settings/simulation output should be saved. The default is True.
recompute_reference :
If results should be taken from a file instead of recomputing
everything each time. The default is True.
output_time :
To print in log the time the calculation took. The default is True.
ref_simulation_output :
For calculation of mismatch factors. The default is None, in which
case the calculation is simply skipped.
Returns
-------
simulation_output : SimulationOutput
Object holding simulation results.
"""
start_time = time.monotonic()
self.init_solver_parameters(accelerator)
simulation_output = self.run(accelerator.elts)
simulation_output.compute_complementary_data(
accelerator.elts, ref_simulation_output
)
if keep_settings:
accelerator.keep_simulation_output(simulation_output, self.id)
accelerator.keep_settings(
simulation_output, exported_phase=self._export_phase
)
end_time = time.monotonic()
delta_t = datetime.timedelta(seconds=end_time - start_time)
if output_time:
logging.info(f"Elapsed time in beam calculation: {delta_t}")
if not recompute_reference:
raise NotImplementedError(
"idea is to take results from file if "
"simulations are too long. will be easy "
"for tracewin."
)
return simulation_output
@property
def cavity_settings_factory(self) -> CavitySettingsFactory:
"""Return the factory with a concise call."""
_list_elts_factory = self.list_of_elements_factory
_instruc_factory = _list_elts_factory.instructions_factory
_element_factory = _instruc_factory.element_factory
_field_map_factory = _element_factory.field_map_factory
cavity_settings_factory = _field_map_factory.cavity_settings_factory
return cavity_settings_factory