CODEWITHSUNDEEP
pythonpytorchdistributiontorch
IPhysResearch
IPhysResearch

Reputation: 129

How to sample from a distribution with constraints in PyTorch?

I got a simple situation like:

Generate samples and log_prob from an uniform distribution for 2-dim variables (m1, m2) which is satisfying m1~U(5, 80), m2~U(5, 80) with constraint m1+m2 < 100.

from torch import distributions
prior = distributions.Uniform(torch.tensor([5.0, 5.0]),
                              torch.tensor([80.0, 80.0]))

I try coding in PyTorch like the above, but I don't know how to construct a torch.distribution with constraint condition. By the way, I see some implementations about torch.distributions.constraints but I can't figure out how to use it.

Upvotes: 0

Views: 1200

Answers (1)

MWB
MWB

Reputation: 12577

This can be achieved using rejection sampling:

d = torch.distributions.Uniform(5, 80)

m1 = 80
m2 = 80

while m1 + m2 >= 100:
    m1 = d.sample()
    m2 = d.sample()
    print(m1, m2)

Example output:

tensor(52.3322) tensor(67.8245)
tensor(68.3232) tensor(40.0983)
tensor(44.7374) tensor(9.9690)

Upvotes: 3

Related Questions