"""
This module contains the simulation result class.
"""
from __future__ import annotations
from enum import IntEnum
from typing import Generic, Optional, Dict, Any, Tuple, TypeVar, Callable
import numpy as np
import scipy.integrate
import scipy.interpolate
from matplotlib.pyplot import Figure
from .. import plotting
from .. elements import Element
from ._advanced_helpers import Device, AdvancedModel
from ._basic_helpers import BasicDevice
from ._radial_dist import boltzmann_radial_potential_linear_density_ebeam
from ..physconst import MINIMAL_N_1D, MINIMAL_KBT
[docs]class Rate(IntEnum):
"""
Enum for conveniently identifying rates produced in advanced simulations
"""
ELECTRON_IONISATION = 101
EI = 101
RADIATIVE_RECOMBINATION = 102
RR = 102
DIELECTRONIC_RECOMBINATION = 103
DR = 103
CHARGE_EXCHANGE = 104
CX = 104
LOSSES_AXIAL_COLLISIONAL = 105
AX_CO = 105
LOSSES_RADIAL_COLLISIONAL = 107
RA_CO = 107
# T_ELECTRON_IONISATION = 201
# T_EI = 201
# T_RADIATIVE_RECOMBINATION = 202
# T_RR = 202
# T_DIELECTRONIC_RECOMBINATION = 203
# T_DR = 203
# T_CHARGE_EXCHANGE = 204
# T_CX = 204
T_LOSSES_AXIAL_COLLISIONAL = 205
T_AX_CO = 205
T_LOSSES_RADIAL_COLLISIONAL = 207
T_RA_CO = 207
T_COLLISIONAL_THERMALISATION = 301
COLLISIONAL_THERMALISATION = 301
T_CT = 301
T_SPITZER_HEATING = 302
SPITZER_HEATING = 302
T_SH = 302
IONISATION_HEAT = 303
# CV = 304
COLLISION_RATE_TOTAL = 401
COLLISION_RATE_SELF = 402
OVERLAP_FACTORS_EBEAM = 501
F_EI = 501
# CHARGE_COMPENSATION = 511
TRAPPING_PARAMETER_AXIAL = 521
W_AX = 521
TRAPPING_PARAMETER_RADIAL = 522
W_RA = 522
TRAP_DEPTH_AXIAL = 523
V_AX = 523
TRAP_DEPTH_RADIAL = 524
V_RA = 524
E_KIN_MEAN = 531
E_KIN_FWHM = 532
_PER_M_PER_S = r"(m$^{-1}$ s$^{-1}$)"
_EV_PER_S = r"(eV s$^{-1}$)"
_PER_S = r"(s$^{-1}$)"
_RATE_LABELS: Dict[Rate, Dict[str, Any]] = {
Rate.EI: dict(
title="EI",
ylabel="Ionisation rate " + _PER_M_PER_S,
),
Rate.RR: dict(
title="RR",
ylabel="Recombination rate " + _PER_M_PER_S,
),
Rate.DR: dict(
title="DR",
ylabel="Recombination rate " + _PER_M_PER_S,
),
Rate.CX: dict(
title="CX",
ylabel="Recombination rate " + _PER_M_PER_S,
),
Rate.AX_CO: dict(
title="Ax. coll. losses",
ylabel="Loss rate " + _PER_M_PER_S,
),
Rate.RA_CO: dict(
title="Rad. coll. losses",
ylabel="Loss rate " + _PER_M_PER_S,
),
# Rate.T_EI: dict(
# title="Electron ionisation"
# ylabel="",
# ),
# Rate.T_RR: dict(
# title="Radiative recombination"
# ylabel="",
# ),
# Rate.T_DR: dict(
# title="Dielectronic recombination"
# ylabel="",
# ),
# Rate.T_CX: dict(
# title="Charge exchange"
# ylabel="",
# ),
Rate.T_AX_CO: dict(
title="Ax. coll. losses",
ylabel="Cooling rate " + _EV_PER_S,
),
Rate.T_RA_CO: dict(
title="Rad. coll. losses",
ylabel="Cooling rate " + _EV_PER_S,
),
Rate.COLLISIONAL_THERMALISATION: dict(
title="Coll. thermalisation",
ylabel="Thermalisation rate " + _EV_PER_S,
),
Rate.SPITZER_HEATING: dict(
title="Spitzer heating",
ylabel="Heating rate " + _EV_PER_S,
),
Rate.IONISATION_HEAT: dict(
title="Ionisation heat",
ylabel="Ionisation heat (eV)",
),
# Rate.CV: dict(
# title="Heat capacity",
# ylabel="$c_V$ (eV/eV)",
# ),
Rate.COLLISION_RATE_TOTAL: dict(
title="Total collision rate",
ylabel=r"$r_i$ " + _PER_S,
),
Rate.COLLISION_RATE_SELF: dict(
title="Self collision rate",
ylabel=r"$r_{ii}$ " + _PER_S,
),
Rate.F_EI: dict(
title="Electron beam overlap",
ylabel=r"$f_{ei}$",
ylim=(0, 1.1),
),
Rate.TRAPPING_PARAMETER_AXIAL: dict(
title="Ax. trapping parameter",
ylabel=r"$\omega_{ax}$",
),
Rate.TRAPPING_PARAMETER_RADIAL: dict(
title="Rad. trapping parameter",
ylabel=r"$\omega_{rad}$",
),
Rate.TRAP_DEPTH_AXIAL: dict(
title="Ax. trapping potential",
ylabel=r"$V_{ax} (V)$",
),
Rate.TRAP_DEPTH_RADIAL: dict(
title="Rad. trapping potential",
ylabel=r"$V_{rad} (V)$",
),
Rate.E_KIN_MEAN: dict(
title="Beam energy mean",
ylabel=r"$E_e$ (eV)",
),
Rate.E_KIN_FWHM: dict(
title="Beam energy FWHM",
ylabel=r"FWHM($E_e$) (eV)",
),
Rate.E_KIN_FWHM: dict(
title="Beam energy FWHM",
ylabel=r"FWHM($E_e$) (eV)",
),
}
# Needed for parametric polymorphism of Result classes
GenericDevice = TypeVar("GenericDevice", Device, BasicDevice)
[docs]class BasicResult(Generic[GenericDevice]):
"""
Instances of this class are containers for the results of ebisim basic_simulations and contain a
variety of convenience methods for simple plot generation etc.
"""
_abundance_label = "Abundance"
def __init__(self, *, t: np.ndarray, N: np.ndarray, device: GenericDevice,
target: Element, res: Optional[Any] = None):
"""
Parameters
----------
t :
An array holding the time step coordinates.
N :
An array holding the occupancy of each charge state at a given time.
device :
If coming from advanced_simulation, this is the machine / electron beam description.
target :
If coming from advanced_simulation, this is the target represented by this Result.
res :
The result object returned by scipy.integrate.solve_ivp. This can contain useful
information about the solver performance etc. Refer to scipy documentation for details.
"""
self.t = t
self.N = N
self.device: GenericDevice = device
self.target = target
self.res = res
@property
def _dense_interpolator(self) -> Optional[Callable[[float], np.ndarray]]:
if self.res and self.res.sol:
return self.res.sol
else:
return None
@property
def _abundance_interpolator(self) -> Callable[[float], np.ndarray]:
return self._dense_interpolator or scipy.interpolate.interp1d(self.t, self.N)
def _check_time_in_domain(self, t: float) -> None:
if t < self.t.min() or t > self.t.max():
raise ValueError("Value for t lies outside the simulated domain.")
def _param_title(self, stub):
"""
Generates a plot title by merging the stub with some general simulation parameters.
Defaults to the stub if no parameters are available
Parameters
----------
stub :
Title stub for the plot
Returns
-------
A LaTeX formatted title string compiled from the stub, current density and fwhm.
"""
title = f"{self.target.latex_isotope()} {stub.lower()}" #: TODO: Fix case
e_kin = self.device.e_kin
j = self.device.j
title = title + f" ($j = {j:0.1f}$ A/cm$^2$, $E_{{e}} = {e_kin:0.1f}$ eV)"
fwhm = self.device.fwhm
if fwhm is not None:
title = title[:-1] + f", FWHM = {fwhm:0.1f} eV)"
return title
[docs] def times_of_highest_abundance(self) -> np.ndarray:
"""
Yields the point of time with the highest abundance for each charge state
Returns
-------
<s>
Array of times.
"""
args = np.argmax(self.N, axis=1)
return self.t[args]
[docs] def abundance_at_time(self, t: float) -> np.ndarray:
"""
Yields the abundance of each charge state at a given time
Parameters
----------
t :
<s>
Point of time to evaluate.
Returns
-------
Abundance of each charge state, array index corresponds to charge state.
"""
self._check_time_in_domain(t)
return self._abundance_interpolator(t)
[docs] def plot_charge_states(self, relative: bool = False, **kwargs: Any) -> Figure:
"""
Plot the charge state evolution of this result object.
Parameters
----------
relative :
Flags whether the absolute numbers or a relative fraction should be plotted at each
timestep, by default False.
kwargs
Keyword arguments are handed down to ebisim.plotting.plot_generic_evolution,
cf. documentation thereof.
If no arguments are provided, reasonable default values are injected.
Returns
-------
Figure handle of the generated plot.
Raises
------
ValueError
If the required data (self.t, self.N) is not available, i.e.
corresponding attributes of this Result instance have not been set correctly.
"""
kwargs.setdefault("xlim", (1e-4, self.t.max()))
kwargs.setdefault("ylim", (0, self.N.sum(axis=0).max()*1.05))
kwargs.setdefault("title", self._param_title("Charge states"))
kwargs.setdefault("yscale", "linear")
kwargs.setdefault("plot_total", True)
if relative:
kwargs["ylim"] = (0, 1.1)
kwargs.setdefault("ylabel", "Relative abundance")
fig = plotting.plot_generic_evolution(self.t, self.N/np.sum(self.N, axis=0), **kwargs)
else:
kwargs.setdefault("ylabel", self._abundance_label)
fig = plotting.plot_generic_evolution(self.t, self.N, **kwargs)
return fig
plot = plot_charge_states #: Alias for plot_charge_states
[docs]class AdvancedResult(BasicResult[Device]):
"""
Instances of this class are containers for the results of ebisim advanced_simulations
and contain a variety of convenience methods for simple plot generation etc.
"""
_abundance_label = "Linear density (m$^{-1}$)"
def __init__(self, *, t: np.ndarray, N: np.ndarray, device: Device, target: Element,
kbT: np.ndarray, model: AdvancedModel, id_: int,
res: Optional[Any] = None, rates: Optional[Dict[Rate, np.ndarray]] = None):
"""
Parameters
----------
t :
An array holding the time step coordinates.
N :
An array holding the occupancy of each charge state at a given time.
device :
If coming from advanced_simulation, this is the machine / electron beam description.
target :
If coming from advanced_simulation, this is the target represented by this Result.
kbT :
An array holding the temperature of each charge state at a given time.
model :
The AdvancedModel instance that was underlying the advanced_simulation
id_ :
The position of this target in the list of all simulated targets.
res :
The result object returned by scipy.integrate.solve_ivp. This can contain useful
information about the solver performance etc. Refer to scipy documentation for details.
rates :
A dictionary containing the different breeding rates in arrays shaped like N.
"""
super().__init__(t=t, N=N, device=device, target=target, res=res)
self.kbT = kbT
self.target = target
self.device = device
self.rates = rates
model = model._replace(
targets=list(model.targets),
bg_gases=list(model.bg_gases),
cxxs_bggas=list(model.cxxs_bggas),
cxxs_trgts=list(model.cxxs_trgts)
)
self.model = model
self.id = id_
@property
def _abundance_interpolator(self) -> Callable[[float], np.ndarray]:
dinterp = self._dense_interpolator
if dinterp is not None:
lower = self.model.lb[self.id]
upper = self.model.ub[self.id]
interp = lambda t: dinterp(t)[lower:upper] # noqa:E731
else:
interp = scipy.interpolate.interp1d(self.t, self.N)
return interp
@property
def _temperature_interpolator(self) -> Callable[[float], np.ndarray]:
dinterp = self._dense_interpolator
if dinterp is not None and self.res is not None:
lower = self.res.y.shape[0]//2 + self.model.lb[self.id]
upper = self.res.y.shape[0]//2 + self.model.ub[self.id]
interp = lambda t: dinterp(t)[lower:upper] # noqa:E731
else:
interp = scipy.interpolate.interp1d(self.t, self.kbT)
return interp
def _density_filter(self, data, threshold):
filtered = data.copy()
filtered[self.N < threshold] = np.nan
return filtered
[docs] def temperature_at_time(self, t: float) -> np.ndarray:
"""
Yields the temperature of each charge state at a given time
Parameters
----------
t :
<s>
Point of time to evaluate.
Returns
-------
<eV>
Temperature of each charge state, array index corresponds to charge state.
"""
self._check_time_in_domain(t)
return self._temperature_interpolator(t)
[docs] def radial_distribution_at_time(self, t: float) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
"""
Yields the radial distribution information at time
Parameters
----------
t :
<s>
Point of time to evaluate.
Returns
-------
phi :
Radial potential
n3d :
On axis 3D density for each charge state.
shapes :
The Boltzmann shape factors for each charge state.
"""
self._check_time_in_domain(t)
interp = self._dense_interpolator
if interp is None:
if self.res is None:
raise ValueError("Required Data not available, "
+ "something went wrong during the creation of this result.")
interp = scipy.interpolate.interp1d(self.t, self.res.y)
y = interp(t)
n = y[:self.model.ub[-1]]
n = np.maximum(n, MINIMAL_N_1D)
kT = y[self.model.ub[-1]:]
kT = np.maximum(kT, MINIMAL_KBT)
phi, n3d, shapes = boltzmann_radial_potential_linear_density_ebeam(
self.device.rad_grid, self.device.current, self.device.r_e, self.device.e_kin,
np.atleast_2d(n).T, np.atleast_2d(kT).T, np.atleast_2d(self.model.q).T,
first_guess=self.device.rad_phi_uncomp,
ldu=(self.device.rad_fd_l, self.device.rad_fd_d, self.device.rad_fd_u)
)
return phi, n3d, shapes
[docs] def plot_radial_distribution_at_time(self, t: float, **kwargs: Any) -> Figure:
"""
Plot the radial ion distribution at time t.
Parameters
----------
t :
<s>
Point of time to evaluate.
kwargs
Keyword arguments are handed down to ebisim.plotting.plot_radial_distribution and
ebisim.plotting.plot_generic_evolution, cf. documentation thereof.
If no arguments are provided, reasonable default values are injected.
Returns
-------
Figure handle of the generated plot.
"""
phi, n3d, shapes = self.radial_distribution_at_time(t)
dens = (n3d * shapes)[self.model.lb[self.id]:self.model.ub[self.id]]
denslim = 10**np.ceil(np.log10(dens.max()))
title = self._param_title(f"Radial distribution at t = {1000*t:.1f} ms")
kwargs.setdefault("xscale", "log")
kwargs.setdefault("yscale", "log")
kwargs.setdefault("title", title)
kwargs.setdefault("xlim", (self.device.r_e/100, self.device.r_dt))
kwargs.setdefault("ylim", (denslim/10**10, denslim))
fig = plotting.plot_radial_distribution(
self.device.rad_grid, dens, phi, self.device.r_e, **kwargs
)
return fig
[docs] def plot_energy_density(self, **kwargs: Any) -> Figure:
"""
Plot the energy density evolution of this result object.
Parameters
----------
kwargs
Keyword arguments are handed down to ebisim.plotting.plot_generic_evolution,
cf. documentation thereof.
If no arguments are provided, reasonable default values are injected.
Returns
-------
Figure handle of the generated plot.
Raises
------
ValueError
If the required data (self.t, self.kbT) is not available, i.e.
corresponding attributes of this Result instance have not been set correctly.
"""
e_den = self.N * self.kbT
ymin = 10**(np.floor(np.log10(e_den[:, 0].sum(axis=0)) - 1))
ymax = 10**(np.ceil(np.log10(e_den.sum(axis=0).max()) + 1))
kwargs.setdefault("xlim", (1e-4, self.t.max()))
kwargs.setdefault("ylim", (ymin, ymax))
kwargs.setdefault("title", self._param_title("Energy density"))
kwargs.setdefault("ylabel", "Linear energy density (eV / m$^{-1}$)")
kwargs.setdefault("plot_total", True)
fig = plotting.plot_generic_evolution(self.t, e_den, **kwargs)
return fig
[docs] def plot_temperature(self, dens_threshold: float = 1000*MINIMAL_N_1D, **kwargs: Any) -> Figure:
"""
Plot the temperature evolution of this result object.
Parameters
----------
dens_threshold :
If given temperatures are only plotted where the particle denisty is larger than
the threshold value.
kwargs
Keyword arguments are handed down to ebisim.plotting.plot_generic_evolution,
cf. documentation thereof.
If no arguments are provided, reasonable default values are injected.
Returns
-------
Figure handle of the generated plot
Raises
------
ValueError
If the required data (self.t, self.kbT) is not available, i.e.
corresponding attributes of this Result instance have not been set correctly.
"""
filtered_kbT = self._density_filter(self.kbT, dens_threshold)
kwargs.setdefault("xlim", (1e-4, self.t.max()))
kwargs.setdefault(
"ylim",
(
10**np.floor(np.nanmin(np.log10(filtered_kbT))),
10**np.ceil(np.nanmax(np.log10(filtered_kbT))),
)
)
kwargs.setdefault("title", self._param_title("Temperature"))
kwargs.setdefault("ylabel", "Temperature (eV)")
fig = plotting.plot_generic_evolution(self.t, filtered_kbT, **kwargs)
return fig
[docs] def plot_rate(self, rate_key: Rate,
dens_threshold: float = 1000*MINIMAL_N_1D, **kwargs) -> Figure:
"""
Plots the requested ionisation- or energy flow rates.
Parameters
----------
rate_key :
The key identifying the rate to be plotted.
See ebisim.simulation.Rate for valid values.
dens_threshold :
If given temperatures are only plotted where the particle denisty is larger than
the threshold value.
kwargs
Keyword arguments are handed down to ebisim.plotting.plot_generic_evolution,
cf. documentation thereof.
If no arguments are provided, reasonable default values are injected.
Returns
-------
Figure handle of the generated plot
Raises
------
ValueError
If the required data (self.rates) is not available, or an invalid key is requested.
"""
if self.rates is None:
raise ValueError("Rates are not available for this result.")
if rate_key not in self.rates:
raise ValueError(
f"The requested rate_key does not exist. Available rates are {self.rates.keys()}."
)
rate = self.rates[rate_key]
if rate.shape[0] != 1: # Don't filter if scalar (i.e. not for indiv charge states)
rate = self._density_filter(rate, dens_threshold)
labels = _RATE_LABELS.get(rate_key, {})
kwargs.setdefault("xlim", (1e-4, self.t.max()))
kwargs.setdefault("ylim", labels.get("ylim", None))
kwargs.setdefault("yscale", labels.get("yscale", "linear"))
kwargs.setdefault("title", self._param_title(labels.get("title", "unkown rate")))
kwargs.setdefault("ylabel", labels.get("ylabel", "Unknown"))
kwargs.setdefault("plot_total", False)
if len(rate.shape) == 1:
rate = rate[np.newaxis, :]
if rate.shape[0] == 1:
kwargs.setdefault("label_lines", False)
kwargs.setdefault("ls", "-")
fig = plotting.plot_generic_evolution(self.t, rate, **kwargs)
return fig