"""Define :class:`Envelope3D`, an envelope solver."""
import logging
from collections.abc import Collection
from pathlib import Path
from lightwin.beam_calculation.beam_calculator import BeamCalculator
from lightwin.beam_calculation.envelope_3d.beam_parameters_factory import (
BeamParametersFactoryEnvelope3D,
)
from lightwin.beam_calculation.envelope_3d.element_envelope3d_parameters_factory import (
ElementEnvelope3DParametersFactory,
)
from lightwin.beam_calculation.envelope_3d.simulation_output_factory import (
SimulationOutputFactoryEnvelope3D,
)
from lightwin.beam_calculation.envelope_3d.transfer_matrix_factory import (
TransferMatrixFactoryEnvelope3D,
)
from lightwin.beam_calculation.envelope_3d.util import ENVELOPE3D_METHODS_T
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 import CavitySettings
from lightwin.core.elements.field_maps.field_map import FieldMap
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.physics.synchronous_phases import (
PHI_S_MODELS,
SYNCHRONOUS_PHASE_FUNCTIONS,
)
from lightwin.util.typing import (
EXPORT_PHASES_T,
REFERENCE_PHASE_POLICY_T,
BeamKwargs,
)
[docs]
class Envelope3D(BeamCalculator):
"""A 3D envelope solver.
As transverse effects are generally not predominant, I do not use this
solver very often and a lot of elements are not implemented. The current
list of explicitly supported elements is:
.. configkeys:: lightwin.beam_calculation.envelope_3d.element_envelope3d_parameters_factory.PARAMETERS_3D
:n_cols: 2
The default behavior when an element in the input ``DAT`` file is not
recognized, is to issue a warning and replace this element by a ``DRIFT``.
Do not hesitate to file an |issue|_ if you need me to implement some
elements.
"""
flag_cython = False
[docs]
def __init__(
self,
*,
reference_phase_policy: REFERENCE_PHASE_POLICY_T,
default_field_map_folder: Path | str,
beam_kwargs: BeamKwargs,
export_phase: EXPORT_PHASES_T,
phi_s_definition: PHI_S_MODELS = "historical",
n_steps_per_cell: int,
method: ENVELOPE3D_METHODS_T = "RK4",
**kwargs,
) -> None:
"""Set the proper motion integration function, according to inputs."""
self.n_steps_per_cell = n_steps_per_cell
self.method: ENVELOPE3D_METHODS_T = method
self._phi_s_definition: PHI_S_MODELS = phi_s_definition
self._phi_s_func = SYNCHRONOUS_PHASE_FUNCTIONS[self._phi_s_definition]
super().__init__(
reference_phase_policy=reference_phase_policy,
default_field_map_folder=default_field_map_folder,
beam_kwargs=beam_kwargs,
export_phase=export_phase,
**kwargs,
)
self.beam_parameters_factory = BeamParametersFactoryEnvelope3D(
self.is_a_3d_simulation,
self.is_a_multiparticle_simulation,
beam_kwargs=self._beam_kwargs,
)
self.transfer_matrix_factory = TransferMatrixFactoryEnvelope3D(
self.is_a_3d_simulation
)
[docs]
def _set_up_specific_factories(self) -> None:
"""Set up the factories specific to the |BC|.
This method is called in the :meth:`.BeamCalculator.__init__`, hence it
appears only in the base |BC|.
"""
self.simulation_output_factory = SimulationOutputFactoryEnvelope3D(
is_multipart=self.is_a_multiparticle_simulation,
beam_calculator_id=self.id,
beam_kwargs=self._beam_kwargs,
)
self.beam_calc_parameters_factory = ElementEnvelope3DParametersFactory(
method=self.method,
n_steps_per_cell=self.n_steps_per_cell,
solver_id=self.id,
beam_kwargs=self._beam_kwargs,
phi_s_definition=self._phi_s_definition,
)
self.list_of_elements_factory = ListOfElementsFactory(
self.is_a_3d_simulation,
self.is_a_multiparticle_simulation,
default_field_map_folder=self.default_field_map_folder,
load_fields=True,
beam_kwargs=self._beam_kwargs,
load_cython_field_maps=False,
field_maps_in_3d=False, # not implemented anyway
elements_to_dump=(),
)
[docs]
def run_with_this(
self,
accelerator_id: str,
set_of_cavity_settings: SetOfCavitySettings,
elts: ListOfElements,
**kwargs,
) -> SimulationOutput:
"""Compute beam propagation with non-nominal settings.
Parameters
----------
accelerator_id :
Associated :attr:`.Accelerator.id`. Looks like:
``0000001_Solution``.
set_of_cavity_settings :
The new cavity settings to try. If it is None, then the cavity
settings are taken from the FieldMap objects.
elts :
List of elements in which the beam must be propagated.
Returns
-------
Holds energy, phase, transfer matrices (among others) packed into a
single object.
"""
single_elts_results = []
w_kin = elts.w_kin_in
phi_abs = elts.phi_abs_in
for elt in elts:
cavity_settings = set_of_cavity_settings.get(elt)
_store_entry_phase_in_settings(phi_abs, cavity_settings)
func = elt.beam_calc_param[self.id].transf_mat_function_wrapper
elt_results = func(w_kin=w_kin, cavity_settings=cavity_settings)
if cavity_settings is not None:
v_cav_mv, phi_s = self._compute_cavity_parameters(elt_results)
cavity_settings.v_cav_mv = v_cav_mv
cavity_settings.phi_s = phi_s
single_elts_results.append(elt_results)
phi_abs += elt_results["phi_rel"][-1]
w_kin = elt_results["w_kin"][-1]
simulation_output = self.simulation_output_factory.create(
accelerator_id=accelerator_id,
elts=elts,
single_elts_results=single_elts_results,
set_of_cavity_settings=set_of_cavity_settings,
)
return simulation_output
[docs]
def init_solver_parameters(self, accelerator: Accelerator) -> None:
"""Create the number of steps, meshing, transfer functions for elts.
The solver parameters are stored in the |E|'s
``beam_calc_param``.
Parameters
----------
accelerator :
Object which |LOE| must be initialized.
"""
elts = accelerator.elts
position = 0.0
index = 0
for elt in elts:
if self.id in elt.beam_calc_param:
logging.debug(
f"Solver already initialized for {elt = }. I will skip "
f"solver param initialisation {elts[0]} to {elts[-1]}"
)
return
solver_param = self.beam_calc_parameters_factory.run(elt)
elt.beam_calc_param[self.id] = solver_param
position, index = solver_param.set_absolute_meshes(position, index)
logging.debug(f"Initialized solver param for {elts[0]} to {elts[-1]}")
return
@property
def is_a_multiparticle_simulation(self) -> bool:
"""Return False."""
return False
@property
def is_a_3d_simulation(self) -> bool:
"""Return True."""
return True
[docs]
def _compute_cavity_parameters(self, results: dict) -> tuple[float, float]:
"""Compute the cavity parameters by calling ``_phi_s_func``.
Parameters
----------
results :
The dictionary of results as returned by the transfer matrix
function wrapper.
Returns
-------
Accelerating voltage in MV and synchronous phase in radians. If the
cavity is failed, two ``np.nan`` are returned.
"""
v_cav_mv, phi_s = self._phi_s_func(**results)
return v_cav_mv, phi_s
[docs]
def _store_entry_phase_in_settings(
phi_bunch_abs: float,
cavity_settings: CavitySettings | Collection[CavitySettings] | None,
) -> None:
"""Set entry phase."""
if cavity_settings is None:
return
if isinstance(cavity_settings, CavitySettings):
cavity_settings = (cavity_settings,)
for settings in cavity_settings:
settings.phi_bunch = phi_bunch_abs
return