# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

import matplotlib.pyplot as plt
import numpy
from scipy.interpolate import interp1d
from typing import Tuple

na = 11037
np = 11034
xa = 139
xp = 239
theta = numpy.linspace(0.2, 0.99, 100)


def profile_likelihood(theta: numpy.ndarray,
                       na: float,
                       np: float,
                       xa: float,
                       xp: float) -> numpy.ndarray:
    ratio = na * theta / (na * theta + np)
    log_like = xa * numpy.log(ratio) + xp * numpy.log(1 - ratio)
    like = numpy.exp(log_like)
    like /= numpy.max(like)
    return like


def likelihood_interval(theta: numpy.ndarray,
                        likelihood: numpy.ndarray,
                        cutoff: float) -> Tuple[float, float]:
    # intersection points occur below and above the maximum likelihood estimate
    mle_index = numpy.argmax(likelihood)
    interp_below_max = interp1d(likelihood[:mle_index], theta[:mle_index])
    interp_above_max = interp1d(likelihood[mle_index:], theta[mle_index :])
    lower_int = numpy.round(interp_below_max(cutoff).flatten()[0], 2)
    upper_int = numpy.round(interp_above_max(cutoff).flatten()[0], 2)
    return (lower_int, upper_int)


def plot_profile_likelihood(
    theta: numpy.ndarray,
    likelihood: numpy.array,
    title='') -> None:
    plt.plot(theta, likelihood)
    plt.axhline(y=0.15, linewidth=1)
    if title.startswith("(b)"):
        plt.axvline(x=1.0, linewidth=1)
    plt.xlabel(r'$\theta$')
    plt.ylabel('Likelihood')
    plt.title(title);


likelihood = profile_likelihood(theta, na, np, xa, xp)
plot_profile_likelihood(theta, likelihood, title='(a) Heart attacks')

ci = likelihood_interval(theta, likelihood, 0.15)
print(f"95% CI = ", ci)

na = 11037
np = 11034
xa = 119
xp = 98
theta = numpy.linspace(0.7, 1.99, 100)

likelihood = profile_likelihood(theta, na, np, xa, xp)
plot_profile_likelihood(theta, likelihood, title='(b) Strokes')

ci = likelihood_interval(theta, likelihood, 0.15)
print(f"95% CI = ", ci)

StarcoderdataJupyterScriptsDedupFiltered