Repository: CUQI-DTU/CUQIpy

Issue:
RV: Variable class with algebra (using AST)
Implement a variable class that can "record" various operations applied to it and allow this information to be used later. For example:
```python
x = Variable()
z = (x+2)**2
```
The variable `z` should know of the operations applied and it should be possible to apply them again to specific values of `x`, say `z(x=1)` should return the value 9.

- [x] Implement abstract syntax tree for a basic variable "node" class
- [x] Support core mathematical operations like (0, -, /, *)
- [x] Support mathematical operations: (abs, pow, neg, matmul)
- [x] Support operations: (getitem)
- [x] Implement repr for displaying operations applied to variable
- [x] Unit tests
- [x] Howto

Discussion:
@nabriis: @jakobsj DoD@jakobsj: @nabriis looks good - I see a howto is added in the PR along with tests of algebra. Please add to DoD separate items for the howto and for unit tests (apart from testing algebra is it possible to add tests of the variable/node implementations)@nabriis: As we discussed if we want stuff like "I remembered to add unit tests" tickbox I want a template to use.@jakobsj: @nabriis Thanks for adding unit tests and howto in the specific DoD here so we can progress this issue. Yes absolutely unit tests will be part of the general template. DoD approved.

Solution patch:
diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml
index 5654f5c11..0cf002c9c 100644
--- a/.github/workflows/docs.yml
+++ b/.github/workflows/docs.yml
@@ -22,6 +22,15 @@ jobs:
       - name: Sphinx build
         run: |
           sphinx-build docs ./public
+
+      # If PR-based build. Upload artifact of the documentation for download and review
+      - name: Upload artifact
+        if: ${{ github.event_name == 'pull_request' }}
+        uses: actions/upload-artifact@v3
+        with:
+          name: docs_build
+          path: ./public
+
       - name: Deploy
         uses: peaceiris/actions-gh-pages@v3
         if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/main' }}
diff --git a/cuqi/experimental/__init__.py b/cuqi/experimental/__init__.py
index 29cf970aa..63d0626ea 100644
--- a/cuqi/experimental/__init__.py
+++ b/cuqi/experimental/__init__.py
@@ -1,2 +1,3 @@
 """ Experimental module for testing new features and ideas. """
 from . import mcmc
+from . import algebra
diff --git a/cuqi/experimental/algebra/__init__.py b/cuqi/experimental/algebra/__init__.py
new file mode 100644
index 000000000..4a9ac77d1
--- /dev/null
+++ b/cuqi/experimental/algebra/__init__.py
@@ -0,0 +1 @@
+from ._ast import VariableNode
diff --git a/cuqi/experimental/algebra/_ast.py b/cuqi/experimental/algebra/_ast.py
new file mode 100644
index 000000000..527c6f536
--- /dev/null
+++ b/cuqi/experimental/algebra/_ast.py
@@ -0,0 +1,325 @@
+"""
+CUQIpy specific implementation of an abstract syntax tree (AST) for algebra on variables.
+
+The AST is used to record the operations applied to variables allowing a delayed evaluation
+of said operations when needed by traversing the tree with the __call__ method.
+
+For example, the following code
+
+    x = VariableNode('x')
+    y = VariableNode('y')
+    z = 2*x + 3*y
+
+will create the following AST:
+
+z = AddNode(
+        MultiplyNode(
+            ValueNode(2),
+            VariableNode('x')
+        ),
+        MultiplyNode(
+            ValueNode(3),
+            VariableNode('y')
+        )
+    )
+
+which can be evaluated by calling the __call__ method:
+
+    z(x=1, y=2) # returns 8
+
+"""
+
+from abc import ABC, abstractmethod
+
+convert_to_node = lambda x: x if isinstance(x, Node) else ValueNode(x)
+""" Converts any non-Node object to a ValueNode object. """
+
+# ====== Base classes for the nodes ======
+
+
+class Node(ABC):
+    """Base class for all nodes in the abstract syntax tree.
+
+    Responsible for building the AST by creating nodes that represent the operations applied to variables.
+
+    Each subclass must implement the __call__ method that will evaluate the node given the input parameters.
+
+    """
+
+    @abstractmethod
+    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 __repr__(self):
+        """String representation of the node. Used for printing the AST."""
+        pass
+
+    def __add__(self, other):
+        return AddNode(self, convert_to_node(other))
+
+    def __radd__(self, other):
+        return AddNode(convert_to_node(other), self)
+
+    def __sub__(self, other):
+        return SubtractNode(self, convert_to_node(other))
+
+    def __rsub__(self, other):
+        return SubtractNode(convert_to_node(other), self)
+
+    def __mul__(self, other):
+        return MultiplyNode(self, convert_to_node(other))
+
+    def __rmul__(self, other):
+        return MultiplyNode(convert_to_node(other), self)
+
+    def __truediv__(self, other):
+        return DivideNode(self, convert_to_node(other))
+
+    def __rtruediv__(self, other):
+        return DivideNode(convert_to_node(other), self)
+
+    def __pow__(self, other):
+        return PowerNode(self, convert_to_node(other))
+
+    def __rpow__(self, other):
+        return PowerNode(convert_to_node(other), self)
+
+    def __neg__(self):
+        return NegateNode(self)
+
+    def __abs__(self):
+        return AbsNode(self)
+
+    def __getitem__(self, i):
+        return GetItemNode(self, convert_to_node(i))
+
+    def __matmul__(self, other):
+        return MatMulNode(self, convert_to_node(other))
+
+    def __rmatmul__(self, other):
+        return MatMulNode(convert_to_node(other), self)
+
+
+class UnaryNode(Node, ABC):
+    """Base class for all unary nodes in the abstract syntax tree.
+
+    Parameters
+    ----------
+    child : Node
+        The direct child node on which the unary operation is performed.
+
+    """
+
+    def __init__(self, child: Node):
+        self.child = child
+
+
+class BinaryNode(Node, ABC):
+    """Base class for all binary nodes in the abstract syntax tree.
+
+    The op_symbol attribute is used for printing the operation in the __repr__ method.
+
+    Parameters
+    ----------
+    left : Node
+        Left child node to the binary operation.
+
+    right : Node
+        Right child node to the binary operation.
+
+    """
+
+    @property
+    @abstractmethod
+    def op_symbol(self):
+        """Symbol used to represent the operation in the __repr__ method."""
+        pass
+
+    def __init__(self, left: Node, right: Node):
+        self.left = left
+        self.right = right
+
+    def __repr__(self):
+        return f"{self.left} {self.op_symbol} {self.right}"
+
+
+class BinaryNodeWithParenthesis(BinaryNode, ABC):
+    """Base class for all binary nodes in the abstract syntax tree that should be printed with parenthesis."""
+
+    def __repr__(self):
+        left = f"({self.left})" if isinstance(self.left, BinaryNode) else str(self.left)
+        right = (
+            f"({self.right})" if isinstance(self.right, BinaryNode) else str(self.right)
+        )
+        return f"{left} {self.op_symbol} {right}"
+
+class BinaryNodeWithParenthesisNoSpace(BinaryNode, ABC):
+    """Base class for all binary nodes in the abstract syntax tree that should be printed with parenthesis but no space."""
+
+    def __repr__(self):
+        left = f"({self.left})" if isinstance(self.left, BinaryNode) else str(self.left)
+        right = (
+            f"({self.right})" if isinstance(self.right, BinaryNode) else str(self.right)
+        )
+        return f"{left}{self.op_symbol}{right}"
+
+
+# ====== Specific implementations of the "leaf" nodes ======
+
+
+class VariableNode(Node):
+    """Node that represents a generic variable, e.g. "x" or "y".
+
+    Parameters
+    ----------
+    name : str
+        Name of the variable. Used for printing and to retrieve the given input value
+        of the variable in the kwargs dictionary when evaluating the tree.
+
+    """
+
+    def __init__(self, name):
+        self.name = name
+
+    def __call__(self, **kwargs):
+        """Retrieves the value of the variable from the passed kwargs. If no value is found, it raises a KeyError."""
+        if not self.name in kwargs:
+            raise KeyError(
+                f"Variable '{self.name}' not found in the given input parameters. Unable to evaluate the expression."
+            )
+        return kwargs[self.name]
+
+    def __repr__(self):
+        return self.name
+
+
+class ValueNode(Node):
+    """Node that represents a constant value. The value can be any python object that is not a Node.
+
+    Parameters
+    ----------
+    value : object
+        The python object that represents the value of the node.
+
+    """
+
+    def __init__(self, value):
+        self.value = value
+
+    def __call__(self, **kwargs):
+        """Returns the value of the node."""
+        return self.value
+
+    def __repr__(self):
+        return str(self.value)
+
+
+# ====== Specific implementations of the "internal" nodes ======
+
+
+class AddNode(BinaryNode):
+    """Node that represents the addition operation."""
+
+    @property
+    def op_symbol(self):
+        return "+"
+
+    def __call__(self, **kwargs):
+        return self.left(**kwargs) + self.right(**kwargs)
+
+
+class SubtractNode(BinaryNode):
+    """Node that represents the subtraction operation."""
+
+    @property
+    def op_symbol(self):
+        return "-"
+
+    def __call__(self, **kwargs):
+        return self.left(**kwargs) - self.right(**kwargs)
+
+
+class MultiplyNode(BinaryNodeWithParenthesis):
+    """Node that represents the multiplication operation."""
+
+    @property
+    def op_symbol(self):
+        return "*"
+
+    def __call__(self, **kwargs):
+        return self.left(**kwargs) * self.right(**kwargs)
+
+
+class DivideNode(BinaryNodeWithParenthesis):
+    """Node that represents the division operation."""
+
+    @property
+    def op_symbol(self):
+        return "/"
+
+    def __call__(self, **kwargs):
+        return self.left(**kwargs) / self.right(**kwargs)
+
+
+class PowerNode(BinaryNodeWithParenthesisNoSpace):
+    """Node that represents the power operation."""
+
+    @property
+    def op_symbol(self):
+        return "^"
+
+    def __call__(self, **kwargs):
+        return self.left(**kwargs) ** self.right(**kwargs)
+
+
+class GetItemNode(BinaryNode):
+    """Node that represents the get item operation. Here the left node is the object and the right node is the index."""
+
+    def __call__(self, **kwargs):
+        return self.left(**kwargs)[self.right(**kwargs)]
+
+    def __repr__(self):
+        left = f"({self.left})" if isinstance(self.left, BinaryNode) else str(self.left)
+        return f"{left}[{self.right}]"
+    
+    @property
+    def op_symbol(self):
+        pass
+
+
+class NegateNode(UnaryNode):
+    """Node that represents the arithmetic negation operation."""
+
+    def __call__(self, **kwargs):
+        return -self.child(**kwargs)
+
+    def __repr__(self):
+        child = (
+            f"({self.child})"
+            if isinstance(self.child, (BinaryNode, UnaryNode))
+            else str(self.child)
+        )
+        return f"-{child}"
+
+
+class AbsNode(UnaryNode):
+    """Node that represents the absolute value operation."""
+
+    def __call__(self, **kwargs):
+        return abs(self.child(**kwargs))
+
+    def __repr__(self):
+        return f"abs({self.child})"
+
+
+class MatMulNode(BinaryNodeWithParenthesis):
+    """Node that represents the matrix multiplication operation."""
+
+    @property
+    def op_symbol(self):
+        return "@"
+
+    def __call__(self, **kwargs):
+        return self.left(**kwargs) @ self.right(**kwargs)
diff --git a/cuqi/experimental/mcmc/_langevin_algorithm.py b/cuqi/experimental/mcmc/_langevin_algorithm.py
index 9c1e9742c..5759a8548 100644
--- a/cuqi/experimental/mcmc/_langevin_algorithm.py
+++ b/cuqi/experimental/mcmc/_langevin_algorithm.py
@@ -3,7 +3,7 @@
 from cuqi.experimental.mcmc import Sampler
 from cuqi.implicitprior import RestorationPrior, MoreauYoshidaPrior
 from cuqi.array import CUQIarray
-from copy import deepcopy
+from copy import copy
 
 class ULA(Sampler): # Refactor to Proposal-based sampler?
     """Unadjusted Langevin algorithm (ULA) (Roberts and Tweedie, 1996)
@@ -316,11 +316,14 @@ def target(self, value):
 
     def _create_smoothed_target(self, value):
         """ Create a smoothed target using a Moreau-Yoshida envelope. """
-        copied_value = deepcopy(value)
+        copied_value = copy(value)
         if isinstance(copied_value.prior, RestorationPrior):
+            # Acceess the prior name
+            name = value.prior.name
             copied_value.prior = MoreauYoshidaPrior(
                 copied_value.prior,
-                self.smoothing_strength)
+                self.smoothing_strength,
+                name=name)
         return copied_value
 
     def validate_target(self):
diff --git a/cuqi/experimental/mcmc/_rto.py b/cuqi/experimental/mcmc/_rto.py
index 9ab5fb90a..9858ed2f2 100644
--- a/cuqi/experimental/mcmc/_rto.py
+++ b/cuqi/experimental/mcmc/_rto.py
@@ -3,7 +3,7 @@
 from scipy.sparse.linalg import LinearOperator as scipyLinearOperator
 import numpy as np
 import cuqi
-from cuqi.solver import CGLS, FISTA
+from cuqi.solver import CGLS, FISTA, ADMM
 from cuqi.experimental.mcmc import Sampler
 
 
@@ -161,6 +161,13 @@ class RegularizedLinearRTO(LinearRTO):
     Regularized Linear RTO (Randomize-Then-Optimize) sampler.
 
     Samples posterior related to the inverse problem with Gaussian likelihood and implicit Gaussian prior, and where the forward model is Linear.
+    The sampler works by repeatedly solving regularized linear least squares problems for perturbed data.
+    The solver for these optimization problems is chosen based on how the regularized is provided in the implicit Gaussian prior.
+    Currently we use the following solvers:
+    FISTA: [1] Beck, Amir, and Marc Teboulle. "A fast iterative shrinkage-thresholding algorithm for linear inverse problems." SIAM journal on imaging sciences 2.1 (2009): 183-202.
+           Used when prior.proximal is callable.
+    ADMM:  [2] Boyd et al. "Distributed optimization and statistical learning via the alternating direction method of multipliers."Foundations and Trends® in Machine learning, 2011.
+           Used when prior.proximal is a list of penalty terms.
 
     Parameters
     ------------
@@ -171,12 +178,19 @@ class RegularizedLinearRTO(LinearRTO):
         Initial point for the sampler. *Optional*.
 
     maxit : int
-        Maximum number of iterations of the inner FISTA solver. *Optional*.
+        Maximum number of iterations of the inner FISTA/ADMM solver. *Optional*.
+
+    inner_max_it : int
+        Maximum number of iterations of the CGLS solver used within the ADMM solver. *Optional*.
         
     stepsize : string or float
         If stepsize is a string and equals either "automatic", then the stepsize is automatically estimated based on the spectral norm.
         If stepsize is a float, then this stepsize is used.
 
+    penalty_parameter : int
+        Penalty parameter of the inner ADMM solver. *Optional*.
+        See [2] or `cuqi.solver.ADMM`
+
     abstol : float
         Absolute tolerance of the inner FISTA solver. *Optional*.
     
@@ -190,7 +204,7 @@ class RegularizedLinearRTO(LinearRTO):
         An example is shown in demos/demo31_callback.py.
         
     """
-    def __init__(self, target=None, initial_point=None, maxit=100, stepsize="automatic", abstol=1e-10, adaptive=True, **kwargs):
+    def __init__(self, target=None, initial_point=None, maxit=100, inner_max_it=10, stepsize="automatic", penalty_parameter=10, abstol=1e-10, adaptive=True, **kwargs):
         
         super().__init__(target=target, initial_point=initial_point, **kwargs)
 
@@ -199,10 +213,13 @@ def __init__(self, target=None, initial_point=None, maxit=100, stepsize="automat
         self.abstol = abstol   
         self.adaptive = adaptive
         self.maxit = maxit
+        self.inner_max_it = inner_max_it
+        self.penalty_parameter = penalty_parameter
 
     def _initialize(self):
         super()._initialize()
-        self._stepsize = self._choose_stepsize()
+        if self._inner_solver == "FISTA":
+            self._stepsize = self._choose_stepsize()
 
     @property
     def proximal(self):
@@ -212,8 +229,7 @@ def validate_target(self):
         super().validate_target()
         if not isinstance(self.target.prior, (cuqi.implicitprior.RegularizedGaussian, cuqi.implicitprior.RegularizedGMRF)):
             raise TypeError("Prior needs to be RegularizedGaussian or RegularizedGMRF")
-        if not callable(self.proximal):
-            raise TypeError("Proximal needs to be callable")
+        self._inner_solver = "FISTA" if callable(self.proximal) else "ADMM"
 
     def _choose_stepsize(self):
         if isinstance(self.stepsize, str):
@@ -237,8 +253,16 @@ def prior(self):
 
     def step(self):
         y = self.b_tild + np.random.randn(len(self.b_tild))
-        sim = FISTA(self.M, y, self.proximal,
-                    self.current_point, maxit = self.maxit, stepsize = self._stepsize, abstol = self.abstol, adaptive = self.adaptive)         
+
+        if self._inner_solver == "FISTA":
+            sim = FISTA(self.M, y, self.proximal,
+                        self.current_point, maxit = self.maxit, stepsize = self._stepsize, abstol = self.abstol, adaptive = self.adaptive)         
+        elif self._inner_solver == "ADMM":
+            sim = ADMM(self.M, y, self.proximal,
+                        self.current_point, self.penalty_parameter, maxit = self.maxit, inner_max_it = self.inner_max_it, adaptive = self.adaptive)  
+        else:
+            raise ValueError("Choice of solver not supported.")
+
         self.current_point, _ = sim.solve()
         acc = 1
         return acc
diff --git a/cuqi/geometry/_geometry.py b/cuqi/geometry/_geometry.py
index 6361d357f..742ccf6ed 100644
--- a/cuqi/geometry/_geometry.py
+++ b/cuqi/geometry/_geometry.py
@@ -522,6 +522,9 @@ class Image2D(Geometry):
     Plotting is handled via matplotlib.pyplot.imshow.
     Colormap is defaulted to grayscale.
 
+    A How-To guide on the use of Image2D as the domain/range geometry of a CUQI 
+    :class:`Model` is available `here <https://cuqi-dtu.github.io/CUQIpy/user/_auto_howtos/myula.html>`_.
+
     Parameters
     -----------
     im_shape : tuple
diff --git a/cuqi/implicitprior/_regularizedGMRF.py b/cuqi/implicitprior/_regularizedGMRF.py
index c6096ed1e..a5c73d5d1 100644
--- a/cuqi/implicitprior/_regularizedGMRF.py
+++ b/cuqi/implicitprior/_regularizedGMRF.py
@@ -63,6 +63,7 @@ def __init__(self, mean=None, prec=None, bc_type='zero', order=1, proximal = Non
             
             # Underlying explicit GMRF
             self._gaussian = GMRF(mean, prec, bc_type=bc_type, order=order, **kwargs)
+            kwargs.pop("geometry", None)
             
             # Init from abstract distribution class
             super(Distribution, self).__init__(**kwargs)
diff --git a/cuqi/implicitprior/_regularizedGaussian.py b/cuqi/implicitprior/_regularizedGaussian.py
index b4d9d2909..f07b364cb 100644
--- a/cuqi/implicitprior/_regularizedGaussian.py
+++ b/cuqi/implicitprior/_regularizedGaussian.py
@@ -1,6 +1,8 @@
 from cuqi.utilities import get_non_default_args
 from cuqi.distribution import Distribution, Gaussian
 from cuqi.solver import ProjectNonnegative, ProjectBox, ProximalL1
+from cuqi.geometry import Continuous1D, Continuous2D, Image2D
+from cuqi.operator import FirstOrderFiniteDifference
 
 import numpy as np
 
@@ -39,17 +41,22 @@ class RegularizedGaussian(Distribution):
     sqrtprec
         See :class:`~cuqi.distribution.Gaussian` for details.
 
-    proximal : callable f(x, scale) or None
-        Euclidean proximal operator f of the regularization function g, that is, a solver for the optimization problem
-        min_z 0.5||x-z||_2^2+scale*g(x).
-
+    proximal : callable f(x, scale), list of tuples (callable proximal operator of f_i, linear operator L_i) or None
+        If callable:
+            Euclidean proximal operator f of the regularization function g, that is, a solver for the optimization problem
+            min_z 0.5||x-z||_2^2+scale*g(x).
+        If list of tuples (callable proximal operator of f_i, linear operator L_i):
+            Each callable proximal operator of f_i accepts two arguments (x, p) and should return the minimizer of p/2||x-z||^2 + f(x) over z for some f.
+            The corresponding regularization takes the form
+                sum_i f_i(L_i x),
+            where the sum ranges from 1 to an arbitrary n.
 
     projector : callable f(x) or None
         Euclidean projection onto the constraint C, that is, a solver for the optimization problem
         min_(z in C) 0.5||x-z||_2^2.
 
     constraint : string or None
-        Preset constraints. Can be set to "nonnegativity" and "box". Required for use in Gibbs.
+        Preset constraints that generate the corresponding proximal parameter. Can be set to "nonnegativity" and "box". Required for use in Gibbs.
         For "box", the following additional parameters can be passed:
             lower_bound : array_like or None
                 Lower bound of box, defaults to zero
@@ -57,10 +64,10 @@ class RegularizedGaussian(Distribution):
                 Upper bound of box, defaults to one
 
     regularization : string or None
-        Preset regularization. Can be set to "l1". Required for use in Gibbs in future update.
-        For "l1", the following additional parameters can be passed:
+        Preset regularization that generate the corresponding proximal parameter. Can be set to "l1" or 'tv'. Required for use in Gibbs in future update.
+        For "l1" or "tv", the following additional parameters can be passed:
             strength : scalar
-                Regularization parameter, i.e., strength*||x||_1 , defaults to one
+                Regularization parameter, i.e., strength*||Lx||_1, defaults to one
 
     """
         
@@ -75,6 +82,7 @@ def __init__(self, mean=None, cov=None, prec=None, sqrtcov=None, sqrtprec=None,
         
         # We init the underlying Gaussian first for geometry and dimensionality handling
         self._gaussian = Gaussian(mean=mean, cov=cov, prec=prec, sqrtcov=sqrtcov, sqrtprec=sqrtprec, **kwargs)
+        kwargs.pop("geometry", None)
 
         # Init from abstract distribution class
         super().__init__(**kwargs)
@@ -88,12 +96,6 @@ def _parse_regularization_input_arguments(self, proximal, projector, constraint,
         if (proximal is not None) + (projector is not None) + (constraint is not None) + (regularization is not None) != 1:
             raise ValueError("Precisely one of proximal, projector, constraint or regularization needs to be provided.")
 
-        if proximal is not None:
-            if not callable(proximal):
-                raise ValueError("Proximal needs to be callable.")
-            if len(get_non_default_args(proximal)) != 2:
-                raise ValueError("Proximal should take 2 arguments.")
-            
         if projector is not None:
             if not callable(projector):
                 raise ValueError("Projector needs to be callable.")
@@ -104,7 +106,8 @@ def _parse_regularization_input_arguments(self, proximal, projector, constraint,
         self._preset = None
         
         if proximal is not None:
-            self._proximal = proximal
+            # No need to generate the proximal and associated information
+            self.proximal = proximal
         elif projector is not None:
             self._proximal = lambda z, gamma: projector(z)
         elif (isinstance(constraint, str) and constraint.lower() == "nonnegativity"):
@@ -113,15 +116,48 @@ def _parse_regularization_input_arguments(self, proximal, projector, constraint,
         elif (isinstance(constraint, str) and constraint.lower() == "box"):
             lower = optional_regularization_parameters["lower_bound"]
             upper = optional_regularization_parameters["upper_bound"]
-            self._proximal = lambda z, gamma: ProjectBox(z, lower, upper)
+            self._proximal = lambda z, _: ProjectBox(z, lower, upper)
             self._preset = "box" # Not supported in Gibbs
         elif (isinstance(regularization, str) and regularization.lower() in ["l1"]):
-            strength = optional_regularization_parameters["strength"]
-            self._proximal = lambda z, gamma: ProximalL1(z, gamma*strength)
+            self._strength = optional_regularization_parameters["strength"]
+            self._proximal = lambda z, gamma: ProximalL1(z, gamma*self._strength)
             self._preset = "l1"
+        elif (isinstance(regularization, str) and regularization.lower() in ["tv"]):
+            self._strength = optional_regularization_parameters["strength"]
+            if isinstance(self.geometry, (Continuous1D, Continuous2D, Image2D)):
+                self._transformation = FirstOrderFiniteDifference(self.geometry.fun_shape, bc_type='neumann')
+            else:
+                raise ValueError("Geometry not supported for total variation")
+            
+            self._regularization_prox = lambda z, gamma: ProximalL1(z, gamma*self._strength)
+            self._regularization_oper = self._transformation
+
+            self._proximal = [(self._regularization_prox, self._regularization_oper)]
+            self._preset = "tv"
         else:
             raise ValueError("Regularization not supported")
 
+    
+    @property
+    def transformation(self):
+        return self._transformation
+    
+    @property
+    def strength(self):
+        return self._strength
+        
+    @strength.setter
+    def strength(self, value):
+        if self._preset not in self.regularization_options():
+            raise TypeError("Strength is only used when the regularization is set to l1 or TV.")
+
+        self._strength = value
+        if self._preset == "tv":
+            self._regularization_prox = lambda z, gamma: ProximalL1(z, gamma*self._strength)
+            self._proximal = [(self._regularization_prox, self._regularization_oper)]
+        elif self._preset == "l1":
+            self._proximal = lambda z, gamma: ProximalL1(z, gamma*self._strength)
+
     # This is a getter only attribute for the underlying Gaussian
     # It also ensures that the name of the underlying Gaussian
     # matches the name of the implicit regularized Gaussian
@@ -135,6 +171,25 @@ def gaussian(self):
     def proximal(self):
         return self._proximal
     
+    @proximal.setter
+    def proximal(self, value):
+        if callable(value):
+            if len(get_non_default_args(value)) != 2:
+                raise ValueError("Proximal should take 2 arguments.")
+        elif isinstance(value, list):
+            for (prox, op) in value:
+                if len(get_non_default_args(prox)) != 2:
+                    raise ValueError("Proximal should take 2 arguments.")
+                if op.shape[1] != self.geometry.par_dim:
+                    raise ValueError("Incorrect shape of linear operator in proximal list.")
+        else:
+            raise ValueError("Proximal needs to be callable or a list. See documentation.")
+        
+        self._proximal = value
+
+        # For all the presets, self._proximal is set directly, 
+        self._preset = None
+            
     @property
     def preset(self):
         return self._preset
@@ -154,7 +209,7 @@ def constraint_options():
 
     @staticmethod
     def regularization_options():
-        return ["l1"]
+        return ["l1", "tv"]
 
 
     # --- Defer behavior of the underlying Gaussian --- #
@@ -206,16 +261,18 @@ def sqrtcov(self):
     def sqrtcov(self, value):
         self.gaussian.sqrtcov = value     
     
-    def get_conditioning_variables(self):
-        return self.gaussian.get_conditioning_variables()
-    
     def get_mutable_variables(self):
-        return self.gaussian.get_mutable_variables()
+        mutable_vars = self.gaussian.get_mutable_variables().copy()
+        if self.preset in self.regularization_options():
+            mutable_vars += ["strength"]
+        return mutable_vars
     
     # Overwrite the condition method such that the underlying Gaussian is conditioned in general, except when conditioning on self.name
     # which means we convert Distribution to Likelihood or EvaluatedDensity.
     def _condition(self, *args, **kwargs):
-
+        if self.preset in self.regularization_options():
+            return super()._condition(*args, **kwargs)
+        
         # Handle positional arguments (similar code as in Distribution._condition)
         cond_vars = self.get_conditioning_variables()
         kwargs = self._parse_args_add_to_kwargs(cond_vars, *args, **kwargs)
@@ -275,7 +332,7 @@ class ConstrainedGaussian(RegularizedGaussian):
         min_(z in C) 0.5||x-z||_2^2.
 
     constraint : string or None
-        Preset constraints. Can be set to "nonnegativity" and "box". Required for use in Gibbs.
+        Preset constraints that generate the corresponding proximal parameter. Can be set to "nonnegativity" and "box". Required for use in Gibbs.
         For "box", the following additional parameters can be passed:
             lower_bound : array_like or None
                 Lower bound of box, defaults to zero
diff --git a/cuqi/solver/_solver.py b/cuqi/solver/_solver.py
index 6fd66f6aa..98a52a389 100644
--- a/cuqi/solver/_solver.py
+++ b/cuqi/solver/_solver.py
@@ -669,7 +669,7 @@ class ADMM(object):
         - flag=2 indicates multiplication of the transpose of A with vector x, that is  A.T @ x.
     b : ndarray.
     penalty_terms : List of tuples (callable proximal operator of f_i, linear operator L_i)
-        Each callable proximal operator f_i accepts two arguments (x, p) and should return the minimizer of p/2||x-z||^2 + f(x) over z for some f.
+        Each callable proximal operator of f_i accepts two arguments (x, p) and should return the minimizer of p/2||x-z||^2 + f(x) over z for some f.
     x0 : ndarray. Initial guess.
     penalty_parameter : Trade-off between linear least squares and regularization term in the solver iterates. Denoted as "rho" in [1].
     maxit : The maximum number of iterations.
diff --git a/demos/howtos/Image2D.py b/demos/howtos/Image2D.py
new file mode 100644
index 000000000..0a1ac893a
--- /dev/null
+++ b/demos/howtos/Image2D.py
@@ -0,0 +1,70 @@
+"""
+The use of Image2D
+==================
+
+This script demonstrates the use of :class:`~cuqi.geometry.Image2D` for domain and range geometries in CUQIpy forward models. The settings of this geometry, in particular the so-called `visual_only` flag, depend on how the forward model input is operated on in the forward model. `visual_only` flag tells the geometry to interpret the variable as a 2D image only when visualizing the variable but keep it as a vector when being operated on in the forward model. Generally, if the forward operator expects an image input, set `visual_only` to `False`, which is the default in CUQIpy. Otherwise, if it expects a vector input, set `visual_only` to `True` in the domain geometry. In the same way, the range geometry `visual_only` setting should be consistent with your forward operator output.
+"""
+
+# %%
+# Import necessary libraries
+import cuqi
+import numpy as np
+import matplotlib.pyplot as plt
+
+# %%
+# Define dimensions
+dim_x, dim_y = 2, 4
+
+# %%
+# Example illustrating the default behavior: `visual_only=False`
+# With `visual_only=False`, the underlying structure of a CUQI array or sample will be reshaped to the
+# shape of the specified geometry, i.e., an image of size (dim_x, dim_y), before being passed to the model.
+# Here we intend to define a forward operator that downsamples the input image by a factor of 2. It's easier
+# to define such a forward operator in terms of an image. We do not need to explicitly set `visual_only` to 
+#`False` because it is the default behavior.
+
+def forward_func_a(image):
+    assert image.shape == (dim_x, dim_y)
+    return image[::2, ::2]
+
+model_A = cuqi.model.LinearModel(forward_func_a,
+                                 domain_geometry=cuqi.geometry.Image2D((dim_x, dim_y)),
+                                 range_geometry=cuqi.geometry.Image2D((dim_x // 2, dim_y // 2)))
+
+print(model_A)
+
+x_sample = cuqi.array.CUQIarray(np.linspace(0, 1, dim_x * dim_y),
+                                geometry=cuqi.geometry.Image2D((dim_x, dim_y)))
+
+# Plot the original sample
+plt.figure()
+x_sample.plot(cmap='Greens')
+plt.title("input image x")
+
+# Apply the model and plot the result
+plt.figure()
+y_sample = model_A @ x_sample
+y_sample.plot(cmap='Greens')
+plt.title("output y=model_A(x)")
+
+# %%
+# Example with `visual_only=True`
+# With `visual_only=True`, the underlying structure of a cuqi array or sample will not be changed before
+# being passed to the model. This is useful when the model expects a vector input. Here we define a 
+# forward operator that reverses the input vector, which can be easily defined in terms of a vector.
+def forward_func_b(x):
+    assert x.shape == (dim_x * dim_y, )
+    return x[::-1]
+
+model_B = cuqi.model.LinearModel(forward_func_b,
+                                 domain_geometry=cuqi.geometry.Image2D((dim_x, dim_y), visual_only=True),
+                                 range_geometry=cuqi.geometry.Image2D((dim_x, dim_y), visual_only=True))
+
+# Note we explicitely set the `visual_only` flag to be `True` here.
+print(model_B)
+
+# Apply the model and plot the result
+plt.figure()
+y_sample = model_B @ x_sample
+y_sample.plot(cmap='Greens')
+plt.title("output y=model_B(x)")
\ No newline at end of file
diff --git a/demos/howtos/algebra.py b/demos/howtos/algebra.py
new file mode 100644
index 000000000..ad30b24b7
--- /dev/null
+++ b/demos/howtos/algebra.py
@@ -0,0 +1,66 @@
+"""
+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.
+
+"""
+
+#%%
+# 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`.
+
+from cuqi.experimental.algebra import VariableNode
+
+x = VariableNode('x')
+y = VariableNode('y')
+
+#%%
+# 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.
+
+print("Basic operations:")
+print(x + 1)
+print(x + y)
+print(x * y)
+print(x / y)
+
+print("\nComplex operations:")
+print(x**2 + 2*x*y + y**2)
+print((x + y)**2)
+
+print("\nProgrammatric operations:")
+print(x[2]+y[3])
+
+# %%
+# Utilizing the Computational Graph
+# ---------------------------------
+# 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.
+
+# Define a mathematical expression
+expr1 = (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)}")
+
+# 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)}")
+
+# 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])}")
+# %%


Test patch:
diff --git a/tests/test_implicit_priors.py b/tests/test_implicit_priors.py
index c4fd363c4..ad617a3c1 100644
--- a/tests/test_implicit_priors.py
+++ b/tests/test_implicit_priors.py
@@ -16,7 +16,7 @@ def test_RegularizedGaussian_guarding_statements():
         cuqi.implicitprior.RegularizedGaussian(np.zeros(5), 1, proximal=lambda s,z: s, constraint="nonnegativity")
 
     # Proximal
-    with pytest.raises(ValueError, match="Proximal needs to be callable"):
+    with pytest.raises(ValueError, match="Proximal needs to be callable or a list. See documentation."):
         cuqi.implicitprior.RegularizedGaussian(np.zeros(5), 1, proximal=1)
 
     with pytest.raises(ValueError, match="Proximal should take 2 arguments"):
@@ -104,3 +104,37 @@ def test_RegularizedUnboundedUniform_is_RegularizedGaussian():
     x = cuqi.implicitprior.RegularizedUnboundedUniform(cuqi.geometry.Continuous1D(5), regularization="l1", strength = 5.0)
     
     assert np.allclose(x.gaussian.sqrtprec, 0.0)
+
+def test_RegularizedGaussian_conditioning_constrained():
+    """ Test that conditioning the implicit regularized Gaussian works as expected """
+    
+    x  = cuqi.implicitprior.RegularizedGMRF(lambda a:a*np.ones(2**2),
+                     prec = lambda b:5*b,
+                     constraint = "nonnegativity",
+                     geometry = cuqi.geometry.Image2D((2,2)))
+    
+    assert x.get_mutable_variables() == ['mean', 'prec']
+    assert x.get_conditioning_variables() == ['a', 'b']
+
+    x = x(a=1, b=2)
+
+    assert np.allclose(x.mean, [1, 1, 1, 1])
+    assert np.allclose(x.prec, 10)
+
+def test_RegularizedGaussian_conditioning_strength():
+    """ Test that conditioning the implicit regularized Gaussian works as expected """
+    
+    x  = cuqi.implicitprior.RegularizedGMRF(lambda a:a*np.ones(2**2),
+                     prec = lambda b:5*b,
+                     regularization = "tv",
+                     strength = lambda c:c*2,
+                     geometry = cuqi.geometry.Image2D((2,2)))
+    
+    assert x.get_mutable_variables() == ['mean', 'prec', 'strength']
+    assert x.get_conditioning_variables() == ['a', 'b', 'c']
+
+    x = x(a=1, b=2, c=3)
+
+    assert np.allclose(x.mean, [1, 1, 1, 1])
+    assert np.allclose(x.prec, 10)
+    assert np.allclose(x.strength, 6)
diff --git a/tests/zexperimental/test_algebra.py b/tests/zexperimental/test_algebra.py
new file mode 100644
index 000000000..06717473d
--- /dev/null
+++ b/tests/zexperimental/test_algebra.py
@@ -0,0 +1,68 @@
+import pytest
+import numpy as np
+from cuqi.experimental.algebra import VariableNode
+
+
+@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_algebra_on_single_variables_is_correct(operations):
+    """ Tests that algebraic operations on variables are recorded correctly by comparing to the lambda expression """
+    X = VariableNode('x')
+    rv = operations(X)
+    val = np.random.randn(2)
+    # Compare random variable recorded operations vs actual operations
+    assert np.allclose(rv(x=val), 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_algebra_on_joint_variables_is_correct(operations):
+    """ Tests that algebraic operations on variables are recorded correctly by comparing to the lambda expression """
+    x = VariableNode('x')
+    y = VariableNode('y')
+    z = VariableNode('z')
+
+    # Define a variable in joint space
+    rv = operations(x, y, z)
+
+    # Check operations work
+    val_x = abs(np.random.randn(2))+1
+    val_y = abs(np.random.randn(2))+1
+    val_z = abs(np.random.randn(2))+1
+    assert np.allclose(rv(x=val_x, y=val_y, z=val_z), operations(val_x, val_y, val_z))
diff --git a/tests/zexperimental/test_mcmc.py b/tests/zexperimental/test_mcmc.py
index ad31b2387..177962dcd 100644
--- a/tests/zexperimental/test_mcmc.py
+++ b/tests/zexperimental/test_mcmc.py
@@ -330,12 +330,15 @@ def create_myula_target(dim=16):
     """Create a target for MYULA."""
     def func(x, restoration_strength=1):
         return x, True
-    likelihood = cuqi.testproblem.Deconvolution1D(
+    y = cuqi.testproblem.Deconvolution1D(
         dim=dim).posterior.likelihood 
-    restoration_prior = cuqi.implicitprior.RestorationPrior(
-        func, geometry=likelihood.model.domain_geometry)
+    x = cuqi.implicitprior.RestorationPrior(
+        func, geometry=y.model.domain_geometry)
+    # access name property to ensure RestorationPrior
+    # name is set 
+    x.name
     posterior = cuqi.distribution.Posterior(
-        likelihood, restoration_prior)
+        y, x)
     return posterior
 
 def create_myula_smoothed_target(dim=16):
@@ -368,6 +371,10 @@ def test_myula():
     assert samples_ula.Ns == 10
     assert samples_myula.Ns == 10
     assert np.allclose(samples_myula.samples, samples_ula.samples)
+    # Assert the name of the MoreauYoshidaPrior created by MYULA is
+    # the same name as the corresponding RestorationPrior (both are named x)
+    assert myula._smoothed_target.prior.name == 'x'
+    assert myula.target.prior.name == 'x'
 
 def test_myula_object_creation_fails_with_target_without_restore_method():
     """ Test that MYULA object creation fails with target that does not