<jupyter_start><jupyter_text>Linear regression for predicting height from weightWe illustrate priors for linear and polynomial regression using the example in sec 4.4  of [Statistical Rethinking ed 2](https://xcelab.net/rm/statistical-rethinking/). The numpyro code is from [Du Phan's site](https://fehiepsi.github.io/rethinking-numpyro/04-geocentric-models.html).<jupyter_code>!pip install -q numpyro@git+https://github.com/pyro-ppl/numpyro
!pip install -q arviz

import numpy as np
np.set_printoptions(precision=3)
import matplotlib.pyplot as plt
import math
import os
import warnings
import pandas as pd

import jax
print("jax version {}".format(jax.__version__))
print("jax backend {}".format(jax.lib.xla_bridge.get_backend().platform))

import jax.numpy as jnp
from jax import random, vmap

rng_key = random.PRNGKey(0)
rng_key, rng_key_ = random.split(rng_key)

import numpyro
import numpyro.distributions as dist
from numpyro.distributions import constraints
from numpyro.distributions.transforms import AffineTransform
from numpyro.diagnostics import hpdi, print_summary
from numpyro.infer import Predictive
from numpyro.infer import MCMC, NUTS
from numpyro.infer import SVI, Trace_ELBO, init_to_value
from numpyro.infer.autoguide import AutoLaplaceApproximation
import numpyro.optim as optim


import arviz as az<jupyter_output>WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)<jupyter_text>DataWe use the "Howell" dataset, which consists of measurements of height, weight, age and sex, of a certain foraging tribe, collected by <NAME>.<jupyter_code>#url = 'https://github.com/fehiepsi/rethinking-numpyro/tree/master/data/Howell1.csv?raw=True'
url = 'https://raw.githubusercontent.com/fehiepsi/rethinking-numpyro/master/data/Howell1.csv'

Howell1 = pd.read_csv(url, sep=';')
d = Howell1
d.info()
d.head()

# get data for adults
d2 = d[d.age >= 18]
az.plot_pair(d2[["weight", "height"]].to_dict(orient="list"))
plt.tight_layout()
plt.savefig('linreg_height_weight_data.pdf', dpi=300)
plt.show()<jupyter_output><empty_output><jupyter_text>Prior predictive distribution<jupyter_code>def plot_prior_predictive_samples(a, b, fname=None):
  plt.subplot(
    xlim=(d2.weight.min(), d2.weight.max()),
    ylim=(-100, 400),
    xlabel="weight",
    ylabel="height",)
  plt.axhline(y=0, c="b", ls='--', label='embryo')
  plt.axhline(y=272, c="r", ls="--",  label='world''s tallest man')
  plt.title("b ~ Normal(0, 10)")
  xbar = d2.weight.mean()
  x = jnp.linspace(d2.weight.min(), d2.weight.max(), 101)
  for i in range(N):
      plt.plot(x, a[i] + b[i] * (x - xbar), "k", alpha=0.2)
  plt.tight_layout()
  if fname:
    plt.savefig(fname, dpi=300)
  plt.show()<jupyter_output><empty_output><jupyter_text>Gaussian prior<jupyter_code>with numpyro.handlers.seed(rng_seed=2971):
    N = 100  # 100 lines
    a = numpyro.sample("a", dist.Normal(178, 20).expand([N]))
    b = numpyro.sample("b", dist.Normal(0, 10).expand([N]))

plot_prior_predictive_samples(a, b, 'linreg_height_weight_gauss_prior.pdf')<jupyter_output><empty_output><jupyter_text>Log-Gaussian prior<jupyter_code>with numpyro.handlers.seed(rng_seed=2971):
    N = 100  # 100 lines
    a = numpyro.sample("a", dist.Normal(178, 28).expand([N]))
    b = numpyro.sample("b", dist.LogNormal(0, 1).expand([N]))

plot_prior_predictive_samples(a, b, 'linreg_height_weight_loggauss_prior.pdf')<jupyter_output><empty_output><jupyter_text>PosteriorWe use the log-gaussian prior.We compute a Laplace approximation to the posterior.<jupyter_code>def model(weight, height):
    a = numpyro.sample("a", dist.Normal(178, 20))
    b = numpyro.sample("b", dist.LogNormal(0, 1))
    sigma = numpyro.sample("sigma", dist.Uniform(0, 50))
    mu = numpyro.deterministic("mu", a + b * (weight - xbar))
    numpyro.sample("height", dist.Normal(mu, sigma), obs=height)

def model2(weight, height=None): # equivalent version
    a = numpyro.sample("a", dist.Normal(178, 20))
    log_b = numpyro.sample("log_b", dist.Normal(0, 1))
    sigma = numpyro.sample("sigma", dist.Uniform(0, 50))
    mu = numpyro.deterministic("mu", a + jnp.exp(log_b) * (weight - xbar))
    numpyro.sample("height", dist.Normal(mu, sigma), obs=height)


m4_3 = AutoLaplaceApproximation(model)
svi = SVI(
    model, m4_3, optim.Adam(1), Trace_ELBO(), weight=d2.weight.values, height=d2.height.values
)
p4_3, losses = svi.run(random.PRNGKey(0), 2000)
post = m4_3.sample_posterior(random.PRNGKey(1), p4_3, (1000,))
{latent: list(post[latent].reshape(-1)[:5]) for latent in post}<jupyter_output><empty_output><jupyter_text>Posterior predictive<jupyter_code># define sequence of weights to compute predictions for
# these values will be on the horizontal axis
weight_seq = jnp.arange(start=25, stop=71, step=1)

# use predictive to compute mu
# for each sample from posterior
# and for each weight in weight_seq
mu = Predictive(m4_3.model, post, return_sites=["mu"])(
    random.PRNGKey(2), weight_seq, None
)["mu"]

print(mu.shape)

# summarize the distribution of mu
mu_mean = jnp.mean(mu, 0)
mu_PI = jnp.percentile(mu, q=(2.5, 97.5), axis=0)
mu_HPDI = hpdi(mu, prob=0.95, axis=0)

# observed output
sim_height = Predictive(m4_3.model, post, return_sites=["height"])(
    random.PRNGKey(2), weight_seq, None
)["height"]
height_PI = jnp.percentile(sim_height, q=(2.5, 97.5), axis=0)


# plot raw data
az.plot_pair(
    d2[["weight", "height"]].to_dict(orient="list"), scatter_kwargs={"alpha": 0.5}
)

# draw MAP line
plt.plot(weight_seq, mu_mean, "k")

# draw HPDI region for line
plt.fill_between(weight_seq, mu_HPDI[0], mu_HPDI[1], color="k", alpha=0.2)

# draw PI region for simulated heights
plt.fill_between(weight_seq, height_PI[0], height_PI[1], color="k", alpha=0.15)
plt.tight_layout()
plt.savefig('linreg_height_weight_loggauss_post.pdf', dpi=300)
plt.show()<jupyter_output>(1000, 46)<jupyter_text>Polynomial regression We will now consider the full dataset, including children. The resulting mapping from weight to height is now nonlinear.  Data<jupyter_code>az.plot_pair(d[["weight", "height"]].to_dict(orient="list")) # d, not d2
plt.tight_layout()
plt.savefig('linreg_height_weight_data_full.pdf', dpi=300)
plt.show()

# standardize
d["weight_s"] = (d.weight - d.weight.mean()) / d.weight.std()

# precompute polynomial terms
d["weight_s2"] = d.weight_s ** 2
d["weight_s3"] = d.weight_s ** 3

ax = az.plot_pair(
    d[["weight_s", "height"]].to_dict(orient="list"), scatter_kwargs={"alpha": 0.5}
)
ax.set(xlabel="weight", ylabel="height", xticks=[])
fig = plt.gcf()

at = jnp.array([-2, -1, 0, 1, 2])
labels = at * d.weight.std() + d.weight.mean()
ax.set_xticks(at)
ax.set_xticklabels([round(label, 1) for label in labels])
#fig

plt.tight_layout()
plt.savefig('linreg_height_weight_data_full_standardized.pdf', dpi=300)
plt.show()<jupyter_output><empty_output><jupyter_text>Fit model<jupyter_code>def fit_predict(model, fname=None):
  guide = AutoLaplaceApproximation(model)
  svi = SVI(
      model,
      guide,
      optim.Adam(0.3),
      Trace_ELBO(),
      weight_s=d.weight_s.values,
      weight_s2=d.weight_s2.values,
      height=d.height.values,
  )
  params, losses = svi.run(random.PRNGKey(0), 3000)

  samples = guide.sample_posterior(random.PRNGKey(1), params, (1000,))
  print_summary({k: v for k, v in samples.items() if k != "mu"}, 0.95, False)

  weight_seq = jnp.linspace(start=-2.2, stop=2, num=30)
  pred_dat = {"weight_s": weight_seq, "weight_s2": weight_seq ** 2}
  post = guide.sample_posterior(random.PRNGKey(1), params, (1000,))
  predictive = Predictive(guide.model, post)
  mu = predictive(random.PRNGKey(2), **pred_dat)["mu"]
  mu_mean = jnp.mean(mu, 0)
  mu_PI = jnp.percentile(mu, q=(2.5, 97.5), axis=0)
  sim_height = predictive(random.PRNGKey(3), **pred_dat)["height"]
  height_PI = jnp.percentile(sim_height, q=(2.5, 97.5), axis=0)

  ax = az.plot_pair(
    d[["weight_s", "height"]].to_dict(orient="list"), scatter_kwargs={"alpha": 0.5})
  plt.plot(weight_seq, mu_mean, "k")
  plt.fill_between(weight_seq, mu_PI[0], mu_PI[1], color="k", alpha=0.2)
  plt.fill_between(weight_seq, height_PI[0], height_PI[1], color="k", alpha=0.15)

  ax.set(xlabel="weight", ylabel="height", xticks=[])
  fig = plt.gcf()
  at = jnp.array([-2, -1, 0, 1, 2])
  labels = at * d.weight.std() + d.weight.mean()
  ax.set_xticks(at)
  ax.set_xticklabels([round(label, 1) for label in labels])

  plt.tight_layout()
  if fname:
    plt.savefig(fname, dpi=300)
  plt.show()<jupyter_output><empty_output><jupyter_text>Linear<jupyter_code>def model_linear(weight_s, weight_s2, height=None):
    a = numpyro.sample("a", dist.Normal(178, 20))
    b1 = numpyro.sample("b1", dist.LogNormal(0, 1)) # Log-Normal prior
    sigma = numpyro.sample("sigma", dist.Uniform(0, 50))
    mu = numpyro.deterministic("mu", a + b1 * weight_s)
    numpyro.sample("height", dist.Normal(mu, sigma), obs=height)


fit_predict(model_linear, 'linreg_height_weight_data_full_linear.pdf')<jupyter_output>100%|██████████| 3000/3000 [00:02<00:00, 1344.90it/s, init loss: 49746.6445, avg. loss [2851-3000]: 2001.7004]<jupyter_text>Quadratic<jupyter_code>def model_quad(weight_s, weight_s2, height=None):
    a = numpyro.sample("a", dist.Normal(178, 20))
    b1 = numpyro.sample("b1", dist.LogNormal(0, 1)) 
    b2 = numpyro.sample("b2", dist.Normal(0, 1)) 
    sigma = numpyro.sample("sigma", dist.Uniform(0, 50))
    mu = numpyro.deterministic("mu", a + b1 * weight_s + b2 * weight_s2)
    numpyro.sample("height", dist.Normal(mu, sigma), obs=height)


fit_predict(model_quad, 'linreg_height_weight_data_full_quad.pdf')
def model_quad_positive_b2(weight_s, weight_s2, height=None):
    a = numpyro.sample("a", dist.Normal(178, 20))
    b1 = numpyro.sample("b1", dist.LogNormal(0, 1))
    #b2 = numpyro.sample("b2", dist.Normal(0, 1)) 
    b2 = numpyro.sample("b2", dist.LogNormal(0, 1))  
    sigma = numpyro.sample("sigma", dist.Uniform(0, 50))
    mu = numpyro.deterministic("mu", a + b1 * weight_s + b2 * weight_s2)
    numpyro.sample("height", dist.Normal(mu, sigma), obs=height)


fit_predict(model_quad_positive_b2, 'linreg_height_weight_data_full_quad_pos_b2.pdf')<jupyter_output>100%|██████████| 3000/3000 [00:02<00:00, 1198.08it/s, init loss: 66975.6406, avg. loss [2851-3000]: 2008.9690]
{'fraction_non_alphanumeric': 0.16121247224635582, 'fraction_numerical': 0.040544454097885894, 'mean_word_length': 5.430788330229671, 'pattern_counts': {'":': 5, '<': 44, '<?xml version=': 0, '>': 43, 'https://': 5, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '11', 'dir': 'jupyter-structured-clean-dedup', 'id': '191136', 'license': 'MIT', 'n_tokens_mistral': 4311, 'n_tokens_neox': 3983, 'n_words': 795, 'path': 'notebooks/linreg_height_weight_numpyro.ipynb', 'repo_name': 'arpitvaghela/probml-notebooks', 'size': '10365'}