Source code for lightwin.core.em_fields.field_helpers

"""Define functions to compute 1D longitudinal electric fields."""

import math
from typing import Any

import numpy as np
from numpy.typing import NDArray

from lightwin.core.em_fields.types import FieldFuncComponent1D, Pos1D


[docs] def null_field_1d(pos: Any) -> float: """Define a null electric/magnetic field.""" return 0.0
[docs] def create_1d_field_func( field_values: NDArray[np.float64], pos_max: float, n_intervals: int, pos_min: float = 0.0, ) -> FieldFuncComponent1D: """Create functions computing field for a given position. Parameters ---------- field_values : Field values on en evenly spaced grid. pos_max : Ending position of the field. n_intervals : Number of intervals in the field map; it is ``n_points - 1``. pos_min : Starting position of the field. Returns ------- FieldFuncComponent1D Function giving a field component at a given 1D position. It also have a ``xp`` and a ``fp`` attributes holding positions and field values. This is used for Cython implementations. .. todo:: Clean this ``xp`` ``fp`` thingy, not clean at all. """ field_values = np.asarray(field_values, dtype=float) total_length = pos_max - pos_min inv_dx = n_intervals / total_length def interp_func(pos: float) -> float: idx_float = (pos - pos_min) * inv_dx idx = int(idx_float) if idx < 0 or idx >= n_intervals: return 0.0 t = idx_float - idx return field_values[idx] * (1.0 - t) + field_values[idx + 1] * t # Attach xp and fp as it is used by Cython interp_func.xp = np.linspace(pos_min, pos_max, n_intervals + 1) interp_func.fp = field_values return interp_func
[docs] def e_1d( pos: Pos1D, e_func: FieldFuncComponent1D, phi: float, amplitude: float, phi_0: float, ) -> float: """Compute normed 1D electric field.""" return amplitude * e_func(pos) * math.cos(phi + phi_0)
[docs] def e_1d_complex( pos: Pos1D, e_func: FieldFuncComponent1D, phi: float, amplitude: float, phi_0: float, ) -> complex: """Compute normed 1D electric field.""" phase = phi + phi_0 return amplitude * e_func(pos) * (math.cos(phase) + 1j * math.sin(phase))
[docs] def shifted_e_spat( e_spat: FieldFuncComponent1D, z_shift: float ) -> FieldFuncComponent1D: """Shift electric field by ``z_shift``.""" def shifted(z_pos: float) -> float: return e_spat(z_pos - z_shift) return shifted
[docs] def rescale_array( array: NDArray[np.float64], norm: float, tol: float = 1e-6 ) -> NDArray[np.float64]: """Rescale given array if ``norm`` is different from unity.""" if abs(norm - 1.0) > tol: array /= norm return array