# An infinite sum of Poisson-weighted convolutions

Kyle Cranmer, Aug 2018

If viewing on GitHub, this looks better with nbviewer: [click here](http://nbviewer.jupyter.org/github/cranmer/poisson-convolution-sum/blob/master/Poisson-weighted-convolutions.ipynb)

Consider a variable $x = \sum_{i=1}^n z_i$, where the $z_i \sim p(z|\theta)$ are iid and $n \sim Pois(\nu)$. For each value of $n$, the distribution $p(x|n)$ is given by repeated convolutions
\begin{equation}
{p(x|n, \theta)} = p(z|\theta)\underbrace{\star \dots \star}_\text{n times} p(z|\theta)
\end{equation}
and for a Poisson mean of $\nu$ the expected distribution for $x$ is  
\begin{equation}
{p(x|\nu, \theta)} = \sum_{n=0}^\infty e^{-\nu} \frac{\nu^n}{n!} p(x|n, \theta)
\end{equation}

The infinite tower of Poisson-weighted convolutions can be computed efficiently with a nice trick documented in *Analytic Confidence Level Calculations using the Likelihood Ratio and Fourier Transform* by Hongbo Hu and Jason Nielsen https://arxiv.org/pdf/physics/9906010.pdf. See also [this old paper](https://arxiv.org/abs/physics/0312050) and this [code](http://phystat.org/phystat/packages/0703002.1.html) for a C++ implementation.

First we take advantage of the convolution theorem relating convolutions to multiplication in the Fourier domain, denoted with a bar.
\begin{equation}
\overline{p(x|n, \theta)} = \overline{p(z|\theta)\underbrace{\star \dots \star}_\text{n times} z|\theta} = \left[\overline{p(z|\theta)} \right]^n 
\end{equation}
and then we can compress the infinite sum into the following exponential
\begin{equation}
\overline{p(x|\nu, \theta)} = \sum_{n=0}^\infty e^{-\nu} \frac{\nu^n}{n!}  \overline{p(x|n)}  =  \exp\left( \nu \left[ \overline{p(z|\theta)}-1 \right] \right)
\end{equation}

In the code below, we will reimplement the technique using pytorch. An interesting feature of implementing this in pytorch is that we can backprop through the inverse FFT, the exponentiation, the multplication, the subtraction, and forward FFT to calculate the gradient
\begin{equation}
\nabla_\theta\, p(x|\nu, \theta)
\end{equation}
and 
\begin{equation}
\nabla_\nu \, p(x|\nu, \theta)
\end{equation}.


```python
import torch
import numpy as np
import matplotlib.pyplot as plt
```

## Classes for Poisson convolutions with complex FFT 


```python
def complex_exp_pois(z, nu):
    real1, imag1 = z.t()
    real1=real1-1.
    real_out = torch.exp(real1*nu)*torch.cos(nu*imag1)
    imag_out = torch.exp(real1*nu)*torch.sin(nu*imag1)
    return torch.stack([real_out, imag_out], dim = -1)
```


```python
# This defines p(z|\theta)
class base_dist():
    def __init__(self):
        self.theta = torch.ones(1,1, requires_grad=True, dtype=torch.double)
        self.sigma = torch.tensor(0.051,requires_grad=True, dtype=torch.double)        
    def forward(self, x):
        max_ = torch.max(x)
        min_ = torch.min(x)
        nbins = x.size()[1]
        norm = torch.tensor(max_/nbins/np.sqrt(2.*np.pi), dtype=torch.double)/self.sigma
        y = torch.empty(1,nbins,2, dtype=torch.double)
        y[0,:,0] = (torch.exp(-(x-self.theta)*(x-self.theta)/2./self.sigma/self.sigma))*norm
        y[0,:,1]= torch.zeros(nbins)
        return y
    def parameters(self):
        return [self.theta, self.sigma]
```


```python
# This does Poisson convolutions of the base distribution p(z|\theta)
class pois_conv():
    def __init__(self):
        self.nu = torch.tensor(3.,requires_grad=True, dtype=torch.double)
        self.base_dist = base_dist()
    def forward(self, x):
        y = self.base_dist.forward(x)
        yfft = torch.fft(y,1,normalized=False)
        y2fft = complex_exp_pois(yfft[0], nu=self.nu).view(-1,nbins,2)
        y2 = torch.ifft(y2fft,1, normalized=False)
        return y2
    def parameters(self):
        temp = list(self.base_dist.parameters())
        temp.append(self.nu)
        return temp
```

## Check with simple Gaussian $p(z|\theta)$ base distribution


```python
min_=0
max_=15
nbins=3000

x = torch.arange(min_, max_,(max_-min_)/nbins, dtype=torch.double)
x=x.view(-1,nbins)
```


```python
y2_ = pois_conv()
y2_.parameters()
y_ = y2_.base_dist
y2_.forward(x)
y2_.parameters()
```




    [tensor([[ 1.]], dtype=torch.float64), tensor(1.00000e-02 *
            5.1000, dtype=torch.float64), tensor(3., dtype=torch.float64)]




```python
#quick helper to for chekcing Poisson probabilities
def myfact(xin):
    return [np.math.factorial(x) for x in xin]
myfact(np.arange(4))
```




    [1, 1, 2, 6]




```python
# Poisson probs
nu = y2_.nu
xint = np.arange(min_, max_)
pois = np.exp(-nu.item())*np.power(nu.item(),xint)/myfact(xint)
scaled_pois = pois
#heights of Gaussian bumps will be lower by width of Gaussian, which scales like sqrt(x()
scaled_pois[1:]= scaled_pois[1:]/np.sqrt(xint[1:])/np.sqrt(2*np.pi)/y_.sigma.item()*max_/nbins
```


```python
with torch.no_grad():
    print(y2_.nu.item())
    #plt.plot(x[0].detach().numpy(),y_.forward(x)[0,:,0].detach().numpy())
    plt.plot(x[0].detach().numpy(),y2_.forward(x)[0,:,0].detach().numpy(), label="computed")
    plt.scatter(xint,scaled_pois, label="predicted peak")
    plt.ylim(0,np.max(1.5*y2_.forward(x)[0,1:,0].detach().numpy()))
    plt.xlabel("x")
    plt.legend()
```

## Do a fit

First make some dummy data


```python
y2_ = pois_conv()
with torch.no_grad():
    y2_.nu+=1.5
    y2_.base_dist.theta+=.1
    y2_.base_dist.sigma+=.1
data = torch.zeros_like(y2_.forward(x))
data.data = y2_.forward(x)
```

Now gradient descent


```python
y2_ = pois_conv()
prefit_prediction = y2_.forward(x)

# do the fit
learning_rate = 0.5
for t in range(10001):
    out = y2_.forward(x)
    loss = (data-out).pow(2).sum()
    if t%500==0: print(f"\nt={t}\tloss={loss}")
    loss.backward()
    with torch.no_grad():
        for param in y2_.parameters():
            if t%500==0: print(f"\tparam ", param.item())
            param -= learning_rate*param.grad
            param.grad.zero_()
```

    
    t=0	loss=0.003986922902300565
    	param  1.0
    	param  0.051
    	param  3.0
    
    t=500	loss=0.0004082261412487773
    	param  1.1069657856157538
    	param  0.1769350213826769
    	param  3.5400895045643637
    
    t=1000	loss=0.00019254864368328299
    	param  1.104256439223183
    	param  0.1674518635502764
    	param  3.769002476383715
    
    t=1500	loss=0.0001108094209823346
    	param  1.1031418506638224
    	param  0.16309889590260718
    	param  3.91099116864832
    
    t=2000	loss=7.033956699939845e-05
    	param  1.1024834462038822
    	param  0.16046594784319412
    	param  4.011160320625154
    
    t=2500	loss=4.735130257572868e-05
    	param  1.102034468115289
    	param  0.15866951450029707
    	param  4.086748654745574
    
    t=3000	loss=3.314869426685245e-05
    	param  1.1017038863350115
    	param  0.15735672225680727
    	param  4.146203055164765
    
    t=3500	loss=2.3863101343512058e-05
    	param  1.1014485810067902
    	param  0.15635343805652147
    	param  4.1942969840307684
    
    t=4000	loss=1.7540769419754663e-05
    	param  1.1012449280000847
    	param  0.15556214994448553
    	param  4.233993118894525
    
    t=4500	loss=1.3103148307666452e-05
    	param  1.1010786615022556
    	param  0.15492333853285223
    	param  4.267256955428157
    
    t=5000	loss=9.914245638328102e-06
    	param  1.1009405412434996
    	param  0.1543983108783803
    	param  4.295459111782398
    
    t=5500	loss=7.57951268948578e-06
    	param  1.1008242563348958
    	param  0.1539606685597788
    	param  4.319593106709681
    
    t=6000	loss=5.844156926365047e-06
    	param  1.1007253141406788
    	param  0.15359169879646456
    	param  4.340401711692256
    
    t=6500	loss=4.538209163846984e-06
    	param  1.1006404082358734
    	param  0.15327771812946442
    	param  4.358454400061147
    
    t=7000	loss=3.545217693286112e-06
    	param  1.1005670380541883
    	param  0.15300845972835223
    	param  4.374196991355995
    
    t=7500	loss=2.7836023312572676e-06
    	param  1.1005032692037167
    	param  0.15277605152118925
    	param  4.387984677341635
    
    t=8000	loss=2.1951300970023443e-06
    	param  1.1004475765038195
    	param  0.15257434506746897
    	param  4.4001046931404355
    
    t=8500	loss=1.7375649295138236e-06
    	param  1.1003987377854716
    	param  0.1523984611320626
    	param  4.410792304345434
    
    t=9000	loss=1.3798512568451279e-06
    	param  1.1003557599943536
    	param  0.15224447378635297
    	param  4.4202423467314675
    
    t=9500	loss=1.0988831508716698e-06
    	param  1.100317826502317
    	param  0.15210918571107085
    	param  4.428617727693528
    
    t=10000	loss=8.772923137453809e-07
    	param  1.1002842587278634
    	param  0.15198996510843235
    	param  4.4360558034693085



```python
with torch.no_grad():
    plt.plot(x[0].detach().numpy(),data[0,:,0].detach().numpy(), label="data")
    plt.plot(x[0].detach().numpy(),y2_.forward(x)[0,:,0].detach().numpy(), label="fit")    
    plt.plot(x[0].detach().numpy(),prefit_prediction[0,:,0].detach().numpy(), alpha=.3, label="pre fit")    
    plt.ylim(0,np.max(1.5*data[0,1:,0].detach().numpy()))
    plt.xlabel('x')
    plt.legend()
```

## It works!

d6086b15d05d0695bc9fb258c4b663b71f9b817a
ipynb
Jupyter Notebook
Poisson-weighted-convolutions.ipynb
cranmer/poisson-convolution-sum
90943b87e7162e64938ba11574e2edaf2c7bd5f5
MIT
2018-08-22T20:13:49.000Z
2020-06-08T03:37:30.000Z
Poisson-weighted-convolutions.ipynb
cranmer/poisson-convolution-sum
90943b87e7162e64938ba11574e2edaf2c7bd5f5
MIT
Poisson-weighted-convolutions.ipynb
cranmer/poisson-convolution-sum
90943b87e7162e64938ba11574e2edaf2c7bd5f5
MIT
2018-08-23T01:17:51.000Z
2018-09-04T19:40:03.000Z
Qwen/Qwen-72B
1. YES
2. YES