From 00caf45ab960759c5bb813762fef08cc6fde2dda Mon Sep 17 00:00:00 2001
From: Chao <chaz@dtu.dk>
Date: Mon, 30 Sep 2024 22:04:59 +0200
Subject: [PATCH 01/21] add (iid) truncated normal

---
 cuqi/distribution/__init__.py          |  1 +
 cuqi/distribution/_truncated_normal.py | 45 ++++++++++++++++++++++++++
 2 files changed, 46 insertions(+)
 create mode 100644 cuqi/distribution/_truncated_normal.py

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/_truncated_normal.py b/cuqi/distribution/_truncated_normal.py
new file mode 100644
index 000000000..89277c39e
--- /dev/null
+++ b/cuqi/distribution/_truncated_normal.py
@@ -0,0 +1,45 @@
+import numpy as np
+from scipy.special import erf
+from cuqi.distribution import Distribution
+
+class TruncatedNormal(Distribution):
+    """
+    Truncated Normal probability distribution. Generates instance of cuqi.distribution.TruncatedNormal.
+    
+    The variables of this distribution are iid.
+
+    
+    Parameters
+    ------------
+    mean: mean of distribution
+    std: standard deviation
+    a: lower bound of the distribution
+    b: upper bound of the distribution
+    """
+    def __init__(self, mean=None, std=None, a=None, b=None, 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.a = a
+        self.b = b
+
+    def logpdf(self, x):
+        # the unnormalized logpdf
+        # check if x falls in the range between np.array a and b
+        if np.any(x < self.a) or np.any(x > self.b):
+            return -np.Inf
+        else:
+             return np.sum(-np.log(self.std*np.sqrt(2*np.pi))-0.5*((x-self.mean)/self.std)**2)
+
+    def gradient(self, x):
+        # check if x falls in the range between np.array a and b
+        if np.any(x < self.a) or np.any(x > self.b):
+            return np.NaN*np.ones_like(x)
+        else:
+            return -(x-self.mean)/(self.std**2)
+
+    def _sample(self,N=1, rng=None):
+        pass
\ No newline at end of file

From 999d57dad113e67b9cfae5a75ea65d699391fa78 Mon Sep 17 00:00:00 2001
From: Chao <chaz@dtu.dk>
Date: Sat, 5 Oct 2024 19:12:53 +0200
Subject: [PATCH 02/21] add sample

---
 cuqi/distribution/_truncated_normal.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/cuqi/distribution/_truncated_normal.py b/cuqi/distribution/_truncated_normal.py
index 89277c39e..c466503a2 100644
--- a/cuqi/distribution/_truncated_normal.py
+++ b/cuqi/distribution/_truncated_normal.py
@@ -16,7 +16,7 @@ class TruncatedNormal(Distribution):
     a: lower bound of the distribution
     b: upper bound of the distribution
     """
-    def __init__(self, mean=None, std=None, a=None, b=None, is_symmetric=False, **kwargs):
+    def __init__(self, mean=None, std=None, a=-1, b=1, is_symmetric=False, **kwargs):
         # Init from abstract distribution class
         super().__init__(is_symmetric=is_symmetric, **kwargs)  
 
@@ -42,4 +42,7 @@ def gradient(self, x):
             return -(x-self.mean)/(self.std**2)
 
     def _sample(self,N=1, rng=None):
-        pass
\ No newline at end of file
+        """
+        Generates random samples from the distribution.
+        """
+        raise NotImplementedError(f"sample is not implemented for {self.__class__.__name__}.")

From 34d5cf21b3719b167533a5b97c1ad43de201d25e Mon Sep 17 00:00:00 2001
From: Chao <chaz@dtu.dk>
Date: Sat, 5 Oct 2024 19:26:40 +0200
Subject: [PATCH 03/21] let truncated normal by default be the same as normal

---
 cuqi/distribution/_truncated_normal.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/cuqi/distribution/_truncated_normal.py b/cuqi/distribution/_truncated_normal.py
index c466503a2..85d8defc7 100644
--- a/cuqi/distribution/_truncated_normal.py
+++ b/cuqi/distribution/_truncated_normal.py
@@ -16,7 +16,7 @@ class TruncatedNormal(Distribution):
     a: lower bound of the distribution
     b: upper bound of the distribution
     """
-    def __init__(self, mean=None, std=None, a=-1, b=1, is_symmetric=False, **kwargs):
+    def __init__(self, mean=None, std=None, a=-np.Inf, b=np.Inf, is_symmetric=False, **kwargs):
         # Init from abstract distribution class
         super().__init__(is_symmetric=is_symmetric, **kwargs)  
 

From c03dbe658aa98096c67d1d7d834165dfb784a6cf Mon Sep 17 00:00:00 2001
From: chaozg <69794131+chaozg@users.noreply.github.com>
Date: Tue, 8 Oct 2024 13:16:07 +0200
Subject: [PATCH 04/21] Update cuqi/distribution/_truncated_normal.py

Co-authored-by: amal-ghamdi <amal.m.alghamdi@gmail.com>
---
 cuqi/distribution/_truncated_normal.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/cuqi/distribution/_truncated_normal.py b/cuqi/distribution/_truncated_normal.py
index 85d8defc7..8b6022ced 100644
--- a/cuqi/distribution/_truncated_normal.py
+++ b/cuqi/distribution/_truncated_normal.py
@@ -4,7 +4,7 @@
 
 class TruncatedNormal(Distribution):
     """
-    Truncated Normal probability distribution. Generates instance of cuqi.distribution.TruncatedNormal.
+    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.
 

From 0eeec67b5c81cb0d56137dc81e3dd6c684b8d679 Mon Sep 17 00:00:00 2001
From: Chao <chaz@dtu.dk>
Date: Wed, 9 Oct 2024 21:16:04 +0200
Subject: [PATCH 05/21] rename a/b to low/high

---
 cuqi/distribution/_truncated_normal.py | 22 +++++++++++++---------
 1 file changed, 13 insertions(+), 9 deletions(-)

diff --git a/cuqi/distribution/_truncated_normal.py b/cuqi/distribution/_truncated_normal.py
index 8b6022ced..af4763288 100644
--- a/cuqi/distribution/_truncated_normal.py
+++ b/cuqi/distribution/_truncated_normal.py
@@ -11,32 +11,36 @@ class TruncatedNormal(Distribution):
     
     Parameters
     ------------
-    mean: mean of distribution
-    std: standard deviation
-    a: lower bound of the distribution
-    b: upper bound of the distribution
+    mean : float or array_like of floats
+        mean of distribution
+    std : float or array_like of floats
+        standard deviation
+    a : float or array_like of floats
+        lower bound of the distribution
+    b : float or array_like of floats
+        upper bound of the distribution
     """
-    def __init__(self, mean=None, std=None, a=-np.Inf, b=np.Inf, is_symmetric=False, **kwargs):
+    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.a = a
-        self.b = b
+        self.low = low
+        self.high = high
 
     def logpdf(self, x):
         # the unnormalized logpdf
         # check if x falls in the range between np.array a and b
-        if np.any(x < self.a) or np.any(x > self.b):
+        if np.any(x < self.low) or np.any(x > self.high):
             return -np.Inf
         else:
              return np.sum(-np.log(self.std*np.sqrt(2*np.pi))-0.5*((x-self.mean)/self.std)**2)
 
     def gradient(self, x):
         # check if x falls in the range between np.array a and b
-        if np.any(x < self.a) or np.any(x > self.b):
+        if np.any(x < self.low) or np.any(x > self.high):
             return np.NaN*np.ones_like(x)
         else:
             return -(x-self.mean)/(self.std**2)

From 683ae553b5aa233690d031cf7e6f80a0b3e49bd0 Mon Sep 17 00:00:00 2001
From: Chao <chaz@dtu.dk>
Date: Wed, 9 Oct 2024 21:19:58 +0200
Subject: [PATCH 06/21] update docstring

---
 cuqi/distribution/_truncated_normal.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/cuqi/distribution/_truncated_normal.py b/cuqi/distribution/_truncated_normal.py
index af4763288..6ddc7d583 100644
--- a/cuqi/distribution/_truncated_normal.py
+++ b/cuqi/distribution/_truncated_normal.py
@@ -15,9 +15,9 @@ class TruncatedNormal(Distribution):
         mean of distribution
     std : float or array_like of floats
         standard deviation
-    a : float or array_like of floats
+    low : float or array_like of floats
         lower bound of the distribution
-    b : float or array_like of floats
+    high : float or array_like of floats
         upper bound of the distribution
     """
     def __init__(self, mean=None, std=None, low=-np.Inf, high=np.Inf, is_symmetric=False, **kwargs):

From 9dadda44a31dfb594058adf2d689c75ee9314da2 Mon Sep 17 00:00:00 2001
From: Chao <chaz@dtu.dk>
Date: Wed, 9 Oct 2024 21:31:32 +0200
Subject: [PATCH 07/21] add gradient to normal and use normal in truncated
 normal

---
 cuqi/distribution/_normal.py           | 3 +++
 cuqi/distribution/_truncated_normal.py | 8 ++++++--
 2 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/cuqi/distribution/_normal.py b/cuqi/distribution/_normal.py
index b101c4788..d3efd384b 100644
--- a/cuqi/distribution/_normal.py
+++ b/cuqi/distribution/_normal.py
@@ -36,6 +36,9 @@ 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, x):
+        return -(x-self.mean)/(self.std**2)
+
     def _sample(self,N=1, rng=None):
 
         """
diff --git a/cuqi/distribution/_truncated_normal.py b/cuqi/distribution/_truncated_normal.py
index 6ddc7d583..a524435ad 100644
--- a/cuqi/distribution/_truncated_normal.py
+++ b/cuqi/distribution/_truncated_normal.py
@@ -1,6 +1,7 @@
 import numpy as np
 from scipy.special import erf
 from cuqi.distribution import Distribution
+from cuqi.distribution import Normal
 
 class TruncatedNormal(Distribution):
     """
@@ -30,20 +31,23 @@ def __init__(self, mean=None, std=None, low=-np.Inf, high=np.Inf, is_symmetric=F
         self.low = low
         self.high = high
 
+        # Init underlying normal distribution
+        self._normal = Normal(self.mean, self.std)
+
     def logpdf(self, 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 np.sum(-np.log(self.std*np.sqrt(2*np.pi))-0.5*((x-self.mean)/self.std)**2)
+             return self._normal.logpdf(x)
 
     def gradient(self, 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 -(x-self.mean)/(self.std**2)
+            return self._normal.gradient(x)
 
     def _sample(self,N=1, rng=None):
         """

From e3789243165f852348b4adcd1ef47113cef2e24f Mon Sep 17 00:00:00 2001
From: Chao <chaz@dtu.dk>
Date: Wed, 9 Oct 2024 22:13:23 +0200
Subject: [PATCH 08/21] add unit tests for logpdf and gradient

---
 tests/test_distribution.py | 44 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 44 insertions(+)

diff --git a/tests/test_distribution.py b/tests/test_distribution.py
index e22427a5e..0d4c82ea1 100644
--- a/tests/test_distribution.py
+++ b/tests/test_distribution.py
@@ -16,6 +16,50 @@ def test_Normal_pdf_mean():
     pX = cuqi.distribution.Normal(0.1,1)
     assert pX.pdf(0.1) == approx(1.0/np.sqrt(2.0*np.pi))
 
+@pytest.mark.parametrize("mean,std,low,high,points",[(np.array([0.0]), 
+                                               np.array([1.0]),
+                                               np.array([-1.0]),
+                                               np.array([1.0]),
+                                               [-1.5, 0.0, 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.0, 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.isinf(x_trun.logpdf(point))
+
+@pytest.mark.parametrize("mean,std,low,high,points",[(np.array([0.0]), 
+                                               np.array([1.0]),
+                                               np.array([-1.0]),
+                                               np.array([1.0]),
+                                               [-1.5, 0.0, 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.0, 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,var,expected",[
                             (2,3.5,[[8.17418321], [3.40055023]]),
                             (3.141592653589793,2.6457513110645907,[[7.80883646],[4.20030911]]),

From 27a0982de6b9f57fd7a96c9e37c7dbbcc84fc0d7 Mon Sep 17 00:00:00 2001
From: Chao <chaz@dtu.dk>
Date: Wed, 9 Oct 2024 22:18:06 +0200
Subject: [PATCH 09/21] tweak parameters abit

---
 tests/test_distribution.py | 26 ++++++++++++++------------
 1 file changed, 14 insertions(+), 12 deletions(-)

diff --git a/tests/test_distribution.py b/tests/test_distribution.py
index 0d4c82ea1..4b2188a73 100644
--- a/tests/test_distribution.py
+++ b/tests/test_distribution.py
@@ -16,16 +16,17 @@ def test_Normal_pdf_mean():
     pX = cuqi.distribution.Normal(0.1,1)
     assert pX.pdf(0.1) == approx(1.0/np.sqrt(2.0*np.pi))
 
-@pytest.mark.parametrize("mean,std,low,high,points",[(np.array([0.0]), 
-                                               np.array([1.0]),
-                                               np.array([-1.0]),
-                                               np.array([1.0]),
-                                               [-1.5, 0.0, 1.5]),
+@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.0, 0.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])]
                                                )])
@@ -38,16 +39,17 @@ def test_TruncatedNormal_logpdf(mean,std,low,high,points):
         else:
             assert np.isinf(x_trun.logpdf(point))
 
-@pytest.mark.parametrize("mean,std,low,high,points",[(np.array([0.0]), 
-                                               np.array([1.0]),
-                                               np.array([-1.0]),
-                                               np.array([1.0]),
-                                               [-1.5, 0.0, 1.5]),
+@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.0, 0.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])]
                                                )])

From 92a65d9f1c5e443fe91ae1bc915f89c2a92b39aa Mon Sep 17 00:00:00 2001
From: Chao <chaz@dtu.dk>
Date: Wed, 9 Oct 2024 22:20:47 +0200
Subject: [PATCH 10/21] update docstring

---
 cuqi/distribution/_truncated_normal.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/cuqi/distribution/_truncated_normal.py b/cuqi/distribution/_truncated_normal.py
index a524435ad..fd3d6c743 100644
--- a/cuqi/distribution/_truncated_normal.py
+++ b/cuqi/distribution/_truncated_normal.py
@@ -5,7 +5,14 @@
 
 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.
+    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.
 

From 6a484b36ecdb9ff313d69950fa3430969e046bf0 Mon Sep 17 00:00:00 2001
From: chaozg <69794131+chaozg@users.noreply.github.com>
Date: Wed, 9 Oct 2024 22:22:08 +0200
Subject: [PATCH 11/21] Update cuqi/distribution/_truncated_normal.py
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: Nicolai André Brogaard Riis <nabr@dtu.dk>
---
 cuqi/distribution/_truncated_normal.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/cuqi/distribution/_truncated_normal.py b/cuqi/distribution/_truncated_normal.py
index fd3d6c743..c3287df6d 100644
--- a/cuqi/distribution/_truncated_normal.py
+++ b/cuqi/distribution/_truncated_normal.py
@@ -56,7 +56,7 @@ def gradient(self, x):
         else:
             return self._normal.gradient(x)
 
-    def _sample(self,N=1, rng=None):
+    def _sample(self, N=1, rng=None):
         """
         Generates random samples from the distribution.
         """

From edec0d7b995d653664e4965998ec8d4863e49eb3 Mon Sep 17 00:00:00 2001
From: Chao <chaz@dtu.dk>
Date: Wed, 9 Oct 2024 22:26:11 +0200
Subject: [PATCH 12/21] update docstring

---
 cuqi/distribution/_truncated_normal.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/cuqi/distribution/_truncated_normal.py b/cuqi/distribution/_truncated_normal.py
index c3287df6d..017a75a8f 100644
--- a/cuqi/distribution/_truncated_normal.py
+++ b/cuqi/distribution/_truncated_normal.py
@@ -27,6 +27,13 @@ class TruncatedNormal(Distribution):
         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 [-1,1]
+        p = cuqi.distribution.TruncatedNormal(mean=0, std=1, low=-1, high=1)
     """
     def __init__(self, mean=None, std=None, low=-np.Inf, high=np.Inf, is_symmetric=False, **kwargs):
         # Init from abstract distribution class

From adac0720584496f59f88632782019b37b3b12dc3 Mon Sep 17 00:00:00 2001
From: Chao <chaz@dtu.dk>
Date: Wed, 9 Oct 2024 22:29:18 +0200
Subject: [PATCH 13/21] update docstring (add code for sampling)

---
 cuqi/distribution/_truncated_normal.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/cuqi/distribution/_truncated_normal.py b/cuqi/distribution/_truncated_normal.py
index 017a75a8f..97f96287d 100644
--- a/cuqi/distribution/_truncated_normal.py
+++ b/cuqi/distribution/_truncated_normal.py
@@ -32,8 +32,13 @@ class TruncatedNormal(Distribution):
     -----------
     .. code-block:: python
 
-        #Generate Normal with mean 0, standard deviation 1 and bounds [-1,1]
-        p = cuqi.distribution.TruncatedNormal(mean=0, std=1, low=-1, high=1)
+        #Generate Normal with mean 0, standard deviation 1 and bounds [-2,2]
+        p = cuqi.distribution.TruncatedNormal(mean=0, std=1, low=-2, high=2)
+        sampler = cuqi.experimental.mcmc.MALA(p, scale=0.1)
+        sampler.sample(10000)
+        samples = sampler.get_samples()
+        plt.figure()
+        samples.plot_trace()
     """
     def __init__(self, mean=None, std=None, low=-np.Inf, high=np.Inf, is_symmetric=False, **kwargs):
         # Init from abstract distribution class

From 28461cf87f1b27c38f40ded55c8128ba7c5cb934 Mon Sep 17 00:00:00 2001
From: Chao <chaz@dtu.dk>
Date: Wed, 9 Oct 2024 23:01:26 +0200
Subject: [PATCH 14/21] implement a naive sampling based on rejection

---
 cuqi/distribution/_truncated_normal.py | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/cuqi/distribution/_truncated_normal.py b/cuqi/distribution/_truncated_normal.py
index 97f96287d..bb2676b75 100644
--- a/cuqi/distribution/_truncated_normal.py
+++ b/cuqi/distribution/_truncated_normal.py
@@ -72,4 +72,15 @@ def _sample(self, N=1, rng=None):
         """
         Generates random samples from the distribution.
         """
-        raise NotImplementedError(f"sample is not implemented for {self.__class__.__name__}.")
+        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()
+            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

From 3b911d7e6034046fc807ad8c8a6c3d8bc568dda8 Mon Sep 17 00:00:00 2001
From: Chao <chaz@dtu.dk>
Date: Wed, 9 Oct 2024 23:11:31 +0200
Subject: [PATCH 15/21] add unit test for sampling

---
 tests/test_distribution.py | 21 +++++++++++++++++----
 1 file changed, 17 insertions(+), 4 deletions(-)

diff --git a/tests/test_distribution.py b/tests/test_distribution.py
index 4b2188a73..e818cc74f 100644
--- a/tests/test_distribution.py
+++ b/tests/test_distribution.py
@@ -28,8 +28,7 @@ def test_Normal_pdf_mean():
                                                [np.array([-0.5, 0.0]),
                                                 np.array([0.5, 0.0]),
                                                 np.array([-2.0, 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)
@@ -51,8 +50,7 @@ def test_TruncatedNormal_logpdf(mean,std,low,high,points):
                                                [np.array([-0.5, 0.0]),
                                                 np.array([0.5, 0.0]),
                                                 np.array([-2.0, 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)
@@ -62,6 +60,21 @@ def test_TruncatedNormal_gradient(mean,std,low,high,points):
         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)
+
 @pytest.mark.parametrize("mean,var,expected",[
                             (2,3.5,[[8.17418321], [3.40055023]]),
                             (3.141592653589793,2.6457513110645907,[[7.80883646],[4.20030911]]),

From 50c256043bd3207adc9efac10f75d134539c1611 Mon Sep 17 00:00:00 2001
From: Chao <chaz@dtu.dk>
Date: Wed, 9 Oct 2024 23:19:45 +0200
Subject: [PATCH 16/21] skip sample test for truncated normal as no easy way to
 carry rng

---
 tests/test_distributions_shape.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/test_distributions_shape.py b/tests/test_distributions_shape.py
index 387f12fc7..abd20211c 100644
--- a/tests/test_distributions_shape.py
+++ b/tests/test_distributions_shape.py
@@ -22,6 +22,7 @@
 skip_sample = [
     cuqi.distribution.Gamma, # Missing force_ndarray
     cuqi.distribution.Lognormal,
+    cuqi.distribution.TruncatedNormal, # Seems no easy way to carry rng
 ]
 skip_gradient = [
     cuqi.distribution.Gamma, # Missing force_ndarray

From 6cc6e6fbfa00421d74305f127e8479777d28ce9a Mon Sep 17 00:00:00 2001
From: Chao <chaz@dtu.dk>
Date: Wed, 9 Oct 2024 23:37:57 +0200
Subject: [PATCH 17/21] add unit test on gradient of normal

---
 tests/test_distribution.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/tests/test_distribution.py b/tests/test_distribution.py
index e818cc74f..da9a071b7 100644
--- a/tests/test_distribution.py
+++ b/tests/test_distribution.py
@@ -87,6 +87,14 @@ 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)
+
 def test_Gaussian_mean():
     mean = np.array([0, 0])
     std = np.array([1, 1])

From 86f7ad91d2b37101ce568a223c0f299974511ee0 Mon Sep 17 00:00:00 2001
From: Chao <chaz@dtu.dk>
Date: Wed, 9 Oct 2024 23:44:53 +0200
Subject: [PATCH 18/21] relocate unit tests on truncated normal

---
 tests/test_distribution.py | 40 +++++++++++++++++++-------------------
 1 file changed, 20 insertions(+), 20 deletions(-)

diff --git a/tests/test_distribution.py b/tests/test_distribution.py
index da9a071b7..78a51f14e 100644
--- a/tests/test_distribution.py
+++ b/tests/test_distribution.py
@@ -16,6 +16,26 @@ def test_Normal_pdf_mean():
     pX = cuqi.distribution.Normal(0.1,1)
     assert pX.pdf(0.1) == approx(1.0/np.sqrt(2.0*np.pi))
 
+@pytest.mark.parametrize("mean,var,expected",[
+                            (2,3.5,[[8.17418321], [3.40055023]]),
+                            (3.141592653589793,2.6457513110645907,[[7.80883646],[4.20030911]]),
+                            (-1e-09, 1000000.0,[[1764052.34596766],[400157.20836722]]),
+                            (1.7724538509055159, 0, [[1.77245385],[1.77245385]])
+                        ])
+def test_Normal_sample_regression(mean,var,expected):
+    rng = np.random.RandomState(0) #Replaces legacy method: np.random.seed(0)
+    samples = cuqi.distribution.Normal(mean,var).sample(2,rng=rng)
+    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,
@@ -75,26 +95,6 @@ def test_TruncatedNormal_sampling(mean,std,low,high):
         sample = samples[:,i]
         assert np.all(sample >= low) and np.all(sample <= high)
 
-@pytest.mark.parametrize("mean,var,expected",[
-                            (2,3.5,[[8.17418321], [3.40055023]]),
-                            (3.141592653589793,2.6457513110645907,[[7.80883646],[4.20030911]]),
-                            (-1e-09, 1000000.0,[[1764052.34596766],[400157.20836722]]),
-                            (1.7724538509055159, 0, [[1.77245385],[1.77245385]])
-                        ])
-def test_Normal_sample_regression(mean,var,expected):
-    rng = np.random.RandomState(0) #Replaces legacy method: np.random.seed(0)
-    samples = cuqi.distribution.Normal(mean,var).sample(2,rng=rng)
-    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)
-
 def test_Gaussian_mean():
     mean = np.array([0, 0])
     std = np.array([1, 1])

From 364db849cbb279d922e0a85f400218ca42ece0b2 Mon Sep 17 00:00:00 2001
From: Chao <chaz@dtu.dk>
Date: Thu, 10 Oct 2024 08:47:54 +0200
Subject: [PATCH 19/21] udate docstring of logpdf and gradient

---
 cuqi/distribution/_truncated_normal.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/cuqi/distribution/_truncated_normal.py b/cuqi/distribution/_truncated_normal.py
index bb2676b75..6422f2c41 100644
--- a/cuqi/distribution/_truncated_normal.py
+++ b/cuqi/distribution/_truncated_normal.py
@@ -54,6 +54,9 @@ def __init__(self, mean=None, std=None, low=-np.Inf, high=np.Inf, is_symmetric=F
         self._normal = Normal(self.mean, self.std)
 
     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):
@@ -62,6 +65,9 @@ def logpdf(self, x):
              return self._normal.logpdf(x)
 
     def gradient(self, x):
+        """
+        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)

From 5f8277f77ddc5bd23d3e1dd6db6cbfcfc1427053 Mon Sep 17 00:00:00 2001
From: Chao <chaz@dtu.dk>
Date: Thu, 10 Oct 2024 09:15:06 +0200
Subject: [PATCH 20/21] update comments

---
 cuqi/distribution/_truncated_normal.py | 1 +
 tests/test_distributions_shape.py      | 2 +-
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/cuqi/distribution/_truncated_normal.py b/cuqi/distribution/_truncated_normal.py
index 6422f2c41..0f18f5945 100644
--- a/cuqi/distribution/_truncated_normal.py
+++ b/cuqi/distribution/_truncated_normal.py
@@ -78,6 +78,7 @@ def _sample(self, N=1, rng=None):
         """
         Generates random samples from the distribution.
         """
+        # FIXME: this implementation does not honor the rng
         max_iter = 1e9 # maximum number of trials to avoid infinite loop
         samples = []
         for i in range(int(max_iter)):
diff --git a/tests/test_distributions_shape.py b/tests/test_distributions_shape.py
index abd20211c..b47979859 100644
--- a/tests/test_distributions_shape.py
+++ b/tests/test_distributions_shape.py
@@ -22,7 +22,7 @@
 skip_sample = [
     cuqi.distribution.Gamma, # Missing force_ndarray
     cuqi.distribution.Lognormal,
-    cuqi.distribution.TruncatedNormal, # Seems no easy way to carry rng
+    cuqi.distribution.TruncatedNormal, # FIXME Missing rng
 ]
 skip_gradient = [
     cuqi.distribution.Gamma, # Missing force_ndarray

From b87545d616b6749fd0b1e636327864205943965d Mon Sep 17 00:00:00 2001
From: Chao <chaz@dtu.dk>
Date: Thu, 10 Oct 2024 09:21:58 +0200
Subject: [PATCH 21/21] update example in docstring (use internal sample
 instead of external sampler)

---
 cuqi/distribution/_truncated_normal.py | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/cuqi/distribution/_truncated_normal.py b/cuqi/distribution/_truncated_normal.py
index 0f18f5945..9886504a9 100644
--- a/cuqi/distribution/_truncated_normal.py
+++ b/cuqi/distribution/_truncated_normal.py
@@ -34,11 +34,7 @@ class TruncatedNormal(Distribution):
 
         #Generate Normal with mean 0, standard deviation 1 and bounds [-2,2]
         p = cuqi.distribution.TruncatedNormal(mean=0, std=1, low=-2, high=2)
-        sampler = cuqi.experimental.mcmc.MALA(p, scale=0.1)
-        sampler.sample(10000)
-        samples = sampler.get_samples()
-        plt.figure()
-        samples.plot_trace()
+        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