Repository: CUQI-DTU/CUQIpy

Issue:
adding ADMM solver

Discussion:

Solution patch:
diff --git a/cuqi/distribution/__init__.py b/cuqi/distribution/__init__.py
index aefebc9c1..3a51f5472 100644
--- a/cuqi/distribution/__init__.py
+++ b/cuqi/distribution/__init__.py
@@ -12,6 +12,7 @@
 from ._smoothed_laplace import SmoothedLaplace
 from ._lognormal import Lognormal
 from ._normal import Normal
+from ._truncated_normal import TruncatedNormal
 from ._posterior import Posterior
 from ._uniform import Uniform
 from ._custom import UserDefinedDistribution, DistributionGallery
diff --git a/cuqi/distribution/_normal.py b/cuqi/distribution/_normal.py
index b101c4788..90074dc2f 100644
--- a/cuqi/distribution/_normal.py
+++ b/cuqi/distribution/_normal.py
@@ -1,5 +1,8 @@
 import numpy as np
+import numbers
 from scipy.special import erf
+from cuqi.geometry import _get_identity_geometries
+from cuqi.utilities import force_ndarray
 from cuqi.distribution import Distribution
 
 class Normal(Distribution):
@@ -27,6 +30,24 @@ def __init__(self, mean=None, std=None, is_symmetric=True, **kwargs):
         self.mean = mean
         self.std = std
 
+    @property
+    def mean(self):
+        """ Mean of the distribution """
+        return self._mean
+
+    @mean.setter
+    def mean(self, value):
+        self._mean = force_ndarray(value, flatten=True)
+
+    @property
+    def std(self):
+        """ Std of the distribution """
+        return self._std
+
+    @std.setter
+    def std(self, value):
+        self._std = force_ndarray(value, flatten=True)
+
     def pdf(self, x):
         return np.prod(1/(self.std*np.sqrt(2*np.pi))*np.exp(-0.5*((x-self.mean)/self.std)**2))
 
@@ -36,6 +57,19 @@ def logpdf(self, x):
     def cdf(self, x):
         return np.prod(0.5*(1 + erf((x-self.mean)/(self.std*np.sqrt(2)))))
 
+    def _gradient(self, val, *args, **kwargs):
+        if not type(self.geometry) in _get_identity_geometries():
+            raise NotImplementedError("Gradient not implemented for distribution {} with geometry {}".format(self,self.geometry))
+        if not callable(self.mean):
+            return -(val-self.mean)/(self.std**2)
+        elif hasattr(self.mean, "gradient"): # for likelihood
+            model = self.mean
+            dev = val - model.forward(*args, **kwargs)
+            print(dev)
+            return model.gradient(1.0/(np.array(self.std)) @ dev, *args, **kwargs)
+        else:
+            raise NotImplementedError("Gradient not implemented for distribution {} with location {}".format(self,self.mean))
+
     def _sample(self,N=1, rng=None):
 
         """
diff --git a/cuqi/distribution/_truncated_normal.py b/cuqi/distribution/_truncated_normal.py
new file mode 100644
index 000000000..ff9d253d7
--- /dev/null
+++ b/cuqi/distribution/_truncated_normal.py
@@ -0,0 +1,129 @@
+import numpy as np
+from scipy.special import erf
+from cuqi.utilities import force_ndarray
+from cuqi.distribution import Distribution
+from cuqi.distribution import Normal
+
+class TruncatedNormal(Distribution):
+    """
+    Truncated Normal probability distribution.
+    
+    Generates instance of cuqi.distribution.TruncatedNormal. 
+    It allows the user to specify upper and lower bounds on random variables 
+    represented by a Normal distribution. This distribution is suitable for a 
+    small dimension setup (e.g. `dim`=3 or 4). Using TruncatedNormal 
+    Distribution with a larger dimension can lead to a high rejection rate when
+    used within MCMC samplers.
+    
+    The variables of this distribution are iid.
+
+    
+    Parameters
+    ------------
+    mean : float or array_like of floats
+        mean of distribution
+    std : float or array_like of floats
+        standard deviation
+    low : float or array_like of floats
+        lower bound of the distribution
+    high : float or array_like of floats
+        upper bound of the distribution
+    
+    Example
+    -----------
+    .. code-block:: python
+
+        #Generate Normal with mean 0, standard deviation 1 and bounds [-2,2]
+        p = cuqi.distribution.TruncatedNormal(mean=0, std=1, low=-2, high=2)
+        samples = p.sample(5000)
+    """
+    def __init__(self, mean=None, std=None, low=-np.Inf, high=np.Inf, is_symmetric=False, **kwargs):
+        # Init from abstract distribution class
+        super().__init__(is_symmetric=is_symmetric, **kwargs)  
+
+        # Init specific to this distribution
+        self.mean = mean
+        self.std = std
+        self.low = low
+        self.high = high
+
+        # Init underlying normal distribution
+        self._normal = Normal(self.mean, self.std, is_symmetric=True, **kwargs)
+
+    @property
+    def mean(self):
+        """ Mean of the distribution """
+        return self._mean
+
+    @mean.setter
+    def mean(self, value):
+        self._mean = force_ndarray(value, flatten=True)
+        if hasattr(self, '_normal'):
+            self._normal.mean = self._mean
+
+    @property
+    def std(self):
+        """ Std of the distribution """
+        return self._std
+
+    @std.setter
+    def std(self, value):
+        self._std = force_ndarray(value, flatten=True)
+        if hasattr(self, '_normal'):
+            self._normal.std = self._std
+
+    @property
+    def low(self):
+        """ Lower bound of the distribution """
+        return self._low
+
+    @low.setter
+    def low(self, value):
+        self._low = force_ndarray(value, flatten=True)
+
+    @property
+    def high(self):
+        """ Higher bound of the distribution """
+        return self._high
+
+    @high.setter
+    def high(self, value):
+        self._high = force_ndarray(value, flatten=True)
+
+    def logpdf(self, x):
+        """
+        Computes the unnormalized logpdf at the given values of x.
+        """
+        # the unnormalized logpdf
+        # check if x falls in the range between np.array a and b
+        if np.any(x < self.low) or np.any(x > self.high):
+            return -np.Inf
+        else:
+             return self._normal.logpdf(x)
+
+    def _gradient(self, x, *args, **kwargs):
+        """
+        Computes the gradient of the unnormalized logpdf at the given values of x.
+        """
+        # check if x falls in the range between np.array a and b
+        if np.any(x < self.low) or np.any(x > self.high):
+            return np.NaN*np.ones_like(x)
+        else:
+            return self._normal.gradient(x, *args, **kwargs)
+
+    def _sample(self, N=1, rng=None):
+        """
+        Generates random samples from the distribution.
+        """
+        max_iter = 1e9 # maximum number of trials to avoid infinite loop
+        samples = []
+        for i in range(int(max_iter)):
+            if len(samples) == N:
+                break
+            sample = self._normal.sample(1,rng)
+            if np.all(sample >= self.low) and np.all(sample <= self.high):
+                samples.append(sample)
+        # raise a error if the number of iterations exceeds max_iter
+        if i == max_iter-1:
+            raise RuntimeError("Failed to generate {} samples within {} iterations".format(N, max_iter))
+        return np.array(samples).T.reshape(-1,N)
\ No newline at end of file
diff --git a/cuqi/experimental/mcmc/_rto.py b/cuqi/experimental/mcmc/_rto.py
index e6e4a2544..55defd043 100644
--- a/cuqi/experimental/mcmc/_rto.py
+++ b/cuqi/experimental/mcmc/_rto.py
@@ -235,8 +235,8 @@ def prior(self):
 
     def step(self):
         y = self.b_tild + np.random.randn(len(self.b_tild))
-        sim = FISTA(self.M, y, self.current_point, self.proximal,
-                    maxit = self.maxit, stepsize = self._stepsize, abstol = self.abstol, adaptive = self.adaptive)         
+        sim = FISTA(self.M, y, self.proximal,
+                    self.current_point, maxit = self.maxit, stepsize = self._stepsize, abstol = self.abstol, adaptive = self.adaptive)         
         self.current_point, _ = sim.solve()
         acc = 1
         return acc
diff --git a/cuqi/sampler/_rto.py b/cuqi/sampler/_rto.py
index c83d6e3d2..6cb549c8b 100644
--- a/cuqi/sampler/_rto.py
+++ b/cuqi/sampler/_rto.py
@@ -267,8 +267,8 @@ def _sample(self, N, Nb):
         samples[:, 0] = self.x0
         for s in range(Ns-1):
             y = self.b_tild + np.random.randn(len(self.b_tild))
-            sim = FISTA(self.M, y, samples[:, s], self.proximal,
-                        maxit = self.maxit, stepsize = _stepsize, abstol = self.abstol, adaptive = self.adaptive)         
+            sim = FISTA(self.M, y, self.proximal,
+                        samples[:, s], maxit = self.maxit, stepsize = _stepsize, abstol = self.abstol, adaptive = self.adaptive)         
             samples[:, s+1], _ = sim.solve()
             
             self._print_progress(s+2,Ns) #s+2 is the sample number, s+1 is index assuming x0 is the first sample
diff --git a/cuqi/solver/__init__.py b/cuqi/solver/__init__.py
index 50cb1fcf8..5b9e11a19 100644
--- a/cuqi/solver/__init__.py
+++ b/cuqi/solver/__init__.py
@@ -7,6 +7,7 @@
     LM,
     PDHG,
     FISTA,
+    ADMM,
     ProjectNonnegative,
     ProjectBox,
     ProximalL1
diff --git a/cuqi/solver/_solver.py b/cuqi/solver/_solver.py
index cb64f4a05..6fd66f6aa 100644
--- a/cuqi/solver/_solver.py
+++ b/cuqi/solver/_solver.py
@@ -584,8 +584,8 @@ class FISTA(object):
     ----------
     A : ndarray or callable f(x,*args).
     b : ndarray.
-    x0 : ndarray. Initial guess.
     proximal : callable f(x, gamma) for proximal mapping.
+    x0 : ndarray. Initial guess.
     maxit : The maximum number of iterations.
     stepsize : The stepsize of the gradient step.
     abstol : The numerical tolerance for convergence checks.
@@ -606,11 +606,11 @@ class FISTA(object):
         b = rng.standard_normal(m)
         stepsize = 0.99/(sp.linalg.interpolative.estimate_spectral_norm(A)**2)
         x0 = np.zeros(n)
-        fista = FISTA(A, b, x0, proximal = ProximalL1, stepsize = stepsize, maxit = 100, abstol=1e-12, adaptive = True)
+        fista = FISTA(A, b, proximal = ProximalL1, x0, stepsize = stepsize, maxit = 100, abstol=1e-12, adaptive = True)
         sol, _ = fista.solve()
 
     """  
-    def __init__(self, A, b, x0, proximal, maxit=100, stepsize=1e0, abstol=1e-14, adaptive = True):
+    def __init__(self, A, b, proximal, x0, maxit=100, stepsize=1e0, abstol=1e-14, adaptive = True):
         
         self.A = A
         self.b = b
@@ -650,8 +650,157 @@ def solve(self):
                 x_new = x_new + ((k-1)/(k+2))*(x_new - x_old)
               
             x = x_new.copy()
+
+class ADMM(object):
+    """Alternating Direction Method of Multipliers for solving regularized linear least squares problems of the form:
+    Minimize ||Ax-b||^2 + sum_i f_i(L_i x),
+    where the sum ranges from 1 to an arbitrary n. See definition of the parameter `penalty_terms` below for more details about f_i and L_i
+
+    Reference:
+    [1] Boyd et al. "Distributed optimization and statistical learning via the alternating direction method of multipliers."Foundations and Trends® in Machine learning, 2011.
+
+    
+    Parameters
+    ----------
+    A : ndarray or callable
+        Represents a matrix or a function that performs matrix-vector multiplications. 
+        When A is a callable, it accepts arguments (x, flag) where:
+        - flag=1 indicates multiplication of A with vector x, that is A @ x.
+        - 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.
+    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.
+    adaptive : Whether to adaptively update the penalty_parameter each iteration such that the primal and dual residual norms are of the same order of magnitude. Based on [1], Subsection 3.4.1
+    
+    Example
+    -----------
+    .. code-block:: python
     
+        from cuqi.solver import ADMM, ProximalL1, ProjectNonnegative
+        import numpy as np
+
+        rng = np.random.default_rng()
+
+        m, n, k = 10, 5, 4
+        A = rng.standard_normal((m, n))
+        b = rng.standard_normal(m)
+        L = rng.standard_normal((k, n))
+
+        x0 = np.zeros(n)
+        admm = ADMM(A, b, x0, penalty_terms = [(ProximalL1, L), (lambda z, _ : ProjectNonnegative(z), np.eye(n))], tradeoff = 10)
+        sol, _ = admm.solve()
+
+    """  
+
+    def __init__(self, A, b, penalty_terms, x0, penalty_parameter = 10, maxit = 100, inner_max_it = 10, adaptive = True):
+
+        self.A = A
+        self.b = b
+        self.x_cur = x0
+
+        dual_len = [penalty[1].shape[0] for penalty in penalty_terms]
+        self.z_cur = [np.zeros(l) for l in dual_len]
+        self.u_cur = [np.zeros(l) for l in dual_len]
+        self.n = penalty_terms[0][1].shape[1]
+        
+        self.rho = penalty_parameter
+        self.maxit = maxit
+        self.inner_max_it = inner_max_it
+        self.adaptive = adaptive
+
+        self.penalty_terms = penalty_terms
+       
+        self.p = len(self.penalty_terms)
+        self._big_matrix = None
+        self._big_vector = None
+
+    def solve(self):
+        """
+        Solves the regularized linear least squares problem using ADMM in scaled form. Based on [1], Subsection 3.1.1
+        """
+        z_new = self.p*[0]
+        u_new = self.p*[0]
+
+        # Iterating
+        for i in range(self.maxit):
+            self._iteration_pre_processing()
+
+            # Main update (Least Squares)
+            solver = CGLS(self._big_matrix, self._big_vector, self.x_cur, self.inner_max_it)
+            x_new, _ = solver.solve()
+        
+            # Regularization update
+            for j, penalty in enumerate(self.penalty_terms):
+                z_new[j] = penalty[0](penalty[1]@x_new + self.u_cur[j], 1.0/self.rho)
+                
+            res_primal = 0.0
+            # Dual update
+            for j, penalty in enumerate(self.penalty_terms):
+                r_partial = penalty[1]@x_new - z_new[j]
+                res_primal += LA.norm(r_partial)**2
+
+                u_new[j] = self.u_cur[j] + r_partial
+            
+            res_dual = 0.0
+            for j, penalty in enumerate(self.penalty_terms):
+                res_dual += LA.norm(penalty[1].T@(z_new[j] - self.z_cur[j]))**2
+
+            # Adaptive approach based on [1], Subsection 3.4.1
+            if self.adaptive:
+                if res_dual > 1e2*res_primal:
+                    self.rho *= 0.5 # More regularization
+                elif res_primal > 1e2*res_dual:
+                    self.rho *= 2.0 # More data fidelity
+
+            self.x_cur, self.z_cur, self.u_cur = x_new, z_new.copy(), u_new
+            
+        return self.x_cur, i
     
+    def _iteration_pre_processing(self):
+            """ Preprocessing
+            Every iteration of ADMM requires solving a linear least squares system of the form
+                minimize 1/(rho) \|Ax-b\|_2^2 + sum_{i=1}^{p} \|penalty[1]x - (y - u)\|_2^2
+            To solve this, all linear least squares terms are combined into a single big term
+            with matrix big_matrix and data big_vector.
+
+            The matrix only needs to be updated when rho changes, i.e., when the adaptive option is used.
+            The data vector needs to be updated every iteration.
+            """
+
+            self._big_vector = np.hstack([np.sqrt(1/self.rho)*self.b] + [self.z_cur[i] - self.u_cur[i] for i in range(self.p)])
+
+            # Check whether matrix needs to be updated
+            if self._big_matrix is not None and not self.adaptive:
+                return
+
+            # Update big_matrix
+            if callable(self.A):
+                def matrix_eval(x, flag):
+                    if flag == 1:
+                        out1 = np.sqrt(1/self.rho)*self.A(x, 1)
+                        out2 = [penalty[1]@x for penalty in self.penalty_terms]
+                        out  = np.hstack([out1] + out2)
+                    elif flag == 2:
+                        idx_start = len(x)
+                        idx_end = len(x)
+                        out1 = np.zeros(self.n)
+                        for _, t in reversed(self.penalty_terms):
+                            idx_start -= t.shape[0]
+                            out1 += t.T@x[idx_start:idx_end]
+                            idx_end = idx_start
+                        out2 = np.sqrt(1/self.rho)*self.A(x[:idx_end], 2)
+                        out  = out1 + out2     
+                    return out
+                self._big_matrix = matrix_eval
+            else:
+                self._big_matrix = np.vstack([np.sqrt(1/self.rho)*self.A] + [penalty[1] for penalty in self.penalty_terms])
+
+
+
+
 def ProjectNonnegative(x):
     """(Euclidean) projection onto the nonnegative orthant.
     
@@ -678,6 +827,22 @@ def ProjectBox(x, lower = None, upper = None):
     
     return np.minimum(np.maximum(x, lower), upper)
 
+def ProjectHalfspace(x, a, b):
+    """(Euclidean) projection onto the halfspace defined {z|<a,z> <= b}.
+    
+    Parameters
+    ----------
+    x : array_like.
+    a : array_like.
+    b : array_like.
+    """  
+
+    ax_b = np.inner(a,x) - b
+    if ax_b <= 0:
+        return x
+    else:
+        return x - (ax_b/np.inner(a,a))*a
+
 def ProximalL1(x, gamma):
     """(Euclidean) proximal operator of the \|x\|_1 norm.
     Also known as the shrinkage or soft thresholding operator.
@@ -687,4 +852,4 @@ def ProximalL1(x, gamma):
     x : array_like.
     gamma : scale parameter.
     """
-    return np.multiply(np.sign(x), np.maximum(np.abs(x)-gamma, 0))
+    return np.multiply(np.sign(x), np.maximum(np.abs(x)-gamma, 0))
\ No newline at end of file
diff --git a/demos/demo38_truncated_normal.py b/demos/demo38_truncated_normal.py
new file mode 100644
index 000000000..7f3b07219
--- /dev/null
+++ b/demos/demo38_truncated_normal.py
@@ -0,0 +1,78 @@
+# %% Initialize and import CUQIpy
+import sys
+sys.path.append("..") 
+import numpy as np
+import cuqi
+from cuqi.utilities import plot_2D_density
+import matplotlib.pyplot as plt
+
+# %% Demo 1
+# This demo shows the use of gradient-based MCMC methods in drawing samples 
+# from a truncated normal distribution. Here we use MALA.
+p = cuqi.distribution.TruncatedNormal(
+    mean=np.array([0, 0]),
+    std=np.array([1, 1]),
+    low=np.array([0, -np.Inf]),
+    high=np.array([np.Inf, 0]))
+
+plt.figure()
+plot_2D_density(p, -5, 5, -5, 5)
+plt.title("Exact PDF")
+
+sampler = cuqi.experimental.mcmc.MALA(p, scale=0.1, initial_point= np.array([1,-1]))
+sampler.sample(10000)
+samples = sampler.get_samples()
+
+plt.figure()
+samples.plot_pair()
+plt.xlim(-5, 5)
+plt.ylim(-5, 5)
+plt.gca().set_aspect('equal')
+plt.title("Samples")
+
+# %% Demo 2
+# This demo shows the use of truncted normal as prior in solving the simplest
+# BIP. Here again we use MALA.
+np.random.seed(0)
+# the forward model
+A_matrix = np.array([[1.0, 1.0]])
+A = cuqi.model.LinearModel(A_matrix)
+x = cuqi.distribution.TruncatedNormal(
+    mean=np.array([0, 0]),
+    std=np.array([1, 1]),
+    low=np.array([0, -np.Inf]),
+    high=np.array([np.Inf, 0]))
+# the data distribution
+b = cuqi.distribution.Gaussian(A@x, 0.1)
+
+# the observed data
+particular_x = np.array([1.5, 1.5])
+b_given_particular_x = b(x=particular_x)
+b_obs = b_given_particular_x.sample()
+print(b_obs)
+
+# the posterior
+joint = cuqi.distribution.JointDistribution(x, b)
+post = joint(b=b_obs)
+
+# MALA sampler
+sampler = cuqi.experimental.mcmc.MALA(post, initial_point=np.array([2.5, -2.5]), scale=0.03)
+
+sampler.warmup(1000)
+sampler.sample(1000)
+samples = sampler.get_samples().burnthin(1000)
+samples.plot_trace()
+
+# plot exact posterior distribution and samples
+# the posterior PDF
+plt.figure()
+plot_2D_density(post, 1, 5, -3, 1)
+plt.title("Exact Posterior PDF")
+
+# samples
+plt.figure()
+samples.plot_pair()
+plt.xlim(1, 5)
+plt.ylim(-3, 1)
+plt.gca().set_aspect('equal')
+plt.title("Samples")
diff --git a/test.ipynb b/test.ipynb
new file mode 100644
index 000000000..cde176af6
--- /dev/null
+++ b/test.ipynb
@@ -0,0 +1,153 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[-3.99513417e-03 -1.32339656e-01 -4.52822633e-02 -7.44973888e-02\n",
+      " -3.35005208e-11]\n"
+     ]
+    }
+   ],
+   "source": [
+    "import numpy as np\n",
+    "import scipy as sp\n",
+    "\n",
+    "from cuqi.solver import CGLS, LM, FISTA, ADMM, ProximalL1, ProjectNonnegative\n",
+    "from scipy.optimize import lsq_linear\n",
+    "\n",
+    "\n",
+    "def test_ADMM_matrix_form():\n",
+    "    # Parameters\n",
+    "    rng = np.random.default_rng(seed = 42)\n",
+    "    m, n = 10, 5\n",
+    "    A = rng.standard_normal((m, n))\n",
+    "    b = rng.standard_normal(m)\n",
+    "    \n",
+    "    k = 4\n",
+    "    L = rng.standard_normal((k, n))\n",
+    "\n",
+    "    x0 = np.zeros(n)\n",
+    "    sol, _ = ADMM(A, b, [(ProximalL1, np.eye(n)), (lambda z, _ : ProjectNonnegative(z), L)],\n",
+    "                   x0, 10, maxit = 100, adaptive = True).solve()\n",
+    "\n",
+    "    print(sol)\n",
+    "    ref_sol = np.array([-3.99513417e-03, -1.32339656e-01, -4.52822633e-02, -7.44973888e-02, -3.35005208e-11])\n",
+    "    # Compare\n",
+    "    assert np.allclose(sol, ref_sol, atol=1e-4)\n",
+    "\n",
+    "test_ADMM_matrix_form()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[-3.99513417e-03 -1.32339656e-01 -4.52822633e-02 -7.44973888e-02\n",
+      " -3.35005152e-11]\n"
+     ]
+    }
+   ],
+   "source": [
+    "\n",
+    "\n",
+    "def test_ADMM_function_form():\n",
+    "    # Parameters\n",
+    "    rng = np.random.default_rng(seed = 42)\n",
+    "    m, n = 10, 5\n",
+    "    A = rng.standard_normal((m, n))\n",
+    "    def A_fun(x, flag):\n",
+    "        if flag == 1:\n",
+    "            return A@x\n",
+    "        if flag == 2:\n",
+    "            return A.T@x\n",
+    "        \n",
+    "    b = rng.standard_normal(m)\n",
+    "    \n",
+    "    k = 4\n",
+    "    L = rng.standard_normal((k, n))\n",
+    "\n",
+    "    x0 = np.zeros(n)\n",
+    "    sol, _ = ADMM(A_fun, b, [(ProximalL1, np.eye(n)), (lambda z, _ : ProjectNonnegative(z), L)],\n",
+    "                   x0, 10, maxit = 100, adaptive = True).solve()\n",
+    "\n",
+    "    print(sol)\n",
+    "    ref_sol = np.array([-3.99513417e-03, -1.32339656e-01, -4.52822633e-02, -7.44973888e-02, -3.35005208e-11])\n",
+    "    # Compare\n",
+    "    assert np.allclose(sol, ref_sol, atol=1e-4)\n",
+    "\n",
+    "test_ADMM_function_form()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[-3.99513417e-03 -1.32339656e-01 -4.52822633e-02 -7.44973888e-02\n",
+      " -3.35005208e-11]\n"
+     ]
+    }
+   ],
+   "source": [
+    "\n",
+    "# Parameters\n",
+    "rng = np.random.default_rng(seed = 42)\n",
+    "m, n = 10, 5\n",
+    "A = rng.standard_normal((m, n))\n",
+    "b = rng.standard_normal(m)\n",
+    "    \n",
+    "k = 4\n",
+    "L = rng.standard_normal((k, n))\n",
+    "\n",
+    "x0 = np.zeros(n)\n",
+    "sol, _ = ADMM(A, b, [(ProximalL1, np.eye(n)), (lambda z, _ : ProjectNonnegative(z), L)],\n",
+    "                   x0, 10, maxit = 100, adaptive = False).solve()\n",
+    "\n",
+    "print(sol)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "base",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}


Test patch:
diff --git a/tests/test_distribution.py b/tests/test_distribution.py
index e22427a5e..f91155895 100644
--- a/tests/test_distribution.py
+++ b/tests/test_distribution.py
@@ -28,6 +28,73 @@ def test_Normal_sample_regression(mean,var,expected):
     target = np.array(expected).T
     assert np.allclose( samples.samples, target)
 
+@pytest.mark.parametrize("mean,std,points,expected",[
+                            (0,1,[-1,0,1],[1,0,-1]),
+   (np.array([0,0]),np.array([1,1]),[[-1,0],[0,0],[0,-1]], [[1,0],[0,0],[0,1]])])
+def test_Normal_gradient(mean,std,points,expected):
+    p = cuqi.distribution.Normal(mean,std)
+    for point, grad in zip(points, expected):
+        assert np.allclose(p.gradient(point), grad)
+
+@pytest.mark.parametrize("mean,std,low,high,points",[(0.0, 
+                                               1.0,
+                                               -1.0,
+                                               1.0,
+                                               [-1.5, -0.5, 0.5, 1.5]),
+                                              (np.array([0.0, 0.0]), 
+                                               np.array([1.0, 1.0]),
+                                               np.array([-1.0, -1.0]),
+                                               np.array([1.0, 1.0]),
+                                               [np.array([-0.5, 0.0]),
+                                                np.array([0.5, 0.0]),
+                                                np.array([-2.0, 0.0]),
+                                                np.array([2.0, 0.0])])])
+def test_TruncatedNormal_logpdf(mean,std,low,high,points):
+    x_trun = cuqi.distribution.TruncatedNormal(mean,std,low=low,high=high)
+    x = cuqi.distribution.Normal(mean,std)
+    for point in points:
+        if np.all(point >= low) and np.all(point <= high):
+            assert x_trun.logpdf(point) == approx(x.logpdf(point))
+        else:
+            assert np.isneginf(x_trun.logpdf(point))
+
+@pytest.mark.parametrize("mean,std,low,high,points",[(0.0, 
+                                               1.0,
+                                               -1.0,
+                                               1.0,
+                                               [-1.5, -0.5, 0.5, 1.5]),
+                                              (np.array([0.0, 0.0]), 
+                                               np.array([1.0, 1.0]),
+                                               np.array([-1.0, -1.0]),
+                                               np.array([1.0, 1.0]),
+                                               [np.array([-0.5, 0.0]),
+                                                np.array([0.5, 0.0]),
+                                                np.array([-2.0, 0.0]),
+                                                np.array([2.0, 0.0])])])
+def test_TruncatedNormal_gradient(mean,std,low,high,points):
+    x_trun = cuqi.distribution.TruncatedNormal(mean,std,low=low,high=high)
+    x = cuqi.distribution.Normal(mean,std)
+    for point in points:
+        if np.all(point >= low) and np.all(point <= high):
+            assert np.all(x_trun.gradient(point) == approx(x.gradient(point)))
+        else:
+            assert np.all(np.isnan(x_trun.gradient(point)))
+
+@pytest.mark.parametrize("mean,std,low,high",[(0.0, 
+                                               1.0,
+                                               -1.0,
+                                               1.0),
+                                              (np.array([0.0, 0.0]), 
+                                               np.array([1.0, 1.0]),
+                                               np.array([-1.0, -1.0]),
+                                               np.array([1.0, 1.0]))])
+def test_TruncatedNormal_sampling(mean,std,low,high):
+    x = cuqi.distribution.TruncatedNormal(mean,std,low=low,high=high)
+    samples = x.sample(10000).samples
+    for i in range(samples.shape[1]):
+        sample = samples[:,i]
+        assert np.all(sample >= low) and np.all(sample <= high)
+
 def test_Gaussian_mean():
     mean = np.array([0, 0])
     std = np.array([1, 1])
diff --git a/tests/test_solver.py b/tests/test_solver.py
index 98177b0dc..8092ee208 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, ProximalL1
+from cuqi.solver import CGLS, LM, FISTA, ADMM, ProximalL1, ProjectNonnegative
 from scipy.optimize import lsq_linear
 
 
@@ -54,8 +54,52 @@ def test_FISTA():
     
     stepsize = 0.99/(sp.linalg.interpolative.estimate_spectral_norm(A)**2)
     x0 = np.zeros(n)
-    sol, _ = FISTA(A, b, x0, proximal = ProximalL1, stepsize = stepsize, maxit = 100, abstol=1e-12, adaptive = True).solve()
+    sol, _ = FISTA(A, b, ProximalL1, x0, stepsize = stepsize, maxit = 100, abstol=1e-12, adaptive = True).solve()
 
     ref_sol = np.array([-1.83273787e-03, -1.72094582e-13,  0.0, -3.35835639e-01, -1.27795593e-01])
     # Compare
-    assert np.allclose(sol, ref_sol, atol=1e-4)
\ No newline at end of file
+    assert np.allclose(sol, ref_sol, atol=1e-4)
+
+def test_ADMM_matrix_form():
+    # Parameters
+    rng = np.random.default_rng(seed = 42)
+    m, n = 10, 5
+    A = rng.standard_normal((m, n))
+    b = rng.standard_normal(m)
+    
+    k = 4
+    L = rng.standard_normal((k, n))
+
+    x0 = np.zeros(n)
+    sol, _ = ADMM(A, b, [(ProximalL1, np.eye(n)), (lambda z, _ : ProjectNonnegative(z), L)],
+                   x0, 10, maxit = 100, adaptive = True).solve()
+
+    ref_sol = np.array([-3.99513417e-03, -1.32339656e-01, -4.52822633e-02, -7.44973888e-02, -3.35005208e-11])
+    # Compare
+    assert np.allclose(sol, ref_sol, atol=1e-4)
+
+
+def test_ADMM_function_form():
+    # Parameters
+    rng = np.random.default_rng(seed = 42)
+    m, n = 10, 5
+    A = rng.standard_normal((m, n))
+    def A_fun(x, flag):
+        if flag == 1:
+            return A@x
+        if flag == 2:
+            return A.T@x
+        
+    b = rng.standard_normal(m)
+    
+    k = 4
+    L = rng.standard_normal((k, n))
+
+    x0 = np.zeros(n)
+    sol, _ = ADMM(A_fun, b, [(ProximalL1, np.eye(n)), (lambda z, _ : ProjectNonnegative(z), L)],
+                   x0, 10, maxit = 100, adaptive = True).solve()
+
+    print(sol)
+    ref_sol = np.array([-3.99513417e-03, -1.32339656e-01, -4.52822633e-02, -7.44973888e-02, -3.35005208e-11])
+    # Compare
+    assert np.allclose(sol, ref_sol, atol=1e-4)