```python
import torch
import numpy as np
from matplotlib import pyplot as plt
from torch.distributions.utils import broadcast_all
from probabll.distributions.utils import mc_kl, mc_entropy
from torch.distributions.kl import kl_divergence
```

Let's import some densities with support over $(0, 1)$


```python
from torch.distributions.uniform import Uniform
from probabll.distributions import RightTruncatedExponential
from probabll.distributions import Kumaraswamy
```

# Mixture of Discrete and Continuous over [0,1)

This distribution is a mixture of a point masses at $0$ and a continuous distribution over the support $(0, 1)$.

The distribution function looks like this:

\begin{equation}
f(x|\alpha, \beta) = \alpha \delta(x) + (1-\alpha) g(x|\beta) \mathbf 1_{(0, 1)}(x)
\end{equation}

where $0 < \alpha < 1$, $g(x|\beta)$ is a (properly normalised) density over $(0,1)$, and $\mathbf 1_{(0, 1)}(x)$ is the indicator function.


```python
from probabll.distributions.mixdc import MixtureD0C01
```


```python
p = MixtureD0C01(
    probs0=torch.full([1], 0.2), 
    cont=RightTruncatedExponential(rate=torch.ones(1), upper=1.))
print(p.p0, p.pc)
```

    tensor([0.2000]) tensor([0.8000])



```python
x = p.rsample(torch.Size([10]))
x
```




    tensor([[0.3865],
            [0.2682],
            [0.5663],
            [0.0000],
            [0.0557],
            [0.4937],
            [0.5800],
            [0.0000],
            [0.3457],
            [0.2662]])




```python
_= plt.hist(p.rsample(torch.Size([1000])).numpy(), bins=50)
```


```python
p2 = MixtureD0C01(
    probs0=torch.full([1], 0.5),
    cont=RightTruncatedExponential(rate=torch.ones(1) * 1., upper=1.))
print(p2.p0, p2.pc)
```

    tensor([0.5000]) tensor([0.5000])



```python
_= plt.hist(p2.rsample(torch.Size([1000])).numpy(), bins=50)
```


```python
mc_kl(p, p2, 1000000)
```




    tensor([0.1929])




```python
kl_divergence(p, p2)
```




    tensor([0.1927])




```python
mc_entropy(p, 1000000)
```




    tensor([0.4673])




```python
p.entropy()
```




    tensor([0.4679])



# Mixture of Discrete and Continuous over [0,1]

This distribution is a mixture of a point masses at $0$, another at $1$, and a continuous distribution over the support $(0, 1)$.

The distribution function looks like this:

\begin{equation}
f(x|\alpha, \beta) = \alpha_0 \delta(x) + \alpha_1 \delta(1 - x) + (1-\alpha_0-\alpha_1) g(x|\beta) \mathbf 1_{(0, 1)}(x)
\end{equation}

where $0 < \alpha_i < 1$, $g(x|\beta)$ is a (properly normalised) density over $(0,1)$, $\mathbf 1_{(0, 1)}(x)$ the indicator function.


```python
from probabll.distributions.mixdc import MixtureD01C01
```


```python
probs = torch.from_numpy(np.array([0.3, 0.5, 0.2])).float()
probs
```




    tensor([0.3000, 0.5000, 0.2000])




```python
p = MixtureD01C01(
    cont=Uniform(torch.zeros([1]), torch.ones([1])),
    probs=probs
)
```


```python
x = p.rsample(torch.Size([10]))
x
```




    tensor([[1.0000],
            [1.0000],
            [1.0000],
            [0.0637],
            [0.0452],
            [1.0000],
            [1.0000],
            [1.0000],
            [0.0000],
            [0.0000]])




```python
_= plt.hist(p.rsample(torch.Size([1000])).numpy(), bins=50)
```


```python
q = MixtureD01C01(
    cont=Uniform(torch.zeros([1]), torch.ones([1])),
    probs=torch.from_numpy(np.array([0.2, 0.5, 0.3])).float()
)
```


```python
p.entropy()
```




    tensor([1.0297])




```python
q.entropy()
```




    tensor([1.0297])




```python
mc_entropy(p, 1000000)
```




    tensor([1.0298])




```python
kl_divergence(p, q)
```




    tensor([0.0405])




```python
mc_kl(p, q, 1000000)
```




    tensor([0.0407])




```python
- torch.log(p.cont.cdf(1) - p.cont.cdf(0))
```




    tensor([-0.])



# Prior and Posterior in VI

Let's specify a prior

\begin{equation}
    \alpha_0 \delta(x) + \alpha_1 \delta(1-x) + (1-\alpha_0-\alpha_1) \mathcal U(x|0, 1) \mathbf 1_{(0,1)}(x)
\end{equation}

and an approximate posterior
\begin{equation}
    \alpha_0 \delta(x) + \alpha_1 \delta(1-x) + (1-\alpha_0-\alpha_1) \text{Exp}_{(0,1)}(x|\lambda) \mathbf 1_{(0,1)}(x)
\end{equation}

where $\text{Exp}_{(0,1)}(\lambda)$ is an RightTruncatedExponential rv.

We provide support to KL for pairs like that, where the densities are mixtures D01C01 (or D0C01) and the prior uses a Uniform density over $(0, 1)$. They are particularly simple due to the uniform prior. Whenever the density used in the approximate posterior admits a closed-form entropy, we have the entire expression in closed form!


```python
p = MixtureD01C01(
    cont=Uniform(torch.zeros([1]), torch.ones([1])),
    probs=torch.from_numpy(np.array([0.2, 0.5, 0.3])).float()
)
```


```python
q1 = MixtureD01C01(
    cont=RightTruncatedExponential(rate=torch.full([1], 10.), upper=1.),
    probs=torch.from_numpy(np.array([0.3, 0.5, 0.2])).float()
)
```


```python
_= plt.hist(q1.rsample(torch.Size([1000])).numpy(), bins=50)
```


```python
kl_divergence(q1, p)
```




    tensor([0.3012])




```python
mc_kl(q1, p, 1000000)
```




    tensor([0.3001])




```python
q1.entropy()
```




    tensor([0.7690])




```python
mc_entropy(q1, 1000000)
```




    tensor([0.7696])



Another possible pair:

\begin{equation}
    \alpha_0 \delta(x) + \alpha_1 \delta(1-x) + (1-\alpha_0-\alpha_1) \mathcal U(x|0, 1) \mathbf 1_{(0,1)}(x)
\end{equation}

and an approximate posterior
\begin{equation}
    \alpha_0 \delta(x) + \alpha_1 \delta(1-x) + (1-\alpha_0-\alpha_1) \text{Kumaraswamy}(x|a, b) \mathbf 1_{(0,1)}(x)
\end{equation}


```python
q2 = MixtureD01C01(
    cont=Kumaraswamy(a=torch.full([1], 5.), b=torch.full([1], 5.)),
    probs=torch.from_numpy(np.array([[0.3, 0.5, 0.2], [0.2, 0.5, 0.3]])).float()
)
```


```python
_= plt.hist(q2.rsample(torch.Size([1000])).numpy(), bins=50)
```


```python
kl_divergence(q2, p)
```




    tensor([0.1590, 0.1777])




```python
mc_kl(q2, p, 1000000)
```




    tensor([0.1637, 0.1774])




```python
q2.entropy()
```




    tensor([0.9112, 0.8520])




```python
mc_entropy(q2, 1000000)
```




    tensor([0.9069, 0.8512])




```python

```

32953beb1d31e15482335ccb16ca016db5aa1392
ipynb
Jupyter Notebook
notebooks/Mixture.ipynb
EelcovdW/dists.pt
994c698c846226ccf1905302f2bab860a2b24ca4
MIT
notebooks/Mixture.ipynb
EelcovdW/dists.pt
994c698c846226ccf1905302f2bab860a2b24ca4
MIT
notebooks/Mixture.ipynb
EelcovdW/dists.pt
994c698c846226ccf1905302f2bab860a2b24ca4
MIT
Qwen/Qwen-72B
1. YES
2. YES