helpinn/solver/base.py - main - bari/HPLSE-poster-code

Репозиторий помещен в архив 16 июня 2026. Вся функциональность доступна только для чтения.
helpinn/solver/base.py
196 строк7 KB
bari
Add time-domain propagation and profile extensions
13 апр 2026, 01:26
13 апр 2026, 01:265ebf062
"""Abstract solver base classes and shared solver numerics."""

from __future__ import annotations

from abc import ABC
from dataclasses import dataclass, field

import numpy as np

from ..plasma.base import PlasmaProfile
from ..utils import safe_divide
from ..pulse.base import Pulse
from .result import Result


def canonical_polarization(polarization: str) -> str:
    """Return the canonical polarization name."""
    if polarization in {"TE", "s"}:
        return "TE"
    if polarization in {"TM", "p"}:
        return "TM"
    raise ValueError("polarization must be TE, TM, s, or p")


def wavelength_to_omega(wavelength: np.ndarray) -> np.ndarray:
    """Convert wavelength to angular frequency in normalized units."""
    return 2.0 * np.pi / wavelength


def spectral_weights(pulse: Pulse) -> np.ndarray:
    """Return normalized spectral intensity weights."""
    power = np.abs(pulse.A) ** 2
    total = power.sum()
    if total == 0:
        return np.full(power.shape, 1.0 / power.size)
    return power / total


def resolve_permittivity(profile: PlasmaProfile, pulse: Pulse) -> np.ndarray:
    """Return the permittivity profile for each wavelength sample."""
    count = pulse.lambda_.size
    if profile.eps is not None:
        return np.broadcast_to(np.asarray(profile.eps, dtype=complex), (count, profile.z.size)).copy()
    omega = wavelength_to_omega(pulse.lambda_)[:, None]
    if profile.w is not None:
        omega_p = np.broadcast_to(np.asarray(profile.w, dtype=float), (count, profile.z.size))
        return np.asarray(profile.plasma_state_equation.permittivity(omega, omega_p), dtype=complex)
    density = np.asarray(profile.n, dtype=float)
    if hasattr(profile.plasma_state_equation, "plasma_frequency_from_density"):
        omega_p_1d = np.asarray(profile.plasma_state_equation.plasma_frequency_from_density(density), dtype=float)
    else:
        omega_p_1d = np.sqrt(np.clip(density, 0.0, None))
    omega_p = omega_p_1d[None, :]
    omega_p = np.broadcast_to(omega_p, (count, profile.z.size))
    return np.asarray(profile.plasma_state_equation.permittivity(omega, omega_p), dtype=complex)


def resolve_eps_derivative(profile: PlasmaProfile, eps: np.ndarray) -> np.ndarray:
    """Return the first spatial derivative of the permittivity."""
    if profile.deps is not None:
        deps = np.asarray(profile.deps, dtype=complex)
        return np.broadcast_to(deps, eps.shape).copy()
    if profile.dw is not None and profile.w is not None and profile.plasma_state_equation is not None:
        return np.gradient(eps, profile.z, axis=1, edge_order=1)
    if profile.dn is not None and profile.n is not None and profile.plasma_state_equation is not None:
        return np.gradient(eps, profile.z, axis=1, edge_order=1)
    return np.gradient(eps, profile.z, axis=1, edge_order=1)


def longitudinal_wavenumber(eps: np.ndarray, pulse: Pulse, eps_in: complex) -> np.ndarray:
    """Return the longitudinal wavenumber on the spatial grid."""
    k0 = wavelength_to_omega(pulse.lambda_)[:, None]
    sin2 = np.sin(pulse.theta_rad)[:, None] ** 2
    return k0 * np.sqrt(eps - eps_in * sin2 + 0j)


def impedance_factor(k: np.ndarray, eps: np.ndarray, polarization: str) -> np.ndarray:
    """Return the polarization-dependent impedance factor."""
    if polarization == "TE":
        return k
    return safe_divide(k, eps.astype(complex))


def local_phase(
    k_left: np.ndarray,
    k_right: np.ndarray,
    dz: float,
    use_midpoint: bool,
    is_jump: bool,
) -> np.ndarray:
    """Return the phase advance across one cell."""
    if use_midpoint and not is_jump:
        return 0.5 * (k_left + k_right) * dz
    return k_left * dz


def magnus_measure(k: np.ndarray, z: np.ndarray) -> float:
    """Return a simple Magnus-condition measure."""
    dz = np.diff(z)
    if dz.size == 0:
        return 0.0
    return float(np.max(np.abs(k[:, :-1] * dz[None, :])))


def combine_fields(
    pulse: Pulse,
    A_modes: np.ndarray,
    B_modes: np.ndarray,
    E_modes: np.ndarray,
    H_modes: np.ndarray,
    reflection_modes: np.ndarray,
    transmission_modes: np.ndarray,
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """Combine wavelength-resolved outputs into one result."""
    weights = spectral_weights(pulse)
    field_weights = pulse.A[:, None]
    A = np.sum(field_weights * A_modes, axis=0)
    B = np.sum(field_weights * B_modes, axis=0)
    E = np.sum(field_weights * E_modes, axis=0)
    H = np.sum(field_weights * H_modes, axis=0)
    reflection = np.sum(weights[:, None] * np.clip(reflection_modes, 0.0, 1.0), axis=0)
    transmission = np.sum(weights[:, None] * np.clip(transmission_modes, 0.0, 1.0), axis=0)
    absorbed = np.clip(1.0 - reflection - transmission, 0.0, 1.0)
    return A, B, E, H, absorbed, reflection, transmission


def build_result(
    profile: PlasmaProfile,
    pulse: Pulse,
    A_modes: np.ndarray,
    B_modes: np.ndarray,
    E_modes: np.ndarray,
    H_modes: np.ndarray,
    reflection_modes: np.ndarray,
    transmission_modes: np.ndarray,
    solver_name: str,
    **extra: object,
) -> Result:
    """Build the public result object from wavelength-resolved arrays."""
    A, B, E, H, absorbed, reflection, transmission = combine_fields(
        pulse,
        A_modes,
        B_modes,
        E_modes,
        H_modes,
        reflection_modes,
        transmission_modes,
    )
    diagnostics = {"solver": solver_name}
    diagnostics.update(extra)
    return Result(
        z=profile.z,
        A=A,
        B=B,
        E=E,
        H=H,
        absorbed_energy=absorbed,
        reflection=reflection,
        transmission=transmission,
        cumulative_R=float(reflection[0]),
        cumulative_T=float(transmission[0]),
        modal_amplitudes=pulse.A,
        modal_wavelengths=pulse.lambda_,
        modal_angles=pulse.theta_rad,
        plane_wave_A=A_modes,
        plane_wave_B=B_modes,
        plane_wave_E=E_modes,
        plane_wave_H=H_modes,
        plane_wave_absorbed_energy=np.clip(1.0 - reflection_modes - transmission_modes, 0.0, 1.0),
        plane_wave_reflection=reflection_modes,
        plane_wave_transmission=transmission_modes,
        diagnostics=diagnostics,
    )


@dataclass
class Solver(ABC):
    """Abstract base class for solvers."""

    plasma_profile: PlasmaProfile
    pulse: Pulse
    check_magnus: bool = True
    results: Result | None = field(init=False, default=None)

    def solve(self) -> Result:
        """Run the solver and store the result."""
        self.results = self._solve()
        return self.results

    def _solve(self) -> Result:
        """Implement the solver-specific calculation."""
        raise NotImplementedError


class SolverBase(Solver):
    """Abstract base class for solvers."""