Add (iid) truncated normal and also add grad to normal #542

chaozg · 2024-09-30T20:07:16Z

fixed #321
fixed #548 (grad of Normal was missing)

With this PR, we can draw samples from truncated iid normal distributions, i.e., truncated version of cuqi.distribution.normal or use them as priors.

demo 1: 2D normal distribution N([0; 0], [1, 0; 0, 1]) truncated in square domain [-2,2]*[-2,2]

import cuqi
import matplotlib.pyplot as plt
import numpy as np
from cuqi.utilities import plot_2D_density

tn = cuqi.distribution.TruncatedNormal(np.array([0,0]),np.array([1,1]),np.array([-2,-2]),np.array([2,2]))

plt.figure()
plot_2D_density(tn, -5, 5, -5, 5)

sampler = cuqi.experimental.mcmc.MALA(tn, scale=0.1, initial_point= np.array([0,0]))
sampler.sample(10000)
samples = sampler.get_samples()
plt.figure()
samples.plot_trace()

demo 2: bottom right quarter of 2D normal distribution N([0; 0], [1, 0; 0, 1])

import cuqi
import matplotlib.pyplot as plt
import numpy as np
from cuqi.utilities import plot_2D_density

tn = cuqi.distribution.TruncatedNormal(np.array([0,0]),np.array([1,1]),np.array([0,-np.Inf]),np.array([np.Inf, 0]))

plt.figure()
plot_2D_density(tn, -5, 5, -5, 5)

sampler = cuqi.experimental.mcmc.MALA(tn, scale=0.1, initial_point= np.array([1,-1]))
sampler.sample(10000)
samples = sampler.get_samples()
plt.figure()
samples.plot_trace()

demo 3: simplest BIP

import numpy as np
import matplotlib.pyplot as plt
import cuqi
from cuqi.utilities import plot_2D_density
np.random.seed(0)

# the forward model
A_matrix = np.array([[1.0, 1.0]])
A = cuqi.model.LinearModel(A_matrix)

# the prior
# x = Gaussian(np.zeros(2), 2.5)
# Bottom right quarter of 2D normal distribution
x = cuqi.distribution.TruncatedNormal(np.array([0, 0]), \
    np.array([1, 1]), np.array([0, -np.Inf]), np.array([np.Inf, 0]))
print(x)

# 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)

# sampling with MALA, ULA and NUTS
sampler = cuqi.experimental.mcmc.MALA(post, initial_point=np.array([2.5, -2.5]), scale=0.03)
# sampler = cuqi.experimental.mcmc.ULA(post, initial_point=np.array([2.5, 2.5]), scale=0.12)
# sampler = cuqi.experimental.mcmc.NUTS(post, initial_point=np.array([2.5, 2.5]))

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")

# samples
plt.figure()
samples.plot_pair()
plt.xlim(1, 5)
plt.ylim(-3, 1)
plt.gca().set_aspect('equal')
plt.title("Posterior Samples")

amal-ghamdi

Thank you @chaozg for adding this very useful distribution! I added some comments, feel free to address them as you see fit.

cuqi/distribution/_truncated_normal.py

Co-authored-by: amal-ghamdi <[email protected]>

nabriis

Really nice @chaozg. I agree with Amals suggestions and I had a few of my own.

nabriis · 2024-10-09T06:45:54Z

cuqi/distribution/_truncated_normal.py

+
+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.


Could you break this line into smaller pieces?

Done as suggested

cuqi/distribution/_truncated_normal.py

nabriis · 2024-10-09T06:53:15Z

cuqi/distribution/_truncated_normal.py

+
+    def _sample(self,N=1, rng=None):
+        """
+        Generates random samples from the distribution.


Perhaps for the direct sampling of the distribution we could use a simple "Rejection" sampling? If asked for "N" samples, we sample a Gaussian, say "N" times, remove out of bounds, sample another "N" times, remote out of bounds. If we have now "N" or more sample we return the first N, else we repeat until we get "N" samples. Could also be used to compare with the samplers you showed in the code.

I just added_sample with a slightly different strategy: I generate samples one by one, check if it is within the bounds, until we get N samples. Do you think strategy is OK?

just a comment here, would it be possible just to pass the rng to sample = self._normal.sample(rng=rng)?

nabriis · 2024-10-09T06:53:52Z

cuqi/distribution/_truncated_normal.py

+from scipy.special import erf
+from cuqi.distribution import Distribution
+
+class TruncatedNormal(Distribution):


You made this really nice showcases in the PR @chaozg. I suggest you add one or two of them as "Examples" in the docstring!

Added a example in the docstring as suggested

Co-authored-by: Nicolai André Brogaard Riis <[email protected]>

chaozg · 2024-10-09T21:50:11Z

Thanks @amal-ghamdi and @nabriis for your review. I think I have addressed your comments so I'm requesting your further review here.

Note that the existing Normal lacks grad() #548 , so I also add it in this PR.

Note that sample() implemented lacks a way to carry rng, and I'm not sure what's the best way to solve this.

For now, I put TruncatedNormal in the skip_sample list so the sampling is not test in test_multivariate_scalar_vars_samplewith a given rng;
Meanwhile, I created a test test_TruncatedNormal_sampling which checks if all samples falls in the bounds, but it's not tested with a fixed rng.

…sampler)

amal-ghamdi

Many thanks @chaozg for the updates! it is nice that the distribution is based on Normal now and that you added the gradient for the Normal distribution. I have only few comments to consider as you see fit.

cuqi/distribution/_truncated_normal.py

amal-ghamdi · 2024-10-10T11:24:43Z

cuqi/distribution/_normal.py

+    def gradient(self, x):
+        return -(x-self.mean)/(self.std**2)
+


This is a nice bonus that we have a gradient for normal now :). Just one suggestion to add two checks:
1- the geometry is not a geometry that needs a chain rule.
2- The distribution is used as prior, not likelihood because we do not account for the chain rule here.

For example, the gradient implementation of cmrf does these two checks

def _gradient(self, val, **kwargs): #Avoid complicated geometries that change the gradient. 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.location): # for prior diff = self._diff_op._matrix @ val return (-2*diff/(diff**2+self.scale**2)) @ self._diff_op._matrix else: warnings.warn('Gradient not implemented for {}'.format(type(self.location)))

amal-ghamdi · 2024-10-10T11:32:25Z

cuqi/distribution/_truncated_normal.py

+
+    def _sample(self,N=1, rng=None):
+        """
+        Generates random samples from the distribution.


just a comment here, would it be possible just to pass the rng to sample = self._normal.sample(rng=rng)?

amal-ghamdi · 2024-10-10T11:37:19Z

tests/test_distribution.py

+        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))


a suggestion to use np.isneginf instead of np.isinf here, to check for negative infinity.

chaozg added 2 commits September 30, 2024 22:04

add (iid) truncated normal

00caf45

add sample

999d57d

chaozg requested review from amal-ghamdi and nabriis October 5, 2024 17:16

let truncated normal by default be the same as normal

34d5cf2

amal-ghamdi requested changes Oct 8, 2024

View reviewed changes

Update cuqi/distribution/_truncated_normal.py

c03dbe6

Co-authored-by: amal-ghamdi <[email protected]>

nabriis requested changes Oct 9, 2024

View reviewed changes

chaozg and others added 13 commits October 9, 2024 21:16

rename a/b to low/high

0eeec67

update docstring

683ae55

add gradient to normal and use normal in truncated normal

9dadda4

add unit tests for logpdf and gradient

e378924

tweak parameters abit

27a0982

update docstring

92a65d9

Update cuqi/distribution/_truncated_normal.py

6a484b3

Co-authored-by: Nicolai André Brogaard Riis <[email protected]>

update docstring

edec0d7

update docstring (add code for sampling)

adac072

implement a naive sampling based on rejection

28461cf

add unit test for sampling

3b911d7

skip sample test for truncated normal as no easy way to carry rng

50c2560

add unit test on gradient of normal

6cc6e6f

chaozg mentioned this pull request Oct 9, 2024

Add gradient to normal #548

Open

3 tasks

relocate unit tests on truncated normal

86f7ad9

chaozg requested review from amal-ghamdi and nabriis October 9, 2024 21:47

udate docstring of logpdf and gradient

364db84

chaozg changed the title ~~Add (iid) truncated normal~~ Add (iid) truncated normal and also add grad to normal Oct 10, 2024

chaozg added 2 commits October 10, 2024 09:15

update comments

5f8277f

update example in docstring (use internal sample instead of external …

b87545d

…sampler)

amal-ghamdi requested changes Oct 10, 2024

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add (iid) truncated normal and also add grad to normal #542

Add (iid) truncated normal and also add grad to normal #542

chaozg commented Sep 30, 2024 •

edited

Loading

amal-ghamdi left a comment

nabriis left a comment

nabriis Oct 9, 2024

chaozg Oct 9, 2024

nabriis Oct 9, 2024

chaozg Oct 9, 2024

amal-ghamdi Oct 10, 2024

nabriis Oct 9, 2024

chaozg Oct 9, 2024

chaozg commented Oct 9, 2024 •

edited

Loading

amal-ghamdi left a comment

amal-ghamdi Oct 10, 2024

amal-ghamdi Oct 10, 2024

amal-ghamdi Oct 10, 2024

Add (iid) truncated normal and also add grad to normal #542

Are you sure you want to change the base?

Add (iid) truncated normal and also add grad to normal #542

Conversation

chaozg commented Sep 30, 2024 • edited Loading

demo 1: 2D normal distribution N([0; 0], [1, 0; 0, 1]) truncated in square domain [-2,2]*[-2,2]

demo 2: bottom right quarter of 2D normal distribution N([0; 0], [1, 0; 0, 1])

demo 3: simplest BIP

amal-ghamdi left a comment

Choose a reason for hiding this comment

nabriis left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

chaozg commented Oct 9, 2024 • edited Loading

amal-ghamdi left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

chaozg commented Sep 30, 2024 •

edited

Loading

chaozg commented Oct 9, 2024 •

edited

Loading