Repository: CUQI-DTU/CUQIpy

Issue:
Bug in plot() of cuqi.samples._samples.Samples
There seems a bug with the `plot` of `cuqi.samples._samples.Samples` when the samples' dimension is 1. 

```python
import cuqi
import matplotlib.pyplot as plt

my_gaussian = cuqi.distribution.Gaussian(mean=1, cov=1)

samples = my_gaussian.sample(10)

samples.plot()
```
The above code displays the following msg and plot:
```bash
[<matplotlib.lines.Line2D at 0x1770b4b60>,
 <matplotlib.lines.Line2D at 0x1770b4b90>,
 <matplotlib.lines.Line2D at 0x1770b4c80>,
 <matplotlib.lines.Line2D at 0x17707bb60>,
 <matplotlib.lines.Line2D at 0x17707bb90>]
```


![Image](https://github.com/user-attachments/assets/096bbad8-f596-41e8-9749-6efe0bed8d56)

Hints from issue comments:

Solution patch:
diff --git a/cuqi/distribution/_joint_distribution.py b/cuqi/distribution/_joint_distribution.py
index aeee61db7..2e32d0555 100644
--- a/cuqi/distribution/_joint_distribution.py
+++ b/cuqi/distribution/_joint_distribution.py
@@ -175,6 +175,10 @@ def _parse_args_add_to_kwargs(self, *args, **kwargs):
                 kwargs[ordered_keys[index]] = arg
         return kwargs
 
+    def _sum_evaluated_densities(self):
+        """ Return the sum of the evaluated densities in the joint distribution """
+        return sum([density.logd() for density in self._evaluated_densities])
+
     def _reduce_to_single_density(self):
         """ Reduce the joint distribution to a single density if possible.
 
@@ -183,7 +187,6 @@ def _reduce_to_single_density(self):
         This method is a hack to allow our current samplers to work with
         the joint distribution. It should be removed in the future.
         """
-
         # Count number of distributions and likelihoods
         n_dist = len(self._distributions)
         n_likelihood = len(self._likelihoods)
@@ -201,16 +204,29 @@ def _reduce_to_single_density(self):
             # Ensure parameter names match, otherwise return the joint distribution
             if set(self._likelihoods[0].get_parameter_names()) != set(self._distributions[0].get_parameter_names()):
                 return self
-            return Posterior(self._likelihoods[0], self._distributions[0])
-        
+            return self._add_constants_to_density(Posterior(self._likelihoods[0], self._distributions[0]))
+
         # If exactly one distribution and no likelihoods its a Distribution
         if n_dist == 1 and n_likelihood == 0:
-            return self._distributions[0]
-
+            return self._add_constants_to_density(self._distributions[0])        
+        
         # If no distributions and exactly one likelihood its a Likelihood
         if n_likelihood == 1 and n_dist == 0:
             return self._likelihoods[0]
 
+        # If only evaluated densities left return joint to ensure logd method is available
+        if n_dist == 0 and n_likelihood == 0:
+            return self
+        
+    def _add_constants_to_density(self, density: Density):
+        """ Add the constants (evaluated densities) to a single density. Used when reducing to single density. """
+
+        if isinstance(density, EvaluatedDensity):
+            raise ValueError("Cannot add the sum of all evaluated densities to an EvaluatedDensity.")
+
+        density._constant += self._sum_evaluated_densities()
+        return density
+
     def _as_stacked(self) -> _StackedJointDistribution:
         """ Return a stacked JointDistribution with the same densities. """
         return _StackedJointDistribution(*self._densities)
diff --git a/cuqi/experimental/mcmc/_gibbs.py b/cuqi/experimental/mcmc/_gibbs.py
index eea58b11d..8e627e30d 100644
--- a/cuqi/experimental/mcmc/_gibbs.py
+++ b/cuqi/experimental/mcmc/_gibbs.py
@@ -6,10 +6,10 @@
 import warnings
 
 try:
-    from progressbar import progressbar
+    from tqdm import tqdm
 except ImportError:
-    def progressbar(iterable, **kwargs):
-        warnings.warn("Module mcmc: Progressbar not found. Install progressbar2 to get sampling progress.")
+    def tqdm(iterable, **kwargs):
+        warnings.warn("Module mcmc: tqdm not found. Install tqdm to get sampling progress.")
         return iterable
 
 # Not subclassed from Sampler as Gibbs handles multiple samplers and samples multiple parameters
@@ -152,13 +152,13 @@ def validate_targets(self):
 
     def sample(self, Ns) -> 'HybridGibbs':
         """ Sample from the joint distribution using Gibbs sampling """
-        for _ in progressbar(range(Ns)):
+        for _ in tqdm(range(Ns)):
             self.step()
             self._store_samples()
 
     def warmup(self, Nb) -> 'HybridGibbs':
         """ Warmup (tune) the Gibbs sampler """
-        for idx in progressbar(range(Nb)):
+        for idx in tqdm(range(Nb)):
             self.step()
             self.tune(idx)
             self._store_samples()
diff --git a/cuqi/experimental/mcmc/_sampler.py b/cuqi/experimental/mcmc/_sampler.py
index 7e803c2e6..3902d4c3f 100644
--- a/cuqi/experimental/mcmc/_sampler.py
+++ b/cuqi/experimental/mcmc/_sampler.py
@@ -7,10 +7,10 @@
 from cuqi.samples import Samples
 
 try:
-    from progressbar import progressbar
+    from tqdm import tqdm
 except ImportError:
-    def progressbar(iterable, **kwargs):
-        warnings.warn("Module mcmc: Progressbar not found. Install progressbar2 to get sampling progress.")
+    def tqdm(iterable, **kwargs):
+        warnings.warn("Module mcmc: tqdm not found. Install tqdm to get sampling progress.")
         return iterable
 
 class Sampler(ABC):
@@ -220,7 +220,7 @@ def sample(self, Ns, batch_size=0, sample_path='./CUQI_samples/') -> 'Sampler':
         if hasattr(self, "_pre_sample"): self._pre_sample()
 
         # Draw samples
-        for _ in progressbar( range(Ns) ):
+        for _ in tqdm( range(Ns) ):
             
             # Perform one step of the sampler
             acc = self.step()
@@ -260,7 +260,7 @@ def warmup(self, Nb, tune_freq=0.1) -> 'Sampler':
         if hasattr(self, "_pre_warmup"): self._pre_warmup()
 
         # Draw warmup samples with tuning
-        for idx in progressbar(range(Nb)):
+        for idx in tqdm(range(Nb)):
 
             # Perform one step of the sampler
             acc = self.step()
diff --git a/cuqi/geometry/_geometry.py b/cuqi/geometry/_geometry.py
index 760fefa57..6361d357f 100644
--- a/cuqi/geometry/_geometry.py
+++ b/cuqi/geometry/_geometry.py
@@ -395,6 +395,8 @@ def grid(self, value):
         self._grid = self._create_dimension(value)
 
     def _plot(self, values, *args, **kwargs):
+        if self.par_dim==1 and "marker" not in kwargs.keys():
+            kwargs["marker"] = "o"
         p = plt.plot(self.grid, values, *args, **kwargs)
         self._plot_config()
         return p
diff --git a/cuqi/model/_model.py b/cuqi/model/_model.py
index d534ba130..365b72f1f 100644
--- a/cuqi/model/_model.py
+++ b/cuqi/model/_model.py
@@ -19,10 +19,10 @@ class Model(object):
         Forward operator.
 
     range_geometry : integer or cuqi.geometry.Geometry
-        If integer is given a _DefaultGeometry is created with dimension of the integer.
+        If integer is given, a cuqi.geometry._DefaultGeometry is created with dimension of the integer.
 
     domain_geometry : integer or cuqi.geometry.Geometry
-        If integer is given a _DefaultGeometry is created with dimension of the integer.
+        If integer is given, a cuqi.geometry._DefaultGeometry is created with dimension of the integer.
 
     gradient : callable function, optional
         The direction-Jacobian product of the forward operator Jacobian with 
@@ -35,14 +35,20 @@ class Model(object):
         The Jacobian of the forward operator with respect to the forward operator input,
         evaluated at a point (`wrt`). The signature of the Jacobian function should be (`wrt`).
         The Jacobian function should return a 2D ndarray of shape (range_dim, domain_dim).
-        The Jacobian function is used to specify the gradient function (vector-Jacobian product)
+        The Jacobian function is used to specify the gradient function by computing the vector-Jacobian
+        product (VJP), here we refer to the vector in the VJP as the `direction` since it is the direction at 
+        which the gradient is computed.
         automatically and thus the gradient function should not be specified when the Jacobian
         function is specified.
 
+
+    :ivar range_geometry: The geometry representing the range.
+    :ivar domain_geometry: The geometry representing the domain.
+
     Example
     -------
 
-    Consider a forward model :math:`F: \mathbb{R}^2 \rightarrow \mathbb{R}` defined by the following forward operator:
+    Consider a forward model :math:`F: \mathbb{R}^2 \\rightarrow \mathbb{R}` defined by the following forward operator:
 
     .. math::
 
@@ -141,11 +147,17 @@ def __init__(self, forward, range_geometry, domain_geometry, gradient=None, jaco
         self._non_default_args = cuqi.utilities.get_non_default_args(self._forward_func)
 
     @property
-    def domain_dim(self): 
+    def domain_dim(self):
+        """
+        The dimension of the domain
+        """
         return self.domain_geometry.par_dim
 
     @property
     def range_dim(self): 
+        """
+        The dimension of the range
+        """
         return self.range_geometry.par_dim
 
     def _2fun(self, x, geometry, is_par):
@@ -466,28 +478,61 @@ class LinearModel(Model):
     forward : 2D ndarray or callable function.
         Forward operator.
 
-    adjoint : 2d ndarray or callable function. (optional if matrix is passed as forward)
+    adjoint : 2D ndarray or callable function. (optional if matrix is passed as forward)
 
     range_geometry : integer or cuqi.geometry.Geometry (optional)
-        If integer is given a _DefaultGeometry is created with dimension of the integer.
+        If integer is given, a cuqi.geometry._DefaultGeometry is created with dimension of the integer.
 
     domain_geometry : integer or cuqi.geometry.Geometry (optional)
-        If integer is given a _DefaultGeometry is created with dimension of the integer.
+        If integer is given, a cuqi.geometry._DefaultGeometry is created with dimension of the integer.
 
-    Attributes
-    -----------
-    range_geometry : cuqi.geometry.Geometry
-        The geometry representing the range.
 
-    domain_geometry : cuqi.geometry.Geometry
-        The geometry representing the domain.
+    :ivar range_geometry: The geometry representing the range.
+    :ivar domain_geometry: The geometry representing the domain.
 
-    Methods
-    -----------
-    :meth:`forward` the forward operator.
-    :meth:`range_dim` the dimension of the range.
-    :meth:`domain_dim` the dimension of the domain.
-    :meth:`get_matrix` returns an ndarray with the matrix representing the forward operator.
+    Example
+    -------
+
+    Consider a linear model represented by a matrix, i.e., :math:`y=Ax` where 
+    :math:`A` is a matrix.
+    
+    We can define such a linear model by passing the matrix :math:`A`:
+
+    .. code-block:: python
+
+        import numpy as np
+        from cuqi.model import LinearModel
+
+        A = np.random.randn(2,3)
+
+        model = LinearModel(A)
+
+    The dimension of the range and domain geometries will be automatically 
+    inferred from the matrix :math:`A`.
+        
+    Meanwhile, such a linear model can also be defined by a forward function 
+    and an adjoint function:
+
+    .. code-block:: python
+
+        import numpy as np
+        from cuqi.model import LinearModel
+
+        A = np.random.randn(2,3)
+
+        def forward(x):
+            return A@x
+
+        def adjoint(y):
+            return A.T@y
+
+        model = LinearModel(forward,
+                            adjoint=adjoint,
+                            range_geometry=2,
+                            domain_geometry=3)
+
+    Note that you would need to specify the range and domain geometries in this
+    case as they cannot be inferred from the forward and adjoint functions.
     """
     # Linear forward model with forward and adjoint (transpose).
     
@@ -552,6 +597,9 @@ def adjoint(self, y, is_par=True):
 
 
     def get_matrix(self):
+        """
+        Returns an ndarray with the matrix representing the forward operator.
+        """
         if self._matrix is not None: #Matrix exists so return it
             return self._matrix
         else:
@@ -594,24 +642,14 @@ class PDEModel(Model):
         Forward operator assembling, solving and observing the pde.
 
     range_geometry : integer or cuqi.geometry.Geometry (optional)
-        If integer is given a _DefaultGeometry is created with dimension of the integer.
+        If integer is given, a cuqi.geometry._DefaultGeometry is created with dimension of the integer.
 
     domain_geometry : integer or cuqi.geometry.Geometry (optional)
-        If integer is given a _DefaultGeometry is created with dimension of the integer.
-
-    Attributes
-    -----------
-    range_geometry : cuqi.geometry.Geometry
-        The geometry representing the range.
+        If integer is given, a cuqi.geometry._DefaultGeometry is created with dimension of the integer.
 
-    domain_geometry : cuqi.geometry.Geometry
-        The geometry representing the domain.
 
-    Methods
-    -----------
-    :meth:`forward` the forward operator.
-    :meth:`range_dim` the dimension of the range.
-    :meth:`domain_dim` the dimension of the domain.
+    :ivar range_geometry: The geometry representing the range.
+    :ivar domain_geometry: The geometry representing the domain.
     """
     def __init__(self, PDE: cuqi.pde.PDE, range_geometry, domain_geometry):
 
diff --git a/cuqi/utilities/__init__.py b/cuqi/utilities/__init__.py
index 3074decf1..b48ecaacc 100644
--- a/cuqi/utilities/__init__.py
+++ b/cuqi/utilities/__init__.py
@@ -11,6 +11,8 @@
     sparse_cholesky,
     approx_derivative,
     check_if_conditional_from_attr,
+    plot_1D_density,
+    plot_2D_density
 )
 
 from ._get_python_variable_name import _get_python_variable_name
diff --git a/cuqi/utilities/_utilities.py b/cuqi/utilities/_utilities.py
index 30ecbce58..2513d42ac 100644
--- a/cuqi/utilities/_utilities.py
+++ b/cuqi/utilities/_utilities.py
@@ -1,4 +1,5 @@
 from cuqi.array import CUQIarray
+from cuqi.density import Density
 import numpy as np
 import inspect
 from numbers import Number
@@ -7,6 +8,7 @@
 from dataclasses import dataclass
 from abc import ABCMeta
 import copy
+import matplotlib.pyplot as plt
 
 
 def force_ndarray(value,flatten=False):
@@ -229,3 +231,117 @@ def approx_gradient(func, x, epsilon= 0.000001):
         eps_vec[i] = 0.0
         
     return FD_gradient
+
+
+# Function for plotting 1D density functions
+def plot_1D_density(density:Density,
+                    v_min, v_max,
+                    N=501, log_scale=False,
+                    **kwargs):
+    """ Plot 1D density function 
+
+    Parameters
+    ----------
+    density : CUQIpy Density
+        The density to be plotted.
+
+    v_min : float
+        Minimum value for the variable.
+    
+    v_max : float
+        Maximum value for the variable.
+
+    N : int
+        Number of grid points for the variable.
+
+    log_scale : bool
+        If True, the density is plotted in log scale.
+    
+    kwargs : dict
+        Additional keyword arguments for the plot that are passed to the
+        underlying plotting method: `matplotlib.pyplot.plot` function 
+        in this case.
+
+    """
+    # Assert that the density is 1D
+    assert density.dim == 1, "The density must be for a scalar variable"
+    ls = np.linspace(v_min, v_max, N)
+
+    # Create a map to evaluate density
+    density_map = (lambda x: x) if log_scale else (lambda x: np.exp(x))
+
+    # Evaluate density on grid
+    y = [density_map(density.logd(grid_point)) for grid_point in ls]
+    p = plt.plot(ls, y, **kwargs)
+    return p
+
+# Function for plotting 2D density functions
+def plot_2D_density(density: Density, 
+                    v1_min, v1_max,
+                    v2_min, v2_max,
+                    N1=201, N2=201,
+                    log_scale=False,
+                    **kwargs):
+    """ Plot 2D density function 
+
+    Parameters
+    ----------
+    density : CUQIpy Density
+        The density to be plotted.
+
+    v1_min : float
+        Minimum value for the first variable.
+    
+    v1_max : float
+        Maximum value for the first variable.
+    
+    v2_min : float
+        Minimum value for the second variable.
+
+    v2_max : float  
+        Maximum value for the second variable.
+
+    N1 : int
+        Number of grid points for the first variable.
+
+    N2 : int
+        Number of grid points for the second variable.
+    
+    log_scale : bool
+        If True, the density is plotted in log scale.
+
+    kwargs : dict
+        Additional keyword arguments for the plot that are passed to the
+        underlying plotting method: `matplotlib.pyplot.imshow` function 
+        in this case.
+
+    """
+    # Assert that the density is 2D
+    assert density.dim == 2,\
+        "The density must be for a two-dimensional variable"
+    # Create grid
+    ls1 = np.linspace(v1_min, v1_max, N1)
+    ls2 = np.linspace(v2_min, v2_max, N2)
+    grid1, grid2 = np.meshgrid(ls1, ls2)
+
+    # Create a map to evaluate density
+    density_map = (lambda x: x) if log_scale else (lambda x: np.exp(x))
+
+    # Evaluate density on grid
+    evaluated_density = np.zeros((N1, N2))
+    for ii in range(N1):
+        for jj in range(N2):
+            evaluated_density[ii,jj] = density_map(
+                density.logd([grid1[ii,jj], grid2[ii,jj]])) 
+
+    # Plot
+    pixelwidth_x = (v1_max-v1_min)/(N2-1)
+    pixelwidth_y = (v2_max-v2_min)/(N2-1)
+
+    hp_x = 0.5*pixelwidth_x
+    hp_y = 0.5*pixelwidth_y
+
+    extent = (v1_min-hp_x, v1_max+hp_x, v2_min-hp_y, v2_max+hp_y)
+
+    im = plt.imshow(evaluated_density, origin='lower', extent=extent, **kwargs)
+    return im
\ No newline at end of file
diff --git a/requirements-dev.txt b/requirements-dev.txt
index 94c44c36f..b10b7c98a 100644
--- a/requirements-dev.txt
+++ b/requirements-dev.txt
@@ -7,4 +7,4 @@ sphinx-panels
 sphinx-copybutton
 sphinx-gallery
 versioneer==0.26
-progressbar2 # For experimental mcmc module
+tqdm # For experimental mcmc module


Test patch:
diff --git a/tests/test_joint_distribution.py b/tests/test_joint_distribution.py
index 456015f38..d0ec32cce 100644
--- a/tests/test_joint_distribution.py
+++ b/tests/test_joint_distribution.py
@@ -310,3 +310,86 @@ def test_MultipleLikelihoodPosterior_should_raise_if_names_do_not_match():
     with pytest.raises(ValueError, match=r"the same parameter name"):
         cuqi.distribution.MultipleLikelihoodPosterior(y.to_likelihood(1), x, z)
 
+@pytest.mark.parametrize("joint, variables", [
+    (
+        cuqi.distribution.JointDistribution(
+            cuqi.distribution.Normal(0, 1, name="x"),
+            cuqi.distribution.Normal(0, 1, name="y")
+        ),
+        {
+            "x": 0,
+            "y": 0
+        }
+    ),
+    (
+        cuqi.distribution.JointDistribution(
+            cuqi.distribution.Uniform(0, 10, name="d"),
+            cuqi.distribution.Uniform(0, 5, name="s"),
+            cuqi.distribution.Gaussian(np.zeros(8), lambda d: d, name="x"),
+            cuqi.distribution.Gaussian(
+                mean=cuqi.testproblem.Deconvolution1D(dim=8).model,
+                cov=lambda s: s,
+                name="y"
+            )
+        ),
+        {
+            "d": 5,
+            "s": 2,
+            "x": np.zeros(8),
+            "y": cuqi.testproblem.Deconvolution1D(dim=8).data
+        }
+    ),
+    (
+        cuqi.distribution.JointDistribution(
+            cuqi.distribution.Uniform(0, 10, name="d"),
+            cuqi.distribution.Uniform(0, 5, name="s"),
+            cuqi.distribution.Gaussian(np.zeros(8), lambda d: d, name="x"),
+            cuqi.distribution.Gaussian(
+                mean=cuqi.testproblem.Deconvolution1D(dim=8).model,
+                cov=lambda s: s,
+                name="y"
+            )
+        ),
+        {
+            "d": 11, # Out of bounds (all should return -Inf)
+            "s": 2,
+            "x": np.zeros(8),
+            "y": cuqi.testproblem.Deconvolution1D(dim=8).data
+        }
+    )
+    ]
+)
+def test_logd_consistency_when_conditioning(joint, variables):
+    """ Test consistency of logd value when conditioning the joint distribution.
+     
+    This ensures we always return the correct value for logd even when reducing to single density.
+    
+    """
+
+    # True value of logd by fully evaluating the joint distribution
+    true_value = joint.logd(**variables)
+
+    # Loop over all variables and evaluate the logd value
+    # where all previously seen variables are conditioned
+    # and not seen variables are given to logd.
+    cond_vars = {}
+    for key, value in variables.items():
+        # Remaining variables are those that are not used for conditioning
+        remaining_vars = {k: v for k, v in variables.items() if k not in cond_vars}
+        
+        # Condition the joint distribution
+        cond_joint = joint(**cond_vars)
+
+        # Evaluate the logd value of conditioned joint distribution
+        # This may be a single density or a joint distribution since
+        # joint can reduce to single density (such as Posterior)
+        logd_value = cond_joint.logd(**remaining_vars)
+
+        # Potential error message if this assert fails:
+        msg = f"Evaluated: \n {joint} \n with variables {variables}.\nFailed at {key}={value} for \n {cond_joint} \n with variables {remaining_vars}."
+        
+        # Assert the logd value matches
+        assert logd_value == pytest.approx(true_value, rel=1e-6), msg
+
+        # Add current variable to the variables that need to be conditioned
+        cond_vars[key] = value
diff --git a/tests/test_samples.py b/tests/test_samples.py
index 8584e10cb..68ab502bc 100644
--- a/tests/test_samples.py
+++ b/tests/test_samples.py
@@ -54,6 +54,22 @@ def test_samples_plot(geom, to_funvals, plot_par):
     with pytest.raises(ValueError,match=r"Cannot pass is_par as a plotting argument"):
         s.plot(is_par=not to_funvals,plot_par=plot_par)
 
+@pytest.mark.parametrize("geom",[
+						(cuqi.geometry.Discrete(1)),
+						(cuqi.geometry.Continuous1D(1))
+						])
+@pytest.mark.parametrize("kwargs",[
+                        ({}),
+                        ({"marker":"*"}),
+                        ])
+def test_1D_samples_plot(geom, kwargs):
+    # Make basic distribution and sample
+    my_gaussian = cuqi.distribution.Gaussian(mean=1, cov=1)
+    samples = my_gaussian.sample(10)
+    samples.geometry = geom
+    fig1 = samples.plot(**kwargs)
+    fig2 = samples.plot(0, **kwargs)
+
 @pytest.mark.parametrize("to_funvals, plot_par", [(True,False),(True,True),(False,False),(False,True)])
 @pytest.mark.parametrize("plot_method, plot_par_supported, arviz",
                          [(Samples.plot_autocorrelation, False, True),
diff --git a/tests/test_utilities.py b/tests/test_utilities.py
index f27364b72..e9103ece5 100644
--- a/tests/test_utilities.py
+++ b/tests/test_utilities.py
@@ -1,8 +1,11 @@
 import pytest
 from scipy.linalg import cholesky
 from scipy.sparse import diags
-from cuqi.utilities import sparse_cholesky
+from cuqi.utilities import sparse_cholesky, plot_1D_density, plot_2D_density
+from cuqi.model import LinearModel
+from cuqi.distribution import Gaussian, Uniform, JointDistribution
 import numpy as np
+import matplotlib.pyplot as plt
 
 
 @pytest.mark.parametrize("P", [
@@ -18,3 +21,48 @@ def test_sparse_cholesky(P):
     L2 = sparse_cholesky(P) 
 
     assert np.allclose(L1, L2.toarray()) # Convert to dense to compare
+
+
+def posterior_2D_components():
+    """A function to create and return 2D Bayesian posterior components"""
+    # Create a CUQIpy model
+    A = LinearModel(np.array([[1.0, 1.0]]))
+    # Create a CUQIpy prior
+    x = Gaussian(mean=np.array([0.0, 0.0]), cov=np.array([1.0, 1.0]))
+    # Create a CUQIpy likelihood
+    y = Gaussian(A@x, cov=0.1)
+    return x, y
+
+
+def likelihood_2D():
+    """A function to create and return a 2D likelihood"""
+    x, y = posterior_2D_components()
+    likelihood = y.to_likelihood(2)
+    return likelihood
+
+
+def posterior_2D():
+    """A function to create and return a 2D posterior"""
+    x, y = posterior_2D_components()
+    joint = JointDistribution(x, y)
+    posterior = joint(y=2)
+    return posterior
+
+
+@pytest.mark.parametrize("density", [
+    Gaussian(mean=0.0, cov=1.0),
+    Uniform(low=0.0, high=1.0),
+    Gaussian(mean=np.array([0.0, 0.0]), cov=np.array([1.0, 1.0])),
+    likelihood_2D(),
+    posterior_2D()
+])
+@pytest.mark.parametrize("log_scale", [True, False])
+def test_plot_density(density, log_scale):
+    """ Test the 1D and 2D density plot. """
+    if density.dim == 1:
+        plot_1D_density(density, -1, 1, color='green', log_scale=log_scale)
+    elif density.dim == 2:
+        plot_2D_density(
+            density, -3, 3, -3, 3, 60, 60, cmap='gray', log_scale=log_scale)
+    else:
+        raise ValueError("Density must be 1D or 2D.")
\ No newline at end of file