"""Classes to load information from FISPACT output JSON file."""
from __future__ import annotations
from typing import TYPE_CHECKING, cast
import io # - needed for dispatch
from functools import singledispatch
from pathlib import Path # - needed for dispatch
import msgspec as ms
import numpy as np
# noinspection PyUnresolvedReferences
from xpypact.time_step import TimeStep # noqa: TC001 - required for Struct field
if TYPE_CHECKING:
from collections.abc import Iterable, Iterator
from xpypact.nuclide import NuclideInfo
from xpypact.xpypact_types import NDArrayFloat
FLOAT_ZERO = 0.0
[docs]
class RunDataCorrected(ms.Struct):
"""Common data for an FISPACT inventory.
Note:
Correction - dose_rate_type and dose_rate_distance are duplicated
the FISPACT time steps. This information is extracted to this header.
"""
timestamp: str
run_name: str
flux_name: str
dose_rate_type: str
dose_rate_distance: float
class InventoryError(ValueError):
"""Base class for inventory exception."""
def __str__(self) -> str:
"""Use the class __doc__ as a message.
Returns:
The __doc__ of the exception class.
"""
return cast("str", self.__class__.__doc__) # pragma: no cover
class InventoryNonMonotonicTimesError(InventoryError):
"""Irradiation and cooling times in FISPACT output should be monotonically increasing."""
[docs]
class Inventory(ms.Struct):
"""Helper class to load FISPACT inventory (output) from JSON file.
Implements list interface over time steps.
"""
run_data: RunData
inventory_data: list[TimeStep]
@property
def meta_info(self) -> RunDataCorrected:
"""Create corrected Inventory header.
Returns:
RunDataCorrected: with common data for the inventory.
"""
ts = self.inventory_data[-1] # The dose_rate in the first timestep can be empty.
return RunDataCorrected(
self.run_data.timestamp,
self.run_data.run_name,
self.run_data.flux_name,
ts.dose_rate.type,
ts.dose_rate.distance,
)
[docs]
def iterate_time_step_nuclides(
self,
) -> Iterator[
tuple[
int,
int,
float,
float,
float,
float,
float,
float,
float,
float,
float,
float,
float,
float,
float,
]
]:
"""Scan the time steps and nuclides in these steps.
Returns:
Iterator over all the nuclides collected in the Inventory.
"""
return (
(
t.number,
n.zai,
n.atoms,
n.grams,
n.activity,
n.alpha_activity,
n.beta_activity,
n.gamma_activity,
n.heat,
n.alpha_heat,
n.beta_heat,
n.gamma_heat,
n.dose,
n.ingestion,
n.inhalation,
)
for t in self.inventory_data
for n in t.nuclides
)
[docs]
def iterate_time_step_gamma(self) -> Iterator[tuple[int, int, float]]:
"""Scan the time steps and gamma spectrum.
Returns:
Iterator over all the gamma bins presented in the Inventory.
"""
return (
(
t.number,
g + 1, # g + 1 corresponds to index of the upper bound of a bin in gbins table
rate,
)
for t in self.inventory_data
if t.gamma_spectrum
for g, rate in enumerate(t.gamma_spectrum.values)
)
[docs]
def __post_init__(self) -> None:
"""Define time steps durations and elapsed time.
Raises:
InventoryNonMonotonicTimesError: if time sequences in JSON are not in order
"""
number = 1
prev_irradiation_time = prev_cooling_time = prev_elapsed_time = 0.0
for ts in self.inventory_data:
duration = ts.irradiation_time - prev_irradiation_time
if duration == FLOAT_ZERO:
duration = ts.cooling_time - prev_cooling_time
if duration < FLOAT_ZERO:
raise InventoryNonMonotonicTimesError # pragma: no cover
ts.duration = duration
prev_elapsed_time = ts.elapsed_time = prev_elapsed_time + duration
if duration == FLOAT_ZERO:
ts.flux = FLOAT_ZERO
ts.number = number
number += 1
prev_irradiation_time, prev_cooling_time = (
ts.irradiation_time,
ts.cooling_time,
)
[docs]
def __iter__(self) -> Iterator[TimeStep]:
"""Iterate over time steps.
Returns:
Iterator over the time steps.
"""
return iter(self.inventory_data)
[docs]
def __len__(self) -> int:
"""Length, delegated to time steps.
Returns:
int: length of the time steps list.
"""
return len(self.inventory_data)
[docs]
def __getitem__(self, item: int) -> TimeStep:
"""List interface delegated to the time steps.
Args:
item: an item index or slice
Returns:
Selected item or items.
"""
return self.inventory_data[item]
def extract_times(time_steps: Iterable[TimeStep]) -> NDArrayFloat:
"""Create vector of elapsed time for all the time steps.
Args:
time_steps: list of time steps
Returns:
Vector with elapsed times.
"""
return np.fromiter((x.elapsed_time for x in time_steps), dtype=float)
# @dispatch(str)
@singledispatch
def from_json(text: str) -> Inventory:
"""Construct Inventory instance from JSON text.
Args:
text: Text source.
Returns:
The loaded Inventory instance.
"""
return ms.json.decode(text, type=Inventory)
@from_json.register
def _(stream: io.IOBase) -> Inventory: # type: ignore[misc]
"""Construct Inventory instance from JSON stream.
Args:
stream: IO stream source.
Returns:
The loaded Inventory instance.
"""
return from_json(stream.read())
@from_json.register
def _(path: Path) -> Inventory: # type: ignore[misc]
"""Construct Inventory instance from JSON path.
Args:
path: path to source.
Returns:
The loaded Inventory instance.
"""
return from_json(path.read_text(encoding="utf8"))
class RunData(ms.Struct, frozen=True, gc=False):
"""FISPACT run title data."""
timestamp: str
run_name: str
flux_name: str
def asdict(self) -> dict[str, str]:
"""Get dict representation."""
return ms.structs.asdict(self)
def astuple(self) -> tuple[str, str, str]:
"""Get tuple representation."""
return cast("tuple[str, str, str]", ms.structs.astuple(self))
@classmethod
def from_json(cls, json_dict: dict[str, str]) -> RunData:
"""Construct RunData instance from JSON.
Args:
json_dict: source dictionary
Returns:
The loaded instance of RunData
"""
return ms.convert(json_dict, cls)