"""Provides attribution functionality for non-linear species."""
from typing import Any, Protocol
import numpy as np
# DEFINE TYPE STYLES
[docs]
class AttributionFunc(Protocol):
"""Defines type protocol for attribution function inputs."""
def __call__(
self, data: dict[str, np.ndarray], /, *, config: dict
) -> dict[str, np.ndarray]:
"""Maps a dict of species arrays to another such dict."""
...
[docs]
def apply_attribution( # noqa: PLR0913, PLR0917
func: AttributionFunc,
diff_func: AttributionFunc,
method: str,
species: str,
sub_dict: dict[str, np.ndarray],
full_dict: dict[str, np.ndarray],
**kwargs: Any,
) -> dict[str, np.ndarray]:
"""Applies attribution methodology using the function ``func``.
Dispatches to the attribution function matching ``method`` for
``sub_dict``.
Args:
func (AttributionFunc): A callable that maps a dict of species arrays to
another such dict. The first argument must be a
dict[str, np.ndarray].
diff_func (AttributionFunc): A callable that calculates the derivative
of ``func`` with respect to the input. Must have the same shape as
``func``. Can be None if no attribution method is used that
requires a derivative.
method (str): Attribution method. Choice of: "none", "residual",
"marginal", "proportional", "differential".
species (str): Name of the species to consider (e.g. 'CO2').
sub_dict (dict[str, np.ndarray]): Dict of species arrays representing a
part of the full dict (e.g. one aircraft identifier).
full_dict (dict[str, np.ndarray]): Dict of species arrays representing
all sub-dicts. Must include ``sub_dict``!
**kwargs: Additional keyword arguments to pass to ``func`` and
``diff_func``.
Raises:
ValueError: If an invalid method is given.
Returns:
dict[str, np.ndarray]: ``func`` value attributable to ``sub_dict``
"""
if method == "none":
return func(sub_dict, **kwargs)
if method == "residual":
return residual_attribution(func, sub_dict, full_dict, species, **kwargs)
if method == "marginal":
return marginal_attribution(diff_func, sub_dict, full_dict, species, **kwargs)
if method == "proportional":
return proportional_attribution(func, sub_dict, full_dict, species, **kwargs)
if method == "differential":
return differential_attribution(
diff_func, sub_dict, full_dict, species, **kwargs
)
raise ValueError(f"Invalid attribution method for species {species}")
[docs]
def residual_attribution(
func: AttributionFunc,
sub_dict: dict[str, np.ndarray],
full_dict: dict[str, np.ndarray],
species: str,
**kwargs: Any,
) -> dict[str, np.ndarray]:
"""Calculates the ``func`` value attributable to ``sub_dict`` (residual).
Uses a residual methodology for species ``species``. The dict
``full_dict``, representing all other sources (i.e. background plus
other aircraft identifiers), must include ``sub_dict``.
Args:
func (AttributionFunc): A callable that maps a dict of species arrays to
another such dict. The first argument must be a
dict[str, np.ndarray].
sub_dict (dict[str, np.ndarray]): Dict of species arrays representing a
part of the full dict (e.g. one aircraft identifier).
full_dict (dict[str, np.ndarray]): Dict of species arrays representing
all sub-dicts. Must include ``sub_dict``!
species (str): Name of the species to consider (e.g. 'CO2').
**kwargs: Additional keyword arguments to pass to ``func``.
Returns:
dict[str, np.ndarray]: ``func`` value attributable to ``sub_dict``
"""
# calculate difference between full and sub dicts
diff_dict = {species: full_dict[species] - sub_dict[species]}
# do attribution
full_res_arr = func(full_dict, **kwargs)[species]
diff_res_arr = func(diff_dict, **kwargs)[species]
return {species: full_res_arr - diff_res_arr}
[docs]
def marginal_attribution(
diff_func: AttributionFunc,
sub_dict: dict[str, np.ndarray],
full_dict: dict[str, np.ndarray],
species: str,
**kwargs: Any,
) -> dict[str, np.ndarray]:
"""Calculates the ``func`` value attributable to ``sub_dict`` (marginal).
Uses a marginal methodology for species ``species``. The dict
``full_dict``, representing all other sources (i.e. background plus
other aircraft identifiers), must include ``sub_dict``.
Args:
diff_func (AttributionFunc): A callable that calculates the derivative
of ``func`` with respect to the input. Must have the same shape as
``func``.
sub_dict (dict[str, np.ndarray]): Dict of species arrays representing a
part of the full dict (e.g. one aircraft identifier).
full_dict (dict[str, np.ndarray]): Dict of species arrays representing
all sub-dicts. Must include ``sub_dict``!
species (str): Name of the species to consider (e.g. 'CO2').
**kwargs: Additional keyword arguments to pass to ``func`` and
``diff_func``.
Returns:
dict[str, np.ndarray]: ``func`` value attributable to ``sub_dict``
"""
# calculate gradient w.r.t. input
deriv = diff_func(full_dict, **kwargs)[species]
return {species: deriv * sub_dict[species]}
[docs]
def proportional_attribution(
func: AttributionFunc,
sub_dict: dict[str, np.ndarray],
full_dict: dict[str, np.ndarray],
species: str,
**kwargs: Any,
) -> dict[str, np.ndarray]:
"""Calculates the ``func`` value attributable to ``sub_dict`` (proportional).
Uses a proportional methodology for species ``species``. The dict
``full_dict``, representing all other sources (i.e. background plus
other aircraft identifiers), must include ``sub_dict``.
Args:
func (AttributionFunc): A callable that maps a dict of species arrays to
another such dict. The first argument must be a
dict[str, np.ndarray].
sub_dict (dict[str, np.ndarray]): Dict of species arrays representing a
part of the full dict (e.g. one aircraft identifier).
full_dict (dict[str, np.ndarray]): Dict of species arrays representing
all sub-dicts. Must include ``sub_dict``!
species (str): Name of the species to consider (e.g. 'CO2').
**kwargs: Additional keyword arguments to pass to ``func``.
Returns:
dict[str, np.ndarray]: ``func`` value attributable to ``sub_dict``
"""
# do attribution
full_res_arr = func(full_dict, **kwargs)[species]
prop = np.zeros_like(full_res_arr, dtype=float)
np.divide(
sub_dict[species],
full_dict[species],
out=prop,
where=full_dict[species] != 0.0,
)
return {species: prop * full_res_arr}
[docs]
def differential_attribution(
diff_func: AttributionFunc,
sub_dict: dict[str, np.ndarray],
full_dict: dict[str, np.ndarray],
species: str,
**kwargs: Any,
) -> dict[str, np.ndarray]:
"""Calculates the ``func`` value attributable to ``sub_dict`` (differential).
Uses a differential methodology for species ``species``. The dict
``full_dict``, representing all other sources (i.e. background plus
other aircraft identifiers), must include ``sub_dict``.
Args:
diff_func (AttributionFunc): A callable that calculates the derivative
of ``func`` with respect to the input. Must have the same shape as
``func``.
sub_dict (dict[str, np.ndarray]): Dict of species arrays representing a
part of the full dict (e.g. one aircraft identifier).
full_dict (dict[str, np.ndarray]): Dict of species arrays representing
all sub-dicts. Must include ``sub_dict``!
species (str): Name of the species to consider (e.g. 'CO2').
**kwargs: Additional keyword arguments to pass to ``func`` and
``diff_func``.
Returns:
dict[str, np.ndarray]: ``func`` value attributable to ``sub_dict``
"""
# calculate gradient w.r.t. input
deriv = diff_func(full_dict, **kwargs)[species]
# do attribution
sub_vals = sub_dict[species]
sub_grad = np.gradient(sub_vals) # derivative (with respect to time)
p = deriv * sub_grad
res_arr = np.cumsum(p, axis=0)
return {species: res_arr}