CausalEstimand class for evaluating causal estimands

[willGraham01]

Partially addresses #37 | Target is #39 since ParameterNodes are needed to get this functionality working, but @mscroggs feel free to hold off the merge until you're happy with the base (since this will need rebasing after any updates to that branch, anyway).

Introduces the CausalProblem class, designed to handle interactions between a Graph (.graph), the causal estimand of interest (.causal_estimand) and the constraints (currently not implemented). Handing of the constraints is not introduced in this pull request, since any comments about how the causal_estimand is being handled would apply to the constraints too, so figured it was best to settle on a design for one before writing the next.

Key Changes

Graph class

Had to introduce two methods related to ParameterNodes - mostly for fetching a tuple of all such nodes (note that this method uses graph.ordered_nodes since it is important that the parameter nodes a consistently given in the same order) and for the setting of the .value attributes of these nodes.

@mscroggs you might want to steal the corresponding changes & their tests and move them into the base PR (#39).

Node class

The sample method for ParameterNodes needs to have the same argument names as the super-class for the "up-the-graph" logic to work without an error, so have removed the _ prefixes and slapped on a # noqa for now, until we find a more elegant solution.

I also took the liberty of telling the linter that unused *args and **kwargs are OK. Since we're going to have a lot of inherited and overwritten methods that may not utilise all the arguments in the signature of the super-class.

CausalProblem

A CausalProblem can be initialised with just a label, and optionally providing a Graph that the instance will store. Said graph will define the relationships between the RVs that the problem considers, and can be set later via the .graph attribute if so desired. Very limited functionality (as in, essentially all you can do is set the .graph attribute) is available for a CausalProblem instance without a set .graph attribute.

Otherwise, the primary way a user is expected to interact with a CausalProblem is to provide the causal estimand function via the set_causal_estimand method. This takes some callable object (sigma) which is assumed to be a function of some RVs, with the RVs to be associated to Nodes in the .graph. At present, we also have to run a few "hacks" due to some design aspects we haven't got round to addressing:

• set_causal_estimand allows one of the arguments to sigma to be a Graph object. This is because our current algorithms module only has methods that support calls via expectation(graph, node_name, ...) etc.
• sigma (and the tests for this method) provides a static rng_key and sample_size, whereas in the long run we probably want these set instance-wise across the whole CausalProblem instance.

Once set, the .causal_estimand method can be called, which will evaluate the given sigma function using the associated Nodes in .graph. This method takes a jax.Array as input, which is assumed to be an array of the parameter values - the order of the values matches that of CausalProblem.graph.parameter_nodes (hence why it is important for this method to always return parameter nodes in the same order). It then calls .graph.set_parameters with the appropriate values extracted from the input array, before running sigma and returning the result. This gets around the issue identified in #37 regarding fixing a single realisation of a RV by accident.

.causal_estimand primarily exists to be sent into a solver / optimiser - I don't envision the user ever wanting to evaluate it manually, and they certainly should not need to evaluate it manually via a parameter vector. If a user does want to evaluate their causal estimand at a particular set of parameters, the recommended way is to run

# Inputs are sent directly to .graph.set_parameters(**)
causal_problem.set_parameter_values(**parameter_values)
# The .parameter_vector property returns the appropriate vector for the
# current parameter values.
# As such, since we just set the parameter values (using their names),
# we can now get the CP instance itself to build the vector, then evaluate
# the causal estimand at that vector.
causal_problem.causal_estimand(causal_problem.parameter_vector)

[mscroggs]

The sample method for ParameterNodes needs to have the same argument names as the super-class for the "up-the-graph" logic to work without an error, so have removed the _ prefixes and slapped on a # noqa for now, until we find a more elegant solution.

I did this exact same change yesterday when I wrote a new test in #42. Yet to think of a more elegant way to do it