The `CausalProblem` class

[willGraham01]

Development

Since there are several moving parts here, I propose to introduce these changes piecemeal, across several PRs.

The two_normal_example integration test (#76) should be updated in each PR so that it works with the latest version of the above classes. Of course, additional tests should also be added for features of these classes too.

Also adopt the ethos of "write TODOs" in these classes when creating new features, and leave these present in the PRs to avoid scope-drifting and account for the inevitable complications we'll encounter going forward. We can convert these TODO statements to more issues and add them to the list below to continue tracking them.

• Opened CausalProblem rework #84 to implement the basic structure outlined above.

Description

Now that we have Graph.model, we should look to decide on how we want to represent and setup (at least in the backend) a causal problem. This was previously the remit of the CausalProblem class, but given our substantial changes to the Graph and ParameterNode classes, it needs to be adapted / reworked.

My vision for the CausalProblem class is something along the lines of...

class CausalProblem:

  _graph: Graph # A reference to the Graph (without conditioning / do operations applied that defines the RVs and structure
  causal_estimand: CausalEstimand # an instance encoding the causal estimand
  constraints: list[Constraint] # a collection of instances encoding the individual constraints

  def __init__(self, graph, ce, *constraints):
    ...

  def lagrangian(self, ...):
    # Returns a callable function L(params, l_multipliers, rng_key) that
    # evaluates the lagrangian of the minimisiation problem

possibly with more methods (like solve?) added as we develop the class further.

The CausalEstimand and Constraint classes should have a common structure:

# define a common base since there are a few subtle differences,
# but largely the same structure / features
class _CPComponent:

    do_with_samples: Callable[..., npt.ArrayLike]
    effect_handlers: tuple[effect_handler, ...]

where the user supplies do_with_samples, a (deterministic) function that maps samples of the random variables (given as a dict[str, ArrayLike]) to the value of the CE/Constraint.

The effect_handlers should be any adaptations to the Graph.model that need to be applied before samples are drawn for the instance. For example, if a constraint involves the do operation or involves conditioning on some value of a RV.