# -*- coding: utf-8 -*-
"""
Stage 2 of the ALBuMS pipeline: instability theories.
Each theory takes the SystemState and the Stage-1 Equilibrium, and returns a
TheoryResult. The actual physics lives in albums.physics so theories.py does
only the orchestration.
"""
from __future__ import annotations
from dataclasses import dataclass, field
from enum import Enum
from typing import TYPE_CHECKING, Protocol, runtime_checkable
import numpy as np
from albums.equilibrium import Capability, Equilibrium
from albums.options import TheoryOptions
from albums.physics.coupled_bunch import dipole_coupled_bunch
from albums.physics.potential import landau_threshold, robinson_frequency
from albums.physics.ptbl import coupled_bunch_mode1, ptbl_alves, ptbl_he
from albums.physics.robinson import (
robinson_coupling,
robinson_glmci,
robinson_no_coupling,
)
from albums.physics.zero_frequency import zero_frequency_instability
if TYPE_CHECKING: # pragma: no cover - hints only, avoids a system.py import cycle
from albums.system import SystemState
[docs]class Category(str, Enum):
"""Physical category a theory belongs to.
The flow-agnostic accessors (Result / ResultArray) resolve "the Robinson
result" / "the PTBL result" by category, so several concrete theories can
share one category. Values are independent of each theory's `name`.
"""
ZERO_FREQUENCY = "zero_frequency"
ROBINSON = "robinson"
HOM = "hom"
PTBL = "ptbl"
[docs]@dataclass
class TheoryResult:
"""Outcome of one instability theory at one operating point."""
name: str
unstable: bool | np.ndarray # scalar flag or per-mode array
growth_rate: float | np.ndarray = 0.0 # Net growth rate [s^-1]
mode_frequency: np.ndarray | None = None # Omega [rad/s]; optional
converged: bool | np.ndarray = True # scalar flag or per-mode array
extra: dict = field(default_factory=dict)
def __repr__(self) -> str:
return f"TheoryResult(name={self.name!r}, unstable={np.any(self.unstable)})"
[docs]@runtime_checkable
class InstabilityTheory(Protocol):
"""Stage-2 strategy interface.
The evaluate methods receives the SystemState, the Stage-1 Equilibrium and
a TheoryOptions bundle, and returns a TheoryResult.
"""
name: str
requires: frozenset[Capability]
self, state: SystemState, eq: Equilibrium, opts: TheoryOptions
) -> TheoryResult: ...
# --------------------------------------------------------------------------- #
# Theories compatible with every solver (need only amplitude form factor). #
# --------------------------------------------------------------------------- #
[docs]class _TheoryRepr:
"""Mixin giving the theory classes a clean repr from their instance config."""
def __repr__(self) -> str:
args = ", ".join(f"{k}={v!r}" for k, v in vars(self).items())
return f"{type(self).__name__}({args})"
[docs]class ZeroFrequency(_TheoryRepr):
"""Zero-frequency (DC Robinson) instability.
Physics: albums.physics.zero_frequency.zero_frequency_instability.
"""
name = "zero_frequency"
category = Category.ZERO_FREQUENCY
requires = frozenset({Capability.AMPLITUDE_FORM_FACTOR})
[docs] def evaluate(
self, state: SystemState, eq: Equilibrium, opts: TheoryOptions
) -> TheoryResult:
ring = state.ring
cavity_list = eq.cavity_list
omega_r = robinson_frequency(ring, cavity_list, eq.F)
unstable = bool(
zero_frequency_instability(
ring,
cavity_list,
eq.I0,
eq.F,
omega_r,
eq.bunch_length,
opts.mode_coupling,
)
)
return TheoryResult(self.name, unstable, growth_rate=1.0 if unstable else 0.0)
[docs]class RobinsonCoupling(_TheoryRepr):
"""Robinson dipole/quadrupole modes with mode coupling.
Physics: albums.physics.robinson.robinson_coupling.
"""
name = "robinson_coupling"
category = Category.ROBINSON
coupled = True
requires = frozenset(
{Capability.AMPLITUDE_FORM_FACTOR, Capability.TAYLOR_POTENTIAL}
)
[docs] def evaluate(
self, state: SystemState, eq: Equilibrium, opts: TheoryOptions
) -> TheoryResult:
ring = state.ring
cavity_list = eq.cavity_list
omega_r = robinson_frequency(ring, cavity_list, eq.F)
landau = landau_threshold(ring, omega_r, eq.c, eq.b)
robinson, Omega, converged, growth = robinson_coupling(
ring, cavity_list, eq.I0, eq.F, omega_r, eq.bunch_length, landau
)
return TheoryResult(
self.name,
robinson,
growth_rate=growth,
mode_frequency=Omega,
converged=converged,
)
[docs]class RobinsonNoCoupling(_TheoryRepr):
"""Robinson dipole to octupole modes without mode coupling.
Physics: albums.physics.robinson.robinson_no_coupling.
"""
name = "robinson_no_coupling"
category = Category.ROBINSON
coupled = False
requires = frozenset(
{Capability.AMPLITUDE_FORM_FACTOR, Capability.TAYLOR_POTENTIAL}
)
[docs] def evaluate(
self, state: SystemState, eq: Equilibrium, opts: TheoryOptions
) -> TheoryResult:
ring = state.ring
cavity_list = eq.cavity_list
omega_r = robinson_frequency(ring, cavity_list, eq.F)
landau = landau_threshold(ring, omega_r, eq.c, eq.b)
robinson, Omega, converged, growth = robinson_no_coupling(
ring, cavity_list, eq.I0, eq.F, omega_r, eq.bunch_length, landau
)
return TheoryResult(
self.name,
robinson,
growth_rate=growth,
mode_frequency=Omega,
converged=converged,
)
[docs]class HOMCoupledBunch(_TheoryRepr):
"""HOM-driven dipole coupled-bunch instability.
Reads the f_HOM / Z_HOM fields of TheoryOptions at evaluation time.
Physics: albums.physics.coupled_bunch.dipole_coupled_bunch.
"""
name = "hom_coupled_bunch"
category = Category.HOM
requires = frozenset(
{Capability.AMPLITUDE_FORM_FACTOR, Capability.TAYLOR_POTENTIAL}
)
[docs] def evaluate(
self, state: SystemState, eq: Equilibrium, opts: TheoryOptions
) -> TheoryResult:
ring = state.ring
cavity_list = eq.cavity_list
omega_r = robinson_frequency(ring, cavity_list, eq.F)
landau = landau_threshold(ring, omega_r, eq.c, eq.b)
HOM = bool(
dipole_coupled_bunch(
ring, eq.I0, opts.f_HOM, opts.Z_HOM, eq.bunch_length, omega_r, landau
)
)
return TheoryResult(self.name, HOM)
[docs]class BoschMode1(_TheoryRepr):
"""Bosch coupled-bunch mode l=1 with Landau threshold.
Physics: albums.physics.ptbl.coupled_bunch_mode1.
"""
name = "bosch_mode1"
category = Category.PTBL
requires = frozenset(
{Capability.AMPLITUDE_FORM_FACTOR, Capability.TAYLOR_POTENTIAL}
)
[docs] def evaluate(
self, state: SystemState, eq: Equilibrium, opts: TheoryOptions
) -> TheoryResult:
ring = state.ring
cavity_list = eq.cavity_list
omega_r = robinson_frequency(ring, cavity_list, eq.F)
landau = landau_threshold(ring, omega_r, eq.c, eq.b)
coupled_mode1, mode1_modes, converged = coupled_bunch_mode1(
ring, cavity_list, eq.I0, eq.F, omega_r, eq.bunch_length, landau
)
return TheoryResult(
self.name,
bool(np.any(coupled_mode1)),
mode_frequency=mode1_modes,
converged=converged,
)
# --------------------------------------------------------------------------- #
# PTBL / FMCI theories restricted by capability. #
# --------------------------------------------------------------------------- #
[docs]class PTBLHe(_TheoryRepr):
"""
Periodic transient beam loading (PTBL) instability via He's perturbation
criterion.
Physics: albums.physics.ptbl.ptbl_he.
"""
name = "ptbl_he"
category = Category.PTBL
requires = frozenset(
{Capability.AMPLITUDE_FORM_FACTOR, Capability.PHASE_FORM_FACTOR}
)
[docs] def evaluate(
self, state: SystemState, eq: Equilibrium, opts: TheoryOptions
) -> TheoryResult:
PTBL = ptbl_he(state.ring, eq.cavity_list, eq.I0, eq.F, eq.PHI)
return TheoryResult(self.name, PTBL)
[docs]class PTBLAlves(_TheoryRepr):
"""
Periodic transient beam loading (PTBL) instability via Alves mode analysis.
Needs the pycolleff LE + full distribution, so Alves only.
Physics: albums.physics.ptbl.ptbl_alves.
"""
name = "ptbl_alves"
category = Category.PTBL
requires = frozenset({Capability.PYCOLLEFF_LE, Capability.FULL_DISTRIBUTION})
[docs] def __init__(self, max_azi: int = 10, max_rad: int = 10):
"""max_azi, max_rad : truncation of the azimuthal / radial mode expansion."""
self.max_azi = max_azi
self.max_rad = max_rad
[docs] def evaluate(
self, state: SystemState, eq: Equilibrium, opts: TheoryOptions
) -> TheoryResult:
PTBL, growth = ptbl_alves(
state.ring, eq.backend_LE, max_azi=self.max_azi, max_rad=self.max_rad
)
return TheoryResult(self.name, PTBL, growth_rate=growth)
[docs]class RobinsonGLMCI(_TheoryRepr):
"""
Robinson instabilities via the Gaussian longitudinal mode-coupling
instability (GLMCI) method. Alves only.
Physics: albums.physics.robinson.robinson_glmci.
"""
name = "robinson_glmci"
category = Category.ROBINSON
coupled = True
requires = frozenset({Capability.PYCOLLEFF_LE, Capability.FULL_DISTRIBUTION})
[docs] def __init__(self, max_azi: int = 2, max_rad: int = 0):
"""max_azi, max_rad : truncation of the azimuthal / radial mode expansion."""
self.max_azi = max_azi
self.max_rad = max_rad
[docs] def evaluate(
self, state: SystemState, eq: Equilibrium, opts: TheoryOptions
) -> TheoryResult:
robinson, Omega, converged, growth = robinson_glmci(
state.ring, eq.backend_LE, max_azi=self.max_azi, max_rad=self.max_rad
)
return TheoryResult(
self.name,
robinson,
growth_rate=growth,
mode_frequency=Omega,
converged=converged,
)
# --------------------------------------------------------------------------- #
# Theory registry — the single source for category groupings. #
# --------------------------------------------------------------------------- #
# Every concrete theory, in canonical precedence order. The order within a
# physical category sets the precedence used by the flow-agnostic category
# accessors (Result / ResultArray pick the first present theory of a category).
ALL_THEORY_CLASSES = (
ZeroFrequency,
RobinsonCoupling,
RobinsonNoCoupling,
HOMCoupledBunch,
BoschMode1,
PTBLHe,
PTBLAlves,
RobinsonGLMCI,
)
[docs]def names_in_category(category: Category) -> tuple[str, ...]:
"""Theory names belonging to a physical category, in precedence order."""
return tuple(t.name for t in ALL_THEORY_CLASSES if t.category == category)
[docs]def coupled_robinson_names() -> tuple[str, ...]:
"""Robinson-family theory names whose plot uses the fast mode-coupling marker."""
return tuple(
t.name
for t in ALL_THEORY_CLASSES
if t.category == Category.ROBINSON and t.coupled
)