**A beta-Poisson model for infectious disease transmission**

This notebook is designed to accompany the upcoming paper of the same title by Joe Hilton and Ian Hall. In the paper, we introduce a branching process model for infectious disease transmission whose offspring distribution is a beta-Poisson mixture distribution. We compare this model's performance as a description of transmission by fitting it to several examples of transmission tree data and comparing these fits to those obtained by the Poisson, geometric, negative binomial, and zero-inflated Poisson (ZIP) models. The code in this notebook includes functions which perform standard probability and likelihood calculations for this model, estimate maximum likelihood parameters, and estimate confidence intervals for these parameters using bootstrapping. We provide a rough outline of the model here, with full derivations left to the paper.

```python
from __future__ import print_function
import math
import numpy as np
import random
import scipy as sp
from scipy import special as spsp
from scipy import sparse as sparse
from scipy import stats as stats
from scipy import optimize as opt
from scipy import interpolate as interpolate
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import matplotlib.path as path
import matplotlib.animation as animation
from ipywidgets import interact, interactive, fixed, interact_manual
import ipywidgets as widgets
import mpmath
import matplotlib as mpl
import matplotlib.cm as cm
import time
cmap = cm.hot
```

**Model description**

The beta-Poisson model describes the person-to-person spread of a pathogen in a population where contact behaviour is homogeneous but transmission behaviour varies from person to person. During their infectious period an infectious individual makes $n$ contacts, drawn from a Poisson distribution with mean $N$. For each case a transmission probability $p$ is chosen from a beta distribution with parameters $(\alpha_1,\alpha_2)$, so that the number of subsequent cases is drawn from a beta-binomial distribution with these parameters and $n$ trials. In our paper we demonstrate that the mean of the resulting distribution is given by
\begin{equation}
\lambda=N\frac{\alpha_1}{\alpha_1+\alpha_2},
\end{equation}
and that if we make the substitution $\Phi=\frac{\alpha_1+\alpha_2}{N}$ the probability mass function can be expressed as:
\begin{equation}
\begin{aligned}
P(x;\lambda,\Phi,N)=\frac{N^x}{\Gamma(x+1)}\frac{\Gamma(x+\Phi \lambda)\Gamma(\Phi N)}{\Gamma(\Phi \lambda)\Gamma(x+\Phi N)}M(x+\Phi \lambda,x+\Phi N,-N).
\end{aligned}
\end{equation}

We demonstrate in the paper that in the limit $N\to\infty$, the beta-Poisson distribution approximates a beta-gamma mixture, i.e. a negative binomial. We will use a slightly unusual parameterisation . The negative binomial with mean $\lambda$ is overdispersed, so that its variance can be expressed as $\lambda(1+\theta)$ for some $\theta>0$. We parameterise the negative binomial in terms of $\lambda$ and $\theta$ so that the pmf is given by
\begin{align}
P(x;\lambda,\theta)=\frac{\Gamma(x+\frac{\lambda}{\theta})}{\Gamma(x+1)\Gamma({\frac{\lambda}{\theta}})}\Bigg(\frac{1}{1+\theta}\Bigg)^{\frac{\lambda}{\theta}}\Bigg(\frac{\theta}{\theta+1}\Bigg)^x.
\end{align}
The negative binomial we obtain in the limit $N\to\infty$ has overdispersion $\theta=\Phi^{-1}$.

```python
def beta_poisson_pmf(x,lmbd,Phi,N):
    if type(x)==int:
        P=spsp.hyp1f1(x+Phi*lmbd,x+Phi*N,-N)
        for n in range(1,x+1): # This loop gives us the N^x/gamma(x+1 term)
            P=(N/n)*P
        for m in range(x): # This loop gives us the term with the two gamma functions in numerator and denominator
            P=((m+Phi*lmbd)/(m+Phi*N))*P
    else:
        P=[]
        for i in range(0,len(x)):
            p=spsp.hyp1f1(x[i]+Phi*lmbd,x[i]+Phi*N,-N)
            for n in range(1,x[i]+1): # This loop gives us the N^x/gamma(x+1 term)
                p=(N/n)*p
            for m in range(x[i]): # This loop gives us the term with the two gamma functions in numerator and denominator
                p=((m+Phi*lmbd)/(m+Phi*N))*p
            P=P+[p]
    return P
```

```python
class beta_poisson_gen:

#    def __init__(self, lmbd, Phi, N):
#        self.lmbd = lmbd
#        self.Phi = Phi
#        self.N = N

    def pmf(self, x, lmbd, Phi, N):
        return np.exp(self.logpmf(x,lmbd,Phi,N))

    def logpmf(self,x,lmbd,Phi,N):
        if type(x)==int:
            log_p = (x*np.log(N) -
                     np.real(spsp.loggamma(x+1)) +
                     np.real(spsp.loggamma(Phi*N)) +
                     np.real(spsp.loggamma(x+Phi*lmbd)) -
                     np.real(spsp.loggamma(x+Phi*N)) -
                     np.real(spsp.loggamma(Phi*lmbd)))
            if x+phi*N<50:
                log_p += np.log(spsp.hyp1f1(x+Phi*lmbd,x+Phi*N,-N))
            else:
                log_p += np.log(float(mpmath.hyp1f1(x+Phi*lmbd,x+Phi*N,-N)))
        else:
            log_p = []
            for i in range(0,len(x)):
                log_pp = (x[i]*np.log(N)-np.real(spsp.loggamma(x[i]+1)) +
                    np.real(spsp.loggamma(Phi*N)) +
                    np.real(spsp.loggamma(x[i]+Phi*lmbd)) -
                    np.real(spsp.loggamma(x[i]+Phi*N)) -
                    np.real(spsp.loggamma(Phi*lmbd)))
                if x[i]+Phi*N<50:
                    log_pp += np.log(spsp.hyp1f1(x[i]+Phi*lmbd,x[i]+Phi*N,-N))
                else:
                    log_pp += np.log(float(mpmath.hyp1f1(x[i]+Phi*lmbd,x[i]+Phi*N,-N)))
                log_p = log_p + [log_pp]
        return log_p

    def pgf(self, s, lmbd, Phi, N):
        if np.size(np.asarray(s))==1:
            G=spsp.hyp1f1(lmbd*Phi,N*Phi,N*(s-1));
        else:
            G=[]
            for i in range(0,np.size(np.asarray(s))):
                G=G+[spsp.hyp1f1(lmbd*Phi,N*phi,N*(s[i]-1))]
        return G

    def extinction_prob(self, lmbd, phi, N ):
        if lmbd<=1:
            return 1
        else:
            def f(s):
                return beta_poisson_pgf(s,lmbd,phi,N)-s
            q=opt.brentq(
                f, 0, 1-1e-4);
            return q

beta_poisson = beta_poisson_gen()
```

```python
print(beta_poisson_pmf([1,2,3,4],2,0.1,5))
print(beta_poisson.pmf([1,2,3,4],2,0.1,5))
```

    [0.11547614919829788, 0.09421813032224745, 0.08866313316657161, 0.07987675522526165]
    [0.11547615 0.09421813 0.08866313 0.07987676]

The next code block generates a widget with which you can plot the pmf of the beta-Poisson distribution and use a set of sliders to control the values of $\lambda$, $\Phi$, and $N$. You should find that when $\lambda$ is close to $N$ or $\Phi$ is large the distribution is more symmetrical, and that decreasing either $\lambda$ or $\Phi$ skews the distribution towards zero.

```python
%matplotlib inline
from ipywidgets import interactive
import matplotlib.pyplot as plt
import numpy as np

def f(lmbd,Phi,N):
    fig, ax=plt.subplots(figsize=(5,5))
    plt.figure(1)
    x = list(range(int(np.ceil(lmbd+4*np.sqrt(lmbd*(1+(N-lmbd)/(Phi*N+1)))))))
    P=beta_poisson_pmf(x,lmbd,Phi,N)
    plt.bar(x,P)
    plt.show()

lmbd_widget=widgets.FloatSlider(min=0.05,max=5,step=0.05,value=1)
Phi_widget=widgets.FloatSlider(min=0.01,max=2,step=0.01,value=0.2)
N_widget=widgets.FloatSlider(min=0.1,max=40,step=0.1,value=5)
def update_lmbd_range(*args):
    lmbd_widget.max=N_widget.value
N_widget.observe(update_lmbd_range,'value')
interactive_plot=interactive(f,lmbd=lmbd_widget,Phi=Phi_widget,N=N_widget)
ourput=interactive_plot.children[-1]
interactive_plot
```

    interactive(children=(FloatSlider(value=1.0, description='lmbd', max=5.0, min=0.05, step=0.05), FloatSlider(va…

**Likelihood calculations**

Suppose that we have a vector $(x_1,...,x_K)$ of secondary case data, i.e. $x_i$ is the number of cases which have case $i$ as their infector (note that in this work we do not address the problem of how to identify these infectors and so generate a dataset). The log-likelihood of parameters $(\lambda,\Phi,N)$ given this data is given by
\begin{align}
\log\mathcal{L}=\sum\limits_{i=1}^Kx_i\log N-\log\Gamma(x_i+1)+\log\Gamma(\Phi N)+\log\Gamma(x_i+\Phi \lambda)-\log\Gamma(x_i+\Phi N)-\log\Gamma(\Phi \lambda)+\log M(x_i+\Phi \lambda,x_i+\Phi N,-N).
\end{align}
For large values of $N$ this formula is numerically unstable, and in this region we will use the log likelihood function of the negative binomial distribution with $\theta=\Phi^{-1}$. One can justify the choice of $N=1,000$ as a cutoff by calculating the beta-Poisson and negative binomial probability mass functions up to a suitable maximum (x=250 seems a generous upper limit for secondary case data) and for a range of $\lambda$ and $\Phi$ values. In an emerging infection context we are interested in $\lambda$ up to about $2$ and, since $\Phi$ is inversely proportional to the level of overdisperion in the data, considering $\Phi$ values below ten is usually sufficient for fitting overdispersed datasets.

In the next code block, `beta_poisson_loglh` calculates the log likelihood of the beta-Poisson parameters $(\lambda,\Phi,N)$ given a set of count data, and `neg_bin_bp_loglh` the log likelihood of the negative binomial parameters $(\lambda,\Phi^{-1})$, or equivalently the beta-Poisson parameters $(\lambda,\Phi,\infty)$. For the purposes of this notebook it is useful for us to define separate functions calculating the negative binomial log likelihood in the $N\to\infty$ extreme of the beta-Poisson model, and as the likelihood function of the negative binomial model in its own right - hence the "`bp`" in the name of the function in the box immediately below.

```python
def beta_poisson_loglh(data,lmbd,phi,N):
    return sum(beta_poisson.logpmf(data,lmbd,phi,N))

def neg_bin_bp_loglh(data,lmbd,phi):
    return sum(stats.nbinom.logpmf(data,lmbd*phi,phi/(phi+1)))
```

Given a vector of count data, we can calculate the maximum likelihood parameters numerically by finding the minimum of the likelihood function. This is a two-dimensional problem, since the MLE of $\lambda$ is just the sample mean. The maximum likelihood estimation is carried out in the function `get_phi_and_N_mles`. Because $N$ can theoretically range up to infinity but its MLE is bounded below by that of $\lambda$, we optimise over $\nu=\frac{1}{N}$ rather than $N$. For values of $\nu$ less than $10^{-3}$ (i.e. values of $N$ greater than $1,000$) the optimiser takes `neg_bin_bp_loglh` as the function to optimise over rather than `beta_poisson_loglh`.

```python
def get_phi_and_N_mles(data,phi_0,N_0):
    def f(params):
        lmbd=np.mean(data)
        phi=params[0]
        if params[1]>0.1e-3:
            N=1/params[1]
            return -beta_poisson_loglh(data,lmbd,phi,N)
        else:
            return -neg_bin_bp_loglh(data,lmbd,phi)
    mle=sp.optimize.minimize(f,[phi_0,N_0],bounds=((1e-6,50),(0,1/np.mean(data))))
    if mle.x[1]<0:
        mle.x[1]=0
    return mle.x[0],mle.x[1]
```

In the supplementary material to the paper we show that the pgf of a beta-Poisson distributed random variable $X$ is
\begin{equation}
    G_X(s)=M(\lambda\Phi,N\Phi,N(s-1)).
\end{equation}
This is calculated using the function `beta_poisson_pgf`. From standard branching process theory, the extinction probability of a beta-Poisson epidemic is the solution to the equation
\begin{equation}
    q=G(q),
\end{equation}
and can be calculated using the function `beta_poisson_extinction_prob`.

These functions are not used to calculate any of the results presented in our paper, but are included here for completeness.

The functions in the following block carry out the equivalent calculations (log likelihood, maximum likelihood estimates, pgf, and extinction probability) for the Poisson, geometric, negative binomial, and ZIP calculations.

The function `nbinom.pmf` in the `scipy.stats` package uses a parameterisation in terms of parameters $n$ and $p$. Under this parameterisation the mean and variance of the distribution are respectively given by $(1-p)n/p$ and $(1-p)n/p^2$. After some rearrangement, one can show that $n=\lambda/\theta$ and $p=1/(\theta+1)$, so that to calculate negative binomial probabilities (and likelihoods) in our $(\lambda,\theta)$ parameterisation we use `stats.nbinom.pmf(x,lmbd/theta,1/(theta+1))`.

```python
def poisson_loglh(data,lmbd):
    return sum(stats.poisson.logpmf(data,lmbd))

def geo_loglh(data,lmbd):
    return sum(stats.geom.logpmf(data,1/(lmbd+1),-1))

def neg_bin_loglh(data,lmbd,theta):
    return sum(stats.nbinom.logpmf(data,lmbd/theta,1/(theta+1)))

def get_theta_mle(data,theta_0):
    def f(theta):
        lmbd=np.mean(data)
        return -neg_bin_loglh(data,lmbd,theta)
    mle=sp.optimize.minimize(f,[theta_0],bounds=((1e-6,50),))
    return mle.x[0]

def zip_pmf(x,lmbd,sigma):
    if type(x)==int:
        return sigma*(x==0)+(1-sigma)*stats.poisson.pmf(x,lmbd)
    else:
        return sigma*np.equal(x,np.zeros(len(x)))+(1-sigma)*stats.poisson.pmf(x,lmbd)

def zip_loglh(data,lmbd,sigma):
    llh=0
    for x in data:
        if x==0:
            llh+=np.log(sigma+(1-sigma)*np.exp(-lmbd))
        else:
            llh+=np.log(1-sigma)+np.log(stats.poisson.pmf(x,lmbd))
    return llh

def get_zip_mles(data,lmbd_0,sigma_0):
    def f(params):
        lmbd=params[0]
        sigma=params[1]
        return -zip_loglh(data,lmbd,sigma)
    mle=sp.optimize.minimize(f,[lmbd_0,sigma_0],bounds=((np.mean(data),50),(0,1-1e-6)))
    return mle.x[0],mle.x[1]
```

**Fits to transmission chain data**

In our paper we study the fitting behaviour of the beta-Poisson model by fitting it (and our other models of interest) to eight transmission chain datasets. We estimate maximum likelihood parameters and calculate the Akaike Information Criterion (AIC) for each fitted model. In the next cell we define a class whose objects contain the parameter fits, log likelihoods, and AICs of each model for a given dataset.

```python
class trans_chain_mles:
    def __init__(self,
                data):

        self.data = data
        self.mean = np.mean(data)
        self.var = np.var(data)

        theta_mle = get_theta_mle(data, self.mean)

        lmbd_mle, sigma_mle = get_zip_mles(data, self.mean, 0.5)

        phi_mle, Nu_mle=get_phi_and_N_mles(data, 1, 1/max(data))

        if Nu_mle>1e-3:
            beta_poi_var = self.mean * (1 + (1 - self.mean*Nu_mle)/(phi_mle + Nu_mle))
        else:
            beta_poi_var = self.mean * (1 + 1/phi_mle)

        poisson_llh = poisson_loglh(data, self.mean)
        geometric_llh = geo_loglh(data, self.mean)
        neg_bin_llh = neg_bin_loglh(data, self.mean, theta_mle)
        zip_llh = zip_loglh(data, lmbd_mle, sigma_mle)

        if Nu_mle>1e-3:
            beta_poi_llh = beta_poisson_loglh(data, self.mean, phi_mle, 1/Nu_mle)
        else:
            beta_poi_llh = neg_bin_loglh(data, self.mean, theta_mle)

        self.poisson = {'lambda' : self.mean,
                        'mean' : self.mean,
                        'var' : self.mean,
                        'llh' : poisson_llh,
                        'AIC' : 2 - 2*poisson_llh}

        self.geometric = {'lambda' : self.mean,
                          'mean' : self.mean,
                          'var' : self.mean * (1+self.mean),
                          'llh' : geometric_llh,
                          'AIC' : 2 - 2*geometric_llh}

        self.neg_bin = {'lambda' : self.mean,
                        'theta' : theta_mle,
                        'mean' : self.mean,
                        'var' : self.mean * (1+theta_mle),
                        'llh' : neg_bin_llh,
                        'AIC' : 4 - 2*neg_bin_llh}

        self.zip = {'lambda' : lmbd_mle,
                    'sigma' : sigma_mle,
                    'mean' : lmbd_mle * (1-sigma_mle),
                    'var' : lmbd_mle * (1-sigma_mle) * (1+lmbd_mle*sigma_mle),
                    'llh' : zip_llh,
                    'AIC' : 4 - 2*zip_llh}

        self.beta_poi = {'lambda' : self.mean,
                         'Phi' : phi_mle,
                         'Nu' : Nu_mle,
                         'mean' : self.mean,
                         'var' : beta_poi_var,
                         'llh' : beta_poi_llh,
                         'AIC' : 6 - 2*beta_poi_llh}

```

The next cell introduces a function which plots the empirical and fitted secondary case distributions.

```python
def plot_chain_data_fits(fits, ax, title, fs=20):

    xVals = range(max(fits.data)+1)

    counts, bins = np.histogram(fits.data, max(fits.data)+1)
    dist = counts / len(fits.data)

    ax.plot(np.where(dist>0)[0], dist[dist>0], 'x', markersize=20)

    poi_line = stats.poisson.pmf(xVals, fits.poisson['lambda'])
    ax.plot(xVals, poi_line, lw=3)

    geom_line = stats.geom.pmf(xVals, 1/(fits.geometric['lambda']+1),-1)
    ax.plot(xVals, geom_line, lw=3)

    neg_bin_line = stats.nbinom.pmf(xVals,
                                  fits.neg_bin['lambda']/fits.neg_bin['theta'],
                                  1/(fits.neg_bin['theta']+1))
    ax.plot(xVals, neg_bin_line, lw=3)

    zip_line = zip_pmf(xVals, fits.zip['lambda'], fits.zip['sigma'])
    ax.plot(xVals, zip_line, lw=3)

    if fits.beta_poi['Nu']>1e-3:
        beta_poi_line = beta_poisson.pmf(xVals,
                                         fits.beta_poi['lambda'],
                                         fits.beta_poi['Phi'],
                                         1/fits.beta_poi['Nu'])
        ax.plot(xVals, beta_poi_line, lw=3)
    else:
        beta_poi_line = neg_bin_line

    ax.axis([-0.5, max(fits.data)+0.5, 0, 1])

    ax.set_aspect((max(fits.data)+1))

    ax.set_xlabel('Secondary cases', fontsize=fs)
    ax.set_ylabel('Probability', fontsize=fs)
    ax.set_title(title, fontsize=fs)
```

In the following cell we define another plotting function, this one plotting the beta distribution which defines the individual-level variation in infectivity.

```python
def plot_underlying_beta(fits):

    x = np.linspace(1e-2,1, 100)
    y = stats.beta.pdf(x,
                       fits.beta_poi['lambda'] * fits.beta_poi['Phi'],
                       (1/fits.beta_poi['Nu'] - fits.beta_poi['lambda']) * fits.beta_poi['Phi'])

    fig,ax = plt.subplots(figsize=(10,10))

    plt.plot(x, y,'r-', lw=5, alpha=0.6, label='beta pdf')

    ax.axis([-0.01, 1.01, 0, 1.01*np.max(y)])

    ax.set_aspect(1.02/(1.01*np.max(y)))

    plt.xlabel('Transmission probability',fontsize=25)

    plt.ylabel('PDF',fontsize=25)

    plt.xticks(fontsize=25)

    plt.yticks(fontsize=25)

    plt.show()
```

**Starting example: Plague**<br/>
From [Gani and Leach 2004](https://www.ncbi.nlm.nih.gov/pmc/articles/PMC3323083/pdf/03-0509.pdf).

To begin with, we calculate the maximum likelihood parameters and associated AIC for each model:

```python
gani_data = [0]*16+[1]*10+[2]*7+[3]*2+[4]*3+[5]+[6]

gani_fits = trans_chain_mles(gani_data)

np.set_printoptions(precision=3)
print('Distribution          | Lambda | Theta | Sigma |Nu=1/N |  Phi  | Log Likelihood |   AIC  |')
print('__________________________________________________________________________________________')
print('Poisson               |', "%.3f"% gani_fits.poisson['lambda'],' |       |       |       |       |   ', "%.3f"% gani_fits.poisson['llh'],'    |', "%.3f"% gani_fits.poisson['AIC'])
print('Geometric             |', "%.3f"% gani_fits.geometric['lambda'],' |       |       |       |       |   ', "%.3f"% gani_fits.geometric['llh'],'    |', "%.3f"% gani_fits.geometric['AIC'])
print('Negative binomial     |', "%.3f"% gani_fits.neg_bin['lambda'],' |', "%.3f"% gani_fits.neg_bin['theta'],'|       |       |       |   ', "%.3f"% gani_fits.neg_bin['llh'],'    |', "%.3f"% gani_fits.neg_bin['AIC'])
print('Zero-inflated Poisson |', "%.3f"% gani_fits.zip['lambda'],' |       |', "%.3f"% gani_fits.zip['sigma'],'|       |       |   ', "%.3f"% gani_fits.zip['llh'],'    |', "%.3f"% gani_fits.zip['AIC'])
print('Beta Poisson          |', "%.3f"% gani_fits.beta_poi['lambda'],' |       |       |', "%.3f"% gani_fits.beta_poi['Nu'],'|', "%.3f"% gani_fits.beta_poi['Phi'],'|   ', "%.3f"% gani_fits.beta_poi['llh'],'    |', "%.3f"% gani_fits.beta_poi['AIC'])
```

    Distribution          | Lambda | Theta | Sigma |Nu=1/N |  Phi  | Log Likelihood |   AIC  |
    __________________________________________________________________________________________
    Poisson               | 1.325  |       |       |       |       |    -67.422     | 136.843
    Geometric             | 1.325  |       |       |       |       |    -63.551     | 129.102
    Negative binomial     | 1.325  | 0.923 |       |       |       |    -63.312     | 130.624
    Zero-inflated Poisson | 1.867  |       | 0.290 |       |       |    -63.945     | 131.890
    Beta Poisson          | 1.325  |       |       | 0.248 | 0.583 |    -63.120     | 132.240

Plotting the data and fitted models demonstrates that the Poisson fails to capture the high incidence of zeros in the data:

```python
fig, ax = plt.subplots(figsize=(10,10))

plot_chain_data_fits(gani_fits, ax, 'Plague')

ax.legend(['Data','Poisson','Geometric','Negative binomial','ZIP','Beta-Poisson'],
          bbox_to_anchor=(1.05,1),
          loc='upper left')

plt.show()
```

We can also visualise the distribution of infectivity across cases in the fitted beta-Poisson by plotting the underlying beta distribution:

```python
plot_underlying_beta(gani_fits)
```

The fitting behaviour of the beta-Poisson model can be visualised by calculating thelikelihood function of each parameter with the other two fixed at their MLEs. The slow decay of the likelihood function associated with increasing values of $\Phi$ is consistent with the relatively high likelihood assigned to the Poisson distribution, which means that models with very low levels of overdispersion (i.e. large values of $\Phi$) can give a good fit.

```python
true_lmbd,true_phi,true_N=np.mean(current_data),phi_mle,1/N_inv_mle
data=current_data
lmbd_lh_array=beta_poisson_loglh(data,np.linspace(0.01,5,500),true_phi,true_N)
phi_lh_array=np.zeros(1000)
for i in range(1000):
    phi_lh_array[i]=beta_poisson_loglh(data,true_lmbd,(i+1)*1e-2,true_N)
n_lh_array=np.zeros(750)
n_lh_array[0]=neg_bin_bp_loglh(data,true_lmbd,true_phi)
for i in range(1,750):
    n_lh_array[i]=beta_poisson_loglh(data,true_lmbd,true_phi,1/(i/1000))

fig,ax=plt.subplots(figsize=(8,8))
ax.plot(np.linspace(0.01,5,500),lmbd_lh_array,'k',linewidth=6)
ax.plot([true_lmbd,true_lmbd],ax.get_ylim(),'--b',linewidth=6)
plt.xlabel('$\lambda$',fontsize=40)
plt.ylabel('Log likelihood',fontsize=40)
plt.xticks(fontsize=40)
plt.yticks(fontsize=40)
plt.show()

fig,ax=plt.subplots(figsize=(8,8))
ax.plot(np.linspace(0.01,10,1000),phi_lh_array,'k',linewidth=6)
ax.plot([true_phi,true_phi],ax.get_ylim(),'--b',linewidth=6)
plt.xlabel('$\Phi$',fontsize=40)
plt.ylabel('Log likelihood',fontsize=40)
plt.xticks(fontsize=40)
plt.yticks(fontsize=40)
plt.show()

fig,ax=plt.subplots(figsize=(8,8))
ax.plot(np.linspace(0.0,0.749,750),n_lh_array,'k',linewidth=6)
ax.plot([1/true_N,1/true_N],ax.get_ylim(),'--b',linewidth=6)
plt.xlabel('$\\nu$',fontsize=40)
plt.ylabel('Log likelihood',fontsize=40)
plt.xticks(fontsize=30)
plt.yticks(fontsize=40)
plt.show()
```

## More datasets

In the cell below we define some more datasets and fit maximum likelihood parameters to each of them for each model.

```python
jezek_data_g1=[0]*114+[1]*23+[2]*8+[3]*1+[5]*1 # First generation cases
jezek_data_g2=[0]*38+[1]*7+[2]*1+[3]*1 # Second generation cases
jezek_data_g3=[0]*9+[1]+[2]
jezek_data_g4=[0]*2+[1]
jezek_data=jezek_data_g1+jezek_data_g2+jezek_data_g3+jezek_data_g4

faye_data=[1,2,2,5,14,1,4,4,1,3,3,8,2,1,1,4,9,9,1,1,17,2,1,1,1,4,3,3,4,2,5,1,2,2,1,9,1,3,1,2,1,1,2]
faye_data=faye_data+[0]*(152-len(faye_data))
fasina_data=[0]*15+[1]*2+[2]+[3]+[12]
leo_data=[0]*162+[1]*19+[2]*8+[3]*7+[7]+[12]+[21]+[23]+[40]
cowling_data=[38,3,2,1,6,81,2,23,2,1,1,1,5,1,1,1,2,1,1,1]
cowling_data=cowling_data+[0]*(166-len(cowling_data))
chowell_data=[0]*13+[1]*5+[2]*4+[3]+[7]
heijne_data=[0]*22+[1]*13+[2]*6+[3]*3+[4]+[5]

jezek_fits = trans_chain_mles(jezek_data)
faye_fits = trans_chain_mles(faye_data)
fasina_fits = trans_chain_mles(fasina_data)
leo_fits = trans_chain_mles(leo_data)
cowling_fits = trans_chain_mles(cowling_data)
chowell_fits = trans_chain_mles(chowell_data)
heijne_fits = trans_chain_mles(heijne_data)
```

```python
print('Dataset | Empirical |      Poisson      |     Geometric      |  Negative binomial  |        ZIP        | Beta-Poisson ')
print('Gani    | (',"%.3f"% gani_fits.mean,',',"%.3f"% gani_fits.var,') | (',"%.3f"% gani_fits.poisson['mean'],',',"%.3f"% gani_fits.poisson['var'],') | (',"%.3f"% gani_fits.geometric['mean'],',',"%.3f"% gani_fits.geometric['var'],')  |  (',"%.3f"% gani_fits.neg_bin['mean'],',',"%.3f"% gani_fits.neg_bin['var'],')  | (',"%.3f"% gani_fits.zip['mean'],',',"%.3f"% gani_fits.zip['var'],') | (',"%.3f"% gani_fits.beta_poi['mean'],',',"%.3f"% gani_fits.beta_poi['var'],') ')
print('Jezek   | (',"%.3f"% jezek_fits.mean,',',"%.3f"% jezek_fits.var,') | (',"%.3f"% jezek_fits.poisson['mean'],',',"%.3f"% jezek_fits.poisson['var'],') | (',"%.3f"% jezek_fits.geometric['mean'],',',"%.3f"% jezek_fits.geometric['var'],')  |  (',"%.3f"% jezek_fits.neg_bin['mean'],',',"%.3f"% jezek_fits.neg_bin['var'],')  | (',"%.3f"% jezek_fits.zip['mean'],',',"%.3f"% jezek_fits.zip['var'],') | (',"%.3f"% jezek_fits.beta_poi['mean'],',',"%.3f"% jezek_fits.beta_poi['var'],') ')
print('Faye    | (',"%.3f"% faye_fits.mean,',',"%.3f"% faye_fits.var,') | (',"%.3f"% faye_fits.poisson['mean'],',',"%.3f"% faye_fits.poisson['var'],') | (',"%.3f"% faye_fits.geometric['mean'],',',"%.3f"% faye_fits.geometric['var'],')  |  (',"%.3f"% faye_fits.neg_bin['mean'],',',"%.3f"% faye_fits.neg_bin['var'],')  | (',"%.3f"% faye_fits.zip['mean'],',',"%.3f"% faye_fits.zip['var'],') | (',"%.3f"% faye_fits.beta_poi['mean'],',',"%.3f"% faye_fits.beta_poi['var'],') ')
print('Fasina  | (',"%.3f"% fasina_fits.mean,',',"%.3f"% fasina_fits.var,') | (',"%.3f"% fasina_fits.poisson['mean'],',',"%.3f"% fasina_fits.poisson['var'],') | (',"%.3f"% fasina_fits.geometric['mean'],',',"%.3f"% fasina_fits.geometric['var'],')  |  (',"%.3f"% fasina_fits.neg_bin['mean'],',',"%.3f"% fasina_fits.neg_bin['var'],')  | (',"%.3f"% fasina_fits.zip['mean'],',',"%.3f"% fasina_fits.zip['var'],') | (',"%.3f"% fasina_fits.beta_poi['mean'],',',"%.3f"% fasina_fits.beta_poi['var'],') ')
print('Leo     | (',"%.3f"% leo_fits.mean,',',"%.3f"% leo_fits.var,') | (',"%.3f"% leo_fits.poisson['mean'],',',"%.3f"% leo_fits.poisson['var'],') | (',"%.3f"% leo_fits.geometric['mean'],',',"%.3f"% leo_fits.geometric['var'],')  |  (',"%.3f"% leo_fits.neg_bin['mean'],',',"%.3f"% leo_fits.neg_bin['var'],')  | (',"%.3f"% leo_fits.zip['mean'],',',"%.3f"% leo_fits.zip['var'],') | (',"%.3f"% leo_fits.beta_poi['mean'],',',"%.3f"% leo_fits.beta_poi['var'],') ')
print('Cowling | (',"%.3f"% cowling_fits.mean,',',"%.3f"% cowling_fits.var,') | (',"%.3f"% cowling_fits.poisson['mean'],',',"%.3f"% cowling_fits.poisson['var'],') | (',"%.3f"% cowling_fits.geometric['mean'],',',"%.3f"% cowling_fits.geometric['var'],') | (',"%.3f"% cowling_fits.neg_bin['mean'],',',"%.3f"% cowling_fits.neg_bin['var'],') | (',"%.3f"% cowling_fits.zip['mean'],',',"%.3f"% cowling_fits.zip['var'],') | (',"%.3f"% cowling_fits.beta_poi['mean'],',',"%.3f"% cowling_fits.beta_poi['var'],') ')
print('Chowell | (',"%.3f"% chowell_fits.mean,',',"%.3f"% chowell_fits.var,') | (',"%.3f"% chowell_fits.poisson['mean'],',',"%.3f"% chowell_fits.poisson['var'],') | (',"%.3f"% chowell_fits.geometric['mean'],',',"%.3f"% chowell_fits.geometric['var'],')  |  (',"%.3f"% chowell_fits.neg_bin['mean'],',',"%.3f"% chowell_fits.neg_bin['var'],')  | (',"%.3f"% chowell_fits.zip['mean'],',',"%.3f"% chowell_fits.zip['var'],') | (',"%.3f"% chowell_fits.beta_poi['mean'],',',"%.3f"% chowell_fits.beta_poi['var'],') ')
print('Heijne  | (',"%.3f"% heijne_fits.mean,',',"%.3f"% heijne_fits.var,') | (',"%.3f"% heijne_fits.poisson['mean'],',',"%.3f"% heijne_fits.poisson['var'],') | (',"%.3f"% heijne_fits.geometric['mean'],',',"%.3f"% heijne_fits.geometric['var'],')  |  (',"%.3f"% heijne_fits.neg_bin['mean'],',',"%.3f"% heijne_fits.neg_bin['var'],')  | (',"%.3f"% heijne_fits.zip['mean'],',',"%.3f"% heijne_fits.zip['var'],') | (',"%.3f"% heijne_fits.beta_poi['mean'],',',"%.3f"% heijne_fits.beta_poi['var'],') ')
```

```python
print('Dataset & Empirical & Poisson & Geometric & Negative binomial & ZIP & Beta-Poisson \')
print(' & Mean & Variance & Mean & Variance & Mean & Variance & Mean & Variance & Mean & Variance & Mean & Variance \')
print('Gani & ',"%.3f"% gani_fits.mean,'&',"%.3f"% gani_fits.var,' & ',"%.3f"% gani_fits.poisson['mean'],' & ',"%.3f"% gani_fits.poisson['var'],' & ',"%.3f"% gani_fits.geometric['mean'],' & ',"%.3f"% gani_fits.geometric['var'],' & ',"%.3f"% gani_fits.neg_bin['mean'],' & ',"%.3f"% gani_fits.neg_bin['var'],' & ',"%.3f"% gani_fits.zip['mean'],' & ',"%.3f"% gani_fits.zip['var'],' & ',"%.3f"% gani_fits.beta_poi['mean'],' & ',"%.3f"% gani_fits.beta_poi['var'],' \')
print('Jezek & ',"%.3f"% jezek_fits.mean,' & ',"%.3f"% jezek_fits.var,' & ',"%.3f"% jezek_fits.poisson['mean'],' & ',"%.3f"% jezek_fits.poisson['var'],' & ',"%.3f"% jezek_fits.geometric['mean'],' & ',"%.3f"% jezek_fits.geometric['var'],' & ',"%.3f"% jezek_fits.neg_bin['mean'],' & ',"%.3f"% jezek_fits.neg_bin['var'],' & ',"%.3f"% jezek_fits.zip['mean'],' & ',"%.3f"% jezek_fits.zip['var'],' & ',"%.3f"% jezek_fits.beta_poi['mean'],' & ',"%.3f"% jezek_fits.beta_poi['var'],' \')
print('Faye & ',"%.3f"% faye_fits.mean,' & ',"%.3f"% faye_fits.var,' & ',"%.3f"% faye_fits.poisson['mean'],' & ',"%.3f"% faye_fits.poisson['var'],' & ',"%.3f"% faye_fits.geometric['mean'],' & ',"%.3f"% faye_fits.geometric['var'],' & ',"%.3f"% faye_fits.neg_bin['mean'],' & ',"%.3f"% faye_fits.neg_bin['var'],' & ',"%.3f"% faye_fits.zip['mean'],' & ',"%.3f"% faye_fits.zip['var'],' & ',"%.3f"% faye_fits.beta_poi['mean'],' & ',"%.3f"% faye_fits.beta_poi['var'],' \')
print('Fasina & ',"%.3f"% fasina_fits.mean,' & ',"%.3f"% fasina_fits.var,' & ',"%.3f"% fasina_fits.poisson['mean'],' & ',"%.3f"% fasina_fits.poisson['var'],' & ',"%.3f"% fasina_fits.geometric['mean'],' & ',"%.3f"% fasina_fits.geometric['var'],' & ',"%.3f"% fasina_fits.neg_bin['mean'],' & ',"%.3f"% fasina_fits.neg_bin['var'],' & ',"%.3f"% fasina_fits.zip['mean'],' & ',"%.3f"% fasina_fits.zip['var'],' & ',"%.3f"% fasina_fits.beta_poi['mean'],' & ',"%.3f"% fasina_fits.beta_poi['var'],' \')
print('Leo & ',"%.3f"% leo_fits.mean,' & ',"%.3f"% leo_fits.var,' & ',"%.3f"% leo_fits.poisson['mean'],' & ',"%.3f"% leo_fits.poisson['var'],' & ',"%.3f"% leo_fits.geometric['mean'],' & ',"%.3f"% leo_fits.geometric['var'],' & ',"%.3f"% leo_fits.neg_bin['mean'],',',"%.3f"% leo_fits.neg_bin['var'],' & ',"%.3f"% leo_fits.zip['mean'],',',"%.3f"% leo_fits.zip['var'],' & ',"%.3f"% leo_fits.beta_poi['mean'],' & ',"%.3f"% leo_fits.beta_poi['var'],' \')
print('Cowling & ',"%.3f"% cowling_fits.mean,' & ',"%.3f"% cowling_fits.var,' & ',"%.3f"% cowling_fits.poisson['mean'],' & ',"%.3f"% cowling_fits.poisson['var'],' & ',"%.3f"% cowling_fits.geometric['mean'],' & ',"%.3f"% cowling_fits.geometric['var'],' & ',"%.3f"% cowling_fits.neg_bin['mean'],' & ',"%.3f"% cowling_fits.neg_bin['var'],' & ',"%.3f"% cowling_fits.zip['mean'],' & ',"%.3f"% cowling_fits.zip['var'],' & ',"%.3f"% cowling_fits.beta_poi['mean'],' & ',"%.3f"% cowling_fits.beta_poi['var'],' \')
print('Chowell & ',"%.3f"% chowell_fits.mean,' & ',"%.3f"% chowell_fits.var,' & ',"%.3f"% chowell_fits.poisson['mean'],' & ',"%.3f"% chowell_fits.poisson['var'],' & ',"%.3f"% chowell_fits.geometric['mean'],' & ',"%.3f"% chowell_fits.geometric['var'],' & ',"%.3f"% chowell_fits.neg_bin['mean'],' & ',"%.3f"% chowell_fits.neg_bin['var'],' & ',"%.3f"% chowell_fits.zip['mean'],' & ',"%.3f"% chowell_fits.zip['var'],' & ',"%.3f"% chowell_fits.beta_poi['mean'],' & ',"%.3f"% chowell_fits.beta_poi['var'],' \')
print('Heijne & ',"%.3f"% heijne_fits.mean,' & ',"%.3f"% heijne_fits.var,' & ',"%.3f"% heijne_fits.poisson['mean'],' & ',"%.3f"% heijne_fits.poisson['var'],' & ',"%.3f"% heijne_fits.geometric['mean'],' & ',"%.3f"% heijne_fits.geometric['var'],' & ',"%.3f"% heijne_fits.neg_bin['mean'],'v',"%.3f"% heijne_fits.neg_bin['var'],' & ',"%.3f"% heijne_fits.zip['mean'],' & ',"%.3f"% heijne_fits.zip['var'],' & ',"%.3f"% heijne_fits.beta_poi['mean'],' & ',"%.3f"% heijne_fits.beta_poi['var'],' \')
```

```python
print('Dataset | Poisson | Geometric | Negative binomial ')
print('__________________________________________________')
print('Gani    | ',"%.3f"% np.exp(gani_fits.poisson['llh'] - gani_fits.beta_poi['llh']),' |  ',"%.3f"% np.exp(gani_fits.geometric['llh'] - gani_fits.beta_poi['llh']),'  |  ',"%.3f"% np.exp(gani_fits.neg_bin['llh'] - gani_fits.beta_poi['llh']))
print('Jezek   | ',"%.3f"% np.exp(jezek_fits.poisson['llh'] - jezek_fits.beta_poi['llh']),' |  ',"%.3f"% np.exp(jezek_fits.geometric['llh'] - jezek_fits.beta_poi['llh']),'  |  ',"%.3f"% np.exp(jezek_fits.neg_bin['llh'] - jezek_fits.beta_poi['llh']))
print('Faye    | ',"%.3f"% np.exp(faye_fits.poisson['llh'] - faye_fits.beta_poi['llh']),' |  ',"%.3f"% np.exp(faye_fits.geometric['llh'] - faye_fits.beta_poi['llh']),'  |  ',"%.3f"% np.exp(faye_fits.neg_bin['llh'] - faye_fits.beta_poi['llh']))
print('Fasina  | ',"%.3f"% np.exp(fasina_fits.poisson['llh'] - fasina_fits.beta_poi['llh']),' |  ',"%.3f"% np.exp(fasina_fits.geometric['llh'] - fasina_fits.beta_poi['llh']),'  |  ',"%.3f"% np.exp(fasina_fits.neg_bin['llh'] - fasina_fits.beta_poi['llh']))
print('Leo     | ',"%.3f"% np.exp(leo_fits.poisson['llh'] - leo_fits.beta_poi['llh']),' |  ',"%.3f"% np.exp(leo_fits.geometric['llh'] - leo_fits.beta_poi['llh']),'  |  ',"%.3f"% np.exp(leo_fits.neg_bin['llh'] - leo_fits.beta_poi['llh']))
print('Cowling | ',"%.3f"% np.exp(cowling_fits.poisson['llh'] - cowling_fits.beta_poi['llh']),' |  ',"%.3f"% np.exp(cowling_fits.geometric['llh'] - cowling_fits.beta_poi['llh']),'  |  ',"%.3f"% np.exp(cowling_fits.neg_bin['llh'] - cowling_fits.beta_poi['llh']))
print('Chowell | ',"%.3f"% np.exp(chowell_fits.poisson['llh'] - chowell_fits.beta_poi['llh']),' |  ',"%.3f"% np.exp(chowell_fits.geometric['llh'] - chowell_fits.beta_poi['llh']),'  |  ',"%.3f"% np.exp(chowell_fits.neg_bin['llh'] - chowell_fits.beta_poi['llh']))
print('Heijne  | ',"%.3f"% np.exp(heijne_fits.poisson['llh'] - heijne_fits.beta_poi['llh']),' |  ',"%.3f"% np.exp(heijne_fits.geometric['llh'] - heijne_fits.beta_poi['llh']),'  |  ',"%.3f"% np.exp(heijne_fits.neg_bin['llh'] - heijne_fits.beta_poi['llh']))
```

```python
print('Dataset | Nu')
print('____________________________')
print('Gani    |', "%.3f"% gani_fits.beta_poi['Nu'],'= 1 /',"%.3f"% (1/gani_fits.beta_poi['Nu']))
print('Jezek   |', "%.3f"% jezek_fits.beta_poi['Nu'],'= 1 /',"%.3f"% (1/jezek_fits.beta_poi['Nu']))
print('Faye    |', "%.3f"% faye_fits.beta_poi['Nu'],'= 1 /',"%.3f"% (1/faye_fits.beta_poi['Nu']))
print('Fasina  |', "%.3f"% fasina_fits.beta_poi['Nu'],'= 1 /',"%.3f"% (1/fasina_fits.beta_poi['Nu']))
print('Leo     |', "%.3f"% leo_fits.beta_poi['Nu'],'= 1 /',"%.3f"% (1/leo_fits.beta_poi['Nu']))
print('Cowling |', "%.3f"% cowling_fits.beta_poi['Nu'],'= 1 /',"%.3f"% (1/cowling_fits.beta_poi['Nu']))
print('Chowell |', "%.3f"% chowell_fits.beta_poi['Nu'],'= 1 /',"%.3f"% (1/chowell_fits.beta_poi['Nu']))
print('Heijne  |', "%.3f"% heijne_fits.beta_poi['Nu'],'= 1 /',"%.3f"% (1/heijne_fits.beta_poi['Nu']))
```

```python
fig = plt.figure(figsize=(50,20))
gani_ax = plt.subplot2grid((2,4),(0,0))
jezek_ax = plt.subplot2grid((2,4),(0,1))
faye_ax = plt.subplot2grid((2,4),(0,2))
fasina_ax = plt.subplot2grid((2,4),(0,3))
leo_ax = plt.subplot2grid((2,4),(1,0))
cowling_ax = plt.subplot2grid((2,4),(1,1))
chowell_ax = plt.subplot2grid((2,4),(1,2))
heijne_ax = plt.subplot2grid((2,4),(1,3))

plot_chain_data_fits(gani_fits, gani_ax, 'Plague, Gani and Leach',30)
plot_chain_data_fits(jezek_fits, jezek_ax, 'Monkeypox, Jezek et al.',30)
plot_chain_data_fits(faye_fits, faye_ax, 'Ebola, Faye et al.',30)
plot_chain_data_fits(fasina_fits, fasina_ax, 'Ebola, Fasina et al.',30)
plot_chain_data_fits(leo_fits, leo_ax, 'SARS, Leo et al.',30)
plot_chain_data_fits(cowling_fits, cowling_ax, 'MERS-CoV, Cowling et al.',30)
plot_chain_data_fits(chowell_fits, chowell_ax, 'MERS-CoV, Chowell et al.',30)
plot_chain_data_fits(heijne_fits, heijne_ax, 'Norovirus, Heijne et al.',30)

fasina_ax.legend(['Data','Poisson','Geometric','Negative binomial','ZIP','Beta-Poisson'],
          bbox_to_anchor=(1.05,1),
          loc='upper left',
          fontsize=30)

plt.tight_layout()
plt.show()
```

## Confidence intervals

The bootstrap sampling used to calculate confidence intervals for the negative binomial and beta-Poisson models is very time consuming when large numbers of samples are used.

```python
def poisson_confidence_intervals(data,interval,points):
    lmbd=np.linspace(interval[0],interval[1],points)
    llh=np.zeros(points)
    for x in data:
        llh += stats.poisson.logpmf(x,lmbd)
    mle_loc=np.argmax(llh)
    lh_normed=np.exp(llh)/np.sum(np.exp(llh))
    current_max=lh_normed[mle_loc]
    interval_weight=current_max
    while interval_weight<0.95:
        max_loc=np.argmax(lh_normed[np.where(lh_normed<current_max)])
        current_max=lh_normed[np.where(lh_normed<current_max)][max_loc]
        interval_weight+=current_max
    ci=[np.min(lmbd[np.where(lh_normed>=current_max)[0]]),np.max(lmbd[np.where(lh_normed>=current_max)[0]])]
    return ci

def geometric_confidence_intervals(data,interval,points):
    lmbd=np.linspace(interval[0],interval[1],points)
    llh=np.zeros(points)
    for x in data:
        llh += stats.geom.logpmf(x,1/(lmbd+1),-1)
    mle_loc=np.argmax(llh)
    lh_normed=np.exp(llh)/np.sum(np.exp(llh))
    current_max=lh_normed[mle_loc]
    interval_weight=current_max
    while interval_weight<0.95:
        max_loc=np.argmax(lh_normed[np.where(lh_normed<current_max)])
        current_max=lh_normed[np.where(lh_normed<current_max)][max_loc]
        interval_weight+=current_max
    ci=[np.min(lmbd[np.where(lh_normed>=current_max)[0]]),np.max(lmbd[np.where(lh_normed>=current_max)[0]])]
    return ci

def zip_confidence_intervals(data,lmbd_interval,lmbd_points,sigma_points):
    lmbd=np.linspace(lmbd_interval[0],lmbd_interval[1],lmbd_points)
    sigma=np.linspace(0,1,sigma_points)

    lmbdgrid,sigmagrid=np.meshgrid(lmbd,sigma)
    mean_grid = np.multiply(lmbdgrid, (1-sigmagrid))
    var_grid = np.multiply(lmbdgrid, np.multiply((1-sigmagrid),(1+np.multiply(lmbdgrid,sigmagrid))))
    llh=np.zeros(lmbdgrid.shape)
    for x in data:
        llh += zip_loglh([x],lmbdgrid,sigmagrid)
    mle_loc=np.unravel_index(np.argmax(llh),llh.shape)
    lh_normed=np.exp(llh)/np.sum(np.exp(llh))
    current_max=lh_normed[mle_loc]
    interval_weight=current_max
    while interval_weight<0.95:
        max_loc=np.unravel_index(np.argmax(lh_normed),lh_normed.shape)
        current_max=lh_normed[max_loc]
        lh_normed[max_loc]=0
        interval_weight+=current_max
    lh_normed=np.exp(llh)/np.sum(np.exp(llh))
    lmbd_ci=[np.min(lmbdgrid[np.where(lh_normed>=current_max)]),np.max(lmbdgrid[np.where(lh_normed>=current_max)])]
    sigma_ci=[np.min(sigmagrid[np.where(lh_normed>=current_max)]),np.max(sigmagrid[np.where(lh_normed>=current_max)])]
    mean_ci=[np.min(mean_grid[np.where(lh_normed>=current_max)]),np.max(mean_grid[np.where(lh_normed>=current_max)])]
    var_ci=[np.min(var_grid[np.where(lh_normed>=current_max)]),np.max(var_grid[np.where(lh_normed>=current_max)])]
    return lmbd_ci,sigma_ci,mean_ci, var_ci
```

```python
def neg_bin_bootstrap(data,no_samples,theta_0):

    sample_size=np.size(data)

    lmbd_samples = []
    theta_samples = []

    print('Now calculating',no_samples,'bootstrap samples.')
    start_time=time.time()
    for i in range(no_samples):
        data_now=random.choices(data,k=sample_size)
        lmbd_samples = lmbd_samples + [np.mean(data_now)]
        theta_samples = theta_samples + [get_theta_mle(data_now,theta_0)]
        if ((i+1)%1000)==0:
            print('Sample',i+1,'of',no_samples,'completed.',time.time()-start_time,'seconds elapsed, approximately',(no_samples-i-1)*(time.time()-start_time)/(i+1),'remaining.')

    lmbd_samples = np.array(lmbd_samples)
    theta_samples = np.array(theta_samples)
    var_samples=np.multiply(lmbd_samples,(1+theta_samples))

    lmbd_ci=[np.percentile(lmbd_samples,2.5),np.percentile(lmbd_samples,97.5)]
    theta_ci=[np.percentile(theta_samples,2.5),np.percentile(theta_samples,97.5)]
    var_ci=[np.percentile(var_samples,2.5),np.percentile(var_samples,97.5)]

    return lmbd_ci,lmbd_samples,theta_ci,theta_samples,var_ci,var_samples
```

```python
def beta_poisson_bootstrap(data,no_samples,phi_0,N_0):

    sample_size=np.size(data)

    lmbd_samples = []
    Phi_samples = []
    Nu_samples = []

    print('Now calculating',no_samples,'bootstrap samples.')
    start_time=time.time()
    for i in range(no_samples):
        data_now=random.choices(data,k=sample_size)
        lmbd_samples = lmbd_samples + [np.mean(data_now)]
        Phi_now, Nu_now=get_phi_and_N_mles(data_now,phi_0,N_0)
        Phi_samples = Phi_samples + [Phi_now]
        Nu_samples = Nu_samples + [Nu_now]
        if ((i+1)%1000)==0:
            print('Sample',i+1,'of',no_samples,'completed.',time.time()-start_time,'seconds elapsed, approximately',(no_samples-i-1)*(time.time()-start_time)/(i+1),'remaining.')

    lmbd_samples = np.array(lmbd_samples)
    Phi_samples = np.array(Phi_samples)
    Nu_samples = np.array(Nu_samples)
    var_samples=lmbd_samples*(1+(1-lmbd_samples*Nu_samples)/(Phi_samples+Nu_samples))

    lmbd_ci=[np.percentile(lmbd_samples,2.5),np.percentile(lmbd_samples,97.5)]
    Phi_ci=[np.percentile(Phi_samples,2.5),np.percentile(Phi_samples,97.5)]
    Nu_ci=[np.percentile(Nu_samples,2.5),np.percentile(Nu_samples,97.5)]
    var_ci=[np.percentile(var_samples,2.5),np.percentile(var_samples,97.5)]

    return lmbd_ci,lmbd_samples,Phi_ci,Phi_samples,Nu_ci,Nu_samples,var_ci,var_samples
```

```python
class confidence_intervals:
    def __init__(self,data):
        self.data = data
        self.mean = np.mean(data)
        self.max = np.max(data)

    def poisson(self,sfs):
        points = 1 + (10**sfs)*self.max
        lmbd_ci = poisson_confidence_intervals(self.data,[0,self.max],points)
        return {'lmbd' : lmbd_ci,
                'mean' : lmbd_ci,
                'var' : lmbd_ci
                      }

    def geometric(self,sfs):
        points = 1 + (10**sfs)*self.max
        lmbd_ci = geometric_confidence_intervals(self.data,[0,self.max],points)
        var_ci = [lmbd_ci[0]*(1+lmbd_ci[0]), lmbd_ci[1]*(1+lmbd_ci[1])]
        return {'lmbd' : lmbd_ci,
                'mean' : lmbd_ci,
                'var' : var_ci
                      }

    def zip(self,sfs):
        lmbd_points = 1 + (10**sfs)*self.max
        sigma_points = 1 + 10**sfs
        lmbd_ci, sigma_ci, mean_ci, var_ci = zip_confidence_intervals(self.data,[0,self.max],lmbd_points,sigma_points)
        return {'lmbd' : lmbd_ci,
                'sigma' : sigma_ci,
                'mean' : mean_ci,
                'var' : var_ci
                      }

    def neg_bin(self,no_samples):
        lmbd_ci,lmbd_samples,theta_ci,theta_samples,var_ci,var_samples = neg_bin_bootstrap(self.data,no_samples,self.mean)
        return {'lmbd' : lmbd_ci,
                'theta' : theta_ci,
                'mean' : lmbd_ci,
                'var' : var_ci}

    def beta_poisson(self,no_samples):
        lmbd_ci,lmbd_samples,Phi_ci,Phi_samples,Nu_ci,Nu_samples,var_ci,var_samples = \
            beta_poisson_bootstrap(self.data,no_samples,1/self.mean,self.max)
        return {'lmbd' : lmbd_ci,
                'Phi' : Phi_ci,
                'Nu' : Nu_ci,
                'mean' : lmbd_ci,
                'var' : var_ci}
```

```python
cis = confidence_intervals(gani_data)
gani_poisson_cis = cis.poisson(3)
gani_geometric_cis = cis.geometric(3)
gani_zip_cis = cis.zip(2)
gani_neg_bin_cis = cis.neg_bin(100)
gani_beta_poisson_cis = cis.beta_poisson(100)

cis = confidence_intervals(jezek_data)
jezek_poisson_cis = cis.poisson(3)
jezek_geometric_cis = cis.geometric(3)
jezek_zip_cis = cis.zip(2)
jezek_neg_bin_cis = cis.neg_bin(100)
jezek_beta_poisson_cis = cis.beta_poisson(100)

cis = confidence_intervals(faye_data)
faye_poisson_cis = cis.poisson(3)
faye_geometric_cis = cis.geometric(3)
faye_zip_cis = cis.zip(2)
faye_neg_bin_cis = cis.neg_bin(100)
faye_beta_poisson_cis = cis.beta_poisson(100)

cis = confidence_intervals(fasina_data)
fasina_poisson_cis = cis.poisson(3)
fasina_geometric_cis = cis.geometric(3)
fasina_zip_cis = cis.zip(2)
fasina_neg_bin_cis = cis.neg_bin(100)
fasina_beta_poisson_cis = cis.beta_poisson(100)

cis = confidence_intervals(leo_data)
leo_poisson_cis = cis.poisson(3)
leo_geometric_cis = cis.geometric(3)
leo_zip_cis = cis.zip(2)
leo_neg_bin_cis = cis.neg_bin(100)
leo_beta_poisson_cis = cis.beta_poisson(100)

cis = confidence_intervals(cowling_data)
cowling_poisson_cis = cis.poisson(3)
cowling_geometric_cis = cis.geometric(3)
cowling_zip_cis = cis.zip(2)
cowling_neg_bin_cis = cis.neg_bin(100)
cowling_beta_poisson_cis = cis.beta_poisson(100)

cis = confidence_intervals(chowell_data)
chowell_poisson_cis = cis.poisson(3)
chowell_geometric_cis = cis.geometric(3)
chowell_zip_cis = cis.zip(2)
chowell_neg_bin_cis = cis.neg_bin(100)
chowell_beta_poisson_cis = cis.beta_poisson(100)

cis = confidence_intervals(heijne_data)
heijne_poisson_cis = cis.poisson(3)
heijne_geometric_cis = cis.geometric(3)
heijne_zip_cis = cis.zip(2)
heijne_neg_bin_cis = cis.neg_bin(100)
heijne_beta_poisson_cis = cis.beta_poisson(100)
```

```python

```

A beta-Poisson model for infectious disease transmission.ipynb


# Introdução ao Firedrake

## Parte 2: Introdução de conceitos básicos de FEM

Ministrante: **Diego Tavares Volpatto (LNCC/ESSS)**

Encontro Acadêmico de Modelagem Computacional - XV EAMC (2022)

<tr>
<td>  </td>
<td>  </td>
</tr>

## Sumário

* Conceitos básicos de FEM utilizando o Firedrake
* Problema de Poisson em 1D e 2D
* Problema de Elasticidade Linear

# Conceitos básicos com Firedrake

## Conceitos básicos de FEM

Aqui veremos conceitos básicos de FEM essenciais para melhor entendimento de como usar o Firedrake. Em suma:

* Formas forte e fraca dos problemas
* Discretização espacial
* Aproximação das soluções (espaços de dimensão finita)

## Forma forte

Antes de mais nada, uma pequena ideia geral sobre FEM. Vamos considerar o seguinte problema-modelo:

Dado um conjunto aberto $\Omega \subset \mathbb{R}^n$ (com $n = 2$ ou $3$), encontrar $p \in C^2(\Omega)$ tal que:

\begin{equation}
\nabla \cdot (-K \nabla p) = f
\end{equation}

sendo $f: \Omega \to \mathbb{R}$ uma função conhecida. Vamos adotar condições de contorno genéricas: $p|_{\Gamma_D} = p_D$ e $(-K \nabla p)\cdot \mathbf{n} = g$ sobre $\Gamma_N$, sendo $\Gamma_D \cap \Gamma_N = \emptyset$ e $\Gamma_D \cup \Gamma_N \equiv \partial \Omega \subset \mathbb{R}^{n-1}$ o contorno de $\Omega$.

Em geral, $K \in \mathbb{R}^{n \times n}$, mas simplificaremos: $K \in \mathbb{R}$.

A formulação do problema, em si, é também denominada de **Forma Forte**, que deve ser satisfeita ponto a ponto em $\Omega$.

## Forma fraca

A derivação clássica de FEM (método de Galerkin) consiste em:

\begin{equation}
\nabla \cdot (-K \nabla p) - f = 0
\end{equation}

\begin{equation}
\int_\Omega \left[\nabla \cdot (-K \nabla p) - f\right] v \,dx = 0
\end{equation}

\begin{equation}
-\int_\Omega v \nabla \cdot (K \nabla p) \,dx = \int_\Omega f\,v \,dx
\end{equation}

Utilizando integração por partes (identidade de Green):

\begin{equation}
\int_\Omega K \nabla p \nabla v \,dx = \int_\Omega f\,v \,dx + \int_{\Gamma_N} g\, v \,ds
\end{equation}

## Forma fraca

Dessa forma,

\begin{equation}
\int_\Omega K \nabla p \nabla v \,dx = \int_\Omega f\,v \,dx + \int_{\Gamma_N} g\, v \,ds
\end{equation}

No entanto, as funções devem ser definidas em algum **espaço de funções**. Ainda no espaço de dimensão infinita, temos que:

* $p \in \{p \in H^1(\Omega): p|_{\Gamma_D} = p_D\} =: \mathcal U$ $\to$ função candidata à solução.
* $v \in \{v \in H^1(\Omega): v|_{\Gamma_D} = 0\} =: \mathcal V$ $\to$ função-teste.

Além disso, $f \in L^2(\Omega)$.

## Formulação variacional abstrata

É interessante definir a formulação variacional a partir de formas bilineares, lineares e até multilineares. Para o Problema de Poisson:

\begin{equation}
a(p, v) := \int_\Omega K \nabla p \nabla v \,dx
\end{equation}

\begin{equation}
L(v) := \int_\Omega f\,v \,dx
\end{equation}

E o produto interno sobre $\Gamma_N$:

\begin{equation}
\langle g, v \rangle_{\Gamma_N} := \int_{\Gamma_N} g\, v \,ds
\end{equation}

Portanto: $a(p, v) = L(v) + \langle g, v \rangle_{\Gamma_N}$.

## Discretização espacial

Precisamos definir a discretização espacial do domínio para então termos onde calcular as soluções numéricas do problema.

Dessa forma, o domínio é particionado em elementos (geométricos e simples) tais que

\begin{equation}
\overline{\Omega} = \bigcup_{e = 1}^{N_e} T_e =: \mathcal{T}_h(\Omega)
\end{equation}

de forma que para $i \ne j$, $T_i \cap T_j$ seja (discretização conforme):

* Vazio;
* Um vértice;
* Uma aresta.

## Aproximação do Espaço de Funções: Espaços de Elementos Finitos

Os espaços $H^1$, $L^2$ e etc são espaços de dimensão infinita! Convenientemente vamos definir espaços de dimensão finita que se aproximem a esses espaços e satisfaçam os requerimentos de regularidade.

Dessa forma, temos espaços definidos por partes a nível de elemento (funções de suporte compacto). Por exemplo, polinômios definidos localmente nos elementos e seus nós.

Poderíamos considerar, por exemplo:

\begin{equation}
\mathcal{U}_h^k := \{p_h: p_h \in \mathcal{P}_k(T_e), \forall T_e \in \mathcal{T}_h(\Omega), p_h|_{\Gamma_D} = p_D\}
\end{equation}
\begin{equation}
\mathcal{P}_k(T) := \{p_h: p_h \text{ é um polinômio de grau} \le k \text{ em } T \in \mathcal{T}_h(\Omega) \}
\end{equation}

Logo, $\mathcal{U}_h^k \subset H^1(\Omega)$ com $k = 1$.

## Em resumo

Será necessário definir, no Firedrake, a famosa tripa da definição genérica de elementos finitos de Ciarlet:

1. A partição do seu domínio (discretização espacial): $\mathcal{T}_h$
2. A natureza das funções definidas a nível de elemento (shape functions, etc): $\mathcal{U}_h^k$, $\mathcal{V}_h^k$, etc
3. Os parâmetros que descrevem as funções discretizadas: os graus de liberdade.

Tudo isso precisaremos definir em código. Além disso, é claro, temos que escrever a forma fraca do problema.

## Leituras adicionais

Fiz apenas breves comentários de conceitos que serão utilizados nos códigos! Caso não tenha uma base teórica confortável em FEM e deseja ampliar, recomendo as seguintes literaturas introdutórias:

* [Numerical Solution of Partial Differential Equations by the Finite Element Method (Claes Johnson)](https://www.google.com.br/books/edition/Numerical_Solution_of_Partial_Differenti/PYXjyoqy5qMC?hl=pt-BR&gbpv=0)
* [Finite Elements, vol 1 (Oden e Carey)](https://www.google.com.br/books/edition/Finite_Elements/k8IXygEACAAJ?hl=pt-BR)
* [The Finite Element Method (Hughes)](https://www.google.com.br/books/edition/Finite_Elements/k8IXygEACAAJ?hl=pt-BR)
* [Understanding and Implementing the Finite Element Method (Gockenbach)](https://www.google.com.br/books/edition/Understanding_and_Implementing_the_Finit/stgyDwAAQBAJ?hl=pt-BR&gbpv=0)
* [(Gratuito e em Português) Introdução ao Método de Elementos Finitos (Mauro Rincon e I-Shih Liu)](http://www.im.ufrj.br/index.php/pt/estrutura/e-books-im/1452-introducao-ao-metodo-de-elementos-finitos)

# Problema de Poisson

## Recapitulando

### Forma forte

Encontrar $p \in C^2(\Omega)$ tal que:

\begin{equation}
\begin{cases}
  \nabla \cdot (-K \nabla p) = f & \text{em } \Omega \\[0.5em]
  p = p_D & \text{sobre } \Gamma_D \\[0.5em]
  (-K \nabla p)\cdot \mathbf{n} = g & \text{sobre } \Gamma_N
\end{cases}
\end{equation}

### Forma fraca

Encontrar $p_h \in \mathcal{U}_h$ para todo $v \in \mathcal{V}_h$ tal que:

\begin{equation}
a(p_h, v) = L(v) + \langle g, v \rangle_{\Gamma_N}
\end{equation}

sendo

\begin{equation}
a(p, v) := \int_\Omega K \nabla p \nabla v \,dx \quad
\end{equation}

\begin{equation}
L(v) := \int_\Omega f\,v \,dx
\end{equation}

## Implementando em Firedrake

Segue basicamente os seguintes passos:

1. Definição (ou leitura) da malha
2. Definição dos Espaços de Elementos Finito sobre a malha
3. Declaração das funções candidatas à solução e funções testes
4. Declaração de qualquer função adicional necessária (termos de fonte, coeficientes, etc)
5. Definição das condições de contorno fortes
6. Escrita da Formulação Variacional
7. Solucionamento do Problema Variacional
8. Coleta (ou visualização) dos resultados

Vamos considerar o caso 1D para compreensão dos conceitos.

### Importando o Firedrake

Antes de mais, é importante que você carregue as ferramentas que o Firedrake provê. Para que tudo fique disponível, simplesmente faça:

```python
import os
os.environ["OMP_NUM_THREADS"] = "1"  # necessário para o Firedrake

from firedrake import *
```

### Definição da malha

Vamos considerar uma malha no intervalo unitário contendo 10 elementos

```python
num_elements = 10
mesh = UnitIntervalMesh(num_elements)
```

Para informações detalhadas, consulte [aqui](https://www.firedrakeproject.org/variational-problems.html#constructing-meshes). Pode consultar as diversas funções convenientes na [documentação das APIs](https://www.firedrakeproject.org/firedrake.html?module-firedrake.mesh#module-firedrake.utility_meshes).

Todos os elementos da [Tabela Periódica de Elementos Finitos](https://www-users.cse.umn.edu/~arnold/femtable/) são suportados.

### Declaração do espaço de funções

Elementos Lagrangeanos Contínuos de grau `p = 1`:

```python
p = 1  # Grau do polinômio
V = FunctionSpace(mesh, "CG", p)
```

Existe uma diversa possibilidade de escolha para os Espaços de Funções, por favor [verifique aqui](https://www.firedrakeproject.org/variational-problems.html#supported-finite-elements) a lista de todas as famílias de elementos disponíveis.

### Declaração das funções candidatas a solução e teste

Feitas simples e diretamente com:

```python
p_h = TrialFunction(V)
v = TestFunction(V)
```

### Termo de fonte

No problema de Poisson, a única função adicional que estamos considerando é o termo de fonte. Vamos escolher um termo de fonte de forma a conhecermos a solução do problema ([Método das Soluções Manufaturadas](https://doi.org/10.1115/1.1436090)):

```python
x, = SpatialCoordinate(mesh)
p_exact = sin(2 * pi * x)
exact_solution = Function(V).interpolate(p_exact)
exact_solution.rename("Exact solution", "label")
```

Adicionalmente, adotaremos $K = 1$ de agora em diante.

```python
K = Constant(1)
```

Coletando os dados da solução para uma rápida visualização:

```python
x_values = mesh.coordinates.vector().dat.data
exact_solution_values = exact_solution.vector().dat.data
```

```python
import matplotlib.pyplot as plt
plt.style.use(['seaborn-poster'])

plt.figure(figsize=(8, 6))

plt.plot(x_values, exact_solution_values)

plt.grid()
plt.xlabel("x")
plt.ylabel("p(x)")

plt.show()
```

Agora, finalmente, um primeiro contato com operadores da [UFL](https://fenics.readthedocs.io/projects/ufl/en/latest/manual/form_language.html):

```python
f_expression = div(-K * grad(p_exact))
f = Function(V).interpolate(f_expression)
```

Dessa forma, definimos nosso termo-fonte a partir da forma forte do problema.

### Condições de Contorno essenciais

Vamos definir condições de contorno que são compatíveis com o termo de fonte:

```python
boundary_value = 0.0
bc = DirichletBC(V, boundary_value, "on_boundary")  # Fronteira 1 (esquerda) e Fronteira 2 (direita)
```

Note que, por simplicidade, toda a fronteira, neste exemplo, está sendo prescrita com condição de Dirichlet.

Caso tenha dúvidas sobre a função de condição de contorno, consulte:

```python
help(DirichletBC)
```

    Help on class DirichletBC in module firedrake.bcs:

    class DirichletBC(BCBase, firedrake.adjoint.dirichletbc.DirichletBCMixin)
     |  DirichletBC(V, g, sub_domain, method=None)
     |
     |  Implementation of a strong Dirichlet boundary condition.
     |
     |  .. note:
     |
     |     This uses facet markers in the domain, so may be used to
     |     applied strong boundary conditions to interior facets (if they
     |     have an appropriate mesh marker). The "on_boundary" string only
     |     applies to the exterior boundaries of the domain.
     |
     |  :arg V: the :class:`.FunctionSpace` on which the boundary condition
     |      should be applied.
     |  :arg g: the boundary condition values. This can be a :class:`.Function` on
     |      ``V``, or a UFL expression that can be interpolated into
     |      ``V``, for example, a :class:`.Constant`, an iterable of
     |      literal constants (converted to a UFL expression), or a
     |      literal constant which can be pointwise evaluated at the nodes
     |      of ``V``.
     |  :arg sub_domain: the integer id(s) of the boundary region over which the
     |      boundary condition should be applied. The string "on_boundary" may be used
     |      to indicate all of the boundaries of the domain. In the case of extrusion
     |      the ``top`` and ``bottom`` strings are used to flag the bcs application on
     |      the top and bottom boundaries of the extruded mesh respectively.
     |  :arg method: the method for determining boundary nodes.
     |      DEPRECATED. The only way boundary nodes are identified is by
     |      topological association.
     |
     |  Method resolution order:
     |      DirichletBC
     |      BCBase
     |      firedrake.adjoint.dirichletbc.DirichletBCMixin
     |      pyadjoint.overloaded_type.FloatingType
     |      pyadjoint.overloaded_type.OverloadedType
     |      builtins.object
     |
     |  Methods defined here:
     |
     |  __init__(self, V, g, sub_domain, method=None)
     |      Initialize self.  See help(type(self)) for accurate signature.
     |
     |  apply(self, r, u=None)
     |      Apply this boundary condition to ``r``.
     |
     |      :arg r: a :class:`.Function` or :class:`.Matrix` to which the
     |          boundary condition should be applied.
     |
     |      :arg u: an optional current state.  If ``u`` is supplied then
     |          ``r`` is taken to be a residual and the boundary condition
     |          nodes are set to the value ``u-bc``.  Supplying ``u`` has
     |          no effect if ``r`` is a :class:`.Matrix` rather than a
     |          :class:`.Function`. If ``u`` is absent, then the boundary
     |          condition nodes of ``r`` are set to the boundary condition
     |          values.
     |
     |
     |      If ``r`` is a :class:`.Matrix`, it will be assembled with a 1
     |      on diagonals where the boundary condition applies and 0 in the
     |      corresponding rows and columns.
     |
     |  dirichlet_bcs(self)
     |
     |  extract_form(self, form_type)
     |
     |  homogenize(self)
     |      Convert this boundary condition into a homogeneous one.
     |
     |      Set the value to zero.
     |
     |  integrals(self)
     |
     |  reconstruct(self, field=None, V=None, g=None, sub_domain=None, use_split=False)
     |
     |  restore(self)
     |      Restore the original value of this boundary condition.
     |
     |      This uses the value passed on instantiation of the object.
     |
     |  set_value(self, val)
     |      Set the value of this boundary condition.
     |
     |      :arg val: The boundary condition values.  See
     |          :class:`.DirichletBC` for valid values.
     |
     |  ----------------------------------------------------------------------
     |  Data descriptors defined here:
     |
     |  function_arg
     |      The value of this boundary condition.
     |
     |  ----------------------------------------------------------------------
     |  Methods inherited from BCBase:
     |
     |  __iter__(self)
     |
     |  as_subspace(self, field, V, use_split)
     |
     |  domain_args(...)
     |      The sub_domain the BC applies to.
     |
     |  extract_forms(self, form_type)
     |
     |  function_space(self)
     |      The :class:`.FunctionSpace` on which this boundary condition should
     |      be applied.
     |
     |  function_space_index(self)
     |
     |  increment_bc_depth(self)
     |
     |  node_set(...)
     |      The subset corresponding to the nodes at which this
     |      boundary condition applies.
     |
     |  nodes(...)
     |      The list of nodes at which this boundary condition applies.
     |
     |  set(self, r, val)
     |      Set the boundary nodes to a prescribed (external) value.
     |      :arg r: the :class:`Function` to which the value should be applied.
     |      :arg val: the prescribed value.
     |
     |  sorted_equation_bcs(self)
     |
     |  zero(self, r)
     |      Zero the boundary condition nodes on ``r``.
     |
     |      :arg r: a :class:`.Function` to which the
     |          boundary condition should be applied.
     |
     |  ----------------------------------------------------------------------
     |  Data descriptors inherited from BCBase:
     |
     |  __dict__
     |      dictionary for instance variables (if defined)
     |
     |  __weakref__
     |      list of weak references to the object (if defined)
     |
     |  ----------------------------------------------------------------------
     |  Methods inherited from pyadjoint.overloaded_type.FloatingType:
     |
     |  create_block_variable(self)
     |
     |  ----------------------------------------------------------------------
     |  Data and other attributes inherited from pyadjoint.overloaded_type.FloatingType:
     |
     |  stop_floating = <class 'pyadjoint.overloaded_type.FloatingType.stop_fl...

### Formulação Variacional

Finalmente, o poder do Firedrake e UFL fica evidente na escrita facilitada da Formulação Variacional. Compare a escrita matemática com a codificação equivalente:

\begin{equation}
a(p_h, v) = L(v) + \langle g, v \rangle_{\Gamma_N}
\end{equation}

sendo

\begin{equation}
a(p, v) := \int_\Omega K \nabla p \nabla v \,dx \quad
\end{equation}

\begin{equation}
L(v) := \int_\Omega f\,v \,dx
\end{equation}

```python
a = inner(K * grad(p_h), grad(v)) * dx

L = f * v * dx
```

### Montando e solucionando o problema

E então, vamos montar o problema variacional e declarar a função onde a solução será armazenada.

\begin{equation}
a(p_h, v) = L(v)
\end{equation}

```python
solution = Function(V)
poisson_problem = LinearVariationalProblem(a, L, solution, bcs=bc)
```

Temos que configurar os parâmetros do solver (`PETSc`). Optaremos pelo LU (método direto). Poderíamos também usar as configs default.

```python
solver_parameters = {
    "ksp_type": "preonly",  # Vamos usar apenas o pré-condicionador
    "pc_type": "lu"  # A escolha do pré-condicionador
}
```

E agora, vamos construir o solver, que monta um sistema do tipo $Ax = b$:

```python
poisson_solver = LinearVariationalSolver(poisson_problem, solver_parameters=solver_parameters)
```

No entanto, isso ainda não resolve o sistema resultante. Para tanto, basta executar:

```python
poisson_solver.solve()
```

A sua solução agora está armazenada em `solution`. Você pode plotar ela com a função conveniente:

```python
plt.figure(figsize=(7, 5))

plot(solution)
plt.plot(x_values, exact_solution_values, "-x", label="Exact solution")

plt.grid()
plt.legend(loc="best")

plt.xlabel("x")
plt.ylabel(r"$p_h$")

plt.show()
```

Se preferir, você também poderia coletar os resultados "na mão" para ter maior controle (como abaixo), mas esse trabalho fica bastante dificultado para problemas que não são 1D.

```python
solution_values = solution.vector().dat.data
```

```python
plt.figure(figsize=(7, 5))

plt.plot(x_values, solution_values, "-", label="Numerical solution")
plt.plot(x_values, exact_solution_values, "-x", label="Exact solution")

plt.grid()
plt.legend(loc="best")

plt.xlabel("x")
plt.ylabel(r"$p_h$")

plt.show()
```

## Resultado final

Por fim, podemos colocar todo o código junto, em um único script! Recomendo você tentar diferentes soluções manufaturadas!

```python
import os
os.environ["OMP_NUM_THREADS"] = "1"  # necessário para o Firedrake

from firedrake import *

# Definição da malha
num_elements = 10
mesh = UnitIntervalMesh(num_elements)

# Declaração do espaço de funções
p = 1  # Grau do polinômio
V = FunctionSpace(mesh, "CG", p)

# Funções trial e test
p_h = TrialFunction(V)
v = TestFunction(V)

# Parâmetros e termo de fonte
x, = SpatialCoordinate(mesh)
p_exact = sin(2 * pi * x)
exact_solution = Function(V).interpolate(p_exact)
exact_solution.rename("Exact solution", "label")
f_expression = div(-grad(p_exact))
f = Function(V).interpolate(f_expression)

# Coletando os valores do termo de fonte e coord. da malha
# em arrays da numpy
x_values = mesh.coordinates.vector().dat.data
exact_solution_values = exact_solution.vector().dat.data

# Condição de contorno de Dirichlet
boundary_value = 0.0
bc = DirichletBC(V, boundary_value, [1, 2])  # Fronteira 1 (esquerda) e Fronteira 2 (direita)

# Formulação variacional
a = inner(K * grad(p_h), grad(v)) * dx
L = f * v * dx

# Montando o problema variacional
solution = Function(V)
poisson_problem = LinearVariationalProblem(a, L, solution, bcs=bc)

# Montando o solver e solucionando o sistema resultante
solver_parameters = {
    "ksp_type": "preonly",  # Vamos usar apenas o pré-condicionador
    "pc_type": "lu"  # A escolha do pré-condicionador
}
poisson_solver = LinearVariationalSolver(poisson_problem, solver_parameters=solver_parameters)
poisson_solver.solve()

# Coletando e plotando os resultados
solution_values = solution.vector().dat.data

plt.figure(figsize=(7, 5))

plt.plot(x_values, solution_values, "-", label="Numerical solution")
plt.plot(x_values, exact_solution_values, "-x", label="Exact solution")

plt.grid()
plt.legend(loc="best")

plt.xlabel("x")
plt.ylabel(r"$p_h$")

plt.show()
```

## Extendendo para o Caso 2D

Algumas modificações/adaptações devem ser feitas. Notavelmente:

* Um domínio/malha 2D
* Uma solução manufaturada equivalente em 2D
* Condições de Contorno

Porém, é necessário redefinir quase tudo para propagar essas modificações.

### Definindo uma malha 2D

Vamos definir uma malha análoga. Note que podemos usar quadriláteros ou triângulos!

```python
N = 10
use_quads = True
mesh = UnitSquareMesh(N, N, quadrilateral=use_quads)

fig, axes = plt.subplots(figsize=(6, 6))
triplot(mesh, axes=axes)
plt.show()
```

### Funções e seus Espaços

Direto ao ponto:

```python
p = 1  # Grau do polinômio
V = FunctionSpace(mesh, "CG", p)

p_h = TrialFunction(V)
v = TestFunction(V)
```

### Definindo a solução manufaturada e o termo de fonte

Aqui, precisamos fazer algumas adaptações:

```python
x, y = SpatialCoordinate(mesh)  # Essa é uma delas
p_exact = sin(2 * pi * x) * sin(2 * pi * y)  # Essa é outra

exact_solution = Function(V).interpolate(p_exact)
exact_solution.rename("Exact solution", "label")
f_expression = div(-grad(p_exact))
f = Function(V).interpolate(f_expression)
```

E mantemos:

```python
K = Constant(1)
```

### Condições de Contorno

Praticamente se mantém da mesma forma, só temos que adicionar os novos contornos:

```python
# Condição de contorno de Dirichlet
boundary_value = 0.0
bc = DirichletBC(V, boundary_value, "on_boundary")  # Fronteira 1 (esquerda), 2 (direita), 3 (inferior) e 4 (superior)
```

### Formulação Variacional, problema e solver

Todos são iguais a como eram no Caso 1D.

```python
# Formulação variacional
a = inner(K * grad(p_h), grad(v)) * dx
L = f * v * dx

# Montando o problema variacional
solution = Function(V)
poisson_problem = LinearVariationalProblem(a, L, solution, bcs=bc)

# Montando o solver e solucionando o sistema resultante
solver_parameters = {
    "ksp_type": "preonly",  # Vamos usar apenas o pré-condicionador
    "pc_type": "lu"  # A escolha do pré-condicionador
}
poisson_solver = LinearVariationalSolver(poisson_problem, solver_parameters=solver_parameters)
poisson_solver.solve()
```

### Visualizando o resultado

Vamos checar visualmente o resultado e comparar com a solução exata.

```python
fig, axes = plt.subplots(figsize=(6, 6))
collection = tripcolor(solution, axes=axes, cmap='coolwarm')
fig.colorbar(collection)
axes.set_xlim([0, 1])
axes.set_ylim([0, 1])
plt.title("Solução Numérica")
plt.xlabel("x")
plt.ylabel("y")
plt.show()
```

```python
fig, axes = plt.subplots(figsize=(6, 6))
collection = tripcolor(exact_solution, axes=axes, cmap='coolwarm')
fig.colorbar(collection)
axes.set_xlim([0, 1])
axes.set_ylim([0, 1])
plt.title("Solução Exata")
plt.xlabel("x")
plt.ylabel("y")
plt.show()
```

notebooks/Parte 2 - Intro de FEM com Firedrake.ipynb


```python
Param=["Voc", "Isc", "Vmp", "Imp", "IrRef", "TRef", "DTIsc", "DTVoc", "a"]
```

```python
Names=["Open circuit voltage [mV]", "Short circuit current [mA]", "Maximum power point voltage [mV]", "Maximum power point current [mA]", "Reference irradiance [W/m^2]", "Reference temperature [°C]", "Isc temperature gradient [mA/°C]", "Voc temperature gradient [mV/°C]", "Diode ideality factor"]
```

```python
Defi=['2512', '501.3', '2244', '485.1', '1367', '28', '0.38', '-6.7', '1.5']
```

```python
print("Definition of solar cell parameters. Default values in parentesis")
print('Press enter in each case to mantain the defaults')
```

    Definition of solar cell parameters. Default values in parentesis
    Press enter in each case to mantain the defaults

```python
New = list()
```

```python
for param, name, defi in zip(Param, Names, Defi):
    print('{n}  ({p}={d})'.format(n=name, p=param, d=defi))
    var = input()
    if var == "":
        New.append(float(defi))
    else:
        New.append(float(var))
    continue
```

    Open circuit voltage [mV]  (Voc=2512)

    Short circuit current [mA]  (Isc=501.3)

    Maximum power point voltage [mV]  (Vmp=2244)

    Maximum power point current [mA]  (Imp=485.1)

    Reference irradiance [W/m^2]  (IrRef=1367)

    Reference temperature [°C]  (TRef=28)

    Isc temperature gradient [mA/°C]  (DTIsc=0.38)

    Voc temperature gradient [mV/°C]  (DTVoc=-6.7)

    Diode ideality factor  (a=1.5)

```python
print(New)

```

    [2512.0, 501.3, 2244.0, 485.1, 1367.0, 28.0, 0.38, -6.7, 1.5]

```python
print(Param)
```

    ['Voc', 'Isc', 'Vmp', 'Imp', 'IrRef', 'TRef', 'DTIsc', 'DTVoc', 'a']

```python
k = 1.38064852e-23
```

```python
type(k)
```

    float

```python
q = 1.6021766208e-19
```

```python
from sympy import *
import sympy as sp
import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import fsolve
```

```python
Rs = Symbol('Rs',real=True)
```

```python
Voc = New[0]
Isc = New[1]
Vmp = New[2]
Imp = New[3]
IrRef = New[4]
TRef = New[5]
DTIsc = New[6]
DTVoc = New[7]
a = New[8]
```

```python
Vt = 1000*k*(TRef+273.15)/q
```

```python
print(Vt)
```

    25.95109031052963

```python
T1 = simplify((a*Vt*Vmp*(2*Imp-Isc))/((Vmp*Isc+Voc*(Imp-Isc))*(Vmp-Imp*Rs)-a*Vt*(Vmp*Isc-Voc*Imp)))
```

```python
print(T1)
```

    -40959057.3860803/(525956480.28*Rs - 2436641598.31791)

```python
T2 = simplify(sp.exp((Vmp+Imp*Rs-Voc)/(a*Vt)))
```

```python
func = T1-T2
```

```python
print(func)
```

    -0.00102327597706717*exp(12.4619041485429*Rs) - 40959057.3860803/(525956480.28*Rs - 2436641598.31791)

```python
Rs = nsolve(func, 0.1)
```

```python
print('Rs=',Rs)
```

    Rs= 0.228661690120130

```python
Rsh=((Vmp-Imp*Rs)*(Vmp-Rs*(Isc-Imp)-a*Vt))/((Vmp-Imp*Rs)*(Isc-Imp)-a*Vt*Imp)
```

```python
print('Rsh=',Rsh)
```

    299.612749506682

```python
I0 = (((Rsh+Rs)*Isc-Voc)/1000)/(Rsh*exp(Voc/(a*Vt))) #Diode dark current
```

```python
print('I0=',I0)
```

    I0= 4.64911624361405e-29

```python
Ipv = ((Rsh+Rs)/(Rsh))*(Isc/1000) #Solar current
```

```python
print('Ipv=',Ipv)
```

    Ipv= 0.501682587541571

```python

```


# Pre-MS Stars in the Small Magellanic Cloud

Models are currently computed with a metallicity of [$m$/H] $= -0.5$, which is higher than observed.

```python
%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np
import scipy.interpolate as scint
```

## Equipartition Magnetic Field Strengths

Surface magnetic field strengths are approximated by assuming the magnetic field is in thermal equipartition with the ambient plasma at an optical depth $\tau_{\rm ross} = 1$. Thermal equipartition implies that
\begin{equation}
    B_{{\rm eq},\, \tau=1} = \left(8\pi P_{\rm gas}\right)^{1/2}.
\end{equation}
Gas pressures for the ambient stellar plasma are estimated using detailed stellar model atmospheres. Combinations of $\log g$ and $T_{\rm eff}$ necessary to determine the surface gas pressure are taken from a standard stellar evolution model isochrones at a given age. Since we're interested in a range of ages for comparison, we'll select multiple ages and compare values of equipartition magnetic field strengths. Provided equipartition values don't change appreciably, predictions from magnetic models computed with a single equipartition magnetic field strength will be robust.

We'll start by importing a series of standard model isochrones at 1, 5, 10, and 15 Myr.

```python
iso_01 = np.genfromtxt('../../evolve/models/smc/std/dmestar_00001.0myr_z-0.50_a+0.00_phx.iso')
iso_05 = np.genfromtxt('../../evolve/models/smc/std/dmestar_00005.0myr_z-0.50_a+0.00_phx.iso')
iso_10 = np.genfromtxt('../../evolve/models/smc/std/dmestar_00010.0myr_z-0.50_a+0.00_phx.iso')
iso_15 = np.genfromtxt('../../evolve/models/smc/std/dmestar_00015.0myr_z-0.50_a+0.00_phx.iso')
```

Next, we'll load the atmospheric boundary condition table used to compute the isochrones. The tables are compiled using `PHOENIX` AMES-COND model atmospheres with a solar composition from Grevesse & Sauval (1998). While interior structure models rely on prescribing the gas pressure and temperature at $\tau_{\rm ross}=10$, we need to choose the table for $\tau_{\rm ross}=1$.

```python
atm = np.genfromtxt('../../evolve/phoenix/ames-cond/tab/Zm0d5.ap0d0_t001.dat')
```

Atmosphere tables are structured so that $T_{\rm eff}$ values are listed in the first column, but $\log g$ is not tabulated; it must be prescribed.

```python
teffs = atm[:,0]
loggs = np.arange(-0.5, 5.6, 0.5)
```

The rest of the table is structured as a set of values (P, T) at each $\log g$. It'll be easier to create new arrays for each of these.

```python
temps = np.empty((len(teffs), len(loggs)))
press = np.empty((len(teffs), len(loggs)))
for i, teff in enumerate(atm[:, 1:]):
    for j, prop in enumerate(teff):
        if j%2 == 0:
            press[i, j/2] = prop
        else:
            temps[i, j/2] = prop
```

With the individual tables formed, we can construct an interpolation surface using a 2D interpolation routine. Note that we only really care about the pressure table, as that sets the equipartition magnetic field strengths. A linear interpolation will suffice for now, but there is structure to the temperature-pressure relation that requires a shape preserving interpolation algorithm.

```python
pgas_surface = scint.interp2d(loggs, teffs, press, kind='linear')
```

We are now ready to compute equipartition magnetic field strengths along each of the isochrones. It'll be easiest to quickly define a generic function.

```python
def computeBeq(iso):
    Beq = np.empty((len(iso)))
    for i in range(len(iso)):
        Beq[i] = np.sqrt(8.0*np.pi*pgas_surface(iso[i, 2], 10**iso[i, 1]))
    return Beq
```

And compute equipartition magnetic field strengths along each isochrone.

```python
Beq_01 = computeBeq(iso_01)
Beq_05 = computeBeq(iso_05)
Beq_10 = computeBeq(iso_10)
Beq_15 = computeBeq(iso_15)
```

Now to compare how equipartition magnetic field strengths vary with age.

```python
fig, ax = plt.subplots(1, 1, figsize=(9., 6.))

ax.set_xlabel('${\\rm Mass}\ (M_{\\odot})$', fontsize=20)
ax.set_ylabel('$\\langle {\\rm B}f \\rangle_{\\rm eq}$', fontsize=20)
ax.grid()

ax.plot(iso_01[:,0], Beq_01, '-', lw=4, color='#333333', alpha=0.7)
ax.plot(iso_05[:,0], Beq_05, '-', lw=4, color='#0099cc', alpha=0.7)
ax.plot(iso_10[:,0], Beq_10, '-', lw=4, color='#b30000', alpha=0.7)
ax.plot(iso_15[:,0], Beq_15, '-', lw=4, color='#e68a00', alpha=0.7)
```

There is significant variation between 1 and 15 Myr at nearly every mass, with the exception of stars in the vicinity of $1 M_{\odot}$. However, if we choose a single set of equipartition magnetic fields (say at 10 Myr), we will know where the effects from magnetic fields will either be over and/or under-estimated.

Let's take, for a start, equipartition magnetic field strengths at 10 Myr.