Constraint creation changes

src/causalprog/causal_problem/causal_estimand.py

@@ -3,6 +3,7 @@
 from collections.abc import Callable
 from typing import Any, Concatenate, TypeAlias
 
+import jax.numpy as jnp
 import numpy.typing as npt
 
 Model: TypeAlias = Callable[..., Any]
@@ -94,13 +95,74 @@ class Constraint(_CPComponent):
     and $\epsilon$ is some tolerance.
     """
 
-    # TODO: (https://github.com/UCL/causalprog/issues/89)
-    # Should explain that Constraint needs more inputs and slightly different
-    # interpretation of the `do_with_samples` object.
-    # Inputs:
-    # - include epsilon as an input (allows constraints to have different tolerances)
-    # - `do_with_samples` should just be $g(\theta)$. Then have the instance build the
-    #   full constraint that will need to be called in the Lagrangian.
-    # - $g$ still needs to be scalar valued? Allow a wrapper function to be applied in
-    #   the event $g$ is vector-valued.
-    # If we do this, will also need to override __call__...
+    data: npt.ArrayLike
+    tolerance: npt.ArrayLike
+    _outer_norm: Callable[[npt.ArrayLike], float]
+
+    def __init__(
+        self,
+        *effect_handlers: ModelMask,
+        model_quantity: Callable[..., npt.ArrayLike],
+        outer_norm: Callable[[npt.ArrayLike], float] | None = None,
+        data: npt.ArrayLike = 0.0,
+        tolerance: float = 1.0e-6,
+    ) -> None:
+        r"""
+        Create a new constraint.
+
+        Constraints have the form
+
+        $$ c(\theta) :=
+        \mathrm{norm}\left( g(\theta)
+        - g_{\mathrm{data}} \right)
+        - \epsilon $$
+
+        where;
+        - $\mathrm{norm}$ is the outer norm of the constraint (`outer_norm`),
+        - $g(\theta)$ is the model quantity involved in the constraint
+            (`model_quantity`),
+        - $g_{\mathrm{data}}$ is the observed data (`data`),
+        - $\epsilon$ is the tolerance in the data (`tolerance`).
+
+        In a causal problem, each constraint appears as the condition $c(\theta)\leq 0$
+        in the minimisation / maximisation (hence the inclusion of the $-\epsilon$
+        term within $c(\theta)$ itself).
+
+        $g$ should be a (possibly vector-valued) function that acts on (a subset of)
+        samples from the random variables of the causal problem. It must accept
+        variable keyword-arguments only, and should access the samples for each random
+        variable by indexing via the RV names (node labels). It should return the
+        model quantity as computed from the samples, that $g_{\mathrm{data}}$ observed.
+
+        $g_{\mathrm{data}}$ should be a fixed value whose shape is broadcast-able with
+        the return shape of $g$. It defaults to $0$ if not explicitly set.
+
+        $\mathrm{norm}$ should be a suitable norm to take on the difference between the
+        model quantity as predicted by the samples ($g$) and the observed data
+        ($g_{\mathrm{data}}$). It must return a scalar value. The default is the 2-norm.
+        """
+        super().__init__(*effect_handlers, do_with_samples=model_quantity)
+
+        if outer_norm is None:
+            self._outer_norm = jnp.linalg.vector_norm
+        else:
+            self._outer_norm = outer_norm
+
+        self.data = data
+        self.tolerance = tolerance
+
+    def __call__(self, samples: dict[str, npt.ArrayLike]) -> npt.ArrayLike:
+        """
+        Evaluate the constraint, given RV samples.
+
+        Args:
+            samples: Mapping of RV (node) labels to drawn samples.
+
+        Returns:
+            Value of the constraint.
+
+        """
+        return (
+            self._outer_norm(self._do_with_samples(**samples) - self.data)
+            - self.tolerance
+        )

tests/test_integration/test_two_normal_example.py

@@ -78,7 +78,9 @@ def test_two_normal_example(
     # Setup the optimisation problem from the graph
     ce = CausalEstimand(do_with_samples=lambda **pv: pv["X"].mean())
     con = Constraint(
-        do_with_samples=lambda **pv: jnp.abs(pv["UX"].mean() - phi_observed) - epsilon
+        model_quantity=lambda **pv: pv["UX"].mean(),
+        data=phi_observed,
+        tolerance=epsilon,
     )
     cp = CausalProblem(
         two_normal_graph(cov=1.0),