CUQIpy 1.2.0.post0.dev314__tar.gz → 1.2.0.post0.dev342__tar.gz

{cuqipy-1.2.0.post0.dev314 → cuqipy-1.2.0.post0.dev342}/CUQIpy.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: CUQIpy
-Version: 1.2.0.post0.dev314
+Version: 1.2.0.post0.dev342
 Summary: Computational Uncertainty Quantification for Inverse problems in Python
 Maintainer-email: "Nicolai A. B. Riis" <nabr@dtu.dk>, "Jakob S. Jørgensen" <jakj@dtu.dk>, "Amal M. Alghamdi" <amaal@dtu.dk>, Chao Zhang <chaz@dtu.dk>
 License:                                  Apache License

{cuqipy-1.2.0.post0.dev314 → cuqipy-1.2.0.post0.dev342}/CUQIpy.egg-info/SOURCES.txt

@@ -45,6 +45,8 @@ cuqi/distribution/_smoothed_laplace.py
 cuqi/distribution/_truncated_normal.py
 cuqi/distribution/_uniform.py
 cuqi/experimental/__init__.py
+cuqi/experimental/algebra/__init__.py
+cuqi/experimental/algebra/_ast.py
 cuqi/experimental/mcmc/__init__.py
 cuqi/experimental/mcmc/_conjugate.py
 cuqi/experimental/mcmc/_conjugate_approx.py

{cuqipy-1.2.0.post0.dev314 → cuqipy-1.2.0.post0.dev342}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: CUQIpy
-Version: 1.2.0.post0.dev314
+Version: 1.2.0.post0.dev342
 Summary: Computational Uncertainty Quantification for Inverse problems in Python
 Maintainer-email: "Nicolai A. B. Riis" <nabr@dtu.dk>, "Jakob S. Jørgensen" <jakj@dtu.dk>, "Amal M. Alghamdi" <amaal@dtu.dk>, Chao Zhang <chaz@dtu.dk>
 License:                                  Apache License

{cuqipy-1.2.0.post0.dev314 → cuqipy-1.2.0.post0.dev342}/cuqi/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2024-11-20T10:51:18+0100",
+ "date": "2024-11-25T09:18:51+0100",
  "dirty": false,
  "error": null,
- "full-revisionid": "42e096c7066472ac4d5701fc5f2de53636246f5d",
- "version": "1.2.0.post0.dev314"
+ "full-revisionid": "63395d14f6c3f964633b20200ada13b8a213da20",
+ "version": "1.2.0.post0.dev342"
 }
 '''  # END VERSION_JSON
 

{cuqipy-1.2.0.post0.dev314 → cuqipy-1.2.0.post0.dev342}/cuqi/experimental/__init__.py

@@ -1,2 +1,3 @@
 """ Experimental module for testing new features and ideas. """
 from . import mcmc
+from . import algebra

cuqipy-1.2.0.post0.dev342/cuqi/experimental/algebra/__init__.py

@@ -0,0 +1 @@
+from ._ast import VariableNode

cuqipy-1.2.0.post0.dev342/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)

{cuqipy-1.2.0.post0.dev314 → cuqipy-1.2.0.post0.dev342}/cuqi/experimental/mcmc/_rto.py

@@ -3,7 +3,7 @@ from scipy.linalg.interpolative import estimate_spectral_norm
 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 @@ class RegularizedLinearRTO(LinearRTO):
         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 @@ class RegularizedLinearRTO(LinearRTO):
         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 @@ class RegularizedLinearRTO(LinearRTO):
 
     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

{cuqipy-1.2.0.post0.dev314 → cuqipy-1.2.0.post0.dev342}/cuqi/implicitprior/_regularizedGMRF.py

@@ -63,6 +63,7 @@ class RegularizedGMRF(RegularizedGaussian):
             
             # 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)

{cuqipy-1.2.0.post0.dev314 → cuqipy-1.2.0.post0.dev342}/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 @@ class RegularizedGaussian(Distribution):
         
         # 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 @@ class RegularizedGaussian(Distribution):
         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 @@ class RegularizedGaussian(Distribution):
         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 @@ class RegularizedGaussian(Distribution):
         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 @@ class RegularizedGaussian(Distribution):
     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 @@ class RegularizedGaussian(Distribution):
 
     @staticmethod
     def regularization_options():
-        return ["l1"]
+        return ["l1", "tv"]
 
 
     # --- Defer behavior of the underlying Gaussian --- #
@@ -206,16 +261,18 @@ class RegularizedGaussian(Distribution):
     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

{cuqipy-1.2.0.post0.dev314 → cuqipy-1.2.0.post0.dev342}/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.

{cuqipy-1.2.0.post0.dev314 → cuqipy-1.2.0.post0.dev342}/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)