Merge pull request #597 from CUQI-DTU/add_577_random_variable_class

Add random variable class

Author: nabriis

cuqi/distribution/_joint_distribution.py

@@ -5,6 +5,7 @@
 from cuqi.distribution import Distribution, Posterior
 from cuqi.likelihood import Likelihood
 from cuqi.geometry import Geometry, _DefaultGeometry1D
+import cuqi
 import numpy as np # for splitting array. Can avoid.
 
 class JointDistribution:
@@ -13,9 +14,11 @@ class JointDistribution:
 
     Parameters
     ----------
-    densities : Density
+    densities : RandomVariable or Density
         The densities to include in the joint distribution.
-        Each density is passed as comma-separated arguments.
+        Each density is passed as comma-separated arguments,
+        and can be either a :class:'Density' such as :class:'Distribution'
+        or :class:`RandomVariable`.
 
     Notes
     -----
@@ -59,7 +62,16 @@ class JointDistribution:
         posterior = joint(y=y_obs)
         
     """
-    def __init__(self, *densities: Density):
+    def __init__(self, *densities: [Density, cuqi.experimental.algebra.RandomVariable]):
+        """ Create a joint distribution from the given densities. """
+
+        # Check if all RandomVariables are simple (not-transformed)
+        for density in densities:
+            if isinstance(density, cuqi.experimental.algebra.RandomVariable) and density.is_transformed:
+                raise ValueError(f"To be used in {self.__class__.__name__}, all RandomVariables must be untransformed.")
+
+        # Convert potential random variables to their underlying distribution
+        densities = [density.distribution if isinstance(density, cuqi.experimental.algebra.RandomVariable) else density for density in densities]
 
         # Ensure all densities have unique names
         names = [density.name for density in densities]

cuqi/experimental/algebra/__init__.py

@@ -1 +1,2 @@
-from ._ast import VariableNode
+from ._ast import VariableNode, Node
+from ._randomvariable import RandomVariable

cuqi/experimental/algebra/_ast.py

@@ -56,6 +56,20 @@ def __repr__(self):
         """String representation of the node. Used for printing the AST."""
         pass
 
+    def get_variables(self, variables=None):
+        """Returns a set with the names of all variables in the sub-tree originated at this node."""
+        if variables is None:
+            variables = set()
+        if isinstance(self, VariableNode):
+            variables.add(self.name)
+        if hasattr(self, "child"):
+            self.child.get_variables(variables)
+        if hasattr(self, "left"):
+            self.left.get_variables(variables)
+        if hasattr(self, "right"):
+            self.right.get_variables(variables)
+        return variables
+
     def __add__(self, other):
         return AddNode(self, convert_to_node(other))
 

cuqi/experimental/algebra/_orderedset.py

@@ -0,0 +1,59 @@
+class _OrderedSet:
+    """A set (i.e. unique elements) that keeps its elements in the order they were added.
+
+    This is a minimal implementation of an ordered set, using a dictionary for storage.
+    """
+    
+    def __init__(self, iterable=None):
+        """Initialize the OrderedSet. 
+
+        If an iterable is provided, add all its elements to the set.
+        """
+        self.dict = dict.fromkeys(iterable if iterable else [])
+
+    def add(self, item):
+        """Add an item to the set. 
+
+        If the item is already in the set, it does nothing.
+        Otherwise, the item is stored as a key in the dictionary, with None as its value.
+        """
+        self.dict[item] = None
+
+    def __contains__(self, item):
+        """Check if an item is in the set. 
+
+        This is equivalent to checking if the item is a key in the dictionary.
+        """
+        return item in self.dict
+
+    def __iter__(self):
+        """Return an iterator over the set. 
+
+        This iterates over the keys in the dictionary.
+        """
+        return iter(self.dict)
+
+    def __len__(self):
+        """Return the number of items in the set."""
+        return len(self.dict)
+
+    def extend(self, other):
+        """Extend the set with the items in another set.
+
+        Raises a TypeError if the other object is not an _OrderedSet.
+        """
+        if not isinstance(other, _OrderedSet):
+            raise TypeError("unsupported operand type(s) for extend: '_OrderedSet' and '{}'".format(type(other).__name__))
+        for item in other:
+            self.add(item)
+
+    def __or__(self, other):
+        """Return a new set that is the union of this set and another set.
+
+        Raises a TypeError if the other object is not an _OrderedSet.
+        """
+        if not isinstance(other, _OrderedSet):
+            raise TypeError("unsupported operand type(s) for |: '_OrderedSet' and '{}'".format(type(other).__name__))
+        new_set = _OrderedSet(self.dict.keys())
+        new_set.extend(other)
+        return new_set