Source code for ogstools.physics.nuclearwasteheat.nuclearwaste

# Copyright (c) 2012-2025, OpenGeoSys Community (http://www.opengeosys.org)
#            Distributed under a Modified BSD License.
#            See accompanying file LICENSE.txt or
#            http://www.opengeosys.org/project/license
#

from dataclasses import dataclass

import numpy as np
from pint.facets.plain import PlainQuantity

from ._unitsetup import Q_




@dataclass
class NuclearWaste:
    "Models the heat generated by one type of nuclear fuel waste."

    name: str
    "Name of this type of waste."
    nuclide_powers: PlainQuantity | np.ndarray
    "Powers for the different leading nuclides."
    decay_consts: PlainQuantity | np.ndarray
    "Decay constants for the different leading nuclides."
    num_bundles: int
    "Number of nuclear waste bundles."
    time_interim: PlainQuantity | float
    "Interim storage time before heat evaluation."
    time_deposit: PlainQuantity | float = 0.0
    "Number of active bundles linearly increases until this time is reached."
    factor: PlainQuantity | float = 1.0
    "Scale the calculated heat by this factor."



    def heat(
        self,
        t: PlainQuantity | float | np.ndarray,
        baseline: bool = False,
        ncl_id: int | None = None,
        time_unit: str = "s",
        power_unit: str = "W",
    ) -> float | np.ndarray:
        """Calculate the heat of a nuclear waste proxy model.

        :param t:           Timevalue(s) at which the heat is calculated.
        :param ncl_id:      If given, only output the heat by nuclide with this
                            id, else sum the heat over all nuclides.
        :param baseline:    If True, evaluate one bundle with no interim or
                            deposition time delay.
        :returns:           Heat generated by the nuclear waste at time t.
        """

        _t = Q_(t).magnitude
        _nuclide_powers = Q_(self.nuclide_powers, power_unit).magnitude
        _decay_consts = Q_(self.decay_consts, f"1/{time_unit}").magnitude
        _num_bundles = self.num_bundles
        _time_interim = Q_(self.time_interim, time_unit).magnitude
        _time_deposit = Q_(self.time_deposit, time_unit).magnitude
        if baseline:
            _num_bundles, _time_interim, _time_deposit = 1, 0.0, 0.0

        t_per_bundle = np.linspace(_t - _time_deposit, _t, _num_bundles)
        t_per_bundle = np.ma.masked_less(t_per_bundle, 0) + _time_interim
        # decay values for each nuclide via matrix multiplicaation
        # results in shape (num_bundles, len(t), len(decay_consts))
        decay = np.exp(-t_per_bundle[..., None] @ _decay_consts[None, ...])
        # sum over all bundles
        res = np.sum(self.factor * _nuclide_powers * decay, axis=0)
        # optionally sum over nuclides
        return np.sum(res, axis=-1) if ncl_id is None else res[..., ncl_id]






@dataclass
class Repository:
    "Models the heat generated by total repository inventory."

    waste: list[NuclearWaste]
    "Waste inventory of the repository."



    def time_deposit(self, time_unit: str = "s") -> float | list[float]:
        "Deposition time for each nuclear waste type."
        if len(self.waste) == 1:
            return Q_(self.waste[0].time_deposit).to(time_unit).magnitude
        return [
            Q_(nw.time_deposit).to(time_unit).magnitude for nw in self.waste
        ]




    def heat(
        self,
        t: PlainQuantity | float | np.ndarray,
        baseline: bool = False,
        ncl_id: int | None = None,
        time_unit: str = "s",
        power_unit: str = "W",
    ) -> float | np.ndarray:
        """Calculate the heat produced by the repository at time t.

        :param t:           Timevalue(s) at which the heat is calculated.
        :param ncl_id:      If given, only output the heat by nuclide with this
                            id, else sum the heat over all nuclides.
        :param baseline:    If True, evaluate one bundle for each waste with no
                            interim or deposition time delay.
        :returns:           Heat generated by the repository at time t.
        """

        result = [
            nw.heat(t, baseline, ncl_id, time_unit, power_unit)
            for nw in self.waste
        ]
        return np.sum(np.array(result), axis=0)