Repository: CUQI-DTU/CUQIpy

Issue:
RV: Sampling on Random Variables
- [x] Add support for sampling non-conditional Random Variables


**The following is skipped:**
 Add support for sampling multiple Random Variables (return as dict)
For example it should be possible to do the following:
```python
x = RandomVariable(Gaussian(0,1))
y = RandomVariable(Gaussian(x, 1))
x.sample(5) # Returns a CUQI samples object with samples of x.
y.sample(10) # Returns dict with 10 samples (in samples object) for "x" and "y"
```

Discussion:
@nabriis: @jakobsj DoD@jakobsj: @nabriis  I'll wait a bit with this one, as I don't yet feel in a position to assess whether the DoD covers the intended work and whether the design around "multiple random variables" and returning samples as a dictionary is how it will look in the end.

Solution patch:
diff --git a/cuqi/experimental/__init__.py b/cuqi/experimental/__init__.py
index 30844cc48..56c4dc7c6 100644
--- a/cuqi/experimental/__init__.py
+++ b/cuqi/experimental/__init__.py
@@ -1,4 +1,4 @@
-""" Experimental module for testing new features and ideas. """
+""" Experimental module for testing new features and ideas.  """
 from . import mcmc
 from . import algebra
 from . import geometry
diff --git a/cuqi/experimental/algebra/_randomvariable.py b/cuqi/experimental/algebra/_randomvariable.py
index b438fa096..35b112a24 100644
--- a/cuqi/experimental/algebra/_randomvariable.py
+++ b/cuqi/experimental/algebra/_randomvariable.py
@@ -6,6 +6,8 @@
 import cuqi
 from cuqi.distribution import Distribution
 from copy import copy, deepcopy
+import numpy as np
+
 
 class RandomVariable:
     """ Random variable defined by a distribution with the option to apply algebraic operations on it.
@@ -165,6 +167,30 @@ def __call__(self, *args, **kwargs) -> Any:
             raise ValueError(f"Expected arguments {self.parameter_names}, got arguments {kwargs}")
 
         return self.tree(**kwargs)
+
+    def sample(self, N=1):
+        """ Sample from the random variable. 
+
+        Parameters
+        ----------
+        N : int, optional
+            Number of samples to draw. Default is 1.
+        """
+
+        if self.is_cond:
+            raise NotImplementedError(
+                "Unable to directly sample from a random variable that has distributions with "
+                "conditioning variables. This is not implemented."
+            )
+
+        if N == 1: return self(**{dist.name: dist.sample() for dist in self.distributions})
+            
+        samples = np.array([
+            self(**{dist.name: dist.sample() for dist in self.distributions})
+            for _ in range(N)
+        ]).reshape(-1, N)  # Ensure correct shape (dim, N)
+
+        return cuqi.samples.Samples(samples)
     
     @property
     def tree(self):
diff --git a/cuqi/solver/_solver.py b/cuqi/solver/_solver.py
index 12d1ceb32..9ceacf025 100644
--- a/cuqi/solver/_solver.py
+++ b/cuqi/solver/_solver.py
@@ -884,4 +884,4 @@ def ProximalL1(x, gamma):
     x : array_like.
     gamma : scale parameter.
     """
-    return np.multiply(np.sign(x), np.maximum(np.abs(x)-gamma, 0))
\ No newline at end of file
+    return np.multiply(np.sign(x), np.maximum(np.abs(x)-gamma, 0))
diff --git a/demos/howtos/algebra.py b/demos/howtos/RandomVariablesAndAlgebra.py
similarity index 67%
rename from demos/howtos/algebra.py
rename to demos/howtos/RandomVariablesAndAlgebra.py
index e7c5956ea..f5150faff 100644
--- a/demos/howtos/algebra.py
+++ b/demos/howtos/RandomVariablesAndAlgebra.py
@@ -121,3 +121,50 @@
 z.condition(x=np.zeros(A.domain_dim))
 
 # %%
+# Sampling from random variables
+# ------------------------------
+# Random variables can be sampled using the `sample` method. The method returns a
+# sample from the distribution of the random variable.
+
+x = RandomVariable(Normal(0, 1))
+
+print(f"Sample from x: {x.sample()}")
+
+# %%
+# This can be combined with algebraic operations to sample from more complex
+# random variables.
+
+z = x + x**2 + 25
+
+print(f"Sample from z: {z.sample()}")
+
+# %%
+# Constructing a Beta distribution using Gamma random variables
+# -------------------------------------------------------------
+# Random variables can also be combined to create new distributions.
+# This is primarily useful for sampling at this stage.
+# For example, a Beta distribution can be constructed from two Gamma distributions:
+# If X ~ Gamma(a1, 1) and Y ~ Gamma(a2, 1), then Z = X / (X + Y) ~ Beta(a1, a2).
+# We illustrate this by comparing samples from a Beta distribution to samples
+# constructed using two Gamma distributions.
+
+from cuqi.distribution import Beta, Gamma
+
+# Define the shape parameters of the Beta distribution
+a1, a2 = 3, 2
+
+# Step 1: Directly define a Beta distribution
+z_ref = RandomVariable(Beta(a1, a2))
+
+# Step 2: Construct the Beta distribution using Gamma random variables
+x = RandomVariable(Gamma(a1, 1))  # X ~ Gamma(a1, 1)
+y = RandomVariable(Gamma(a2, 1))  # Y ~ Gamma(a2, 1)
+z = x / (x + y)                   # Z ~ Beta(a1, a2)
+
+# Step 3: Sample from both distributions
+z_samples = z.sample(10000)       # Samples from constructed Beta distribution
+z_ref_samples = z_ref.sample(10000)  # Samples from direct Beta distribution
+
+# Step 4: Plot histograms of the samples for comparison
+z_samples.hist_chain([0], bins=100)
+z_ref_samples.hist_chain([0], bins=100)


Test patch:
diff --git a/tests/zexperimental/test_randomvariable.py b/tests/zexperimental/test_randomvariable.py
index 3399c1acb..46ada8368 100644
--- a/tests/zexperimental/test_randomvariable.py
+++ b/tests/zexperimental/test_randomvariable.py
@@ -148,7 +148,6 @@ def test_randomvariable_algebra_works_on_joint_space(operations):
     val_z = abs(np.random.randn(z.dim))+1
     assert np.allclose(rv(x=val_x, y=val_y, z=val_z), operations(val_x, val_y, val_z))
 
-@pytest.mark.xfail(reason="sampling method for random variables is not yet implemented")
 @pytest.mark.parametrize("operations", [
     lambda x, y, z: x + y,
     lambda x, y, z: x * y,
@@ -328,3 +327,30 @@ def test_RV_condition_maintains_parameter_name_order():
     assert z.condition(s=1).parameter_names == ['x', 'y']
     assert z.condition(d=1).parameter_names == ['x', 'y']
     assert z.condition(d=1, s=1).parameter_names == ['x', 'y']
+
+def test_RV_sampling_unable_if_conditioning_variables_from_lambda():
+    x = RandomVariable(cuqi.distribution.Gaussian(0, lambda s: s))
+
+    with pytest.raises(NotImplementedError, match=r"Unable to directly sample from a random variable that has distributions with conditioning variables"):
+        x.sample()
+
+    x.condition(s=1).sample() # This should work
+
+def test_RV_sampling_unable_if_conditioning_variables_from_RV():
+
+    x = RandomVariable(cuqi.distribution.Gaussian(0, 1))
+    y = RandomVariable(cuqi.distribution.Gaussian(x, 1))
+
+    with pytest.raises(NotImplementedError, match=r"Unable to directly sample from a random variable that has distributions with conditioning variables"):
+        y.sample() # One might expect this to work, but it is not implemented at this time.
+
+def test_RV_sample_against_distribution_sample():
+    x = RandomVariable(cuqi.distribution.Gaussian(np.zeros(2), 1))
+
+    np.random.seed(0)
+    rv_samples = x.sample(10)
+
+    np.random.seed(0)
+    dist_samples = x.distribution.sample(10)
+
+    assert np.allclose(rv_samples.samples, dist_samples.samples)