```python
import matplotlib.pyplot as plt
import numpy as np
import scipy.special
from scipy.optimize import minimize
from scipy.stats import geom
from typing import Tuple
```


```python
cycles = np.array([i for i in range(1, 14)])
smokers = np.array([29, 16, 17, 4, 3, 9, 4, 5, 1, 1, 1, 3, 7])
non_smokers = np.array([198, 107, 55, 38, 18, 22, 7, 9, 5, 3, 6, 6, 12])
```

(a) Fitting a geometric model to the group of smokers and the group of non-smokers we can compare the estimated probability of pregnancy per cycle.


```python
def plot_pregnancy_per_cycle(
    cyles: np.ndarray,
    pregnancy_counts: np.ndarray,
    title_label: str) -> None:
    p = pregnancy_counts[0] / np.sum(pregnancy_counts)
    prob = geom.pmf(cycles, p)
    plt.bar(cycles, prob)
    plt.xlabel('Number of cycles')
    plt.ylabel('Probability p')
    plt.title(f'Probability of pregnancy per cycle for {title_label}');
```


```python
plot_pregnancy_per_cycle(cycles, smokers, 'smokers')
```


```python
plot_pregnancy_per_cycle(cycles, non_smokers, 'non-smokers')
```

From the plots it seems that the probability of getting pregnant in less cycles is higher for non-smokers. 

(b) To check the adequacy of the geometric model we follow the Example (4.6) by plotting a Geometric plot. The idea is to plot the number of cycles k against a function of the counts $n_k$ that is expected to be a straight line for Geometric data. Since $p_k = P(X = k) = (1 - p)^{k - 1}p$, then

$$
\log p_k = k \log(1 - p) + \log \Big( \frac{p}{1 - p} \Big)
$$

indicating that we should plot $\log p_k$ versus $k$, where $p_k$ is estimated by $n_k \big/ \sum n_k$.


```python
def geometric_plot(cycles: np.ndarray, counts: np.ndarray) -> None:
    p_k = counts / np.sum(counts)
    plt.plot(cycles, np.log(p_k), 'o')
    plt.xlabel('Number of cycles k')
    plt.ylabel(r'Log($p_k$)');
    plt.title('(a) Geometric plot')
```


```python
geometric_plot(cycles, smokers)
```


```python
geometric_plot(cycles, non_smokers)
```

It is clear that for both smokers and non-smokers the data does not follow a Geometric prescription.

(c) We are told that conditional on $p$ to let the pregancy counts $X_p \sim geom(p)$, where the probability of pregnancy per cycle $p$ itself is random with beta distribution with density:

$$
f(p) = \frac{1}{B(\alpha, \beta)}p^{\alpha - 1} (1 - p)^{\beta - 1}
$$

The beta-geometric probabilities are given by

\begin{align}
P(X = k) & = E \big[ P (X_p = k | p) \big]  \\
& = E \big[ (1 - p)^{k-1}p \big] \\
& = \int_{0}^{1} (1 - p)^{k-1} p f(p) dp \\
& = \frac{1}{B(\alpha, \beta)} \int_{0}^{1} p^{\alpha - 1} (1 - p)^{\beta - 1} (1 - p)^{k - 1} p dp \\
& = \frac{1}{B(\alpha, \beta)} \int_{0}^{1} p^{(\alpha + 1) - 1} (1 - p)^{(\beta + k - 1) - 1} dp \\
& = \frac{B(\alpha + 1, \beta + k - 1)}{B(\alpha, \beta)}
\end{align}

for $k = 0, 1, ...$ where the definition of the [beta function](https://en.wikipedia.org/wiki/Beta_function) is used to evaluate the integral.

(d) To fit the beta-geometric model to the data, the log-likelihood of $(\alpha, \beta)$ is:

$$
\log L(\alpha, \beta) = \sum_{k=1}^{12}n_k \log P(X=k) + n_{[>12]} \log\{1 - P(X > 12)\}
$$

where $P(X = k)$ is given in (b) above. Numerical optimization of the log-likelihood yields.


```python
def neg_log_likelihood(x: np.ndarray, cycles: np.ndarray, counts: np.ndarray) -> float:
    alpha = x[0]
    beta = x[1]
    n_k = counts
    p_k = beta_geometric_probability(alpha, beta, cycles)
    neg_log_like = -np.sum(n_k[:-1] * p_k[:-1]) - n_k[-1] * np.log(1 - p_k[-1])
    return neg_log_like

def beta_geometric_probability(alpha: float, beta: float, k: np.ndarray) -> np.ndarray:
    prob = scipy.special.beta(alpha + 1, beta + k - 1) / scipy.special.beta(alpha, beta)
    return prob
```


```python
def fit_beta_geometric_model(
    cycles: np.ndarray, counts: np.ndarray) -> Tuple[float, float]:
    # Yikes the results are sensitive to these initial conditions! The max value is
    # consistent but occurs at different alpha and beta values.
    x0 = np.array([1.0, 1.0]) 
    bounds = ((1e-10, None), (1e-10, None))
    result = minimize(neg_log_likelihood, x0, args=(cycles, counts), bounds=bounds)
    #print(result)
    return result.x[0], result.x[1]

def print_mle(alpha_mle: float, beta_mle: float, title: str) -> None:
    (print(title + " alpha-hat = ", np.round(alpha_mle, 2), ", ", 
           "beta-hat = ", np.round(beta_mle, 7)))
```


```python
alpha_hat, beta_hat = fit_beta_geometric_model(cycles, smokers)
print_mle(alpha_hat, beta_hat, 'Smokers')
```


```python
def expected_frequencies(mle_probs: np.ndarray, total_count: float) -> np.ndarray:
    freqs = mle_probs * total_count
    return freqs

def chi_squared(obs_freqs: np.ndarray, exp_freqs: np.ndarray) -> float:
    chi_sq = np.sum((obs_freqs - exp_freqs)**2 / exp_freqs)
    return chi_sq

def print_chi_squared(chi_sq: float) -> None:
    print("chi-squared = ", np.round(chi_sq, 1))
```


```python
N = np.sum(smokers)
mle_probs = beta_geometric_probability(alpha_hat, beta_hat, cycles)
print(mle_probs, np.sum(mle_probs))
exp_freqs = expected_frequencies(mle_probs, N)
chi_sq = chi_squared(smokers, exp_freqs)
print_chi_squared(chi_sq)
```

The MLE probabilities are clearly wrong as all the probability mass is put on the first cyle. The same happens for the non-smokers below. The problem seems to stem from the MLE values but I am not sure where the bug is in the code (my code or scipy) or possibly in my math. Anyway I think the general workflow is correct in this notebook.


```python
alpha_hat, beta_hat = fit_beta_geometric_model(cycles, non_smokers)
print_mle(alpha_hat, beta_hat, 'Non-smokers')
```


```python
N = np.sum(non_smokers)
mle_probs = beta_geometric_probability(alpha_hat, beta_hat, cycles)
print(mle_probs, np.sum(mle_probs))
exp_freqs = expected_frequencies(mle_probs, N)
chi_sq = chi_squared(non_smokers, exp_freqs)
print_chi_squared(chi_sq)
```

c816c82d99243daedd8a175d8dbc798d5fedb953
ipynb
Jupyter Notebook
python/chapter-4/exercises/EX4-6.ipynb
covuworie/in-all-likelihood
6638bec8bb4dde7271adb5941d1c66e7fbe12526
MIT
python/chapter-4/exercises/EX4-6.ipynb
covuworie/in-all-likelihood
6638bec8bb4dde7271adb5941d1c66e7fbe12526
MIT
2020-03-24T17:53:04.000Z
2021-08-23T20:16:17.000Z
python/chapter-4/exercises/EX4-6.ipynb
covuworie/in-all-likelihood
6638bec8bb4dde7271adb5941d1c66e7fbe12526
MIT
2021-11-21T10:24:59.000Z
2021-11-21T10:24:59.000Z
Qwen/Qwen-72B
1. YES
2. YES