OpenFisca-Core-test 46.0.0__py3-none-any.whl

openfisca_core/tracers/__init__.py

@@ -0,0 +1,40 @@
+# Transitional imports to ensure non-breaking changes.
+# Could be deprecated in the next major release.
+#
+# How imports are being used today:
+#
+#   >>> from openfisca_core.module import symbol
+#
+# The previous example provokes cyclic dependency problems
+# that prevent us from modularizing the different components
+# of the library so to make them easier to test and to maintain.
+#
+# How could them be used after the next major release:
+#
+#   >>> from openfisca_core import module
+#   >>> module.symbol()
+#
+# And for classes:
+#
+#   >>> from openfisca_core import module
+#   >>> module.Symbol()
+#
+# See: https://www.python.org/dev/peps/pep-0008/#imports
+
+from .computation_log import ComputationLog
+from .flat_trace import FlatTrace
+from .full_tracer import FullTracer
+from .performance_log import PerformanceLog
+from .simple_tracer import SimpleTracer
+from .trace_node import TraceNode
+from .tracing_parameter_node_at_instant import TracingParameterNodeAtInstant
+
+__all__ = [
+    "ComputationLog",
+    "FlatTrace",
+    "FullTracer",
+    "PerformanceLog",
+    "SimpleTracer",
+    "TraceNode",
+    "TracingParameterNodeAtInstant",
+]

openfisca_core/tracers/computation_log.py

@@ -0,0 +1,196 @@
+from __future__ import annotations
+
+import sys
+
+import numpy
+
+from openfisca_core import types as t
+from openfisca_core.indexed_enums import Enum, EnumArray
+
+
+class ComputationLog:
+    _full_tracer: t.FullTracer
+
+    def __init__(self, full_tracer: t.FullTracer) -> None:
+        self._full_tracer = full_tracer
+
+    def lines(
+        self,
+        aggregate: bool = False,
+        max_depth: int = sys.maxsize,
+        ignore_default: bool = False,
+        tax_benefit_system: t.TaxBenefitSystem | None = None,
+    ) -> list[str]:
+        """Generate lines for the computation log.
+
+        Args:
+            aggregate: Show aggregated statistics instead of individual values.
+            max_depth: Maximum depth to display in the computation tree.
+            ignore_default: Hide variables with default values and their children.
+            tax_benefit_system: Used to get default values for variables.
+
+        Returns:
+            List of strings representing the computation log lines.
+
+        """
+        depth = 1
+
+        lines_by_tree = [
+            self._get_node_log(
+                node, depth, aggregate, max_depth, ignore_default, tax_benefit_system
+            )
+            for node in self._full_tracer.trees
+        ]
+
+        return self._flatten(lines_by_tree)
+
+    def print_log(
+        self,
+        aggregate: bool = False,
+        max_depth: int = sys.maxsize,
+        ignore_default: bool = False,
+        tax_benefit_system: t.TaxBenefitSystem | None = None,
+    ) -> 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.
+
+        If ``max_depth`` is ``sys.maxsize`` (default), print the entire
+        computation.
+
+        If ``max_depth`` is set, for example to ``3``, only print computed
+        vectors up to a depth of ``max_depth``.
+
+        If ``ignore_default`` is ``True``, do not display variables that have
+        their default value. When a variable with its default value is hidden,
+        its children are also hidden, even if they have non-default values.
+        """
+        for line in self.lines(
+            aggregate, max_depth, ignore_default, tax_benefit_system
+        ):
+            print(line)
+
+    def _get_node_log(
+        self,
+        node: t.TraceNode,
+        depth: int,
+        aggregate: bool,
+        max_depth: int = sys.maxsize,
+        ignore_default: bool = False,
+        tax_benefit_system: t.TaxBenefitSystem | None = None,
+    ) -> list[str]:
+        if depth > max_depth:
+            return []
+
+        # Check if we should ignore this node because it has its default value
+        if ignore_default and self._is_default_value(node, tax_benefit_system):
+            # Don't display this node or its children
+            return []
+
+        node_log = [self._print_line(depth, node, aggregate)]
+
+        children_logs = [
+            self._get_node_log(
+                child,
+                depth + 1,
+                aggregate,
+                max_depth,
+                ignore_default,
+                tax_benefit_system,
+            )
+            for child in node.children
+        ]
+
+        return node_log + self._flatten(children_logs)
+
+    def _is_default_value(
+        self, node: t.TraceNode, tax_benefit_system: t.TaxBenefitSystem | None
+    ) -> bool:
+        """Check if a node's value is the default value for its variable.
+
+        Args:
+            node: The trace node to check
+            tax_benefit_system: The tax benefit system to get variable information
+
+        Returns:
+            True if the value is the default value, False otherwise
+        """
+        if tax_benefit_system is None or node.value is None:
+            return False
+
+        try:
+            variable = tax_benefit_system.get_variable(node.name)
+            if variable is None:
+                return False
+
+            default_value = variable.default_value
+            if default_value is None:
+                return False
+
+            # For arrays, check if all values equal the default
+            if isinstance(node.value, numpy.ndarray):
+                # Handle EnumArray specially
+                if isinstance(node.value, EnumArray):
+                    # For Enum, compare the enum values
+                    if variable.value_type == Enum:
+                        # Check if all values in the array are the default enum value
+                        return numpy.all(node.value == default_value)
+                    return False
+
+                # For numeric/boolean arrays, check if all values equal default
+                try:
+                    return numpy.all(node.value == default_value)
+                except (ValueError, TypeError):
+                    # If comparison fails, assume not default
+                    return False
+
+            # For scalar values, direct comparison
+            return node.value == default_value
+
+        except Exception:
+            # If we can't determine, assume not default to be safe
+            return False
+
+    def _print_line(self, depth: int, node: t.TraceNode, aggregate: bool) -> str:
+        indent = "  " * depth
+        value = node.value
+
+        if value is None:
+            formatted_value = "{'avg': '?', 'max': '?', 'min': '?'}"
+
+        elif aggregate:
+            try:
+                formatted_value = str(  # pyright: ignore[reportCallIssue]
+                    {
+                        "avg": numpy.mean(value),  # pyright: ignore[reportArgumentType,reportCallIssue]
+                        "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}"
+
+    @staticmethod
+    def display(value: t.VarArray, max_depth: int = sys.maxsize) -> str:
+        if isinstance(value, EnumArray):
+            value = value.decode_to_str()
+        return numpy.array2string(value, max_line_width=max_depth)
+
+    @staticmethod
+    def _flatten(lists: list[list[str]]) -> list[str]:
+        return [item for list_ in lists for item in list_]
+
+
+__all__ = ["ComputationLog"]

openfisca_core/tracers/flat_trace.py

@@ -0,0 +1,87 @@
+from __future__ import annotations
+
+import numpy
+
+from openfisca_core import types as t
+from openfisca_core.indexed_enums import EnumArray
+
+
+class FlatTrace:
+    _full_tracer: t.FullTracer
+
+    def __init__(self, full_tracer: t.FullTracer) -> None:
+        self._full_tracer = full_tracer
+
+    def get_trace(self) -> t.FlatNodeMap:
+        trace: t.FlatNodeMap = {}
+
+        for node in self._full_tracer.browse_trace():
+            # We don't want cache read to overwrite data about the initial
+            # calculation.
+            #
+            # We therefore use a non-overwriting update.
+            trace.update(
+                {
+                    key: node_trace
+                    for key, node_trace in self._get_flat_trace(node).items()
+                    if key not in trace
+                },
+            )
+
+        return trace
+
+    def get_serialized_trace(self) -> t.SerializedNodeMap:
+        return {
+            key: {**flat_trace, "value": self.serialize(flat_trace["value"])}
+            for key, flat_trace in self.get_trace().items()
+        }
+
+    def _get_flat_trace(
+        self,
+        node: t.TraceNode,
+    ) -> t.FlatNodeMap:
+        key = self.key(node)
+
+        return {
+            key: {
+                "dependencies": [self.key(child) for child in node.children],
+                "parameters": {
+                    self.key(parameter): self.serialize(parameter.value)
+                    for parameter in node.parameters
+                },
+                "value": node.value,
+                "calculation_time": node.calculation_time(),
+                "formula_time": node.formula_time(),
+            },
+        }
+
+    @staticmethod
+    def key(node: t.TraceNode) -> t.NodeKey:
+        """Return the key of a node."""
+        name = node.name
+        period = node.period
+        return t.NodeKey(f"{name}<{period}>")
+
+    @staticmethod
+    def serialize(
+        value: None | t.VarArray | t.ArrayLike[object],
+    ) -> None | t.ArrayLike[object]:
+        if value is None:
+            return None
+
+        if isinstance(value, EnumArray):
+            return value.decode_to_str().tolist()
+
+        if isinstance(value, numpy.ndarray) and numpy.issubdtype(
+            value.dtype,
+            numpy.dtype(bytes),
+        ):
+            return value.astype(numpy.dtype(str)).tolist()
+
+        if isinstance(value, numpy.ndarray):
+            return value.tolist()
+
+        return value
+
+
+__all__ = ["FlatTrace"]

openfisca_core/tracers/full_tracer.py

@@ -0,0 +1,172 @@
+from __future__ import annotations
+
+import sys
+import time
+from collections.abc import Iterator
+
+from openfisca_core import types as t
+
+from .computation_log import ComputationLog
+from .flat_trace import FlatTrace
+from .performance_log import PerformanceLog
+from .simple_tracer import SimpleTracer
+from .trace_node import TraceNode
+
+
+class FullTracer:
+    _simple_tracer: t.SimpleTracer
+    _trees: list[t.TraceNode]
+    _current_node: None | t.TraceNode
+
+    def __init__(self) -> None:
+        self._simple_tracer = SimpleTracer()
+        self._trees = []
+        self._current_node = None
+
+    @property
+    def stack(self) -> t.SimpleStack:
+        """Return the stack of traces."""
+        return self._simple_tracer.stack
+
+    @property
+    def trees(self) -> list[t.TraceNode]:
+        """Return the tree of traces."""
+        return self._trees
+
+    @property
+    def computation_log(self) -> t.ComputationLog:
+        """Return the computation log."""
+        return ComputationLog(self)
+
+    @property
+    def performance_log(self) -> t.PerformanceLog:
+        """Return the performance log."""
+        return PerformanceLog(self)
+
+    @property
+    def flat_trace(self) -> t.FlatTrace:
+        """Return the flat trace."""
+        return FlatTrace(self)
+
+    def record_calculation_start(
+        self,
+        variable: t.VariableName,
+        period: t.PeriodInt | t.Period,
+    ) -> None:
+        self._simple_tracer.record_calculation_start(variable, period)
+        self._enter_calculation(variable, period)
+        self._record_start_time()
+
+    def record_parameter_access(
+        self,
+        parameter: str,
+        period: t.Period,
+        value: t.VarArray,
+    ) -> None:
+        if self._current_node is not None:
+            self._current_node.parameters.append(
+                TraceNode(name=parameter, period=period, value=value),
+            )
+
+    def record_calculation_result(self, value: t.VarArray) -> None:
+        if self._current_node is not None:
+            self._current_node.value = value
+
+    def record_calculation_end(self) -> None:
+        self._simple_tracer.record_calculation_end()
+        self._record_end_time()
+        self._exit_calculation()
+
+    def print_computation_log(
+        self,
+        aggregate: bool = False,
+        max_depth: int = sys.maxsize,
+        ignore_default: bool = False,
+        tax_benefit_system: t.TaxBenefitSystem | None = None,
+    ) -> None:
+        self.computation_log.print_log(
+            aggregate, max_depth, ignore_default, tax_benefit_system
+        )
+
+    def generate_performance_graph(self, dir_path: str) -> None:
+        self.performance_log.generate_graph(dir_path)
+
+    def generate_performance_tables(self, dir_path: str) -> None:
+        self.performance_log.generate_performance_tables(dir_path)
+
+    def get_nb_requests(self, variable: str) -> int:
+        return sum(self._get_nb_requests(tree, variable) for tree in self.trees)
+
+    def get_flat_trace(self) -> t.FlatNodeMap:
+        return self.flat_trace.get_trace()
+
+    def get_serialized_flat_trace(self) -> t.SerializedNodeMap:
+        return self.flat_trace.get_serialized_trace()
+
+    def browse_trace(self) -> Iterator[t.TraceNode]:
+        def _browse_node(node: t.TraceNode) -> Iterator[t.TraceNode]:
+            yield node
+
+            for child in node.children:
+                yield from _browse_node(child)
+
+        for node in self._trees:
+            yield from _browse_node(node)
+
+    def _enter_calculation(
+        self,
+        variable: t.VariableName,
+        period: t.PeriodInt | t.Period,
+    ) -> None:
+        new_node = TraceNode(
+            name=variable,
+            period=period,
+            parent=self._current_node,
+        )
+
+        if self._current_node is None:
+            self._trees.append(new_node)
+
+        else:
+            self._current_node.append_child(new_node)
+
+        self._current_node = new_node
+
+    def _record_start_time(
+        self,
+        time_in_s: float | None = None,
+    ) -> None:
+        if time_in_s is None:
+            time_in_s = self._get_time_in_sec()
+
+        if self._current_node is not None:
+            self._current_node.start = time_in_s
+
+    def _record_end_time(
+        self,
+        time_in_s: None | t.Time = None,
+    ) -> None:
+        if time_in_s is None:
+            time_in_s = self._get_time_in_sec()
+
+        if self._current_node is not None:
+            self._current_node.end = time_in_s
+
+    def _exit_calculation(self) -> None:
+        if self._current_node is not None:
+            self._current_node = self._current_node.parent
+
+    def _get_nb_requests(self, tree: t.TraceNode, variable: str) -> int:
+        tree_call = tree.name == variable
+        children_calls = sum(
+            self._get_nb_requests(child, variable) for child in tree.children
+        )
+
+        return tree_call + children_calls
+
+    @staticmethod
+    def _get_time_in_sec() -> t.Time:
+        return time.time_ns() / (10**9)
+
+
+__all__ = ["FullTracer"]

openfisca_core/tracers/performance_log.py

@@ -0,0 +1,133 @@
+from __future__ import annotations
+
+import csv
+import importlib.resources
+import itertools
+import json
+import os
+import typing
+
+from openfisca_core import tracers
+
+if typing.TYPE_CHECKING:
+    Trace = dict[str, dict]
+    Calculation = tuple[str, dict]
+    SortedTrace = list[Calculation]
+
+
+class PerformanceLog:
+    def __init__(self, full_tracer: tracers.FullTracer) -> None:
+        self._full_tracer = full_tracer
+
+    def generate_graph(self, dir_path: str) -> None:
+        with open(os.path.join(dir_path, "performance_graph.html"), "w") as f:
+            template = importlib.resources.read_text(
+                "openfisca_core.scripts.assets",
+                "index.html",
+            )
+
+            perf_graph_html = template.replace(
+                "{{data}}",
+                json.dumps(self._json()),
+            )
+
+            f.write(perf_graph_html)
+
+    def generate_performance_tables(self, dir_path: str) -> None:
+        flat_trace = self._full_tracer.get_flat_trace()
+
+        csv_rows = [
+            {
+                "name": key,
+                "calculation_time": trace["calculation_time"],
+                "formula_time": trace["formula_time"],
+            }
+            for key, trace in flat_trace.items()
+        ]
+
+        self._write_csv(
+            os.path.join(dir_path, "performance_table.csv"),
+            csv_rows,
+        )
+
+        aggregated_csv_rows = [
+            {"name": key, **aggregated_time}
+            for key, aggregated_time in self.aggregate_calculation_times(
+                flat_trace,
+            ).items()
+        ]
+
+        self._write_csv(
+            os.path.join(dir_path, "aggregated_performance_table.csv"),
+            aggregated_csv_rows,
+        )
+
+    def aggregate_calculation_times(
+        self,
+        flat_trace: Trace,
+    ) -> dict[str, dict]:
+        def _aggregate_calculations(calculations: list) -> dict:
+            calculation_count = len(calculations)
+
+            calculation_time = sum(
+                calculation[1]["calculation_time"] for calculation in calculations
+            )
+
+            formula_time = sum(
+                calculation[1]["formula_time"] for calculation in calculations
+            )
+
+            return {
+                "calculation_count": calculation_count,
+                "calculation_time": tracers.TraceNode.round(calculation_time),
+                "formula_time": tracers.TraceNode.round(formula_time),
+                "avg_calculation_time": tracers.TraceNode.round(
+                    calculation_time / calculation_count,
+                ),
+                "avg_formula_time": tracers.TraceNode.round(
+                    formula_time / calculation_count,
+                ),
+            }
+
+        def _groupby(calculation: Calculation) -> str:
+            return calculation[0].split("<")[0]
+
+        all_calculations: SortedTrace = sorted(flat_trace.items())
+
+        return {
+            variable_name: _aggregate_calculations(list(calculations))
+            for variable_name, calculations in itertools.groupby(
+                all_calculations,
+                _groupby,
+            )
+        }
+
+    def _json(self) -> dict:
+        children = [self._json_tree(tree) for tree in self._full_tracer.trees]
+        calculations_total_time = sum(child["value"] for child in children)
+
+        return {
+            "name": "All calculations",
+            "value": calculations_total_time,
+            "children": children,
+        }
+
+    def _json_tree(self, tree: tracers.TraceNode) -> dict:
+        calculation_total_time = tree.calculation_time()
+        children = [self._json_tree(child) for child in tree.children]
+
+        return {
+            "name": f"{tree.name}<{tree.period}>",
+            "value": calculation_total_time,
+            "children": children,
+        }
+
+    def _write_csv(self, path: str, rows: list[dict]) -> None:
+        fieldnames = list(rows[0].keys())
+
+        with open(path, "w") as csv_file:
+            writer = csv.DictWriter(csv_file, fieldnames=fieldnames)
+            writer.writeheader()
+
+            for row in rows:
+                writer.writerow(row)

openfisca_core/tracers/simple_tracer.py

@@ -0,0 +1,64 @@
+from __future__ import annotations
+
+from openfisca_core import types as t
+
+
+class SimpleTracer:
+    """A simple tracer that records a stack of traces."""
+
+    #: The stack of traces.
+    _stack: t.SimpleStack
+
+    def __init__(self) -> None:
+        self._stack = []
+
+    @property
+    def stack(self) -> t.SimpleStack:
+        """Return the stack of traces."""
+        return self._stack
+
+    def record_calculation_start(
+        self, variable: t.VariableName, period: t.PeriodInt | t.Period
+    ) -> None:
+        """Record the start of a calculation.
+
+        Args:
+            variable: The variable being calculated.
+            period: The period for which the variable is being calculated.
+
+        Examples:
+            >>> from openfisca_core import tracers
+
+            >>> tracer = tracers.SimpleTracer()
+            >>> tracer.record_calculation_start("variable", 2020)
+            >>> tracer.stack
+            [{'name': 'variable', 'period': 2020}]
+
+        """
+        self.stack.append({"name": variable, "period": period})
+
+    def record_calculation_result(self, value: t.ArrayLike[object]) -> None:
+        """Ignore calculation result."""
+
+    def record_parameter_access(
+        self, parameter: str, period: t.Period, value: t.ArrayLike[object]
+    ) -> None:
+        """Ignore parameter access."""
+
+    def record_calculation_end(self) -> None:
+        """Record the end of a calculation.
+
+        Examples:
+            >>> from openfisca_core import tracers
+
+            >>> tracer = tracers.SimpleTracer()
+            >>> tracer.record_calculation_start("variable", 2020)
+            >>> tracer.record_calculation_end()
+            >>> tracer.stack
+            []
+
+        """
+        self.stack.pop()
+
+
+__all__ = ["SimpleTracer"]