Source code for openfisca_core.tracers.computation_log
from __future__ import annotations
import typing
from typing import List, Optional, Union
import numpy
from .. import tracers
from openfisca_core.indexed_enums import EnumArray
if typing.TYPE_CHECKING:
from numpy.typing import ArrayLike
Array = Union[EnumArray, ArrayLike]
[docs]class ComputationLog:
_full_tracer: tracers.FullTracer
def __init__(self, full_tracer: tracers.FullTracer) -> None:
self._full_tracer = full_tracer
def display(
self,
value: Optional[Array],
) -> str:
if isinstance(value, EnumArray):
value = value.decode_to_str()
return numpy.array2string(value, max_line_width = float("inf"))
def _get_node_log(
self,
node: tracers.TraceNode,
depth: int,
aggregate: bool,
) -> List[str]:
def print_line(depth: int, node: tracers.TraceNode) -> str:
indent = ' ' * depth
value = node.value
if value is None:
formatted_value = "{'avg': '?', 'max': '?', 'min': '?'}"
elif aggregate:
try:
formatted_value = str({
'avg': numpy.mean(value),
'max': numpy.max(value),
'min': numpy.min(value),
})
except TypeError:
formatted_value = "{'avg': '?', 'max': '?', 'min': '?'}"
else:
formatted_value = self.display(value)
return f"{indent}{node.name}<{node.period}> >> {formatted_value}"
node_log = [print_line(depth, node)]
children_logs = [
self._get_node_log(child, depth + 1, aggregate)
for child
in node.children
]
return node_log + self._flatten(children_logs)
def _flatten(
self,
list_of_lists: List[List[str]],
) -> List[str]:
return [item for _list in list_of_lists for item in _list]
def lines(self, aggregate: bool = False) -> List[str]:
depth = 1
lines_by_tree = [
self._get_node_log(node, depth, aggregate)
for node
in self._full_tracer.trees
]
return self._flatten(lines_by_tree)
[docs] def print_log(self, aggregate = False) -> None:
"""
Print the computation log of a simulation.
If ``aggregate`` is ``False`` (default), print the value of each
computed vector.
If ``aggregate`` is ``True``, only print the minimum, maximum, and
average value of each computed vector.
This mode is more suited for simulations on a large population.
"""
for line in self.lines(aggregate):
print(line) # noqa T001