Add fail/warn mechanism to ColumnObjects and ParamFunctions (#16)

- As discussed in [GETTSIM
#1055](ttsim-dev/gettsim#1055), we sometimes
want to fail or issue a warning if a certain node is visited in the
graph. This PR adds a mechanism for that along with corresponding tests
- Implementation: 
  - All `ColumnObjects` and `ParamFunctions` are endowed with attributes
`warn_msg_if_included: str | None = None` and `fail_msg_if_included: str
| None = None`.
  - These are checked by a corresponding fail/warn function each, which
depends on the TT DAG and the specialised environment without tree logic
and with derived functions (guaranteed to catch the correct type of
object)
  - For warn functions, only unique messages are displayed. E.g., in the
GETTSIM example triggering this, we have at least three functions, which
all have the same message.
- As we visited that code, we realised that `tt_dag` is a better name
for `tax_transfer_dag` as we are using that abbreviation much more
extensively that 'tax_transfer' by now. Renamed, same for
`tax_transfer_function` -> `tt_function`.

Author: hmgaudecker

CHANGES.md

@@ -6,6 +6,9 @@ releases are available on [Anaconda.org](https://anaconda.org/conda-forge/ttsim)
 
 ## v1.0a2 — 2025-07-xx
 
+- {gh}`16` Add fail/warn mechanism to ColumnObjects and ParamFunctions.
+  ({ghuser}`hmgaudecker`)
+
 - {gh}`15` Do not call len() on unsized arrays. ({ghuser}`hmgaudecker`)
 
 - {gh}`14` Do not loop over the attributes of Jax arrays in

src/ttsim/interface_dag_elements/fail_if.py

@@ -38,7 +38,7 @@
 )
 
 if TYPE_CHECKING:
-    from collections.abc import Callable
+    from collections.abc import Callable, Iterable
 
     from ttsim.interface_dag_elements.input_data import FlatData
     from ttsim.typing import (
@@ -704,9 +704,42 @@ def backend_has_changed(
         )
 
 
+@fail_function()
+def tt_dag_includes_function_with_fail_msg_if_included_set(
+    specialized_environment__without_tree_logic_and_with_derived_functions: SpecEnvWithoutTreeLogicAndWithDerivedFunctions,
+    specialized_environment__tt_dag: nx.DiGraph,
+    labels__processed_data_columns: UnorderedQNames,
+) -> None:
+    """Fail if the TT DAG includes functions with `fail_msg_if_included` set."""
+
+    env = specialized_environment__without_tree_logic_and_with_derived_functions
+    issues = ""
+    for node in specialized_environment__tt_dag:
+        if (
+            # This may run before 'fail_if.root_nodes_are_missing'
+            node not in env
+            or
+            # ColumnObjects overridden by data are fine
+            (
+                not isinstance(env[node], PolicyInput)
+                and node in labels__processed_data_columns
+            )
+        ):
+            continue
+        # Check because ParamObjects can be overridden by ColumnObjects down the road.
+        if hasattr(env[node], "fail_msg_if_included"):  # noqa: SIM102
+            if msg := env[node].fail_msg_if_included:
+                issues += f"{node}:\n\n{msg}\n\n\n"
+    if issues:
+        raise ValueError(
+            "The TT DAG includes the following functions with `fail_msg_if_included` "
+            f"set.\n\n{issues}"
+        )
+
+
 @fail_function()
 def tt_root_nodes_are_missing(
-    specialized_environment__tax_transfer_dag: nx.DiGraph,
+    specialized_environment__tt_dag: nx.DiGraph,
     specialized_environment__with_partialled_params_and_scalars: SpecEnvWithPartialledParamsAndScalars,
     processed_data: QNameData,
     labels__grouping_levels: OrderedQNames,
@@ -715,7 +748,7 @@ def tt_root_nodes_are_missing(
 
     Parameters
     ----------
-    specialized_environment__tax_transfer_dag
+    specialized_environment__tt_dag
         The DAG of taxes and transfers functions.
     specialized_environment__with_partialled_params_and_scalars
         The specialized environment with partialled params and scalars.
@@ -738,9 +771,8 @@ def tt_root_nodes_are_missing(
         )
     # Obtain root nodes
     root_nodes = nx.subgraph_view(
-        specialized_environment__tax_transfer_dag,
-        filter_node=lambda n: specialized_environment__tax_transfer_dag.in_degree(n)
-        == 0,
+        specialized_environment__tt_dag,
+        filter_node=lambda n: specialized_environment__tt_dag.in_degree(n) == 0,
     ).nodes
 
     missing_nodes = [
@@ -848,7 +880,7 @@ def format_errors_and_warnings(text: str, width: int = 79) -> str:
     return "\n\n".join(wrapped_paragraphs)
 
 
-def format_list_linewise(some_list: list[Any]) -> str:  # type: ignore[type-arg, unused-ignore]
+def format_list_linewise(some_list: Iterable[Any]) -> str:  # type: ignore[type-arg, unused-ignore]
     formatted_list = '",\n    "'.join(some_list)
     return textwrap.dedent(
         """

src/ttsim/interface_dag_elements/labels.py

@@ -142,14 +142,14 @@ def input_columns(
 
 @interface_function()
 def root_nodes(
-    specialized_environment__tax_transfer_dag: nx.DiGraph,
+    specialized_environment__tt_dag: nx.DiGraph,
     processed_data_columns: UnorderedQNames,
 ) -> UnorderedQNames:
     """Names of the columns in `processed_data` required for the tax transfer function.
 
     Parameters
     ----------
-    specialized_environment__tax_transfer_dag:
+    specialized_environment__tt_dag:
         The tax transfer DAG.
     processed_data:
         The processed data.
@@ -161,9 +161,8 @@ def root_nodes(
     """
     # Obtain root nodes
     root_nodes = nx.subgraph_view(
-        specialized_environment__tax_transfer_dag,
-        filter_node=lambda n: specialized_environment__tax_transfer_dag.in_degree(n)
-        == 0,
+        specialized_environment__tt_dag,
+        filter_node=lambda n: specialized_environment__tt_dag.in_degree(n) == 0,
     ).nodes
 
     # Restrict the passed data to the subset that is actually used.

src/ttsim/interface_dag_elements/raw_results.py

@@ -19,9 +19,9 @@
 def columns(
     labels__root_nodes: UnorderedQNames,
     processed_data: QNameData,
-    specialized_environment__tax_transfer_function: Callable[[QNameData], QNameData],
+    specialized_environment__tt_function: Callable[[QNameData], QNameData],
 ) -> QNameData:
-    return specialized_environment__tax_transfer_function(
+    return specialized_environment__tt_function(
         {k: v for k, v in processed_data.items() if k in labels__root_nodes},
     )
 

src/ttsim/interface_dag_elements/specialized_environment.py

@@ -340,7 +340,7 @@ def _apply_rounding(element: ColumnFunction, xnp: ModuleType) -> ColumnFunction:
 
 
 @interface_function()
-def tax_transfer_dag(
+def tt_dag(
     with_partialled_params_and_scalars: SpecEnvWithPartialledParamsAndScalars,
     labels__column_targets: OrderedQNames,
 ) -> nx.DiGraph:
@@ -361,16 +361,16 @@ def tt_function_set_annotations() -> bool:
 
 
 @interface_function()
-def tax_transfer_function(
-    tax_transfer_dag: nx.DiGraph,
+def tt_function(
+    tt_dag: nx.DiGraph,
     with_partialled_params_and_scalars: SpecEnvWithPartialledParamsAndScalars,
     labels__column_targets: OrderedQNames,
     backend: Literal["numpy", "jax"],
     tt_function_set_annotations: bool,
 ) -> Callable[[QNameData], QNameData]:
     """Returns a function that takes a dictionary of arrays and unpacks them as keyword arguments."""
     ttf_with_keyword_args = concatenate_functions(
-        dag=tax_transfer_dag,
+        dag=tt_dag,
         functions=with_partialled_params_and_scalars,
         targets=list(labels__column_targets),
         return_type="dict",

src/ttsim/interface_dag_elements/warn_if.py

@@ -15,9 +15,12 @@
 if TYPE_CHECKING:
     import datetime
 
+    import networkx as nx
+
     from ttsim.typing import (
         PolicyEnvironment,
         QNameData,
+        SpecEnvWithoutTreeLogicAndWithDerivedFunctions,
         UnorderedQNames,
     )
 
@@ -114,3 +117,37 @@ def evaluation_date_set_in_multiple_places(
 `evaluation_year`.
 """
         warnings.warn(UserWarning(msg), stacklevel=2)
+
+
+@warn_function()
+def tt_dag_includes_function_with_warn_msg_if_included_set(
+    specialized_environment__without_tree_logic_and_with_derived_functions: SpecEnvWithoutTreeLogicAndWithDerivedFunctions,  # noqa: E501
+    specialized_environment__tt_dag: nx.DiGraph,
+    labels__processed_data_columns: UnorderedQNames,
+) -> None:
+    """Warn if the TT DAG includes functions with `warn_msg_if_included` set."""
+
+    env = specialized_environment__without_tree_logic_and_with_derived_functions
+    my_warnings = set()
+    for node in specialized_environment__tt_dag:
+        if (
+            # This may run before 'fail_if.root_nodes_are_missing'
+            node not in env
+            or
+            # ColumnObjects overridden by data are fine
+            (
+                not isinstance(env[node], PolicyInput)
+                and node in labels__processed_data_columns
+            )
+        ):
+            continue
+        # Check because ParamObjects can be overridden by ColumnObjects down the road.
+        if hasattr(env[node], "fail_msg_if_included"):  # noqa: SIM102
+            if msg := env[node].warn_msg_if_included:
+                my_warnings |= {f"{msg}\n\n\n"}
+    if my_warnings:
+        msg = (
+            "The TT DAG includes elements with `warn_msg_if_included` set to the "
+            f"following values:\n\n{format_list_linewise(my_warnings)}"
+        )
+        warnings.warn(UserWarning(msg), stacklevel=2)

src/ttsim/main_args.py

@@ -166,5 +166,5 @@ class SpecializedEnvironment(MainArg):
     with_partialled_params_and_scalars: SpecEnvWithPartialledParamsAndScalars | None = (
         None
     )
-    tax_transfer_dag: nx.DiGraph | None = None
-    tax_transfer_function: Callable[[QNameData], QNameData] | None = None
+    tt_dag: nx.DiGraph | None = None
+    tt_function: Callable[[QNameData], QNameData] | None = None

src/ttsim/main_target.py

@@ -26,6 +26,9 @@ class WarnIf(MainTargetABC):
     evaluation_date_set_in_multiple_places: str = (
         "warn_if__evaluation_date_set_in_multiple_places"
     )
+    tt_dag_includes_function_with_warn_msg_if_included_set: str = (
+        "warn_if__tt_dag_includes_function_with_warn_msg_if_included_set"
+    )
 
 
 @dataclass(frozen=True)
@@ -66,6 +69,9 @@ class FailIf(MainTargetABC):
         "fail_if__targets_are_not_in_specialized_environment_or_data"
     )
     targets_tree_is_invalid: str = "fail_if__targets_tree_is_invalid"
+    tt_dag_includes_function_with_fail_msg_if_included_set: str = (
+        "fail_if__tt_dag_includes_function_with_fail_msg_if_included_set"
+    )
 
 
 @dataclass(frozen=True)
@@ -94,8 +100,8 @@ class SpecializedEnvironment(MainTargetABC):
     with_partialled_params_and_scalars: str = (
         "specialized_environment__with_partialled_params_and_scalars"
     )
-    tax_transfer_dag: str = "specialized_environment__tax_transfer_dag"
-    tax_transfer_function: str = "specialized_environment__tax_transfer_function"
+    tt_dag: str = "specialized_environment__tt_dag"
+    tt_function: str = "specialized_environment__tt_function"
 
 
 @dataclass(frozen=True)