# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: 'Python 3.7.6 64-bit (''venv'': venv)'
#     name: python_defaultSpec_1598173300831
# ---

# +
import math

import arviz as az
import matplotlib.pyplot as plt
import pandas as pd
from IPython.display import set_matplotlib_formats

import jax.numpy as jnp
from jax import lax, random
from jax.scipy.special import expit, logit

import numpyro
import numpyro.distributions as dist
import numpyro.optim as optim
from numpyro.diagnostics import print_summary
from numpyro.distributions.transforms import OrderedTransform
from numpyro.infer import ELBO, MCMC, NUTS, SVI, Predictive, init_to_value
from numpyro.infer.autoguide import AutoLaplaceApproximation

# %config InlineBackend.figure_formats = ["svg"]
az.style.use("arviz-darkgrid")
numpyro.set_host_device_count(4)

# +
ratings_counts = jnp.array([12, 36, 7, 41])

pr_r = ratings_counts / ratings_counts.sum()
cpr_r = jnp.cumsum(pr_r, -1)
cpr_r
# -

lco_r = logit(cpr_r)
lco_r

plt.plot(range(1, 5), cpr_r, "--o")
plt.gca().set(xlabel="rank", ylabel="cumulative proportion", ylim=(-0.1, 1.1))
plt.show()

# + tags=[]
hurricanes = pd.read_csv("../data/Hurricanes.csv", sep=";")
d = hurricanes
print(d.shape)
display(d.sample(3))
d.describe()
# -

d.category.value_counts().sort_index()

#
# ## Hurricane fatalities and gender of names
#
# ### Description
#
# Data used in Jung et al 2014 analysis of effect of gender of name on hurricane fatalities. Note that hurricanes Katrina (2005) and Audrey (1957) were removed from the data.
#
# ### Format
#
# name : Given name of hurricane  
# year : Year of hurricane  
# deaths : number of deaths  
# category : Severity code for storm  
# min_pressure : Minimum pressure, a measure of storm strength; low is stronger  
# damage_norm : Normalized estimate of damage in dollars  
# female : Indicator variable for female name  
# femininity : 1-11 scale from totally masculine (1) to totally feminine (11) for name. Average of 9 scores from 9 raters.  
#

import seaborn as sns
sns.pairplot(d[['femininity', 'category', 'min_pressure', 'damage_norm', 'deaths']])


# + tags=[]
def model1(deaths=None, femininity=None):
    a = numpyro.sample("a", dist.Normal(0, .3))
    bF = numpyro.sample("bF", dist.Normal(0.2, .25))
    lambda_ = numpyro.deterministic("lambda", jnp.exp(a + bF * femininity))
    numpyro.sample("deaths", dist.Poisson(lambda_), obs=deaths)


# +
x = jnp.linspace(d.femininity.min(), d.femininity.max(), 100)
lmu = Predictive(model1, num_samples=100, return_sites=["lambda"]).get_samples(
    random.PRNGKey(53),
    femininity=x,
)["lambda"]
plt.axhline(y=d.deaths.min(), c="k", ls="--")
plt.axhline(y=d.deaths.max(), c="k", ls="--")#, lw=0.5)
plt.title("prior pred")
plt.xlabel('femininity')

for i in range(40):
    plt.plot(x, lmu[i], "k", alpha=0.3)

# + tags=[]
mcmc = MCMC(NUTS(model1), 500, 500)
mcmc.run(random.PRNGKey(0), femininity=d.femininity.values, deaths=d.deaths.values)
mcmc.print_summary()
post1 = mcmc.get_samples()

# +
x = jnp.linspace(d.femininity.min(), d.femininity.max(), 100)
pred = Predictive(model1, post1, return_sites=["lambda"]).get_samples(random.PRNGKey(33), femininity=x)
lmu = jnp.mean(pred['lambda'], 0)
lci = jnp.percentile(pred['lambda'], q=(5.5, 94.5), axis=0)

plt.scatter(d.femininity.values, d.deaths.values)
plt.plot(x, lmu, 'k')
plt.fill_between(x, lci[0], lci[1], color='k', alpha=0.2)


# -

def model2(deaths=None):
    a = numpyro.sample("a", dist.Normal(0, .3))
    lambda_ = numpyro.deterministic("lambda", jnp.exp(a))
    numpyro.sample("deaths", dist.Poisson(lambda_), obs=deaths)


# + tags=[]
mcmc = MCMC(NUTS(model2), 500, 500)
mcmc.run(random.PRNGKey(0), deaths=d.deaths.values)
mcmc.print_summary()
post2 = mcmc.get_samples()


# -

def model3(deaths=None, femininity=None):
    a = numpyro.sample("a", dist.Normal(0, 5))
    bF = numpyro.sample("bF", dist.Normal(0.2, .25))    
    phi = numpyro.sample("phi", dist.Exponential(1))
    lambda_ = numpyro.deterministic("lambda", jnp.exp(a + bF * femininity))
    numpyro.sample("deaths", dist.GammaPoisson(lambda_ / phi, 1 / phi), obs=deaths)


# + tags=[]
mcmc = MCMC(NUTS(model3), 500, 500)
mcmc.run(random.PRNGKey(0), femininity=d.femininity.values, deaths=d.deaths.values)
mcmc.print_summary()
post3 = mcmc.get_samples()

# +
x = jnp.linspace(d.femininity.min(), d.femininity.max(), 100)
pred = Predictive(model3, post3, return_sites=["lambda"]).get_samples(random.PRNGKey(33), femininity=x)
lmu = jnp.mean(pred['lambda'], 0)
lci = jnp.percentile(pred['lambda'], q=(5.5, 94.5), axis=0)

plt.scatter(d.femininity.values, d.deaths.values)
plt.plot(x, lmu, 'k')
plt.fill_between(x, lci[0], lci[1], color='k', alpha=0.2)
# -

deaths_pred = Predictive(model3, post3)(random.PRNGKey(123), femininity=d.femininity.values)["deaths"]
plt.scatter(d.femininity.values, d.deaths.values)
plt.errorbar(
    d.femininity.values,
    jnp.mean(deaths_pred, 0),
    jnp.std(deaths_pred, 0) / 2,
    fmt="o",
    c="k",
    mfc="none",
    ms=7,
    elinewidth=1,
)
plt.plot(d.femininity.values, jnp.percentile(deaths_pred, 5.5, 0), "k+")
plt.plot(d.femininity.values, jnp.percentile(deaths_pred, 94.5, 0), "k+")



StarcoderdataJupyterScriptsDedupFiltered