Repository: CUQI-DTU/CUQIpy

Issue:
fix documentation build (1-line fix)
change the version of setuptools  to be the latest in the documentation build

Hints from issue comments:
@amal-ghamdi: duplicated, see #180

Solution patch:
diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml
index 0cf002c9c..9a0eb63c9 100644
--- a/.github/workflows/docs.yml
+++ b/.github/workflows/docs.yml
@@ -16,7 +16,7 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          pip install --upgrade pip setuptools==65.7.0
+          pip install --upgrade pip setuptools
           pip install -r requirements.txt
           pip install -r requirements-dev.txt
       - name: Sphinx build
diff --git a/cuqi/distribution/_joint_distribution.py b/cuqi/distribution/_joint_distribution.py
index 2e32d0555..682e96273 100644
--- a/cuqi/distribution/_joint_distribution.py
+++ b/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]
diff --git a/cuqi/experimental/__init__.py b/cuqi/experimental/__init__.py
index 63d0626ea..30844cc48 100644
--- a/cuqi/experimental/__init__.py
+++ b/cuqi/experimental/__init__.py
@@ -1,3 +1,4 @@
 """ Experimental module for testing new features and ideas. """
 from . import mcmc
 from . import algebra
+from . import geometry
diff --git a/cuqi/experimental/algebra/__init__.py b/cuqi/experimental/algebra/__init__.py
index 4a9ac77d1..dbeb4da1c 100644
--- a/cuqi/experimental/algebra/__init__.py
+++ b/cuqi/experimental/algebra/__init__.py
@@ -1 +1,2 @@
-from ._ast import VariableNode
+from ._ast import VariableNode, Node
+from ._randomvariable import RandomVariable
diff --git a/cuqi/experimental/algebra/_ast.py b/cuqi/experimental/algebra/_ast.py
index 527c6f536..65267a294 100644
--- a/cuqi/experimental/algebra/_ast.py
+++ b/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))
 
diff --git a/cuqi/experimental/algebra/_orderedset.py b/cuqi/experimental/algebra/_orderedset.py
new file mode 100644
index 000000000..b91fedac1
--- /dev/null
+++ b/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
diff --git a/cuqi/experimental/algebra/_randomvariable.py b/cuqi/experimental/algebra/_randomvariable.py
new file mode 100644
index 000000000..b2b771c65
--- /dev/null
+++ b/cuqi/experimental/algebra/_randomvariable.py
@@ -0,0 +1,360 @@
+from __future__ import annotations
+from typing import List, Any, Union
+from ._ast import VariableNode, Node
+from ._orderedset import _OrderedSet
+import operator
+import cuqi
+from cuqi.distribution import Distribution
+from copy import copy
+
+
+class RandomVariable:
+    """ Random variable defined by a distribution with the option to apply algebraic operations on it.
+
+    Random variables allow for the definition of Bayesian Problems in a natural way. In the context
+    of code, the random variable can be viewed as a lazily evaluated variable/array. It records
+    operations applied to it and acts as a function that, when called, evaluates the operations
+    and returns the result.
+
+    In CUQIpy, random variables can be in two forms: (1) a 'primal' random variable that is directly
+    defined by a distribution, e.g. x ~ N(0, 1), or (2) a 'transformed' random variable that is defined by
+    applying algebraic operations on one or more random variables, e.g. y = x + 1.
+
+    This distinction is purely for the purpose of the implementation in CUQIpy, as mathematically both
+    x ~ N(0, 1) and y = x + 1 ~ N(1, 1) are random variables. The distinction is useful for the
+    code implementation. In the future some operations like the above may allow primal random variables
+    that are transformed if the distribution can be analytically described.
+
+    Parameters
+    ----------
+    distributions : Distribution or list of Distributions
+        The distribution from which the random variable originates. If multiple distributions are
+        provided, the random variable is defined by the passed abstract syntax `tree` representing the
+        algebraic operations applied to one or more random variables.
+
+    tree : Node, optional
+        The tree, represented by the syntax tree nodes, that contain the algebraic operations applied to the random variable.
+        Specifically, the root of the tree should be provided.
+
+    name : str, optional
+        Name of the random variable. If not provided, the name is extracted from either the distribution provided
+        or from the variable name in the code. The name provided must match the parameter name of the distribution.
+
+    Example
+    -------
+
+    Basic usage:
+
+    .. code-block:: python
+
+        from cuqi.distribution import Gaussian
+
+        x = RandomVariable(Gaussian(0, 1))
+
+    Defining Bayesian problem using random variables:
+
+    .. code-block:: python
+
+        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
+        A, y_obs, info = Deconvolution1D().get_components()
+
+        # Bayesian problem
+        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))
+
+        BP = BayesianProblem(y, x, s, d)
+        BP.set_data(y=y_obs)
+        BP.UQ()
+
+    Defining random variable from multiple distributions:
+
+    .. code-block:: python
+
+        from cuqi.distribution import Gaussian, Gamma
+        from cuqi.experimental.algebra import RandomVariable, VariableNode
+
+        # Define the variables
+        x = VariableNode('x')
+        y = VariableNode('y')
+
+        # Define the distributions (names must match variables)
+        dist_x = Gaussian(0, 1, name='x')
+        dist_y = Gamma(1, 1e-4, name='y')
+
+        # Define the tree (this is the algebra that defines the random variable along with the distributions)
+        tree = x + y
+
+        # Define random variable from 2 distributions with relation x+y
+        rv = RandomVariable([dist_x, dist_y], tree)
+
+    """
+
+
+    def __init__(self, distributions: Union['Distribution', List['Distribution']], tree: 'Node' = None, name: str = None):
+        """ Create random variable from distribution """
+
+        if isinstance(distributions, Distribution):
+            distributions = [distributions]
+        
+        if not  isinstance(distributions, list) and not isinstance(distributions, _OrderedSet):
+            raise ValueError("Expected a distribution or a list of distributions")
+
+        # Convert single distribution(s) to internal datastructure _OrderedSet.
+        # We use ordered set to ensure that the order of the distributions is preserved.
+        # which in turn ensures that the parameter names are always in the same order.
+        if not isinstance(distributions, _OrderedSet):
+            distributions = _OrderedSet(distributions)
+
+        # If tree is provided, check it is consistent with the given distributions
+        if tree:
+            tree_var_names = tree.get_variables()
+            dist_par_names = {dist._name for dist in distributions}
+
+            if len(tree_var_names) != len(distributions):
+                raise ValueError(
+                    f"There are {len(tree_var_names)} variables in the tree, but {len(distributions)} distributions are provided. "
+                    "This may be due to passing multiple distributions with the same parameter name. "
+                    f"The tree variables are {tree_var_names} and the distribution parameter names are {dist_par_names}."
+                )
+            
+            if not all(var_name in dist_par_names for var_name in tree_var_names):
+                raise ValueError(
+                    f"Variable names in the tree {tree_var_names} do not match the parameter names in the distributions {dist_par_names}. "
+                    "Ensure the name is inferred from the variable or explicitly provide it using name='var_name' in the distribution."
+                )
+
+        # Match random variable name with distribution parameter name (for single distribution)
+        if len(distributions) == 1 and tree is None:
+            dist = next(iter(distributions))
+            dist_par_name = dist._name
+            if dist_par_name is not None:
+                if name is not None and dist_par_name != name:
+                    raise ValueError(f"Parameter name '{dist_par_name}' of the distribution does not match the input name '{name}' for the random variable.")
+                name = dist_par_name
+    
+        self._distributions = distributions
+        """ The distribution from which the random variable originates. """
+        
+        self._tree = tree
+        """ The tree representation of the random variable. """
+
+        self._original_variable = None
+        """ Stores the original variable if this is a conditioned copy"""
+
+        self._name = name
+        """ Name of the random variable. """
+
+
+    def __call__(self, *args, **kwargs) -> Any:
+        """ Evaluate random variable at a given parameter value. For example, for random variable `X`, `X(1)` gives `1` and `(X+1)(1)` gives `2` """
+
+        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)
+
+        # Check if kwargs match parameter names using a all compare
+        if not all([name in kwargs for name in self.parameter_names]) or not all([name in self.parameter_names for name in kwargs]):
+            raise ValueError(f"Expected arguments {self.parameter_names}, got arguments {kwargs}")
+
+        return self.tree(**kwargs)
+    
+    @property
+    def tree(self):
+        if self._tree is None:
+            if len(self._distributions) > 1:
+                raise ValueError("Tree for multiple distributions can not be created automatically and need to be passed as an argument to the {} initializer.".format(type(self).__name__))
+            self._tree = VariableNode(self.name)
+        return self._tree
+
+    @property
+    def name(self):
+        """ Name of the random variable. If not provided, the name is extracted from the variable name in the code. """
+        if self._is_copy: # Extract the original variable name if this is a copy
+            return self._original_variable.name
+        if self._name is None: # If None extract the name from the stack
+            self._name = cuqi.utilities._get_python_variable_name(self)
+        if self._name is not None:
+            self._inject_name_into_distribution(self._name)
+        return self._name
+    
+    @name.setter
+    def name(self, name):
+        if self._is_copy:
+            raise ValueError("This random variable is derived from the conditional random variable named "+self._original_variable.name+". The name of the derived random variable cannot be set, but follows the name of the original random variable.")
+        self._name = name
+        
+    @property
+    def distribution(self) -> cuqi.distribution.Distribution:
+        """ Distribution from which the random variable originates. """
+        if len(self._distributions) > 1:
+            raise ValueError("Cannot get distribution from random variable defined by multiple distributions")
+        self._inject_name_into_distribution()
+        return next(iter(self._distributions))
+    
+    @property
+    def distributions(self) -> set:
+        """ Distributions from which the random variable originates. """
+        self._inject_name_into_distribution()
+        return self._distributions
+        
+    @property
+    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
+
+    @property
+    def dim(self):
+        if self.is_transformed:
+            raise NotImplementedError("Dimension not implemented for transformed random variables")
+        return self.distribution.dim
+
+    @property
+    def geometry(self):
+        if self.is_transformed:
+            raise NotImplementedError("Geometry not implemented for transformed random variables")
+        return self.distribution.geometry
+
+    @geometry.setter
+    def geometry(self, geometry):
+        if self.is_transformed:
+            raise NotImplementedError("Geometry not implemented for transformed random variables")
+        self.distribution.geometry = geometry
+
+    @property
+    def expression(self):
+        """ Expression (formula) of the random variable. """
+        return str(self.tree)
+
+    @property
+    def is_transformed(self):
+        """ Returns True if the random variable is transformed. """
+        return not isinstance(self.tree, VariableNode)
+    
+    @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 _inject_name_into_distribution(self, name=None):
+        if len(self._distributions) == 1:
+            dist = next(iter(self._distributions))
+            if dist._name is None:
+                if name is None:
+                    name = self.name
+                dist._name = name
+    
+    def _parse_args_add_to_kwargs(self, args, kwargs) -> dict:
+        """ Parse args and add to kwargs if any. Arguments follow self.parameter_names order. """
+        if len(args) != len(self.parameter_names):
+            raise ValueError(f"Expected {len(self.parameter_names)} arguments, got {len(args)}. Parameters are: {self.parameter_names}")
+        
+        # Add args to kwargs
+        for arg, name in zip(args, self.parameter_names):
+            kwargs[name] = arg
+            
+        return kwargs
+
+    def __repr__(self):
+        # Create strings for parameter name ~ distribution pairs
+        parameter_strings = [f"{name} ~ {distribution}" for name, distribution in zip(self.parameter_names, self.distributions)]
+        # Join strings with newlines
+        parameter_strings = "\n".join(parameter_strings)
+        # Add initial newline and indentations
+        parameter_strings = "\n".join(["\t"+line for line in parameter_strings.split("\n")])
+        # Print parameter strings with newlines
+        if self.is_transformed:
+            title = f"Transformed Random Variable"
+        else:
+            title = f""
+        if self.is_transformed:
+            body = (
+                f"\n"
+                f"Expression: {self.tree}\n"
+                f"Components: \n{parameter_strings}"
+                )
+        else:
+            body = parameter_strings.replace("\t","")
+        return title+body
+   
+    @property
+    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. """
+        new_variable = copy(self)
+        new_variable._distributions = copy(self.distributions)
+        new_variable._tree = copy(self._tree)
+        new_variable._original_variable = self
+        return new_variable
+    
+    def _apply_operation(self, operation, other=None) -> 'RandomVariable':
+        """
+        Apply a specified operation to this RandomVariable.
+        """
+        if isinstance(other, cuqi.distribution.Distribution):
+            raise ValueError("Cannot apply operation to distribution. Use .rv to create random variable first.")
+        if other is None: # unary operation case
+            return RandomVariable(self.distributions, operation(self.tree))
+        elif isinstance(other, RandomVariable): # binary operation case with another random variable that has distributions
+            return RandomVariable(self.distributions | other.distributions, operation(self.tree, other.tree))
+        return RandomVariable(self.distributions, operation(self.tree, other)) # binary operation case with any other object (constant)
+
+    def __add__(self, other) -> 'RandomVariable':
+        return self._apply_operation(operator.add, other)
+
+    def __radd__(self, other) -> 'RandomVariable':
+        return self.__add__(other)
+
+    def __sub__(self, other) -> 'RandomVariable':
+        return self._apply_operation(operator.sub, other)
+
+    def __rsub__(self, other) -> 'RandomVariable':
+        return self._apply_operation(lambda x, y: operator.sub(y, x), other)
+
+    def __mul__(self, other) -> 'RandomVariable':
+        return self._apply_operation(operator.mul, other)
+
+    def __rmul__(self, other) -> 'RandomVariable':
+        return self.__mul__(other)
+
+    def __truediv__(self, other) -> 'RandomVariable':
+        return self._apply_operation(operator.truediv, other)
+
+    def __rtruediv__(self, other) -> 'RandomVariable':
+        return self._apply_operation(lambda x, y: operator.truediv(y, x), other)
+
+    def __matmul__(self, other) -> 'RandomVariable':
+        if isinstance(other, cuqi.model.Model) and not isinstance(other, cuqi.model.LinearModel):
+            raise TypeError("Cannot apply matmul to non-linear models")
+        return self._apply_operation(operator.matmul, other)
+
+    def __rmatmul__(self, other) -> 'RandomVariable':
+        if isinstance(other, cuqi.model.Model) and not isinstance(other, cuqi.model.LinearModel):
+            raise TypeError("Cannot apply matmul to non-linear models")
+        return self._apply_operation(lambda x, y: operator.matmul(y, x), other)
+
+    def __neg__(self) -> 'RandomVariable':
+        return self._apply_operation(operator.neg)
+
+    def __abs__(self) -> 'RandomVariable':
+        return self._apply_operation(abs)
+
+    def __pow__(self, other) -> 'RandomVariable':
+        return self._apply_operation(operator.pow, other)
+
+    def __getitem__(self, other) -> 'RandomVariable':
+        return self._apply_operation(operator.getitem, other)
diff --git a/cuqi/experimental/geometry/__init__.py b/cuqi/experimental/geometry/__init__.py
new file mode 100644
index 000000000..3264de1a4
--- /dev/null
+++ b/cuqi/experimental/geometry/__init__.py
@@ -0,0 +1 @@
+from ._productgeometry import _ProductGeometry
diff --git a/cuqi/experimental/geometry/_productgeometry.py b/cuqi/experimental/geometry/_productgeometry.py
new file mode 100644
index 000000000..1ff377327
--- /dev/null
+++ b/cuqi/experimental/geometry/_productgeometry.py
@@ -0,0 +1,181 @@
+from cuqi.geometry import Geometry
+import numpy as np
+
+class _ProductGeometry(Geometry):
+    """ A class for representing a product geometry. A product geometry
+    represents the product space of multiple geometries of type :class:`Geometry`.
+    See the example below for a product geometry of two geometries.
+
+    Parameters
+    ----------
+    \*geometries : cuqi.geometry.Geometry
+        The geometries to be combined into a product geometry. Each geometry
+        is passed as a comma-separated argument.
+
+    Example
+    -------
+    .. code-block:: python
+        import numpy as np
+        from cuqi.geometry import Continuous1D, Discrete
+        from cuqi.experimental.geometry import _ProductGeometry
+        geometry1 = Continuous1D(np.linspace(0, 1, 100))
+        geometry2 = Discrete(["sound_speed"])
+        product_geometry = _ProductGeometry(geometry1, geometry2)
+    """
+
+    def __init__(self, *geometries):
+        self.geometries = geometries
+
+    @property
+    def geometries(self):
+        """List of geometries that are combined to form the product geometry."""
+        return self._geometries
+    
+    @geometries.setter
+    def geometries(self, geometries):
+        # Check if all geometries are of type Geometry.
+        for g in geometries:
+            if not isinstance(g, Geometry):
+                raise TypeError(
+                    "All geometries must be of type Geometry. "
+                    "Received: {}".format(type(g))
+                )
+        self._geometries = geometries
+
+    @property
+    def fun_shape(self):
+        """Shape of the function representation. Returns a tuple, where
+        each element of the tuple is the shape of the function 
+        representation of each geometry."""
+        return tuple([g.fun_shape for g in self.geometries])
+    
+    @property
+    def fun_dim(self):
+        """Dimension of the function representation which is the sum of
+        the function representation dimensions of each geometry."""
+        return sum([g.fun_dim for g in self.geometries])
+
+    @property
+    def par_shape(self):
+        """Shape of the parameter representation. Returns a tuple, where
+        each element of the tuple is the shape of the parameter
+        representation of each geometry."""
+        return tuple([g.par_shape for g in self.geometries])
+    
+    @property
+    def par_dim(self):
+        """Dimension of the parameter representation which is the sum of
+        the parameter representation dimensions of each geometry."""
+        return sum(self.par_dim_list)
+    
+    @property
+    def par_dim_list(self):
+        """List of the parameter representation dimensions of each
+        geometry. This property is useful for indexing a stacked parameter
+        vector."""
+        return [g.par_dim for g in self.geometries]
+    
+    @property
+    def stacked_par_split_indices(self):
+        """Indices at which the stacked parameter vector should be split
+        to obtain the parameter vectors for each geometry. For example, if
+        the stacked parameter vector is [1, 2, 3, 4, 5, 6] and the parameter
+        vectors for each geometry are [1, 2], [3, 4], and [5, 6], then the
+        split indices are [2, 4]"""
+        return np.cumsum(self.par_dim_list[:-1])
+
+    @property
+    def number_of_geometries(self):
+        """Number of geometries in the product geometry."""
+        return len(self.geometries)
+    
+    def _split_par(self, par):
+        """Splits a stacked parameter vector into parameter vectors for each
+        geometry."""
+        return tuple(np.split(par, self.stacked_par_split_indices))
+    
+    def _plot(self, values, **kwargs):
+        """Plotting function for the product geometry."""
+        raise NotImplementedError(
+            f"Plotting not implemented for {self.__class__.__name__}.")
+
+    def par2fun(self, *pars):
+        """Converts parameter vector(s) into function values for each 
+        geometry. The parameter vector can be stacked (all parameters are
+        in one vector) or unstacked (one parameter vector corresponds to
+        each geometry). In all cases, the order of the parameter vectors
+        should follow the order of the geometries in the product, i.e., the
+        first parameter vector corresponds to the first geometry and so on."""
+
+        # If one argument is passed, then it is assumed that the parameter
+        # vector is stacked and split it.
+        # No effect if the parameter vector is already split and corresponds
+        # to one geometry.
+        if len(pars) == 1:
+            pars = self._split_par(pars[0])
+
+        # Convert parameter vectors to function values for each geometry.
+        funvals = []
+        for i, g in enumerate(self.geometries):
+            funval_i = g.par2fun(pars[i])
+            funvals.append(funval_i)
+        return tuple(funvals)
+
+    def fun2par(self, *funvals, stacked=False):
+        """Converts (multiple) function values into the corresponding
+        parameter vectors. If the flag stacked is set to True, then the
+        parameter vectors are stacked into one vector. Otherwise, the
+        parameter vectors are returned as a tuple. The order of function
+        values should follow the order of the geometries in the product,
+        i.e., the first function value corresponds to the first geometry
+        and so on."""
+
+        pars = []
+        for i, g in enumerate(self.geometries):
+            par_i = g.fun2par(funvals[i])
+            pars.append(par_i)
+
+        # stack parameters:
+        if stacked:
+            # if single sample
+            if len(pars[0].shape) == 1:
+                stacked_val = np.hstack(pars)
+            elif len(pars[0].shape) == 2:
+                stacked_val = np.vstack(pars)
+            else:
+                raise ValueError(
+                    "Cannot stack parameter vectors with more than 2 dimensions."
+                    )
+
+        return  stacked_val if stacked else tuple(pars)
+
+    def vec2fun(self, *funvecs):
+        """Maps function vector representation, if available, to function 
+        values. The order of the function vectors should follow the order of
+        the geometries in the product, i.e., the first function vector
+        corresponds to the first geometry and so on."""
+        funvals = []  
+        for i, g in enumerate(self.geometries):
+            funvals.append(g.vec2fun(funvecs[i]))
+
+        return tuple(funvals)
+    
+    def fun2vec(self, *funvals):
+        """Maps function values to a vector representation of the function
+        values, if available. The order of the function values should follow
+        the order of the geometries in the product, i.e., the first function
+        value corresponds to the first geometry and so on."""
+        funvecs = []
+        for i, g in enumerate(self.geometries):
+            funvecs.append(g.fun2vec(funvals[i]))
+        
+        return tuple(funvecs)
+        
+    
+    def __repr__(self) -> str:
+        """Representation of the product geometry."""
+        string = "{}(".format(self.__class__.__name__) + "\n"
+        for g in self.geometries:
+            string += "\t{}\n".format(g.__repr__())
+        string += ")"
+        return string
\ No newline at end of file
diff --git a/cuqi/model/_model.py b/cuqi/model/_model.py
index 0b9abf265..a811eaa6f 100644
--- a/cuqi/model/_model.py
+++ b/cuqi/model/_model.py
@@ -351,6 +351,17 @@ def forward(self, *args, is_par=True, **kwargs):
             new_model._non_default_args = [x.name] # Defaults to x if distribution had no name
             return new_model
 
+        # If input is a random variable, we handle it separately
+        if isinstance(x, cuqi.experimental.algebra.RandomVariable):
+            return self._handle_random_variable(x)
+        
+        # If input is a Node from internal abstract syntax tree, we let the Node handle the operation
+        # We use NotImplemented to indicate that the operation is not supported from the Model class
+        # in case of operations such as "@" that can be interpreted as both __matmul__ and __rmatmul__
+        # the operation may be delegated to the Node class.
+        if isinstance(x, cuqi.experimental.algebra.Node):
+            return NotImplemented
+
         # Else we apply the forward operator
         return self._apply_func(self._forward_func,
                                 self.range_geometry,
@@ -463,6 +474,20 @@ def _check_gradient_can_be_computed(self, direction, wrt):
             not type(self.domain_geometry) in _get_identity_geometries():
             raise NotImplementedError("Gradient not implemented for model {} with domain geometry {}".format(self,self.domain_geometry))
 
+    def _handle_random_variable(self, x):
+        """ Private function that handles the case of the input being a random variable. """
+        # If random variable is not a leaf-type node (e.g. internal node) we return NotImplemented
+        if not isinstance(x.tree, cuqi.experimental.algebra.VariableNode):
+            return NotImplemented        
+        
+        # In leaf-type node case we simply change the parameter name of model to match the random variable name
+        dist = x.distribution
+        if dist.dim != self.domain_dim:
+            raise ValueError("Attempting to match parameter name of Model with given random variable, but random variable dimension does not match model domain dimension.")
+        
+        new_model = copy(self)
+        new_model._non_default_args = [dist.name]
+        return new_model
 
     def __len__(self):
         return self.range_dim
diff --git a/cuqi/problem/_problem.py b/cuqi/problem/_problem.py
index 824a2d878..455fe7704 100644
--- a/cuqi/problem/_problem.py
+++ b/cuqi/problem/_problem.py
@@ -771,11 +771,11 @@ def gradfunc(x): return -density.gradient(x)
         if self._check_posterior(self, CMRF, must_have_gradient=True): # Use L-BFGS-B for CMRF prior as it has better performance for this multi-modal posterior
             if disp: print(f"Using scipy.optimize.L_BFGS_B on negative log of {density.__class__.__name__}")
             if disp: print("x0: ones vector")
-            solver = cuqi.solver.L_BFGS_B(func, x0, gradfunc=gradfunc)
+            solver = cuqi.solver.ScipyLBFGSB(func, x0, gradfunc=gradfunc)
         else:
             if disp: print(f"Using scipy.optimize.minimize on negative log of {density.__class__.__name__}")
             if disp: print("x0: ones vector")
-            solver = cuqi.solver.minimize(func, x0, gradfunc=gradfunc)
+            solver = cuqi.solver.ScipyMinimizer(func, x0, gradfunc=gradfunc)
 
         x_MAP, solver_info = solver.solve()
 
diff --git a/cuqi/solver/__init__.py b/cuqi/solver/__init__.py
index 5b9e11a19..c3851ae65 100644
--- a/cuqi/solver/__init__.py
+++ b/cuqi/solver/__init__.py
@@ -1,8 +1,8 @@
 from ._solver import (
-    L_BFGS_B,
-    minimize,
-    maximize,
-    LS,
+    ScipyLBFGSB,
+    ScipyMinimizer,
+    ScipyMaximizer,
+    ScipyLeastSquares,
     CGLS,
     LM,
     PDHG,
diff --git a/cuqi/solver/_solver.py b/cuqi/solver/_solver.py
index 98a52a389..5c7868c2d 100644
--- a/cuqi/solver/_solver.py
+++ b/cuqi/solver/_solver.py
@@ -15,7 +15,7 @@
     has_cholmod = False
 
 
-class L_BFGS_B(object):
+class ScipyLBFGSB(object):
     """Wrapper for :meth:`scipy.optimize.fmin_l_bfgs_b`.
 
     Minimize a function func using the L-BFGS-B algorithm.
@@ -30,14 +30,10 @@ class L_BFGS_B(object):
         Initial guess.
     gradfunc : callable f(x,*args), optional
         The gradient of func. 
-        If None, then the solver approximates the gradient.
+        If None, the solver approximates the gradient with a finite difference scheme.
     kwargs : keyword arguments passed to scipy's L-BFGS-B algorithm. See documentation for scipy.optimize.minimize
-
-    Methods
-    ----------
-    :meth:`solve`: Runs the solver and returns the solution and info about the optimization.
     """
-    def __init__(self,func,x0, gradfunc = None, **kwargs):
+    def __init__(self, func, x0, gradfunc = None, **kwargs):
         self.func= func
         self.x0 = x0
         self.gradfunc = gradfunc
@@ -83,7 +79,7 @@ def solve(self):
                 "nfev": solution[2]['funcalls']}
         return solution[0], info
 
-class minimize(object):
+class ScipyMinimizer(object):
     """Wrapper for :meth:`scipy.optimize.minimize`.
 
     Minimize a function func using scipy's optimize.minimize module.
@@ -115,12 +111,8 @@ class minimize(object):
         ‘trust-krylov’ 
         If not given, chosen to be one of BFGS, L-BFGS-B, SLSQP, depending if the problem has constraints or bounds.
     kwargs : keyword arguments passed to scipy's minimizer. See documentation for scipy.optimize.minimize
-
-    Methods
-    ----------
-    :meth:`solve`: Runs the solver and returns the solution and info about the optimization.
     """
-    def __init__(self,func,x0, gradfunc = None, method = None, **kwargs):
+    def __init__(self, func, x0, gradfunc = '2-point', method = None, **kwargs):
         self.func= func
         self.x0 = x0
         self.method = method
@@ -147,18 +139,20 @@ def solve(self):
         info = {"success": solution['success'],
                 "message": solution['message'],
                 "func": solution['fun'],
-                "grad": solution['jac'],
                 "nit": solution['nit'], 
                 "nfev": solution['nfev']}
+        # if gradfunc is callable, record the gradient in the info dict
+        if 'jac' in solution.keys():
+            info['grad'] = solution['jac']
         if isinstance(self.x0,CUQIarray):
             sol = CUQIarray(solution['x'],geometry=self.x0.geometry)
         else:
             sol = solution['x']
         return sol, info
 
-class maximize(minimize):
-    """Simply calls ::class:: cuqi.solver.minimize with -func."""
-    def __init__(self,func,x0, gradfunc = None, method = None, **kwargs):
+class ScipyMaximizer(ScipyMinimizer):
+    """Simply calls ::class:: cuqi.solver.ScipyMinimizer with -func."""
+    def __init__(self, func, x0, gradfunc = None, method = None, **kwargs):
         def nfunc(*args,**kwargs):
             return -func(*args,**kwargs)
         if gradfunc is not None:
@@ -170,7 +164,7 @@ def ngradfunc(*args,**kwargs):
 
 
 
-class LS(object):
+class ScipyLeastSquares(object):
     """Wrapper for :meth:`scipy.optimize.least_squares`.
 
     Solve nonlinear least-squares problems with bounds:
@@ -189,7 +183,7 @@ class LS(object):
         Initial guess.
     Jac : callable f(x,*args), optional
         The Jacobian of func. 
-        If None, then the solver approximates the Jacobian.
+        If not specified, the solver approximates the Jacobian with a finite difference scheme.
     loss: callable rho(x,*args)
         Determines the loss function
         'linear' : rho(z) = z. Gives a standard least-squares problem.
@@ -203,7 +197,7 @@ class LS(object):
         'dogbox', dogleg algorithm with rectangular trust regions, for small problems with bounds.
         'lm', Levenberg-Marquardt algorithm as implemented in MINPACK. Doesn't handle bounds and sparse Jacobians.
     """
-    def __init__(self, func, x0, jacfun=None, method='trf', loss='linear', tol=1e-6, maxit=1e4):
+    def __init__(self, func, x0, jacfun='2-point', method='trf', loss='linear', tol=1e-6, maxit=1e4):
         self.func = func
         self.x0 = x0
         self.jacfun = jacfun
diff --git a/cuqi/utilities/_get_python_variable_name.py b/cuqi/utilities/_get_python_variable_name.py
index 1ca9445a5..9651d5c15 100644
--- a/cuqi/utilities/_get_python_variable_name.py
+++ b/cuqi/utilities/_get_python_variable_name.py
@@ -9,7 +9,7 @@
 def _get_python_variable_name(var):
     """ Retrieve the Python variable name of an object. Takes the first variable name appearing on the stack that is not in the ignore list. """
 
-    ignored_var_names = ["self", "cls", "obj", "var", "_"]
+    ignored_var_names = ["self", "cls", "obj", "var", "_", "result", "args", "kwargs", "par_name", "name", "distribution", "dist"]
 
     # First get the stack size and loop (in reverse) through the stack
     # It can be a bit slow to loop through stack size so we limit the levels
@@ -29,7 +29,7 @@ def _get_python_variable_name(var):
             if len(var_names) > 0:
                 return var_names[0]
 
-    warnings.warn("Could not automatically find variable name for object: {}. Use keyword `name` when defining distribution to specify a name. If code runs slowly and variable name is not needed set config.MAX_STACK_SEARCH_DEPTH to 0.".format(var))
+    warnings.warn("Could not automatically find variable name for object. Did you assign (=) the object to a python variable? Alternatively, use keyword `name` when defining distribution to specify a name. If code runs slowly and variable name is not needed set config.MAX_STACK_SEARCH_DEPTH to 0. These names are reserved {} and should not be used as object name.".format(ignored_var_names))
 
     return None
 
diff --git a/demos/demo12_gradients.py b/demos/demo12_gradients.py
index afb9e74c4..9940df3ff 100644
--- a/demos/demo12_gradients.py
+++ b/demos/demo12_gradients.py
@@ -33,7 +33,7 @@
     # Solve posterior problem using BFGS
     def f(x): return -TP.posterior.logd(x)
     def gradf(x): return -TP.posterior.gradient(x)
-    solver = cuqi.solver.L_BFGS_B(f, x0, gradf)
+    solver = cuqi.solver.ScipyLBFGSB(f, x0, gradf)
     x_MAP, solution_info = solver.solve()
 print('relative error MAP:', np.linalg.norm(x_MAP-x_true)/np.linalg.norm(x_true))
 
diff --git a/demos/howtos/algebra.py b/demos/howtos/algebra.py
index ad30b24b7..43deaf078 100644
--- a/demos/howtos/algebra.py
+++ b/demos/howtos/algebra.py
@@ -1,66 +1,85 @@
 """
-Algebra in CUQIpy
-=================
+Random Variables and Algebra in CUQIpy
+======================================
 
 CUQIpy provides a simple algebraic framework for defining and manipulating
-variables.
-
-In this example, we will demonstrate how to define a simple algebraic structure
-and perform some basic algebraic operations.
+random variables.
 
+In this example, we demonstrate how to define random variables, apply algebraic
+operations on them, and finally use them in Bayesian Problems.
 """
 
 #%%
-# Defining Variables
-# ------------------
-# To utilize the algebraic framework in CUQIpy, we first need to define some
-# variables to apply the algebraic operations on. In this example, we will
-# define variables `x` and `y`.
+# Defining Random Variables
+# -------------------------
+# Random variables can be defined using the RandomVariable class. The RandomVariable
+# class requires a distribution object to be passed as an argument. The distribution
+# object can be any distribution object from the `cuqi.distribution` module.
 
-from cuqi.experimental.algebra import VariableNode
+from cuqi.distribution import Normal
+from cuqi.experimental.algebra import RandomVariable
 
-x = VariableNode('x')
-y = VariableNode('y')
 
-#%%
+x = RandomVariable(Normal(0, 1))
+y = RandomVariable(Normal(0, 1))
+
+# %%
 # Recording Algebraic Operations
-# --------------------
-# We can now perform some algebraic operations on the variables `x` and `y`.
-# The algebraic operations are recorded in a computational graph (abstract syntax tree).
-# The operations are recoding the correct ordering and adhering to the rules of
-# algebra.
+# ------------------------------
+# We can now perform some algebraic operations on the random variables. The
+# operations are recorded in a computational graph, which can be evaluated
+# later.
 
-print("Basic operations:")
-print(x + 1)
-print(x + y)
-print(x * y)
-print(x / y)
+print("Basic operations: \n")
+print(f"x + y yields:\n{x + y}\n")
+print(f"x - y yields:\n{x - y}\n")
+print(f"x * y yields:\n{x * y}\n")
+print(f"x / y yields:\n{x / y}\n")
 
-print("\nComplex operations:")
-print(x**2 + 2*x*y + y**2)
-print((x + y)**2)
+# %%
+print("Complex operations: \n")
+print(f"x**2 + 2*x*y + y**2 yields:\n{x**2 + 2*x*y + y**2}\n")
+print(f"(x + y)**2 yields\n{(x + y)**2}\n")
 
-print("\nProgrammatric operations:")
-print(x[2]+y[3])
+# %%
+print("Array operations: \n")
+print(f"x[0] + y[1] yields:\n{x[0] + y[1]}\n")
 
 # %%
-# Utilizing the Computational Graph
+# Utilizing the recorded operations
 # ---------------------------------
-# The computational graph can be utilized to evaluate the algebraic expressions
-# when desired. This means we can define mathematical expressions without
-# needing to evaluate them immediately.
+# We can evaluate the recorded operations by calling the random variable object
+# with the desired values for the random variables.
 
-# Define a mathematical expression
-expr1 = (x + y)**2
+# Define a new random variable 'z'
+z = (x + y)**2
 
 # Evaluate the expression (using the __call__ method)
-print(f"Expression {expr1} evaluated at x=2, y=3 yields {expr1(x=2, y=3)}")
+print(f"z={z.expression} evaluated at x=1, y=2 yields: {z(x=1, y=2)}")
+# %%
+# Building Bayesian Problems
+# --------------------------
+# Random variables can be used to define Bayesian problems. In this example we build
+# an example Bayesian problem using the Deconvolution1D test problem.
+
+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().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))
 
-# Another example
-expr2 = x**2 + 2*x*y + y**2 + 16
-print(f"Expression {expr2} evaluated at x=2, y=3 yields {expr2(x=2, y=3)}")
+# Combine into a Bayesian Problem and perform UQ
+BP = BayesianProblem(y, x, s, d)
+BP.set_data(y=y_obs)
+BP.UQ(exact={"x": info.exactSolution})
 
-# Another example utilizing array indexing
-expr3 = x[1] + y[2]
-print(f"Expression {expr3} evaluated at x=[1,2,3], y=[4,5,6] yields {expr3(x=[1,2,3], y=[4,5,6])}")
 # %%
diff --git a/demos/try09.py b/demos/try09.py
index 7051e1ca3..1b6d1c6e5 100644
--- a/demos/try09.py
+++ b/demos/try09.py
@@ -48,28 +48,28 @@ def posterior_logpdf_grad(x):
 print('relative error exact MAP:', np.linalg.norm(x_MAP_exact-x_true)/np.linalg.norm(x_true))
 
 #%% L_BFGS_B MAP
-solver = cuqi.solver.L_BFGS_B(posterior_logpdf, x0, gradfunc = posterior_logpdf_grad)
+solver = cuqi.solver.ScipyLBFGSB(posterior_logpdf, x0, gradfunc = posterior_logpdf_grad)
 x_MAP_LBFGS, info_MAP_LBFGS = solver.solve()
 print('relative error L-BFGS MAP:', np.linalg.norm(x_MAP_LBFGS-x_true)/np.linalg.norm(x_true))
 
-#%% L_BFGS_B MAP using minimize
-solver = cuqi.solver.minimize(posterior_logpdf, x0, gradfunc = posterior_logpdf_grad, method = "L-BFGS-B")
+#%% L_BFGS_B MAP using ScipyMinimizer
+solver = cuqi.solver.ScipyMinimizer(posterior_logpdf, x0, gradfunc = posterior_logpdf_grad, method = "L-BFGS-B")
 x_MAP_LBFGS2, info_MAP_LBFGS2 = solver.solve()
 print('relative error L-BFGS MAP:', np.linalg.norm(x_MAP_LBFGS2-x_true)/np.linalg.norm(x_true))
 
 #%% BFGS MAP without gradient input
-solver = cuqi.solver.minimize(posterior_logpdf, x0)
+solver = cuqi.solver.ScipyMinimizer(posterior_logpdf, x0)
 x_MAP_BFGS, info_MAP_BFGS = solver.solve()
 print('relative error BFGS MAP:', np.linalg.norm(x_MAP_BFGS-x_true)/np.linalg.norm(x_true))
 
 #%% BFGS MAP with gradient input
-solver = cuqi.solver.minimize(posterior_logpdf, x0, gradfunc = posterior_logpdf_grad)
+solver = cuqi.solver.ScipyMinimizer(posterior_logpdf, x0, gradfunc = posterior_logpdf_grad)
 x_MAP_BFGSgrad, info_MAP_BFGSgrad = solver.solve()
 print('relative error BFGS MAP:', np.linalg.norm(x_MAP_BFGSgrad-x_true)/np.linalg.norm(x_true))
 
 
 #%% SLSQP MAP
-solver = cuqi.solver.minimize(posterior_logpdf, x0, method = 'SLSQP')
+solver = cuqi.solver.ScipyMinimizer(posterior_logpdf, x0, method = 'SLSQP')
 x_MAP_SLSQP, info_MAP_SLSQP = solver.solve()
 print('relative error SLSQP MAP:', np.linalg.norm(x_MAP_SLSQP-x_true)/np.linalg.norm(x_true))
 
@@ -95,12 +95,12 @@ def likelihood_logpdf_grad(x):
     return grad
 
 # L_BFGS_B MAP
-solver = cuqi.solver.L_BFGS_B(likelihood_logpdf, x0, gradfunc = likelihood_logpdf_grad)
+solver = cuqi.solver.ScipyLBFGSB(likelihood_logpdf, x0, gradfunc = likelihood_logpdf_grad)
 x_ML_LBFGS, info_ML_LBFGS = solver.solve()
 print('relative error L-BFGS ML:', np.linalg.norm(x_MAP_LBFGS-x_true)/np.linalg.norm(x_true))
 
 # BFGS MAP
-solver = cuqi.solver.minimize(likelihood_logpdf, x0)
+solver = cuqi.solver.ScipyMinimizer(likelihood_logpdf, x0)
 x_ML_BFGS, info_ML_BFGS = solver.solve()
 print('relative error BFGS ML:', np.linalg.norm(x_MAP_BFGS-x_true)/np.linalg.norm(x_true))
 


Test patch:
diff --git a/tests/test_solver.py b/tests/test_solver.py
index 8092ee208..bdbb5ba3e 100644
--- a/tests/test_solver.py
+++ b/tests/test_solver.py
@@ -1,7 +1,7 @@
 import numpy as np
 import scipy as sp
 
-from cuqi.solver import CGLS, LM, FISTA, ADMM, ProximalL1, ProjectNonnegative
+from cuqi.solver import ScipyLBFGSB, ScipyMinimizer, ScipyLeastSquares, CGLS, LM, FISTA, ADMM, ProximalL1, ProjectNonnegative
 from scipy.optimize import lsq_linear
 
 
@@ -20,6 +20,61 @@ def test_CGLS():
     # Compare
     assert np.allclose(sol, ref_sol, rtol=1e-3)
 
+def test_ScipyMinimizer_without_gradient():
+    solver = ScipyMinimizer(sp.optimize.rosen,
+                            np.array([1.3, 0.7, 0.8, 1.9, 1.2]),
+                            method='Nelder-Mead',
+                            tol=1e-6)
+    sol, _ = solver.solve()
+    sol_ref = np.array([1.0, 1.0, 1.0, 1.0, 1.0])
+    assert np.allclose(sol, sol_ref)
+
+def test_ScipyMinimizer_with_gradient():
+    solver = ScipyMinimizer(sp.optimize.rosen,
+                            np.array([1.3, 0.7, 0.8, 1.9, 1.2]),
+                            gradfunc=sp.optimize.rosen_der,
+                            method='BFGS',
+                            tol=1e-6)
+    sol, _ = solver.solve()
+    sol_ref = np.array([1.0, 1.0, 1.0, 1.0, 1.0])
+    assert np.allclose(sol, sol_ref)
+
+def test_ScipyLBFGSB_without_gradient():
+    solver = ScipyLBFGSB(sp.optimize.rosen,
+                         np.array([1.3, 0.7, 0.8, 1.9, 1.2]))
+    sol, _ = solver.solve()
+    sol_ref = np.array([1.0, 1.0, 1.0, 1.0, 1.0])
+    assert np.allclose(sol, sol_ref)
+
+def test_ScipyLBFGSB_with_gradient():
+    solver = ScipyLBFGSB(sp.optimize.rosen,
+                         np.array([1.3, 0.7, 0.8, 1.9, 1.2]),
+                         gradfunc=sp.optimize.rosen_der)
+    sol, _ = solver.solve()
+    sol_ref = np.array([1.0, 1.0, 1.0, 1.0, 1.0])
+    assert np.allclose(sol, sol_ref)
+
+def test_ScipyLeastSquares_without_Jac():
+    def fun_rosenbrock(x):
+        return np.array([10 * (x[1] - x[0]**2), (1 - x[0])])
+    x0 = np.array([2, 2])
+    solver = ScipyLeastSquares(fun_rosenbrock, x0)
+    sol, _ = solver.solve()
+    sol_ref = np.array([1, 1])
+    assert np.allclose(sol, sol_ref)
+
+def test_ScipyLeastSquares_with_Jac():
+    def fun_rosenbrock(x):
+        return np.array([10 * (x[1] - x[0]**2), (1 - x[0])])
+    def jac_rosenbrock(x):
+        return np.array([
+            [-20 * x[0], 10],
+            [-1, 0]])
+    x0 = np.array([2, 2])
+    solver = ScipyLeastSquares(fun_rosenbrock, x0, jacfun=jac_rosenbrock)
+    sol, _ = solver.solve()
+    sol_ref = np.array([1, 1])
+    assert np.allclose(sol, sol_ref)
 
 def test_LM():
     # compare to MATLAB's original code solution
diff --git a/tests/zexperimental/test_productgeometry.py b/tests/zexperimental/test_productgeometry.py
new file mode 100644
index 000000000..cb0b711cf
--- /dev/null
+++ b/tests/zexperimental/test_productgeometry.py
@@ -0,0 +1,165 @@
+# %%
+import cuqi
+import pytest
+import numpy as np
+from cuqi.geometry import Continuous1D, Discrete, MappedGeometry, Continuous2D
+from cuqi.experimental.geometry import _ProductGeometry
+
+@pytest.fixture
+def product_geometry():
+    """fixture function to create a product geometry"""
+    geometry1 = Continuous1D(np.linspace(0, 1, 100))
+    geometry2 = MappedGeometry(
+        Discrete(["sound_speed"]), map=lambda x: x**2, imap=lambda x: np.sqrt(x)
+    )
+    geometry3 = Continuous2D((3,3))
+    product_geometry = _ProductGeometry(geometry1, geometry2, geometry3)
+    return product_geometry
+
+@pytest.fixture
+def par1(product_geometry):
+    """fixture function to create a parameter vector"""
+    geometry1, geometry2, geometry3 = product_geometry.geometries
+
+    g1_dim = geometry1.par_dim
+    g2_dim = geometry2.par_dim
+    g3_dim = geometry3.par_dim
+
+    par1 = np.ones(g1_dim+g2_dim+g3_dim)
+
+    par1[g1_dim:g1_dim+g2_dim] = 2
+    par1[g1_dim+g2_dim:] = 3
+    return par1
+
+@pytest.fixture
+def par2(product_geometry):
+    """fixture function to create a parameter vector"""
+    geometry1, geometry2, geometry3 = product_geometry.geometries
+
+    par2 = (np.ones(geometry1.par_dim),
+            np.ones(geometry2.par_dim)*2,
+            np.ones(geometry3.par_dim)*3)
+    return par2
+
+@pytest.fixture
+def par3(product_geometry):
+    """fixture function to create a parameter vector"""
+    geometry1, geometry2, geometry3 = product_geometry.geometries
+
+    g1_dim = geometry1.par_dim
+    g2_dim = geometry2.par_dim
+    g3_dim = geometry3.par_dim
+    
+    Ns = 5
+    par3 = np.ones((g1_dim+g2_dim+g3_dim, Ns))
+    par3[g1_dim:,:] = 2
+    par3[g1_dim+g2_dim:,:] = 3
+
+    return par3
+
+@pytest.fixture
+def par4(product_geometry):
+    """fixture function to create a parameter vector"""
+    geometry1, geometry2, geometry3 = product_geometry.geometries
+
+    Ns = 5
+    par4 = (np.ones((geometry1.par_dim, Ns)),
+            2*np.ones((geometry2.par_dim, Ns)),
+            3*np.ones((geometry3.par_dim, Ns)))
+    return par4
+
+
+def test_par2fun_single_sample(product_geometry, par1, par2):
+    """Test par2fun for single sample"""
+    geometry1, geometry2, geometry3 = product_geometry.geometries
+
+    fun1 = product_geometry.par2fun(par1)
+    fun2 = product_geometry.par2fun(par2[0], par2[1], par2[2])
+    fun2_a = geometry1.par2fun(par2[0])
+    fun2_b = geometry2.par2fun(par2[1])
+    fun2_c = geometry3.par2fun(par2[2])
+
+    assert np.allclose(len(fun1), 3)
+    assert np.allclose(len(fun1), len(fun2))
+    assert np.allclose(fun1[0], fun2[0])
+    assert np.allclose(fun1[1], fun2[1])
+    assert np.allclose(fun1[2], fun2[2])
+    assert np.allclose(fun2[0], fun2_a)
+    assert np.allclose(fun2[1], fun2_b)
+    assert np.allclose(fun2[2], fun2_c)
+
+def test_par2fun_multiple_samples(product_geometry, par3, par4):
+    """Test par2fun for multiple samples"""
+    geometry1, geometry2, geometry3 = product_geometry.geometries
+
+    fun3 = product_geometry.par2fun(par3)
+    fun4 = product_geometry.par2fun(par4[0], par4[1], par4[2])
+    fun4_a = geometry1.par2fun(par4[0])
+    fun4_b = geometry2.par2fun(par4[1])
+    fun4_c = geometry3.par2fun(par4[2])
+
+    assert(np.allclose(len(fun3), len(fun4)))
+    assert(np.allclose(len(fun3), 3))
+    assert(np.allclose(fun3[0], fun4[0]))
+    assert(np.allclose(fun3[1], fun4[1]))
+    assert(np.allclose(fun3[2], fun4[2]))
+    assert(np.allclose(fun4[0], fun4_a))
+    assert(np.allclose(fun4[1], fun4_b))
+    assert(np.allclose(fun4[2], fun4_c))
+
+def test_fun2par_single_sample(product_geometry, par2):
+    """Test fun2par for single sample"""
+
+    fun2 = product_geometry.par2fun(par2[0], par2[1], par2[2])
+    par_same_as_par2 = product_geometry.fun2par(fun2[0], fun2[1], fun2[2])
+
+    assert np.allclose(len(par2), len(par_same_as_par2))
+    assert np.allclose(len(par2), 3)
+    assert np.allclose(par2[0], par_same_as_par2[0])
+    assert np.allclose(par2[1], par_same_as_par2[1])
+    assert np.allclose(par2[2], par_same_as_par2[2])
+
+def test_fun2par_single_sample_stacked(product_geometry, par1):
+    """Test fun2par for single sample stacked"""
+
+    fun1 = product_geometry.par2fun(par1)
+    par_same_as_par1 = product_geometry.fun2par(fun1[0], fun1[1], fun1[2], stacked=True)
+
+    assert np.allclose(par1.shape, par_same_as_par1.shape)
+    assert np.allclose(par1, par_same_as_par1)
+
+def test_fun2par_multiple_samples(product_geometry, par4):
+    """Test fun2par for multiple samples"""
+
+    fun4 = product_geometry.par2fun(par4[0], par4[1], par4[2])
+    par_same_as_par4 = product_geometry.fun2par(fun4[0], fun4[1], fun4[2])
+
+    assert np.allclose(len(par4), len(par_same_as_par4))
+    assert np.allclose(len(par4), 3)
+    assert np.allclose(par4[0], par_same_as_par4[0])
+    assert np.allclose(par4[1], par_same_as_par4[1])
+    assert np.allclose(par4[2], par_same_as_par4[2])
+
+def test_fun2par_multiple_samples_stacked(product_geometry, par3):
+    """Test fun2par for multiple samples stacked"""
+
+    fun3 = product_geometry.par2fun(par3)
+
+    par_same_as_par3 = product_geometry.fun2par(fun3[0], fun3[1], fun3[2], stacked=True)
+
+    assert np.allclose(par3.shape, par_same_as_par3.shape)
+    assert np.allclose(par3, par_same_as_par3)
+
+def test_number_of_geometries(product_geometry):
+    """Test number of geometries"""
+    assert np.allclose(product_geometry.number_of_geometries, 3)
+
+def test_stacked_par_split_indices(product_geometry):
+    """Test stacked_par_split_indices are correct"""
+    assert np.allclose(product_geometry.stacked_par_split_indices, [100, 101])
+
+def test_input_not_geometry_fails():
+    """Test that passing input other than Geometry fails"""
+    with pytest.raises(TypeError,
+                       match="All geometries must be of type Geometry"):
+        _ProductGeometry(1, 2, 3)
\ No newline at end of file
diff --git a/tests/zexperimental/test_randomvariable.py b/tests/zexperimental/test_randomvariable.py
new file mode 100644
index 000000000..3c82d200b
--- /dev/null
+++ b/tests/zexperimental/test_randomvariable.py
@@ -0,0 +1,318 @@
+import pytest
+import cuqi
+import numpy as np
+from cuqi.experimental.algebra._ast import VariableNode
+from cuqi.experimental.algebra import RandomVariable
+
+def test_randomvariable_name_consistency():
+
+    # Case 1: distribution with explicit name
+    x_rv = RandomVariable(cuqi.distribution.Gaussian(0, 1, name="x"))
+
+    assert x_rv.name == "x"
+
+    # Case 2: distribution with implicit name
+    x = RandomVariable(cuqi.distribution.Gaussian(0, 1))
+
+    assert x.name == "x"
+
+
+def test_algrabraic_operations_on_distribution_should_create_randomvariable():
+
+    x = RandomVariable(cuqi.distribution.Gaussian(0, 1))
+
+    assert isinstance(x+1, RandomVariable)
+    assert isinstance(1+x, RandomVariable)
+
+    assert isinstance(x-1, RandomVariable)
+    assert isinstance(1-x, RandomVariable)
+
+    assert isinstance(x*1, RandomVariable)
+    assert isinstance(1*x, RandomVariable)
+
+    assert isinstance(x/1, RandomVariable)
+    assert isinstance(1/x, RandomVariable)
+
+    assert isinstance(x**1, RandomVariable)
+
+    assert isinstance(x@1, RandomVariable)
+    assert isinstance(1@x, RandomVariable)
+
+    assert isinstance(-x, RandomVariable)
+
+    assert isinstance(abs(x), RandomVariable)
+
+    assert isinstance(x[0], RandomVariable)
+
+
+@pytest.mark.parametrize("operations", [
+    lambda x: x+1,
+    lambda x: x**2,
+    lambda x: -x,
+    lambda x: abs(x),
+    lambda x: x/1,
+    lambda x: 1/x,
+    lambda x: x*1,
+    lambda x: 1*x,
+    lambda x: 3/abs((x+1)**2-10),
+    lambda x: 9/(x+1)**2,
+    lambda x: x/3+4,
+    lambda x: x[0],
+    lambda x: x[0:1],
+    lambda x: cuqi.model.LinearModel(np.ones((2,2)))@x,
+])
+def test_algebra_on_randomvariables_can_be_computed_and_is_correct(operations):
+    X = RandomVariable(cuqi.distribution.Gaussian(np.zeros(2), 1))
+    rv = operations(X)
+    val = np.random.randn(2)
+    # Compare random variable recorded operations vs actual operations
+    assert np.allclose(rv(val), operations(val))
+
+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,
+    lambda x: -x,
+    lambda x: abs(x),
+    lambda x: x/1,
+    lambda x: 1/x,
+    lambda x: x*1,
+    lambda x: 1*x,
+    lambda x: 3/abs((x+1)**2-10),
+    lambda x: 9/(x+1)**2,
+    lambda x: x/3+4,
+    lambda x: x[0],
+    lambda x: x[0:1],
+])
+def test_randomvariable_works_with_distribution_conditioning(operations):
+
+    # Define x and y | x (y conditioned on x) with some algebraic operations
+    x = RandomVariable(cuqi.distribution.Gaussian(np.zeros(2), 1))
+    y = RandomVariable(cuqi.distribution.Gaussian(np.zeros(2), operations(x)))
+
+    # Condition y on a random fixed value of x
+    val = np.random.randn(2)
+    y_cond_x = y.condition(x=val)
+
+    # Check basic classes are correct
+    assert isinstance(y, RandomVariable)
+    assert isinstance(y.distribution.cov, RandomVariable)
+
+    # Check conditioning works and provides the expected result
+    assert np.allclose(y_cond_x.distribution.cov, operations(val))
+
+@pytest.mark.parametrize("operations", [
+    lambda x, y, z: x + y - z,
+    lambda x, y, z: x * y - z,
+    lambda x, y, z: x / y - z,
+    lambda x, y, z: x ** y - z,
+    lambda x, y, z: x @ y - z,
+    lambda x, y, z: x + y + z,
+    lambda x, y, z: x * y * z,
+    lambda x, y, z: x / y / z,
+    lambda x, y, z: x ** y ** z,
+    lambda x, y, z: x + y * z,
+    lambda x, y, z: x * y + z,
+    lambda x, y, z: x + y @ z,
+    lambda x, y, z: x @ y + z,
+    lambda x, y, z: x[0]+z+y,
+    lambda x, y, z: abs(x[0])+z-y,
+    lambda x, y, z: (x + y) * z,
+    lambda x, y, z: (x - y) / z,
+    lambda x, y, z: x ** (y + z),
+    lambda x, y, z: abs(x) + abs(y * z),
+    lambda x, y, z: (x + y) * (z + 1) - x * y,
+    lambda x, y, z: x / (y ** z) + (x * y) ** z,
+    lambda x, y, z: abs(x) * (abs(y) + abs(z)),
+    lambda x, y, z: abs(x * y) + abs(y * z) + abs(z * x),
+    lambda x, y, z: x * y + y * z + z * x,
+])
+def test_randomvariable_algebra_works_on_joint_space(operations):
+    """ Test that algebraic operations on random variables work in joint space """
+    x = RandomVariable(cuqi.distribution.Gaussian(0, 1))
+    y = RandomVariable(cuqi.distribution.Gaussian(0, 1))
+    z = RandomVariable(cuqi.distribution.Gaussian(0, 1))
+
+    # Define a random variable in joint space
+    rv = operations(x, y, z)
+
+    # Check that the random variable is a random variable
+    assert isinstance(rv, RandomVariable)
+
+    # Check operations work
+    val_x = abs(np.random.randn(x.dim))+1
+    val_y = abs(np.random.randn(y.dim))+1
+    val_z = abs(np.random.randn(z.dim))+1
+    assert np.allclose(rv(x=val_x, y=val_y, z=val_z), operations(val_x, val_y, val_z))
+
+@pytest.mark.xfail(reason="sampling method for random variables is not yet implemented")
+@pytest.mark.parametrize("operations", [
+    lambda x, y, z: x + y,
+    lambda x, y, z: x * y,
+    lambda x, y, z: x / y,
+    lambda x, y, z: x ** y,
+    lambda x, y, z: x + y + z,
+    lambda x, y, z: x * y * z,
+    lambda x, y, z: x / y / z,
+    lambda x, y, z: x ** y ** z,
+    lambda x, y, z: x + y * z,
+    lambda x, y, z: x * y + z,
+    lambda x, y, z: (x + y) * z,
+    lambda x, y, z: (x - y) / z,
+    lambda x, y, z: x ** (y + z),
+    lambda x, y, z: abs(x) + abs(y * z),
+    lambda x, y, z: (x + y) * (z + 1) - x * y,
+    lambda x, y, z: x / (y ** z) + (x * y) ** z,
+    lambda x, y, z: abs(x) * (abs(y) + abs(z)),
+    lambda x, y, z: abs(x * y) + abs(y * z) + abs(z * x),
+    lambda x, y, z: x * y + y * z + z * x,
+])
+def test_randomvariable_sample(operations):
+    """ Test that random variable sampling works """
+    x = RandomVariable(cuqi.distribution.Gaussian(0, 1))
+    y = RandomVariable(cuqi.distribution.Gaussian(0, 1))
+    z = RandomVariable(cuqi.distribution.Gaussian(0, 1))
+
+    # Define a random variable in joint space
+    rv = operations(x, y, z)
+
+    # Fix rng and sample
+    np.random.seed(0)
+    result = rv.sample()
+
+    # Check that the result compares to the expected result
+    np.random.seed(0)
+    expected_result = operations(x.sample(), y.sample(), z.sample())
+
+    assert np.allclose(result, expected_result)
+
+@pytest.mark.xfail(reason="logd method for random variable not yet implemented")
+def test_logp_conditional():
+    """ This tests logp evaluation for conditional random variables """
+    # Base example logp value
+    true_val = cuqi.distribution.Gaussian(3, 7).logd(13)
+
+    # Distribution with no specified parameters
+    x = RandomVariable(cuqi.distribution.Gaussian(cov=lambda s:s, geometry=1))
+
+    # Test logp evaluates correctly in various cases
+    assert x.logd(mean=3, s=7, x=13) == true_val
+    assert x.condition(x=13).logd(mean=3, s=7) == true_val
+    assert x.condition(x=13, mean=3).logd(s=7) == true_val
+    assert x.condition(x=13, mean=3, s=7).logd() == true_val
+    assert x.condition(mean=3).logd(s=7, x=13) == true_val
+    assert x.condition(mean=3, s=7).logd(x=13) == true_val
+    assert x.condition(mean=3, x=13).logd(s=7) == true_val
+
+def test_rv_attributes():
+    """ Test various attributes of random variable"""
+
+    x = RandomVariable(cuqi.distribution.Gaussian(0, 1, geometry=10))
+
+    assert x.dim == 10
+    assert isinstance(x.geometry, cuqi.geometry._DefaultGeometry1D)
+    assert x.name == "x"
+    assert isinstance(x.tree, VariableNode)
+    assert x.tree.name == "x"
+    assert isinstance(x.distribution, cuqi.distribution.Gaussian)
+
+def test_rv_variable_name_detection():
+    """Test that the rv name is detected correctly at different levels of the python stack. """
+
+    # Test that the density variable name is detected correctly at current level.
+    x = RandomVariable(cuqi.distribution.Gaussian(geometry=1))
+    assert x.name == 'x'
+
+    # Test that variable name is detected correctly 1 level deep.
+    def inner_name():
+        y = RandomVariable(cuqi.distribution.Gaussian(geometry=1))
+        assert y.name == 'y'
+    inner_name()
+
+    # Test variable name is detected correctly at n levels deep.
+    class recursive_name:
+        def __init__(self, max_recursion=10):
+            self.max_recursion = max_recursion
+        def __call__(self, current_recursion=0):
+            if current_recursion == self.max_recursion:
+                z = RandomVariable(cuqi.distribution.Gaussian(geometry=1))
+                assert z.name == 'z'
+            else:
+                self(current_recursion + 1)
+    recursive_name()()
+
+def test_variable_name_accross_frames():
+    """ Test variable name can be inferred across multiple stack frames. """
+
+    h = RandomVariable(cuqi.distribution.Gaussian(geometry=1)) # Name should be 'h'
+
+    def recursive_return_rv(rv, recursions):
+        if recursions == 0:
+            assert rv.name == 'h' # h was defined many frames above, and name should be inferred correctly.
+        else:
+            recursive_return_rv(rv, recursions - 1)
+    
+    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))
+    x2 = x.condition(mean=1)
+    x3 = x2.condition(cov=1)
+
+    # Names should be the same as the original density.  
+    assert x3.name == 'x'
+    assert x2.name == 'x' 
+    assert x.name == 'x'
+
+    # Ensure that the name cannot be changed for conditioned densities.
+    with pytest.raises(ValueError, match=r"This random variable is derived from the conditional random variable named x"):
+        x2.name = 'y'
+
+    x.name = 'y'
+
+    # Ensure that the name is changed for the other conditioned densities.
+    assert x2.name == 'y'
+    assert x3.name == 'y'
+
+def test_RV_should_catch_non_linear_model_used_as_linear_model():
+    A = cuqi.testproblem.Poisson1D().model
+    x = RandomVariable(cuqi.distribution.Gaussian(0, 1, geometry=A.domain_geometry))
+
+    with pytest.raises(TypeError, match=r"Cannot apply matmul to non-linear models"):
+        y=A@x
+
+def test_ensure_that_RV_evaluation_requires_all_parameters():
+    x = RandomVariable(cuqi.distribution.Gaussian(0, 1))
+
+    # raise ValueError(f"Expected {self.parameter_names} arguments, got {kwargs}")
+    with pytest.raises(ValueError, match=r"Expected arguments \['x'\], got arguments"):
+        x()
+
+    with pytest.raises(ValueError, match=r"Expected arguments \['x'\], got arguments \{'y': 1\}"):
+        x(y=1)
+
+    with pytest.raises(ValueError, match=r"Expected arguments \['x'\], got arguments \{'x': 1, 'y': 1\}"):
+        x(x=1, y=1)
+
+    y = RandomVariable(cuqi.distribution.Gaussian(0, 1))
+    z = (x+y)**2
+
+    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)
+
+    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.