"""High-level trajectory integration orchestration.
This module glues together the low-level physics kernels, image-charge
generation, and optional self-consistency loop. The public functions provide
the main entry points for running the modern LW integrator programmatically.
"""
from __future__ import annotations
from typing import Optional, Tuple
import numpy as np
from .equations import retarded_equations_of_motion
from .images import generate_conducting_image, generate_switching_image
from .self_consistency import SelfConsistencyConfig, self_consistent_step
from .types import IntegratorConfig, ParticleState, SimulationType, Trajectory
[docs]def retarded_integrator(
steps: int,
h_step: float,
wall_z: float,
aperture_radius: float,
sim_type: SimulationType,
init_rider: ParticleState,
init_driver: Optional[ParticleState],
mean: float,
cav_spacing: float,
z_cutoff: float,
self_consistency: Optional[SelfConsistencyConfig] = None,
) -> Tuple[Trajectory, Trajectory]:
"""Run the retarded-field integrator for rider and driver trajectories.
Parameters
----------
steps:
Total number of integration updates to compute.
h_step:
Temporal step between states (``Δτ`` in the covariant formulation).
wall_z:
Conducting wall location for boundary-condition simulations.
aperture_radius:
Aperture radius used by the wall/image generators.
sim_type:
Boundaries and interaction type encoded as :class:`SimulationType`.
init_rider:
Initial state of the primary bunch.
init_driver:
Optional initial state of the opposing bunch (for ``BUNCH_TO_BUNCH``).
mean:
Historical bunch separation parameter retained for compatibility.
cav_spacing:
Longitudinal spacing between cavities when using switching walls.
z_cutoff:
Threshold beyond which the switching wall no longer mirrors charges.
self_consistency:
Optional :class:`SelfConsistencyConfig` to iterate each step until the
Lorentz factor converges.
Returns
-------
tuple[Trajectory, Trajectory]
Two trajectories: the rider (primary bunch) and the driver (image or
opposing bunch), each represented as a list of particle states.
"""
trajectory: Trajectory = [{} for _ in range(steps)]
trajectory_drv: Trajectory = [{} for _ in range(steps)]
for i in range(steps):
if i == 0:
trajectory[i] = init_rider
if sim_type == SimulationType.CONDUCTING_WALL:
trajectory_drv[i] = generate_conducting_image(
init_rider, wall_z, aperture_radius
)
elif sim_type == SimulationType.SWITCHING_WALL:
trajectory_drv[i] = generate_switching_image(
init_rider, wall_z, aperture_radius, z_cutoff
)
elif sim_type == SimulationType.BUNCH_TO_BUNCH:
if init_driver is None:
raise ValueError(
"SimulationType.BUNCH_TO_BUNCH requires init_driver state"
)
trajectory_drv[i] = init_driver
else:
trajectory[i] = self_consistent_step(
retarded_equations_of_motion,
h_step,
trajectory,
trajectory_drv,
i - 1,
aperture_radius,
sim_type,
self_consistency,
)
if sim_type == SimulationType.SWITCHING_WALL:
trajectory_drv[i] = generate_switching_image(
trajectory[i], wall_z, aperture_radius, z_cutoff
)
if np.mean(trajectory[i]["z"]) > z_cutoff:
z_cutoff += cav_spacing
wall_z += cav_spacing
elif sim_type == SimulationType.CONDUCTING_WALL:
trajectory_drv[i] = generate_conducting_image(
trajectory[i], wall_z, aperture_radius
)
elif sim_type == SimulationType.BUNCH_TO_BUNCH:
if init_driver is None:
raise ValueError(
"SimulationType.BUNCH_TO_BUNCH requires init_driver state"
)
trajectory_drv[i] = self_consistent_step(
retarded_equations_of_motion,
h_step,
trajectory_drv,
trajectory,
i - 1,
aperture_radius,
sim_type,
self_consistency,
)
return trajectory, trajectory_drv
[docs]def run_integrator(
config: IntegratorConfig,
init_rider: ParticleState,
init_driver: Optional[ParticleState],
) -> Tuple[Trajectory, Trajectory]:
"""Convenience wrapper using :class:`IntegratorConfig`.
All parameters are supplied via ``config`` which mirrors the keyword
arguments accepted by :func:`retarded_integrator`.
"""
return retarded_integrator(
steps=config.steps,
h_step=config.time_step,
wall_z=config.wall_position,
aperture_radius=config.aperture_radius,
sim_type=config.simulation_type,
init_rider=init_rider,
init_driver=init_driver,
mean=config.bunch_mean,
cav_spacing=config.cavity_spacing,
z_cutoff=config.z_cutoff,
)
__all__ = [
"retarded_integrator",
"run_integrator",
]