Lagrangian respects effect handlers (#94)

* Reinstate check after renaming caused bugs in sampling

* Remove old CausalProblem class

* Add barebones classes to be populated later

* Update two normal example test to use new infrastructure

* Rework so that the lagrangian can be passed model parameters and the multiplier values as separate args

* ruffing

* Refactor out g.model argument from the Lagrangian call

* Make TODOs obvious so I don't forget to do them

* Add docstrings and more TODOs

* Todo resolution and addition

* Make _CPConstraint callable

* Hide _CPComponent attributes that we don't expect to change

* Test __call__ for _CPComponents

* Add note about __call__ in docstring

* Fix bug in how handlers are applied

* Write tests for features

* Edit Constraint so it is created in pieces

* Rework Constraint.__init__ and docstring to match new format

* Update two_normal_example integration test

* Remove todo note

* Create wrapper class to make handlers easier. __eq__ placeholder for now

* Tidy eq docstring

* Write necessary comparison methods

* Write model association method and implement in Lagrangian

* Docstirngs and breakout HandlerToApply class to submodule

* Docstirngs and breakout HandlerToApply class to submodule

* Tests for HandlerToApply class

* Tests for can_use_same_model

* Tests for associating models to components of the CP

* Reorganise classes now that structure is somewhat rigid

* Add module-level import for user-facing functions

* Update src/causalprog/causal_problem/causal_problem.py

Co-authored-by: Matthew Scroggs <matthew.w.scroggs@gmail.com>

* Reinstate if not else fix

---------

Co-authored-by: Matthew Scroggs <matthew.w.scroggs@gmail.com>

Author: willGraham01

src/causalprog/causal_problem/__init__.py

@@ -1 +1,7 @@
 """Classes for defining causal problems."""
+
+from .causal_problem import CausalProblem
+from .components import CausalEstimand, Constraint
+from .handlers import HandlerToApply
+
+__all__ = ("CausalEstimand", "CausalProblem", "Constraint", "HandlerToApply")

src/causalprog/causal_problem/_base_component.py

@@ -0,0 +1,86 @@
+"""Base class for components of causal problems."""
+
+from collections.abc import Callable
+from typing import Any
+
+import numpy.typing as npt
+
+from causalprog.causal_problem.handlers import EffectHandler, HandlerToApply, Model
+
+
+class _CPComponent:
+    """
+    Base class for components of a Causal Problem.
+
+    A _CPComponent has an attached method that it can apply to samples
+    (`do_with_samples`), which will be passed sample values of the RVs
+    during solution of a Causal Problem and used to evaluate the causal
+    estimand or constraint the instance represents.
+
+    It also has a sequence of effect handlers that need to be applied
+    to the sampling model before samples can be drawn to evaluate this
+    component. For example, if a component requires conditioning on the
+    value of a RV, the `condition` handler needs to be applied to the
+    underlying model, before generating samples to pass to the
+    `do_with_sample` method. `effect_handlers` will be applied to the model
+    in the order they are given.
+    """
+
+    do_with_samples: Callable[..., npt.ArrayLike]
+    effect_handlers: tuple[HandlerToApply, ...]
+
+    @property
+    def requires_model_adaption(self) -> bool:
+        """Return True if effect handlers need to be applied to model."""
+        return len(self.effect_handlers) > 0
+
+    def __call__(self, samples: dict[str, npt.ArrayLike]) -> npt.ArrayLike:
+        """
+        Evaluate the estimand or constraint, given sample values.
+
+        Args:
+            samples: Mapping of RV (node) labels to samples of that RV.
+
+        Returns:
+            Value of the estimand or constraint, given the samples.
+
+        """
+        return self._do_with_samples(**samples)
+
+    def __init__(
+        self,
+        *effect_handlers: HandlerToApply | tuple[EffectHandler, dict[str, Any]],
+        do_with_samples: Callable[..., npt.ArrayLike],
+    ) -> None:
+        self.effect_handlers = tuple(
+            h if isinstance(h, HandlerToApply) else HandlerToApply.from_pair(h)
+            for h in effect_handlers
+        )
+        self._do_with_samples = do_with_samples
+
+    def apply_effects(self, model: Model) -> Model:
+        """Apply any necessary effect handlers prior to evaluating."""
+        adapted_model = model
+        for handler in self.effect_handlers:
+            adapted_model = handler.handler(adapted_model, handler.options)
+        return adapted_model
+
+    def can_use_same_model_as(self, other: "_CPComponent") -> bool:
+        """
+        Determine if two components use the same (predictive) model.
+
+        Two components rely on the same model if they apply the same handlers
+        to the model, which occurs if and only if `self.effect_handlers` and
+        `other.effect_handlers` contain identical entries, in the same order.
+        """
+        if (not isinstance(other, _CPComponent)) or (
+            len(self.effect_handlers) != len(other.effect_handlers)
+        ):
+            return False
+
+        return all(
+            my_handler == their_handler
+            for my_handler, their_handler in zip(
+                self.effect_handlers, other.effect_handlers, strict=True
+            )
+        )

src/causalprog/causal_problem/causal_problem.py

@@ -6,7 +6,11 @@
 import numpy.typing as npt
 from numpyro.infer import Predictive
 
-from causalprog.causal_problem.causal_estimand import CausalEstimand, Constraint
+from causalprog.causal_problem._base_component import _CPComponent
+from causalprog.causal_problem.components import (
+    CausalEstimand,
+    Constraint,
+)
 from causalprog.graph import Graph
 
 
@@ -40,6 +44,11 @@ class CausalProblem:
     causal_estimand: CausalEstimand
     constraints: list[Constraint]
 
+    @property
+    def _ordered_components(self) -> list[_CPComponent]:
+        """Internal ordering for components of the causal problem."""
+        return [*self.constraints, self.causal_estimand]
+
     def __init__(
         self,
         graph: Graph,
@@ -51,6 +60,83 @@ def __init__(
         self.causal_estimand = causal_estimand
         self.constraints = list(constraints)
 
+    def _associate_models_to_components(
+        self, n_samples: int
+    ) -> tuple[list[Predictive], list[int]]:
+        """
+        Create models to be used by components of the problem.
+
+        Depending on how many constraints (and the causal estimand) require effect
+        handlers to wrap `self._underlying_graph.model`, we will need to create several
+        predictive models to sample from. However, we also want to minimise the number
+        of such models we have to make, in order to minimise the time we spend
+        actually computing samples.
+
+        As such, in this method we determine:
+        - How many models we will need to build, by grouping the constraints and the
+          causal estimand by the handlers they use.
+        - Build these models, returning them in a list called `models`.
+        - Build another list that maps the index of components in
+          `self._ordered_components` to the index of the model in `models` that they
+          use. The causal estimand is by convention the component at index -1 of this
+          returned list.
+
+        Args:
+            n_samples: Value to be passed to `numpyro.Predictive`'s `num_samples`
+                argument for each of the models that are constructed from the underlying
+                graph.
+
+        Returns:
+            list[Predictive]: List of Predictive models, whose elements contain all the
+                models needed by the components.
+            list[int]: Mapping of component indexes (as per `self_ordered_components`)
+                to the index of the model in the first return argument that the
+                component uses.
+
+        """
+        models: list[Predictive] = []
+        grouped_component_indexes: list[list[int]] = []
+        for index, component in enumerate(self._ordered_components):
+            # Determine if this constraint uses the same handlers as those of any of
+            # the other sets.
+            belongs_to_existing_group = False
+            for group in grouped_component_indexes:
+                # Pull any element from the group to compare models to.
+                # Items in a group are known to have the same model, so we can just
+                # pull out the first one.
+                group_element = self._ordered_components[group[0]]
+                # Check if the current constraint can also use this model.
+                if component.can_use_same_model_as(group_element):
+                    group.append(index)
+                    belongs_to_existing_group = True
+                    break
+
+            # If the component does not fit into any existing group, create a new
+            # group for it. And add the model corresponding to the group to the
+            # list of models.
+            if not belongs_to_existing_group:
+                grouped_component_indexes.append([index])
+
+                models.append(
+                    Predictive(
+                        component.apply_effects(self._underlying_graph.model),
+                        num_samples=n_samples,
+                    )
+                )
+
+        # Now "invert" the grouping, creating a mapping that maps the index of a
+        # component to the (index of the) model it uses.
+        component_index_to_model_index = []
+        for index in range(len(self._ordered_components)):
+            for group_index, group in enumerate(grouped_component_indexes):
+                if index in group:
+                    component_index_to_model_index.append(group_index)
+                    break
+        # All indexes should belong to at least one group (worst case scenario,
+        # their own individual group). Thus, it is safe to do the above to create
+        # the mapping from component index -> model (group) index.
+        return models, component_index_to_model_index
+
     def lagrangian(
         self, n_samples: int = 1000, *, maximum_problem: bool = False
     ) -> Callable[[dict[str, npt.ArrayLike], npt.ArrayLike, jax.Array], npt.ArrayLike]:
@@ -89,26 +175,28 @@ def lagrangian(
         """
         maximisation_prefactor = -1.0 if maximum_problem else 1.0
 
+        # Build association between self.constraints and the model-samples that each
+        # one needs to use. We do this here, since once it is constructed, it is
+        # fixed, and doesn't need to be done each time we call the Lagrangian.
+        models, component_to_index_mapping = self._associate_models_to_components(
+            n_samples
+        )
+
         def _inner(
             parameter_values: dict[str, npt.ArrayLike],
             l_mult: jax.Array,
             rng_key: jax.Array,
         ) -> npt.ArrayLike:
-            # In general, we will need to check which of our CE/CONs require masking,
-            # and do multiple predictive models to account for this...
-            # We can always pre-build the predictive models too, so we should replace
-            # the "model" input with something that can map the right predictive models
-            # to the CE/CONS that need them.
-            # TODO: https://github.com/UCL/causalprog/issues/90
-            predictive_model = Predictive(
-                model=self._underlying_graph.model, num_samples=n_samples
+            # Draw samples from all models
+            all_samples = tuple(
+                sample_model(model, rng_key, parameter_values) for model in models
             )
-            all_samples = sample_model(predictive_model, rng_key, parameter_values)
 
-            value = maximisation_prefactor * self.causal_estimand(all_samples)
+            value = maximisation_prefactor * self.causal_estimand(all_samples[-1])
             # TODO: https://github.com/UCL/causalprog/issues/87
             value += sum(
-                l_mult[i] * c(all_samples) for i, c in enumerate(self.constraints)
+                l_mult[i] * c(all_samples[component_to_index_mapping[i]])
+                for i, c in enumerate(self.constraints)
             )
             return value
 

src/causalprog/causal_problem/causal_estimand.py renamed to src/causalprog/causal_problem/components.py

@@ -6,66 +6,13 @@
 import jax.numpy as jnp
 import numpy.typing as npt
 
+from causalprog.causal_problem._base_component import _CPComponent
+
 Model: TypeAlias = Callable[..., Any]
 EffectHandler: TypeAlias = Callable[Concatenate[Model, ...], Model]
 ModelMask: TypeAlias = tuple[EffectHandler, dict]
 
 
-class _CPComponent:
-    """
-    Base class for components of a Causal Problem.
-
-    A _CPComponent has an attached method that it can apply to samples
-    (`do_with_samples`), which will be passed sample values of the RVs
-    during solution of a Causal Problem and used to evaluate the causal
-    estimand or constraint the instance represents.
-
-    It also has a sequence of effect handlers that need to be applied
-    to the sampling model before samples can be drawn to evaluate this
-    component. For example, if a component requires conditioning on the
-    value of a RV, the `condition` handler needs to be applied to the
-    underlying model, before generating samples to pass to the
-    `do_with_sample` method. `effect_handlers` will be applied to the model
-    in the order they are given.
-    """
-
-    do_with_samples: Callable[..., npt.ArrayLike]
-    effect_handlers: tuple[ModelMask, ...]
-
-    @property
-    def requires_model_adaption(self) -> bool:
-        """Return True if effect handlers need to be applied to model."""
-        return len(self._effect_handlers) > 0
-
-    def __call__(self, samples: dict[str, npt.ArrayLike]) -> npt.ArrayLike:
-        """
-        Evaluate the estimand or constraint, given sample values.
-
-        Args:
-            samples: Mapping of RV (node) labels to samples of that RV.
-
-        Returns:
-            Value of the estimand or constraint, given the samples.
-
-        """
-        return self._do_with_samples(**samples)
-
-    def __init__(
-        self,
-        *effect_handlers: ModelMask,
-        do_with_samples: Callable[..., npt.ArrayLike],
-    ) -> None:
-        self._effect_handlers = tuple(effect_handlers)
-        self._do_with_samples = do_with_samples
-
-    def apply_effects(self, model: Model) -> Model:
-        """Apply any necessary effect handlers prior to evaluating."""
-        adapted_model = model
-        for handler, handler_options in self._effect_handlers:
-            adapted_model = handler(adapted_model, handler_options)
-        return adapted_model
-
-
 class CausalEstimand(_CPComponent):
     """
     A Causal Estimand.