Merge pull request #600 from CUQI-DTU/add_conditioning_RV_578

Add support for conditioning on random variables

Author: nabriis

cuqi/experimental/algebra/_ast.py

@@ -51,6 +51,11 @@ def __call__(self, **kwargs):
         """Evaluate node at a given parameter value. This will traverse the sub-tree originated at this node and evaluate it given the recorded operations."""
         pass
 
+    @abstractmethod
+    def condition(self, **kwargs):
+        """ Conditions the tree by replacing any VariableNode with a ValueNode if the variable is in the kwargs dictionary. """
+        pass
+
     @abstractmethod
     def __repr__(self):
         """String representation of the node. Used for printing the AST."""
@@ -129,6 +134,9 @@ class UnaryNode(Node, ABC):
     def __init__(self, child: Node):
         self.child = child
 
+    def condition(self, **kwargs):
+        return self.__class__(self.child.condition(**kwargs))
+
 
 class BinaryNode(Node, ABC):
     """Base class for all binary nodes in the abstract syntax tree.
@@ -155,6 +163,9 @@ def __init__(self, left: Node, right: Node):
         self.left = left
         self.right = right
 
+    def condition(self, **kwargs):
+        return self.__class__(self.left.condition(**kwargs), self.right.condition(**kwargs))
+
     def __repr__(self):
         return f"{self.left} {self.op_symbol} {self.right}"
 
@@ -205,6 +216,11 @@ def __call__(self, **kwargs):
             )
         return kwargs[self.name]
 
+    def condition(self, **kwargs):
+        if self.name in kwargs:
+            return ValueNode(kwargs[self.name])
+        return self
+
     def __repr__(self):
         return self.name
 
@@ -226,6 +242,9 @@ def __call__(self, **kwargs):
         """Returns the value of the node."""
         return self.value
 
+    def condition(self, **kwargs):
+        return self
+
     def __repr__(self):
         return str(self.value)
 

cuqi/experimental/algebra/_orderedset.py

@@ -19,6 +19,13 @@ def add(self, item):
         """
         self.dict[item] = None
 
+    def remove(self, item):
+        """Remove an item from the set. 
+
+        If the item is not in the set, it raises a KeyError.
+        """
+        del self.dict[item]
+
     def __contains__(self, item):
         """Check if an item is in the set. 
 
@@ -47,6 +54,18 @@ def extend(self, other):
         for item in other:
             self.add(item)
 
+    def replace(self, old_item, new_item):
+        """Replace old_item with new_item at the same position, preserving order."""
+        if old_item not in self.dict:
+            raise KeyError(f"{old_item} not in set")
+        
+        items = list(self.dict.keys())  # Preserve order
+        index = items.index(old_item)  # Find position
+        items[index] = new_item  # Replace at the same position
+
+        # Reconstruct the ordered set with the new item in place
+        self.dict = dict.fromkeys(items)
+
     def __or__(self, other):
         """Return a new set that is the union of this set and another set.
 
@@ -57,3 +76,7 @@ def __or__(self, other):
         new_set = _OrderedSet(self.dict.keys())
         new_set.extend(other)
         return new_set
+
+    def __repr__(self):
+        """Return a string representation of the set."""
+        return "_OrderedSet({})".format(list(self.dict.keys()))

cuqi/experimental/algebra/_randomvariable.py

@@ -5,8 +5,7 @@
 import operator
 import cuqi
 from cuqi.distribution import Distribution
-from copy import copy
-
+from copy import copy, deepcopy
 
 class RandomVariable:
     """ Random variable defined by a distribution with the option to apply algebraic operations on it.
@@ -210,7 +209,7 @@ def distributions(self) -> set:
     def parameter_names(self) -> str:
         """ Name of the parameter that the random variable can be evaluated at. """
         self._inject_name_into_distribution()
-        return [distribution.name for distribution in self.distributions] # Consider renaming .name to .par_name for distributions
+        return [distribution._name for distribution in self.distributions] # Consider renaming .name to .par_name for distributions
 
     @property
     def dim(self):
@@ -239,21 +238,89 @@ def expression(self):
     def is_transformed(self):
         """ Returns True if the random variable is transformed. """
         return not isinstance(self.tree, VariableNode)
-    
+
+    @property
+    def is_cond(self):
+        """ Returns True if the random variable is a conditional random variable. """
+        return any(dist.is_cond for dist in self.distributions)
+
+    def condition(self, *args, **kwargs):
+        """Condition the random variable on a given value. Only one of either positional or keyword arguments can be passed.
+        
+        Parameters
+        ----------
+        *args : Any
+            Positional arguments to condition the random variable on. The order of the arguments must match the order of the parameter names.
+
+        **kwargs : Any
+            Keyword arguments to condition the random variable on. The keys must match the parameter names.
+        
+        """
+
+        # Before conditioning, capture repr to ensure all variable names are injected
+        self.__repr__()
+
+        if args and kwargs:
+            raise ValueError("Cannot pass both positional and keyword arguments to RandomVariable")
+        
+        if args:
+            kwargs = self._parse_args_add_to_kwargs(args, kwargs)
+
+        # Create a deep copy of the random variable to ensure the original tree is not modified
+        new_variable = self._make_copy(deep=True)
+
+        for kwargs_name in list(kwargs.keys()):
+            value = kwargs.pop(kwargs_name)
+
+            # Condition the tree turning the variable into a constant
+            if kwargs_name in self.parameter_names:
+                new_variable._tree = new_variable.tree.condition(**{kwargs_name: value})
+            
+            # Condition the random variable on both the distribution parameter name and distribution conditioning variables
+            for dist in self.distributions:
+                if kwargs_name == dist.name:
+                    new_variable._remove_distribution(dist.name)
+                elif kwargs_name in dist.get_conditioning_variables():
+                    new_variable._replace_distribution(dist.name, dist(**{kwargs_name: value}))
+
+        # Check if any kwargs are left unprocessed
+        if kwargs:
+            raise ValueError(f"Conditioning variables {list(kwargs.keys())} not found in the random variable {self}")
+
+        return new_variable
+
     @property
     def _non_default_args(self) -> List[str]:
         """List of non-default arguments to distribution. This is used to return the correct
         arguments when evaluating the random variable.
         """
         return self.parameter_names
 
+    def _replace_distribution(self, name, new_distribution):
+        """ Replace distribution with a given name with a new distribution in the same position of the ordered set. """
+        for dist in self.distributions:
+            if dist._name == name:
+                self._distributions.replace(dist, new_distribution)
+                break
+
+    def _remove_distribution(self, name):
+        """ Remove distribution with a given name from the set of distributions. """
+        for dist in self.distributions:
+            if dist._name == name:
+                self._distributions.remove(dist)
+                break
+
     def _inject_name_into_distribution(self, name=None):
         if len(self._distributions) == 1:
             dist = next(iter(self._distributions))
+
+            if dist._is_copy:
+                dist = dist._original_density
+
             if dist._name is None:
                 if name is None:
                     name = self.name
-                dist._name = name
+                dist.name = name # Inject using setter
 
     def _parse_args_add_to_kwargs(self, args, kwargs) -> dict:
         """ Parse args and add to kwargs if any. Arguments follow self.parameter_names order. """
@@ -293,8 +360,12 @@ def _is_copy(self):
         """ Returns True if this is a copy of another random variable, e.g. by conditioning. """
         return hasattr(self, '_original_variable') and self._original_variable is not None
 
-    def _make_copy(self):
-        """ Returns a shallow copy of the density keeping a pointer to the original. """
+    def _make_copy(self, deep=False) -> 'RandomVariable':
+        """ Returns a copy of the density keeping a pointer to the original. """
+        if deep:
+            new_variable = deepcopy(self)
+            new_variable._original_variable = self
+            return new_variable
         new_variable = copy(self)
         new_variable._distributions = copy(self.distributions)
         new_variable._tree = copy(self._tree)

cuqi/likelihood/_likelihood.py

@@ -43,6 +43,14 @@ def name(self):
     def name(self, value):
         self.distribution.name = value
 
+    @property
+    def _name(self):
+        return self.distribution._name
+
+    @_name.setter
+    def _name(self, value):
+        self.distribution._name = value
+
     @property
     def FD_enabled(self):
         """ Return FD_enabled of the likelihood from the underlying distribution """

demos/howtos/algebra.py

@@ -83,3 +83,41 @@
 BP.UQ(exact={"x": info.exactSolution})
 
 # %%
+# Conditioning on random variables (example 1)
+from cuqi.distribution import Gaussian
+from cuqi.experimental.algebra import RandomVariable
+
+x = RandomVariable(Gaussian(0, lambda s: s))
+y = RandomVariable(Gaussian(0, lambda d: d))
+
+z = x+y
+
+z.condition(x=1)
+
+# %%
+# Or conditioning on the variables s, or d
+z.condition(s=1)
+
+# %%
+# Conditioning on random variables (example 2)
+from cuqi.testproblem import Deconvolution1D
+from cuqi.distribution import Gaussian, Gamma, GMRF
+from cuqi.experimental.algebra import RandomVariable
+from cuqi.problem import BayesianProblem
+import numpy as np
+
+# Forward model
+A, y_obs, info = Deconvolution1D(dim=4).get_components()
+
+# Bayesian Problem (defined using Random Variables)
+d = RandomVariable(Gamma(1, 1e-4))
+s = RandomVariable(Gamma(1, 1e-4))
+x = RandomVariable(GMRF(np.zeros(A.domain_dim), d))
+y = RandomVariable(Gaussian(A @ x, 1/s))
+
+
+z = x+y
+
+z.condition(x=np.zeros(A.domain_dim))
+
+# %%

tests/zexperimental/test_randomvariable.py

@@ -72,7 +72,6 @@ def test_randomvariable_returns_correct_parameter_name():
     z = RandomVariable(cuqi.distribution.Gaussian(0, 1))
     assert cuqi.utilities.get_non_default_args(z) == ["z"]
 
-@pytest.mark.xfail(reason="Conditional random variables are not yet implemented")
 @pytest.mark.parametrize("operations", [
     lambda x: x+1,
     lambda x: x**2,
@@ -258,7 +257,6 @@ def recursive_return_rv(rv, recursions):
 
     recursive_return_rv(h, 10)
 
-@pytest.mark.xfail(reason="Conditional random variables are not yet implemented")
 def test_rv_name_consistency():
 
     x = RandomVariable(cuqi.distribution.Gaussian(geometry=1))
@@ -306,13 +304,27 @@ def test_ensure_that_RV_evaluation_requires_all_parameters():
     with pytest.raises(ValueError, match=r"Expected arguments \['x', 'y'\], got arguments \{'x': 1\}"):
         z(x=1)
 
-@pytest.mark.xfail(reason="Conditional random variables are not yet implemented")
 def test_RV_sets_name_of_internal_conditional_density_if_par_name_not_set_and_does_not_set_original_density():
-    Z_s = cuqi.distribution.Gaussian(0, lambda s: s)
-    Z = Z_s(s=3)
-    z = RandomVariable(Z)
+    # Case 1
+    z = RandomVariable(cuqi.distribution.Gaussian(0, lambda s: s))
 
     assert z.name == 'z'
-    assert z.distribution.par_name == "z"
-    assert Z_s.par_name is None # Should not be set for the original density.
-    assert Z.par_name is None # Should not be set for the conditioned density.
+    assert z.distribution.name == "z"
+
+    # Case 2 (conditioned density. Should not be able to set name here)
+    z = RandomVariable(cuqi.distribution.Gaussian(0, lambda s: s)(s=3))
+
+    assert z.name == 'z'
+    assert z.distribution.name == "z"
+
+def test_RV_condition_maintains_parameter_name_order():
+
+    x = RandomVariable(cuqi.distribution.Gaussian(0, lambda s: s))
+    y = RandomVariable(cuqi.distribution.Gaussian(0, lambda d: d))
+
+    z = x+y
+
+    assert z.parameter_names == ['x', 'y']
+    assert z.condition(s=1).parameter_names == ['x', 'y']
+    assert z.condition(d=1).parameter_names == ['x', 'y']
+    assert z.condition(d=1, s=1).parameter_names == ['x', 'y']