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