Codecov Report

❌ Patch coverage is 89.28571% with 21 lines in your changes missing coverage. Please review.

📢 Thoughts on this report? Let us know!

@JuergenWiemers This PR should fix all the stuff we talked about in the last couple of days. I hope I didn't miss any corner case. If you find the time, feel free to try it out and complain if I missed something!

Thanks! I looked through most of it but don't quite get the exact workings yet.

In particular, could you explain in the docstrings or here how the issue we had with respect to providing some input data and requesting all targets is resolved in specialized_envrionment_from_policy_inputs?

Oh, forgot to voice the feeling I had that maybe we can simplify more in specialized_environment now? That is, didn't we include what we factor out now as a special case?

Thanks a lot for this!!!

Oh I hoped it becomes clear through the module docstring, but here is a longer explanation:

The issues were mostly about the handling of derived functions. A few weeks ago we added a hack to be able to create templates / the TT DAG without requiring input data: We set the policy input names as input qnames if no input (or processed) data was provided. This way, we could use our regular machinery and if no input data was provided, everything was fine.

Now, this PR generalises this by extending the same logic (if some input data is missing, just assume it is there if the goal is to build a template or plot a DAG) to cases where users provided partial input data (e.g. one root node that they want to override in the template or the plot).

Let's go through it function by function:

1. without_tree_logic_and_with_derived_functions

• In specialized_environment this builds derived functions based on i) input data and ii) column functions. I.e. wealth_fam is created if either i) wealth is in the input data or ii) there is a column function called wealth.
• In specialized_environment_from_policy_inputs we also build derived functions based on policy inputs. That is, wealth_fam is created if either i) wealth is in the input data, ii) there is a column function called wealth or iii) there is a policy input called wealth.

In the two applications we look at here (plotting TT DAGs and creating templates from DAG root nodes), we typically don't provide input data. Hence, when using specialized_environment instead of specialized_environment_from_policy_inputs, we would never create the function wealth_fam. wealth_fam is then a root node of the DAG, which is incorrect.

2. without_processed_data_nodes_with_dummy_callables

This one is new in specialized_environment_from_policy_inputs. It creates dummy callables from policy inputs and parameters such that they are part of the DAG.

3. complete_dag

This one is new in specialized_environment_from_policy_inputs. Here we build the DAG without partialling parameters into the functions. This function yields the "complete" DAG, i.e. the DAG that contains all params, all column functions and all policy inputs. It still depends on input data though, i.e. if users overwrite columns via input data or feed input data that implies different kinds of time conversions (_w instead of _m for example), this will show up here.

4. with_processed_params_and_scalars / with_partialled_params_and_scalars

These are imports from specialized_environment. We need them because we need SpecEnvWithPartialledParamsAndScalars to create the input templates.

Oh, forgot to voice the feeling I had that maybe we can simplify more in specialized_environment now? That is, didn't we include what we factor out now as a special case?

I don't think we have to adjust specialized_environment, we handles this special case in main.py and labels.py. But you're right, I forgot to remove those!

Oh, forgot to voice the feeling I had that maybe we can simplify more in specialized_environment now? That is, didn't we include what we factor out now as a special case?

I don't think we have to adjust specialized_environment, we handles this special case in main.py and labels.py. But you're right, I forgot to remove those!

Yeah, I'd think the fail/warn logic should need some adjustment then. That is, when requesting something from the standard specialized environment without providing data, we should always see an error, right? (did not check now whether you already implemented it, at least I don't remember from my earlier look)

And thanks for the explainer, very clear! Apologies I missed that docstring, will have another look tomorrow.

"Pruning" the ancestors-DAG based on user inputs works as advertised:

from __future__ import annotations
import numpy as np
from ttsim import InputData
from ttsim.plot import dag
import mettsim.middle_earth as middle_earth

fig = dag.tt(
    root=middle_earth.ROOT_PATH,
    policy_date_str="2000-01-01",
    node_selector=dag.NodeSelector(
        type="ancestors",
        node_paths=[("housing_benefits", "amount_m_fam")],
    ),
    input_data=InputData.tree(
        tree={
            "p_id": np.array([0]),
            "housing_benefits": {
                "income": {"amount_m_fam": np.array([0]),},
                "eligibility": {"number_of_adults_fam": np.array([0]),
                }
            },
        },
    ),    
    include_params=False,
    show_node_description=True,
    output_path="METTSIM_housing_benefits_pruned.html",
)

gives

But I have two minor quibbles (you told me to complain! 😉)

• The output contains

  Your data provides the column:
  
  [
      "housing_benefits__eligibility__number_of_adults_fam",
  ]
  
  This is already present among the hard-coded functions of the taxes and
  transfers system. ...
  
  Turn off warnings by setting `include_warn_nodes=False` in `main`. If you want
  to be selective about warnings, include these among the `main_targets`.

  • Shouldn't it also show housing_benefits__income__amount_m_fam?
  • I cannot set include_warn_nodes=False with ttsim.plot.dag.tt. I guess the idea is to move the plotting functionality to something like main(main_target=MainTarget.plot.dag.tt, ...) at a later point?
• It feels a bit awkward to have to provide a "dummy dataset" as input_data just to get a pruned graph. I mean, the data don't really matter, what matters is just the nodes, right? Something like

  fig = dag.tt(
      root=middle_earth.ROOT_PATH,
      policy_date_str="2000-01-01",
      node_selector=dag.NodeSelector(
          type="ancestors",
          node_paths=[("housing_benefits", "amount_m_fam")],
      ),
      input_node_paths=[
          ("housing_benefits", "income", "amount_m_fam"),
          ("housing_benefits", "eligibility", "number_of_adults_fam"),
      ],
      include_params=False,
      show_node_description=True,
      output_path="METTSIM_housing_benefits_pruned.html",
  )

  (because it mirrors the node_paths syntax; a nested dict instead of the tuples would also be fine) would feel more natural to me. (Maybe these input_node_paths could be converted to input_data behind the scenes?) To be honest, I already thought the same with respect to MainTarget.templates.input_data_dtypes ("Why do I have to provide data to get a pruned input template?").

Or is there a deeper reason why input_data needs to be given for "pruned" templates/DAG plots?

Yeah, I'd think the fail/warn logic should need some adjustment then. That is, when requesting something from the standard specialized environment without providing data, we should always see an error, right? (did not check now whether you already implemented it, at least I don't remember from my earlier look)

I drafted a version that doesn't rely on the hack. The idea is that we use input dependent interface functions; whenever there is input data, use it. If there is no input data, just use empty sets / dicts instead.

dags doesn't work well with functions that don't have any dependencies (they are not part of the DAG), so I added a fake dependency to one function that should always be satisfied:

@input_dependent_interface_function(
    include_if_no_input_present=[
        "input_data__df_and_mapper__df",
        "input_data__df_and_mapper__mapper",
        "input_data__df_with_nested_columns",
        "input_data__tree",
    ],
    leaf_name="processed_data_or_empty_dict",
)
def processed_data_or_empty_dict_no_input_data_provided(
    labels__grouping_levels: OrderedQNames,  # fake input
) -> QNameData:
    """The processed data or an empty dict.

    When computing nodes from this namespace (helpers to create templates or plot the
    DAG), we don't necessarily need to have processed data, as processed data only
    customizes the targets.
    """
    return {}

Tbh I find the current implementation a bit confusing (creating processed data that may or may not be more than an empty dict), but I don't see another way. LMK in case you have ideas.

Edit: Now I understand where you were coming from with the processed_data dummy data structure. Needed to think about Jürgen's comment to find out you were already thinking miles ahead. This makes things much easier I believe.

Shouldn't it also show housing_benefits__income__amount_m_fam?

No, because housing_benefits__income__amount_m_fam is a derived function (not what I would consider hard coded). One could throw a different warning there of course.

I cannot set include_warn_nodes=False with ttsim.plot.dag.tt. I guess the idea is to move the plotting functionality to something like main(main_target=MainTarget.plot.dag.tt, ...) at a later point?

Good point, that should be False by default in that case.

We want to keep the plotting infrastructure outside of main. Allows for a bit more flexibility when providing plotly kwargs (not implemented yet, but the idea is that users can customize the plot as they please) and it would bloat the interface DAG.

It feels a bit awkward to have to provide a "dummy dataset" as input_data just to get a pruned graph. I mean, the data don't really matter, what matters is just the nodes, right?

I agree. I thought it might be more convenient for users to pass input data using the same syntax as with main, but it does more harm than good here, I believe.

It feels a bit awkward to have to provide a "dummy dataset" as input_data just to get a pruned graph. I mean, the data don't really matter, what matters is just the nodes, right?

I agree. I thought it might be more convenient for users to pass input data using the same syntax as with main, but it does more harm than good here, I believe.

Darn, had I not destroyed my phone hitting a wasp crawling towards my kids' soft-drinks, I would have mentioned the connection earlier and saved you some work... We really just need the input_data__columns there, which is exactly what @JuergenWiemers suggested as input_node_paths above. Whether we create these from a dataset or a user supplies them directly should not matter.

While you are at it, let's get the interface of the dag-plotters right. Three components:

1. ttsim-plot-specific options
2. elements of main
3. arguments passed to plotly.

1 seems obvious and 3. should just be the residual (**kwargs). Probably we should make sure that for 2., we allow everything that could be useful? That is, orig_policy_objects, input_data, processed_data, labels, tt_targets (always all elements)?

Edit: I realised that I use the word target two describe two things that are a bit different. 1) target in the sense of end-points of the DAG and 2) target as the root of the TTSIM DAG plot (the middle point if neighbors are plotted, the root node if descendants are plotted, the end node if ancestors are plotted). Definitely need to improve naming in the plotting interface in that regard.

I think I finally have a solution that works. The basic problem is that we

• have to select all available nodes as targets for the DAG in order to allow for plotting of neighbors and descendants (in other words, we have to be agnostic about what the end nodes of the DAG are).
• but when adding input data, we want to get rid of sets of nodes overridden by one input column. This is not trivial because we have to select some of these nodes as targets (see above)!

The solution is to first create a DAG setting all nodes as targets (but removing functions that are in the data) and then removing all nodes that are not ancestors or descendants of any of the targets.

This comes a bit out of the blue, so here is an example:

• The user is interested in "c". We need to accommodate all possible types of plots: i) a plot of the entire DAG, ii) ancestors of "c", iii) descendents of "c", iv) neighbors of "c". That's why the first step is to create the entire DAG shown above.
• If no input data is given, its easy: Just take this DAG and then choose the nodes to be plotted starting from the target node.
• If "b" is overwritten by input data, we are left with this DAG:

• Because we want to get rid of "d" in the final output, we eliminate it by searching for the ancestors /descendants of "c". Then, we can, as in the case with no input data, take this DAG and select the nodes we want to plot

Apart from that, I removed the NodeSelector dataclass. This way, users can set targets as they do when computing taxes and transfers, no need to introduce a new concept.

This one is finished from my POV. @JuergenWiemers feel free to try it out here if you find the time.

You can pass input data as regular InputData object, or you can just pass the labels of the input columns. Here is an example using GETTSIM:

from gettsim import plot

plot.dag.tt(
    policy_date_str="2025-01-01",
    primary_nodes=["einkommensteuer__betrag_y_sn"],
    selection_type="ancestors",
    labels={"input_columns": ["einnahmen__renten__gesetzliche_y"]},
)

There are some orphaned nodes when plotting the entire GETTSIM DAG. This is not a plotting issue, we just have some objects in the policy environment we don't use currently (geburtstag, stuff related to the Nachhaltigkeitsfaktor in the pension formula, etc).

This all looks very good to me! Being able to specify input/"known" nodes for templates/DAG plots through labels is a big improvement for the interface!

I noticed just a tiny asymmetry in the DAG plots:

• In ancestors plots, nodes given through the labels option are included in the plot
• In descendants plots, these nodes are not included.

Here is an example for the latter:

    fig_p_id_parent_pruned = dag.tt(
        root=middle_earth.ROOT_PATH,
        policy_date_str="2000-01-01",
        primary_nodes={"p_id_parent_1"},
        labels=Labels(input_columns={"wealth_fam"}),
        selection_type="descendants",
        include_params=False,

wealth_fam does not show in the plot.

I noticed just a tiny asymmetry in the DAG plots:

* In ancestors plots, nodes given through the `labels` option are _included_ in the plot

* In descendants plots, these nodes are _not included_.

I did not check the example, but it seems like this might be entirely expected? That is, if you are asking for the descendents of some node y and pass x as the data, x cannot possibly be a descendent of y (even though the policy function x might have y as an input, but passing x as an input means this is overridden).

Thanks for trying it out! @JuergenWiemers

I noticed just a tiny asymmetry in the DAG plots

I agree with HM here. The DAG plot is supposed to show what GETTSIM calculates eventually. Including the input node in the plot when we use selection_type='ancestors' makes sense because the input node is a root node of the DAG.

With selection_type='descendants', the input node is completely independent of its ancestors; it would be confusing to indicate a relationship between them by including it in the DAG.

Maybe I am expecting something unreasonable. Let me try again:

• First an ancestors plot of orc_hunting_bounty__amount, where an immediate ancestor (orc_hunting_bounty__amount_without_topup) is given as "known":

  fig = dag.tt(
      root=middle_earth.ROOT_PATH,
      policy_date_str="2000-01-01",
      primary_nodes={"orc_hunting_bounty__amount"},
      labels=Labels(input_columns={"orc_hunting_bounty__amount_without_topup"}),
      selection_type="ancestors",
      include_params=False,
      title="Ancestors Plot: orc_hunting_bounty__amount"
  )

The node in labels is shown in the plot. The immediate ancestors of orc_hunting_bounty__amount_without_topup are cut off, but not orc_hunting_bounty__amount_without_topup itself, as expected.

• Now let's reverse the roles in the descendants plot:

  fig = dag.tt(
      root=middle_earth.ROOT_PATH,
      policy_date_str="2000-01-01",
      primary_nodes={"orc_hunting_bounty__amount_without_topup"},
      labels=Labels(input_columns={"orc_hunting_bounty__amount"}),
      selection_type="descendants",
      include_params=False,
      title="Descendants Plot: orc_hunting_bounty__amount_without_topup"
  )

  
  I would have expected orc_hunting_bounty__amount to show in the last plot.

I thought the idea of the descendants plot with additional "known nodes" was to prune the immediate descendants of the known nodes, but the known nodes itself should be included. This is what I mean by "asymmetry in the plots".

It's probably a tiny thing that is not worth the hassle...

Thanks for the explanation!

I thought the idea of the descendants plot with additional "known nodes" was to prune the immediate descendants of the known nodes, but the known nodes itself should be included. This is what I mean by "asymmetry in the plots".

The behaviour could be changed in this way, ofc, but I do not think it would be helpful in a broader set of use cases. Think about GETTSIM and you are simply throwing your input data at the plotting function (the column names of the data end up being labels__input_columns). If you asked for the descendants of a single node, you would still get all other input columns as orphaned nodes.

Maybe a different way to view it: There is no way to recur to the structure of the graph that would obtain if you had not passed this particular input data. Well, maybe we could achieve that, but I am pretty sure that would be like opening Pandora's box...

In order to get the plot you want, you could just pass orc_hunting_bounty__amount as a primary node, too.