# Code 6: Time Series#

import matplotlib.pyplot as plt
import arviz as az
import pandas as pd
import numpy as np

import tensorflow as tf
import tensorflow_probability as tfp

tfd = tfp.distributions
tfb = tfp.bijectors
root = tfd.JointDistributionCoroutine.Root

import datetime
print(f"Last Run {datetime.datetime.now()}")
az.style.use("arviz-grayscale")
plt.rcParams["figure.dpi"] = 300

## Time Series Analysis as a Regression Problem#

### Code 6.1#

co2_by_month["date_month"] = pd.to_datetime(co2_by_month["date_month"])
co2_by_month["CO2"] = co2_by_month["CO2"].astype(np.float32)
co2_by_month.set_index("date_month", drop=True, inplace=True)

num_forecast_steps = 12 * 10  # Forecast the final ten years, given previous data
co2_by_month_training_data = co2_by_month[:-num_forecast_steps]
co2_by_month_testing_data = co2_by_month[-num_forecast_steps:]

### Figure 6.1#

def plot_co2_data(fig_ax=None):
if not fig_ax:
fig, ax = plt.subplots(1, 1, figsize=(10, 5))
else:
fig, ax = fig_ax
ax.plot(co2_by_month_training_data, label="training data")
ax.plot(co2_by_month_testing_data, color="C4", label="testing data")
ax.legend()
ax.set(
ylabel="Atmospheric CO₂ concentration (ppm)",
xlabel="Year"
)
ax.text(0.99, .02,
"""Source: Scripps Institute for Oceanography CO₂ program
http://scrippsco2.ucsd.edu/data/atmospheric_co2/primary_mlo_co2_record""",
transform=ax.transAxes,
horizontalalignment="right",
alpha=0.5)
fig.autofmt_xdate()
return fig, ax

_ = plot_co2_data()
plt.savefig("img/chp06/fig1_co2_by_month.png")
/var/folders/7p/srk5qjp563l5f9mrjtp44bh800jqsw/T/ipykernel_14469/1698526664.py:19: UserWarning: This figure was using constrained_layout, but that is incompatible with subplots_adjust and/or tight_layout; disabling constrained_layout.
fig.autofmt_xdate()

### Code 6.2 and Figure 6.2#

trend_all = np.linspace(0., 1., len(co2_by_month))[..., None]
trend_all = trend_all.astype(np.float32)
trend = trend_all[:-num_forecast_steps, :]

seasonality_all = pd.get_dummies(
co2_by_month.index.month).values.astype(np.float32)
seasonality = seasonality_all[:-num_forecast_steps, :]

fig, ax = plt.subplots(figsize=(10, 4))
X_subset = np.concatenate([trend, seasonality], axis=-1)[-50:]
im = ax.imshow(X_subset.T, cmap="cet_gray_r")

label_loc = np.arange(1, 50, 12)
ax.set_xticks(label_loc)
ax.set_yticks([])
ax.set_xticklabels(co2_by_month.index.year[-50:][label_loc])
fig.colorbar(im, ax=ax, orientation="horizontal", shrink=.6)

plt.savefig("img/chp06/fig2_sparse_design_matrix.png")

### Code 6.3#

tfd = tfp.distributions
root = tfd.JointDistributionCoroutine.Root

@tfd.JointDistributionCoroutine
def ts_regression_model():
intercept = yield root(tfd.Normal(0., 100., name='intercept'))
trend_coeff = yield root(tfd.Normal(0., 10., name='trend_coeff'))
seasonality_coeff = yield root(
tfd.Sample(tfd.Normal(0., 1.),
sample_shape=seasonality.shape[-1],
name='seasonality_coeff'))
noise = yield root(tfd.HalfCauchy(loc=0., scale=5., name='noise_sigma'))
y_hat = (intercept[..., None] +
tf.einsum('ij,...->...i', trend, trend_coeff) +
tf.einsum('ij,...j->...i', seasonality, seasonality_coeff))
observed = yield tfd.Independent(
tfd.Normal(y_hat, noise[..., None]),
reinterpreted_batch_ndims=1,
name='observed')

# If we remove the decorator @tfd.JointDistributionCoroutine above:
# ts_regression_model = tfd.JointDistributionCoroutine(ts_regression_model)

# check shape
ts_regression_model.log_prob_parts(ts_regression_model.sample(4))
StructTuple(
intercept=<tf.Tensor: shape=(4,), dtype=float32, numpy=array([-7.068921, -6.320923, -5.556099, -6.88211 ], dtype=float32)>,
trend_coeff=<tf.Tensor: shape=(4,), dtype=float32, numpy=array([-3.8205433, -3.7595298, -3.3554578, -3.2555058], dtype=float32)>,
seasonality_coeff=<tf.Tensor: shape=(4,), dtype=float32, numpy=array([-13.463518, -15.407171, -14.842403, -16.500769], dtype=float32)>,
noise_sigma=<tf.Tensor: shape=(4,), dtype=float32, numpy=array([-5.227767 , -2.787994 , -2.855589 , -2.3863304], dtype=float32)>,
observed=<tf.Tensor: shape=(4,), dtype=float32, numpy=array([-2344.5051, -1561.3213, -1644.2692, -1306.9598], dtype=float32)>
)

### Code 6.4 and Figure 6.3#

# Draw 100 prior and prior predictive samples
prior_samples = ts_regression_model.sample(100)
prior_predictive_timeseries = prior_samples[-1]

fig, ax = plt.subplots(figsize=(10, 5))
ax.plot(co2_by_month.index[:-num_forecast_steps],
tf.transpose(prior_predictive_timeseries), alpha=.5)
ax.set_xlabel("Year")
fig.autofmt_xdate()

plt.savefig("img/chp06/fig3_prior_predictive1.png")
/var/folders/7p/srk5qjp563l5f9mrjtp44bh800jqsw/T/ipykernel_14469/2060325595.py:9: UserWarning: This figure was using constrained_layout, but that is incompatible with subplots_adjust and/or tight_layout; disabling constrained_layout.
fig.autofmt_xdate()

### Code 6.5#

run_mcmc = tf.function(
autograph=False, jit_compile=True)
%%time
mcmc_samples, sampler_stats = run_mcmc(
observed=co2_by_month_training_data["CO2"].values[None, ...])
WARNING:tensorflow:From /opt/miniconda3/envs/bmcp/lib/python3.9/site-packages/tensorflow_probability/python/distributions/distribution.py:342: calling MultivariateNormalDiag.__init__ (from tensorflow_probability.python.distributions.mvn_diag) with scale_identity_multiplier is deprecated and will be removed after 2020-01-01.
Instructions for updating:
scale_identity_multiplier is deprecated; please combine it into scale_diag directly instead.
2021-11-19 17:31:24.773325: W tensorflow/compiler/tf2xla/kernels/random_ops.cc:102] Warning: Using tf.random.uniform with XLA compilation will ignore seeds; consider using tf.random.stateless_uniform instead if reproducible behavior is desired. sanitize_seed/seed
2021-11-19 17:31:24.921521: W tensorflow/compiler/tf2xla/kernels/assert_op.cc:38] Ignoring Assert operator mcmc_retry_init/assert_equal_1/Assert/AssertGuard/Assert
CPU times: user 24.4 s, sys: 596 ms, total: 25 s
Wall time: 25.2 s
regression_idata = az.from_dict(
posterior={
k:np.swapaxes(v.numpy(), 1, 0)
for k, v in mcmc_samples._asdict().items()},
sample_stats={
k:np.swapaxes(sampler_stats[k], 1, 0)
for k in ["target_log_prob", "diverging", "accept_ratio", "n_steps"]}
)

az.summary(regression_idata)
mean sd hdi_3% hdi_97% mcse_mean mcse_sd ess_bulk ess_tail r_hat
intercept 317.417 0.323 316.828 318.046 0.014 0.010 559.0 1101.0 1.0
trend_coeff 80.648 0.299 80.096 81.225 0.005 0.003 3699.0 2994.0 1.0
seasonality_coeff[0] 0.095 0.369 -0.569 0.800 0.014 0.010 693.0 1495.0 1.0
seasonality_coeff[1] 0.647 0.366 -0.055 1.309 0.013 0.009 789.0 1580.0 1.0
seasonality_coeff[2] 1.310 0.370 0.626 2.002 0.014 0.010 712.0 1514.0 1.0
seasonality_coeff[3] 2.443 0.369 1.751 3.140 0.014 0.010 741.0 1697.0 1.0
seasonality_coeff[4] 2.861 0.366 2.168 3.549 0.014 0.010 714.0 1637.0 1.0
seasonality_coeff[5] 2.217 0.371 1.511 2.878 0.014 0.010 694.0 1493.0 1.0
seasonality_coeff[6] 0.712 0.368 0.015 1.406 0.014 0.010 710.0 1584.0 1.0
seasonality_coeff[7] -1.296 0.365 -1.946 -0.575 0.013 0.009 749.0 1748.0 1.0
seasonality_coeff[8] -3.041 0.365 -3.709 -2.364 0.013 0.009 783.0 1541.0 1.0
seasonality_coeff[9] -3.157 0.365 -3.860 -2.491 0.013 0.009 793.0 1559.0 1.0
seasonality_coeff[10] -1.987 0.362 -2.665 -1.326 0.013 0.009 729.0 1310.0 1.0
seasonality_coeff[11] -0.858 0.364 -1.559 -0.180 0.013 0.009 747.0 1482.0 1.0
noise_sigma 1.581 0.050 1.489 1.675 0.001 0.001 4282.0 3167.0 1.0
axes = az.plot_trace(regression_idata, compact=True);

### Code 6.6#

# We can draw posterior predictive sample with jd.sample_distributions()
posterior_dist, posterior_and_predictive = ts_regression_model.sample_distributions(
value=mcmc_samples)
posterior_predictive_samples = posterior_and_predictive[-1]
posterior_predictive_dist = posterior_dist[-1]

# Since we want to also plot the posterior predictive distribution for
# each components, conditioned on both training and testing data, we
# construct the posterior predictive distribution as below:
nchains = regression_idata.posterior.dims['chain']

trend_posterior = mcmc_samples.intercept + \
tf.einsum('ij,...->i...', trend_all, mcmc_samples.trend_coeff)
seasonality_posterior = tf.einsum(
'ij,...j->i...', seasonality_all, mcmc_samples.seasonality_coeff)

y_hat = trend_posterior + seasonality_posterior
posterior_predictive_dist = tfd.Normal(y_hat, mcmc_samples.noise_sigma)
posterior_predictive_samples = posterior_predictive_dist.sample()

### Figure 6.4#

# plot components
fig, ax = plt.subplots(2, 1, figsize=(10, 5), sharex=True)

for i in range(nchains):
ax[0].plot(co2_by_month.index[:-num_forecast_steps],
trend_posterior[:-num_forecast_steps, -100:, i], alpha=.05)
ax[1].plot(co2_by_month.index[:-num_forecast_steps],
seasonality_posterior[:-num_forecast_steps, -100:, i], alpha=.05)

ax[0].set_title('Trend (Linear)')
ax[1].set_title('Seasonality (Month of the year effect)')
ax[1].set_xlabel("Year")
fig.autofmt_xdate()

plt.savefig("img/chp06/fig4_posterior_predictive_components1.png")
/var/folders/7p/srk5qjp563l5f9mrjtp44bh800jqsw/T/ipykernel_14469/865649276.py:13: UserWarning: This figure was using constrained_layout, but that is incompatible with subplots_adjust and/or tight_layout; disabling constrained_layout.
fig.autofmt_xdate()

### Figure 6.5#

fig, ax = plt.subplots(1, 1, figsize=(10, 5))
sample_shape = posterior_predictive_samples.shape[1:]
ax.plot(co2_by_month.index,
tf.reshape(posterior_predictive_samples,
[-1, tf.math.reduce_prod(sample_shape)])[:, :500],
color='gray', alpha=.01)

plot_co2_data((fig, ax))
plt.savefig("img/chp06/fig5_posterior_predictive1.png")
/var/folders/7p/srk5qjp563l5f9mrjtp44bh800jqsw/T/ipykernel_14469/1698526664.py:19: UserWarning: This figure was using constrained_layout, but that is incompatible with subplots_adjust and/or tight_layout; disabling constrained_layout.
fig.autofmt_xdate()

### Code 6.7#

n_changepoints = 8
n_tp = 500
t = np.linspace(0, 1, n_tp)
s = np.linspace(0, 1, n_changepoints + 2)[1:-1]
A = (t[:, None] > s)

k, m = 2.5, 40
delta = np.random.laplace(.1, size=n_changepoints)
growth = (k + A @ delta) * t
offset = m + A @ (-s * delta)
trend = growth + offset
_, ax = plt.subplots(4, 1, figsize=(10, 10))
ax[0].imshow(A.T, cmap="cet_gray_r", aspect='auto', interpolation='none');
ax[0].axis('off')
ax[0].set_title(r'$\mathbf{A}$')
ax[1].plot(t, growth, lw=2)
ax[1].set_title(r'$(k + \mathbf{A}\delta) t$')
ax[2].plot(t, offset, lw=2)
ax[2].set_title(r'$m + \mathbf{A} \gamma$')
ax[3].plot(t, trend, lw=2);
ax[3].set_title('Step linear function as trend');

lines = [np.where(t > s_)[0][0] for s_ in s]
for ax_ in ax[1:]:
ax_.vlines(t[lines], *ax_.get_ylim(), linestyles='--');

plt.savefig("img/chp06/fig6_step_linear_function.png");

### Code 6.8 and Figure 6.7#

def gen_fourier_basis(t, p=365.25, n=3):
x = 2 * np.pi * (np.arange(n) + 1) * t[:, None] / p
return np.concatenate((np.cos(x), np.sin(x)), axis=1)

n_tp = 500
p = 12
t_monthly = np.asarray([i % p for i in range(n_tp)])
monthly_X = gen_fourier_basis(t_monthly, p=p, n=3)

fig, ax = plt.subplots(figsize=(10, 3))
ax.plot(monthly_X[:p*2, 0])
ax.plot(monthly_X[:p*2, 1:], alpha=.25)

plt.savefig("img/chp06/fig7_fourier_basis.png")

### Code 6.9#

# Generate trend design matrix
n_changepoints = 12

n_tp = seasonality_all.shape[0]
t = np.linspace(0, 1, n_tp, dtype=np.float32)
s = np.linspace(0, max(t), n_changepoints + 2, dtype=np.float32)[1: -1]
A = (t[:, None] > s).astype(np.float32)

# Set n=6 here so that there are 12 columns, which is the same as seasonality_all
X_pred = gen_fourier_basis(np.where(seasonality_all)[1],
p=seasonality_all.shape[-1],
n=6)
n_pred = X_pred.shape[-1]

def gen_gam_jd(training=True):
@tfd.JointDistributionCoroutine
def gam():
beta = yield root(tfd.Sample(
tfd.Normal(0., 1.),
sample_shape=n_pred,
name='beta'))
seasonality = tf.einsum('ij,...j->...i', X_pred, beta)

k = yield root(tfd.HalfNormal(10., name='k'))
m = yield root(tfd.Normal(
co2_by_month_training_data['CO2'].mean(), scale=5., name='m'))
tau = yield root(tfd.HalfNormal(10., name='tau'))
delta = yield tfd.Sample(tfd.Laplace(0., tau),
sample_shape=n_changepoints,
name='delta')

growth_rate = k[..., None] + tf.einsum('ij,...j->...i', A, delta)
offset = m[..., None] + tf.einsum('ij,...j->...i', A, -s * delta)
trend = growth_rate * t + offset

y_hat = seasonality + trend
if training:
y_hat = y_hat[..., :co2_by_month_training_data.shape[0]]

noise_sigma = yield root(tfd.HalfNormal(scale=5., name='noise_sigma'))
observed = yield tfd.Independent(
tfd.Normal(y_hat, noise_sigma[..., None]),
reinterpreted_batch_ndims=1,
name='observed'
)

return gam

gam = gen_gam_jd()

## Autoregressive Models#

### Figure 6.8#

prior_samples = gam.sample(100)
prior_predictive_timeseries = prior_samples[-1]

fig, ax = plt.subplots(figsize=(10, 5))
ax.plot(co2_by_month.index[:-num_forecast_steps],
tf.transpose(prior_predictive_timeseries), alpha=.5)
ax.set_xlabel("Year")
fig.autofmt_xdate()

plt.savefig("img/chp06/fig8_prior_predictive2.png");
/var/folders/7p/srk5qjp563l5f9mrjtp44bh800jqsw/T/ipykernel_14469/970221687.py:8: UserWarning: This figure was using constrained_layout, but that is incompatible with subplots_adjust and/or tight_layout; disabling constrained_layout.
fig.autofmt_xdate()
%%time
mcmc_samples, sampler_stats = run_mcmc(
seed=tf.constant([-12341, 62345], dtype=tf.int32),
observed=co2_by_month_training_data.T)
2021-11-19 17:32:18.333839: W tensorflow/compiler/tf2xla/kernels/assert_op.cc:38] Ignoring Assert operator mcmc_retry_init/assert_equal_1/Assert/AssertGuard/Assert
CPU times: user 1min 4s, sys: 860 ms, total: 1min 5s
Wall time: 1min 5s
gam_idata = az.from_dict(
posterior={
k:np.swapaxes(v.numpy(), 1, 0)
for k, v in mcmc_samples._asdict().items()},
sample_stats={
k:np.swapaxes(sampler_stats[k], 1, 0)
for k in ["target_log_prob", "diverging", "accept_ratio", "n_steps"]}
)

axes = az.plot_trace(gam_idata, compact=True);
gam_full = gen_gam_jd(training=False)
posterior_dists, _ = gam_full.sample_distributions(value=mcmc_samples)
# plot components
_, ax = plt.subplots(2, 1, figsize=(10, 5))

k, m, tau, delta = mcmc_samples[1:5]
growth_rate = k[..., None] + tf.einsum('ij,...j->...i', A, delta)
offset = m[..., None] + tf.einsum('ij,...j->...i', A, -s * delta)
trend_posterior = growth_rate * t + offset
seasonality_posterior = tf.einsum('ij,...j->...i', X_pred, mcmc_samples[0])

for i in range(nchains):
ax[0].plot(co2_by_month.index,
trend_posterior[-100:, i, :].numpy().T, alpha=.05)
ax[1].plot(co2_by_month.index,
seasonality_posterior[-100:, i, :].numpy().T, alpha=.05)

### Figure 6.9#

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

# fitted_with_forecast = posterior_dists[-1].distribution.mean().numpy()
fitted_with_forecast = posterior_dists[-1].distribution.sample().numpy()

ax.plot(co2_by_month.index, fitted_with_forecast[-100:, 0, :].T, color='gray', alpha=.1);
ax.plot(co2_by_month.index, fitted_with_forecast[-100:, 1, :].T, color='gray', alpha=.1);

plot_co2_data((fig, ax));

average_forecast = np.mean(fitted_with_forecast, axis=(0, 1)).T
ax.plot(co2_by_month.index, average_forecast, ls='--', label='GAM forecast', alpha=.5);

plt.savefig("img/chp06/fig9_posterior_predictive2.png");
/var/folders/7p/srk5qjp563l5f9mrjtp44bh800jqsw/T/ipykernel_14469/1698526664.py:19: UserWarning: This figure was using constrained_layout, but that is incompatible with subplots_adjust and/or tight_layout; disabling constrained_layout.
fig.autofmt_xdate()

### Code 6.10 and Figure 6.10#

n_t = 200

@tfd.JointDistributionCoroutine
def ar1_with_forloop():
sigma = yield root(tfd.HalfNormal(1.))
rho = yield root(tfd.Uniform(-1., 1.))
x0 = yield tfd.Normal(0., sigma)
x = [x0]
for i in range(1, n_t):
x_i = yield tfd.Normal(x[i-1] * rho, sigma)
x.append(x_i)

nplot = 4
fig, axes = plt.subplots(nplot, 1)
for ax, rho in zip(axes, np.linspace(-1.01, 1.01, nplot)):
test_samples = ar1_with_forloop.sample(value=(1., rho))
ar1_samples = tf.stack(test_samples[2:])
ax.plot(ar1_samples, alpha=.5, label=r'$\rho$=%.2f' % rho)
ax.legend(bbox_to_anchor=(1, 1), loc='upper left',
ax.get_xaxis().set_visible(False)
fig.suptitle(r'AR(1) process with varies autoregressive coefficient ($\rho$)')

plt.savefig("img/chp06/fig10_ar1_process.png")

### Code 6.11#

@tfd.JointDistributionCoroutine
def ar1_without_forloop():
sigma = yield root(tfd.HalfNormal(1.))
rho = yield root(tfd.Uniform(-1., 1.))

def ar1_fun(x):
# We apply the backshift operation here
x_tm1 = tf.concat([tf.zeros_like(x[..., :1]), x[..., :-1]], axis=-1)
loc = x_tm1 * rho[..., None]
return tfd.Independent(tfd.Normal(loc=loc, scale=sigma[..., None]),
reinterpreted_batch_ndims=1)

dist = yield tfd.Autoregressive(
distribution_fn=ar1_fun,
sample0=tf.zeros([n_t], dtype=rho.dtype),
num_steps=n_t)

seed = [1000, 5234]
_, ax = plt.subplots(figsize=(10, 5))
rho = np.linspace(-1.01, 1.01, 5)
sigma = np.ones(5)
test_samples = ar1_without_forloop.sample(value=(sigma, rho), seed=seed)
ar1_samples = tf.transpose(test_samples[-1])
ax.plot(ar1_samples, alpha=.5)
ax.set_title('AR(1) process with varies autoregressive coefficient (rho)');

A more general way to implement this is to use a Lag operator $$B$$

B = np.diag(np.ones(5 - 1), -1)
B @ B
array([[0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0.],
[1., 0., 0., 0., 0.],
[0., 1., 0., 0., 0.],
[0., 0., 1., 0., 0.]])
B = np.diag(np.ones(n_t - 1), -1)

@tfd.JointDistributionCoroutine
def ar1_lag_operator():
sigma = yield root(tfd.HalfNormal(1., name='sigma'))
rho = yield root(tfd.Uniform(-1., 1., name='rho'))

def ar1_fun(x):
loc = tf.einsum('ij,...j->...i', B, x) * rho[..., None]
return tfd.Independent(tfd.Normal(loc=loc, scale=sigma[..., None]),
reinterpreted_batch_ndims=1)

dist = yield tfd.Autoregressive(
distribution_fn=ar1_fun,
sample0=tf.zeros([n_t], dtype=rho.dtype),
num_steps=n_t,
name="ar1")

_, ax = plt.subplots(figsize=(10, 5))
rho = np.linspace(-1.01, 1.01, 5)
sigma = np.ones(5)
test_samples = ar1_lag_operator.sample(value=(sigma, rho), seed=seed)
ar1_samples = tf.transpose(test_samples[-1])
ax.plot(ar1_samples, alpha=.5)
ax.set_title('AR(1) process with varies autoregressive coefficient (rho)');

Note that since we are using a stateless RNG seeding, we got the same result (yay!)

### Code 6.12#

def gam_trend_seasonality():
beta = yield root(tfd.Sample(
tfd.Normal(0., 1.),
sample_shape=n_pred,
name='beta'))
seasonality = tf.einsum('ij,...j->...i', X_pred, beta)

k = yield root(tfd.HalfNormal(10., name='k'))
m = yield root(tfd.Normal(
co2_by_month_training_data['CO2'].mean(), scale=5., name='m'))
tau = yield root(tfd.HalfNormal(10., name='tau'))
delta = yield tfd.Sample(tfd.Laplace(0., tau),
sample_shape=n_changepoints,
name='delta')

growth_rate = k[..., None] + tf.einsum('ij,...j->...i', A, delta)
offset = m[..., None] + tf.einsum('ij,...j->...i', A, -s * delta)
trend = growth_rate * t + offset
noise_sigma = yield root(tfd.HalfNormal(scale=5., name='noise_sigma'))
return seasonality, trend, noise_sigma

def generate_gam(training=True):

@tfd.JointDistributionCoroutine
def gam():
seasonality, trend, noise_sigma = yield from gam_trend_seasonality()
y_hat = seasonality + trend
if training:
y_hat = y_hat[..., :co2_by_month_training_data.shape[0]]

# likelihood
observed = yield tfd.Independent(
tfd.Normal(y_hat, noise_sigma[..., None]),
reinterpreted_batch_ndims=1,
name='observed'
)

return gam

gam = generate_gam()
def generate_gam_ar_likelihood(training=True):

@tfd.JointDistributionCoroutine
def gam_with_ar_likelihood():
seasonality, trend, noise_sigma = yield from gam_trend_seasonality()
y_hat = seasonality + trend
if training:
y_hat = y_hat[..., :co2_by_month_training_data.shape[0]]

# Likelihood
rho = yield root(tfd.Uniform(-1., 1., name='rho'))

def ar_fun(y):
loc = tf.concat([tf.zeros_like(y[..., :1]), y[..., :-1]],
axis=-1) * rho[..., None] + y_hat
return tfd.Independent(
tfd.Normal(loc=loc, scale=noise_sigma[..., None]),
reinterpreted_batch_ndims=1)

observed = yield tfd.Autoregressive(
distribution_fn=ar_fun,
sample0=tf.zeros_like(y_hat),
num_steps=1,
name='observed'
)

return gam_with_ar_likelihood

gam_with_ar_likelihood = generate_gam_ar_likelihood()

plt.plot(tf.transpose(gam_with_ar_likelihood.sample(50)[-1]));
%%time
mcmc_samples, sampler_stats = run_mcmc(
seed=tf.constant([-234272345, 73564234], dtype=tf.int32),
observed=co2_by_month_training_data.T)
2021-11-19 17:34:00.601657: W tensorflow/compiler/tf2xla/kernels/assert_op.cc:38] Ignoring Assert operator mcmc_retry_init/assert_equal_1/Assert/AssertGuard/Assert
CPU times: user 1min 19s, sys: 916 ms, total: 1min 20s
Wall time: 1min 21s

### Code 6.13#

gam_ar_likelihood_idata = az.from_dict(
posterior={
k:np.swapaxes(v.numpy(), 1, 0)
for k, v in mcmc_samples._asdict().items()},
sample_stats={
k:np.swapaxes(sampler_stats[k], 1, 0)
for k in ["target_log_prob", "diverging", "accept_ratio", "n_steps"]}
)

axes = az.plot_trace(gam_ar_likelihood_idata, compact=True);
# plot components
_, ax = plt.subplots(2, 1, figsize=(10, 5))

k, m, tau, delta = mcmc_samples[1:5]
growth_rate = k[..., None] + tf.einsum('ij,...j->...i', A, delta)
offset = m[..., None] + tf.einsum('ij,...j->...i', A, -s * delta)
trend_posterior = growth_rate * t + offset
seasonality_posterior = tf.einsum('ij,...j->...i', X_pred, mcmc_samples[0])

for i in range(nchains):
ax[0].plot(co2_by_month.index,
trend_posterior[-100:, i, :].numpy().T, alpha=.05)
ax[1].plot(co2_by_month.index,
seasonality_posterior[-100:, i, :].numpy().T, alpha=.05)
gam_with_ar_likelihood_full = generate_gam_ar_likelihood(training=False)
_, values = gam_with_ar_likelihood_full.sample_distributions(value=mcmc_samples)
fitted_with_forecast = values[-1].numpy()

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

ax.plot(co2_by_month.index, fitted_with_forecast[-100:, 0, :].T, color='gray', alpha=.1);
ax.plot(co2_by_month.index, fitted_with_forecast[-100:, 1, :].T, color='gray', alpha=.1);

plot_co2_data((fig, ax));

average_forecast = np.mean(fitted_with_forecast, axis=(0, 1)).T
ax.plot(co2_by_month.index, average_forecast, ls='--', label='GAM forecast', alpha=.5);
/var/folders/7p/srk5qjp563l5f9mrjtp44bh800jqsw/T/ipykernel_14469/1698526664.py:19: UserWarning: This figure was using constrained_layout, but that is incompatible with subplots_adjust and/or tight_layout; disabling constrained_layout.
fig.autofmt_xdate()

### Figure 6.11#

fig, axes = plt.subplots(1, 3, figsize=(4*3, 4))
az.plot_posterior(gam_idata, var_names=['noise_sigma'], alpha=.5, lw=2.5, ax=axes[0]);
axes[0].set_title(r'$\sigma_{noise}$ (Normal)')

az.plot_posterior(gam_ar_likelihood_idata, var_names=['noise_sigma', 'rho'], alpha=.5, lw=2.5, ax=axes[1:]);
axes[1].set_title(r'$\sigma_{noise}$ (AR(1))')
axes[2].set_title(r'$\rho$')

plt.savefig("img/chp06/fig11_ar1_likelihood_rho.png");

### Code 6.14#

def generate_gam_ar_latent(training=True):

@tfd.JointDistributionCoroutine
def gam_with_latent_ar():
seasonality, trend, noise_sigma = yield from gam_trend_seasonality()

# Latent AR(1)
ar_sigma = yield root(tfd.HalfNormal(.1, name='ar_sigma'))
rho = yield root(tfd.Uniform(-1., 1., name='rho'))
def ar_fun(y):
loc = tf.concat([tf.zeros_like(y[..., :1]), y[..., :-1]],
axis=-1) * rho[..., None]
return tfd.Independent(
tfd.Normal(loc=loc, scale=ar_sigma[..., None]),
reinterpreted_batch_ndims=1)
temporal_error = yield tfd.Autoregressive(
distribution_fn=ar_fun,
sample0=tf.zeros_like(trend),
num_steps=trend.shape[-1],
name='temporal_error')

# Linear prediction
y_hat = seasonality + trend + temporal_error
if training:
y_hat = y_hat[..., :co2_by_month_training_data.shape[0]]

# Likelihood
observed = yield tfd.Independent(
tfd.Normal(y_hat, noise_sigma[..., None]),
reinterpreted_batch_ndims=1,
name='observed'
)

return gam_with_latent_ar

gam_with_latent_ar = generate_gam_ar_latent()

plt.plot(tf.transpose(gam_with_latent_ar.sample(50)[-1]));
%%time
mcmc_samples, sampler_stats = run_mcmc(
seed=tf.constant([36245, 734565], dtype=tf.int32),
observed=co2_by_month_training_data.T)
2021-11-19 17:50:22.041962: W tensorflow/compiler/tf2xla/kernels/assert_op.cc:38] Ignoring Assert operator mcmc_retry_init/assert_equal_1/Assert/AssertGuard/Assert
CPU times: user 3min 9s, sys: 2.9 s, total: 3min 12s
Wall time: 3min 13s
nuts_trace_ar_latent = az.from_dict(
posterior={
k:np.swapaxes(v.numpy(), 1, 0)
for k, v in mcmc_samples._asdict().items()},
sample_stats={
k:np.swapaxes(sampler_stats[k], 1, 0)
for k in ["target_log_prob", "diverging", "accept_ratio", "n_steps"]}
)

axes = az.plot_trace(
nuts_trace_ar_latent,
var_names=['beta', 'tau', 'ar_sigma', 'rho', 'noise_sigma'],
compact=True);
gam_with_latent_ar_full = generate_gam_ar_latent(training=False)
posterior_dists, ppc_samples = gam_with_latent_ar_full.sample_distributions(value=mcmc_samples)

### Figure 6.12#

# plot components
fig, ax = plt.subplots(3, 1, figsize=(10, 7.5), sharex=True)

beta, k, m, tau, delta = mcmc_samples[:5]
growth_rate = k[..., None] + tf.einsum('ij,...j->...i', A, delta)
offset = m[..., None] + tf.einsum('ij,...j->...i', A, -s * delta)
trend_posterior = growth_rate * t + offset
seasonality_posterior = tf.einsum('ij,...j->...i', X_pred, beta)
temporal_error = mcmc_samples[-1]
# temporal_error_ = mcmc_samples[7]
# temporal_error = tf.concat([tf.zeros_like(temporal_error_[..., :1]),
#                             temporal_error_], axis=-1)

for i in range(nchains):
ax[0].plot(co2_by_month.index, trend_posterior[-100:, i, :].numpy().T, alpha=.05);
ax[1].plot(co2_by_month.index, seasonality_posterior[-100:, i, :].numpy().T, alpha=.05);
ax[2].plot(co2_by_month.index, temporal_error[-100:, i, :].numpy().T, alpha=.05);

ax[0].set_title('Trend (Step Linear)')
ax[1].set_title('Seasonality (Month of the year effect)')
ax[2].set_title('AR(1)')
ax[2].set_xlabel("Year")
fig.autofmt_xdate()

plt.savefig("img/chp06/fig12_posterior_predictive_ar1.png");
/var/folders/7p/srk5qjp563l5f9mrjtp44bh800jqsw/T/ipykernel_14469/4037595605.py:23: UserWarning: This figure was using constrained_layout, but that is incompatible with subplots_adjust and/or tight_layout; disabling constrained_layout.
fig.autofmt_xdate()
fig, ax = plt.subplots(1, 1, figsize=(10, 5))

# fitted_with_forecast = posterior_dists[-1].distribution.mean().numpy()
# fitted_with_forecast = posterior_dists[-1].distribution.sample().numpy()
fitted_with_forecast = ppc_samples[-1].numpy()

ax.plot(co2_by_month.index, fitted_with_forecast[-100:, 0, :].T, color='gray', alpha=.1);
ax.plot(co2_by_month.index, fitted_with_forecast[-100:, 1, :].T, color='gray', alpha=.1);

plot_co2_data((fig, ax));

average_forecast = np.mean(fitted_with_forecast, axis=(0, 1)).T
ax.plot(co2_by_month.index, average_forecast, ls='--', label='GAM forecast', alpha=.5);
/var/folders/7p/srk5qjp563l5f9mrjtp44bh800jqsw/T/ipykernel_14469/1698526664.py:19: UserWarning: This figure was using constrained_layout, but that is incompatible with subplots_adjust and/or tight_layout; disabling constrained_layout.
fig.autofmt_xdate()

### Figure 6.13#

axes = az.plot_posterior(
nuts_trace_ar_latent,
var_names=['noise_sigma', 'ar_sigma', 'rho'],
alpha=.5, lw=2.5,
figsize=(4*3, 4));
axes[0].set_title(r'$\sigma_{noise}$')
axes[1].set_title(r'$\sigma_{AR}$')
axes[2].set_title(r'$\rho$')

plt.savefig("img/chp06/fig13_ar1_likelihood_rho2.png");

#### Reflection on autoregressive and smoothing#

num_steps = 100

x = np.linspace(0, 50, num_steps)
f = np.exp(1.0 + np.power(x, 0.5) - np.exp(x/15.0))
y = f + np.random.normal(scale=1.0, size=x.shape)

plt.plot(x, y, 'ok', label='Observed')
plt.plot(x, f, 'r', label='f(x)')
plt.legend()
plt.xlabel('x');
plt.ylabel('y');

### Code 6.15#

@tfd.JointDistributionCoroutine
def smoothing_grw():
alpha = yield root(tfd.Beta(5, 1.))
variance = yield root(tfd.HalfNormal(10.))
sigma0 = tf.sqrt(variance * alpha)
sigma1 = tf.sqrt(variance * (1. - alpha))
z = yield tfd.Sample(tfd.Normal(0., sigma0), num_steps)
observed = yield tfd.Independent(
tfd.Normal(tf.math.cumsum(z, axis=-1), sigma1[..., None]),
name='observed'
)
%%time
mcmc_samples, sampler_stats = run_mcmc(
observed=tf.constant(y[None, ...], dtype=tf.float32))
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:From /opt/miniconda3/envs/bmcp/lib/python3.9/site-packages/tensorflow_probability/python/distributions/distribution.py:342: calling Independent.__init__ (from tensorflow_probability.python.distributions.independent) with reinterpreted_batch_ndims=None is deprecated and will be removed after 2022-03-01.
Instructions for updating:
Please pass an integer value for reinterpreted_batch_ndims. The current behavior corresponds to reinterpreted_batch_ndims=tf.size(distribution.batch_shape_tensor()) - 1.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
2021-11-19 17:56:39.831415: W tensorflow/core/framework/op_kernel.cc:1745] OP_REQUIRES failed at functional_ops.cc:373 : INTERNAL: No function library
2021-11-19 17:56:39.831472: W tensorflow/core/framework/op_kernel.cc:1745] OP_REQUIRES failed at functional_ops.cc:373 : INTERNAL: No function library
2021-11-19 17:56:39.852216: W tensorflow/core/framework/op_kernel.cc:1745] OP_REQUIRES failed at functional_ops.cc:373 : INTERNAL: No function library
2021-11-19 17:56:39.852258: W tensorflow/core/framework/op_kernel.cc:1745] OP_REQUIRES failed at functional_ops.cc:373 : INTERNAL: No function library
2021-11-19 17:56:39.874546: W tensorflow/core/framework/op_kernel.cc:1745] OP_REQUIRES failed at functional_ops.cc:373 : INTERNAL: No function library
2021-11-19 17:56:39.874599: W tensorflow/core/framework/op_kernel.cc:1745] OP_REQUIRES failed at functional_ops.cc:373 : INTERNAL: No function library
2021-11-19 17:56:39.885913: W tensorflow/core/framework/op_kernel.cc:1745] OP_REQUIRES failed at functional_ops.cc:373 : INTERNAL: No function library
2021-11-19 17:56:39.885967: W tensorflow/core/framework/op_kernel.cc:1745] OP_REQUIRES failed at functional_ops.cc:373 : INTERNAL: No function library
2021-11-19 17:56:39.920939: W tensorflow/core/framework/op_kernel.cc:1745] OP_REQUIRES failed at functional_ops.cc:373 : INTERNAL: No function library
2021-11-19 17:56:39.920993: W tensorflow/core/framework/op_kernel.cc:1745] OP_REQUIRES failed at functional_ops.cc:373 : INTERNAL: No function library
2021-11-19 17:56:40.343209: W tensorflow/core/framework/op_kernel.cc:1745] OP_REQUIRES failed at functional_ops.cc:373 : INTERNAL: No function library
2021-11-19 17:56:40.343251: W tensorflow/core/framework/op_kernel.cc:1745] OP_REQUIRES failed at functional_ops.cc:373 : INTERNAL: No function library
2021-11-19 17:56:40.350681: W tensorflow/core/framework/op_kernel.cc:1745] OP_REQUIRES failed at functional_ops.cc:373 : INTERNAL: No function library
2021-11-19 17:56:40.350794: W tensorflow/core/framework/op_kernel.cc:1745] OP_REQUIRES failed at functional_ops.cc:373 : INTERNAL: No function library
2021-11-19 17:56:40.357544: W tensorflow/core/framework/op_kernel.cc:1745] OP_REQUIRES failed at functional_ops.cc:373 : INTERNAL: No function library
2021-11-19 17:56:40.357614: W tensorflow/core/framework/op_kernel.cc:1745] OP_REQUIRES failed at functional_ops.cc:373 : INTERNAL: No function library
2021-11-19 17:56:40.366483: W tensorflow/core/framework/op_kernel.cc:1745] OP_REQUIRES failed at functional_ops.cc:373 : INTERNAL: No function library
2021-11-19 17:56:40.366524: W tensorflow/core/framework/op_kernel.cc:1745] OP_REQUIRES failed at functional_ops.cc:373 : INTERNAL: No function library
2021-11-19 17:56:40.611639: W tensorflow/compiler/tf2xla/kernels/assert_op.cc:38] Ignoring Assert operator mcmc_retry_init/assert_equal_1/Assert/AssertGuard/Assert
WARNING:tensorflow:5 out of the last 9 calls to <function windowed_adaptive_nuts at 0x7fae934af550> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has experimental_relax_shapes=True option that relaxes argument shapes that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.
CPU times: user 36.1 s, sys: 342 ms, total: 36.4 s
Wall time: 36.5 s
_, ax = plt.subplots(2, 1, figsize=(10, 5))
ax[0].plot(mcmc_samples[0], alpha=.5)
ax[1].plot(mcmc_samples[1], alpha=.5);

### Figure 6.14#

nsample, nchain = mcmc_samples[-1].shape[:2]

z = tf.reshape(tf.math.cumsum(mcmc_samples[-1], axis=-1), [nsample*nchain, -1])
lower, upper = np.percentile(z, [5, 95], axis=0)

_, ax = plt.subplots(figsize=(10, 4))

ax.plot(x, y, 'o', label='Observed')
ax.plot(x, f, label='f(x)')
ax.fill_between(x, lower, upper, color='C1', alpha=.25)
ax.plot(x, tf.reduce_mean(z, axis=0), color='C4', ls='--', label='z')
ax.legend()
ax.set_xlabel('x');
ax.set_ylabel('y');

plt.savefig("img/chp06/fig14_smoothing_with_gw.png");

### Code 6.16 and Figure 6.15#

us_monthly_birth["date_month"] = pd.to_datetime(us_monthly_birth["date_month"])
us_monthly_birth.set_index("date_month", drop=True, inplace=True)

def plot_birth_data(fig_ax=None):
if not fig_ax:
fig, ax = plt.subplots(1, 1, figsize=(10, 5))
else:
fig, ax = fig_ax

ax.plot(us_monthly_birth, lw=2)
ax.set_ylabel("Birth (thousands)")
ax.set_xlabel("Year")
fig.suptitle("Monthly live birth U.S.A",
fontsize=15)
ax.text(0.99, .02,
"Source: Stoffer D (2019). “astsa: Applied Statistical Time Series Analysis.”",
transform=ax.transAxes,
horizontalalignment="right",
alpha=0.5)
fig.autofmt_xdate()
return fig, ax

_ = plot_birth_data()

plt.savefig("img/chp06/fig15_birth_by_month.png")
/var/folders/7p/srk5qjp563l5f9mrjtp44bh800jqsw/T/ipykernel_14469/1306182099.py:21: UserWarning: This figure was using constrained_layout, but that is incompatible with subplots_adjust and/or tight_layout; disabling constrained_layout.
fig.autofmt_xdate()

### Code 6.17#

# y ~ Sarima(1,1,1)(1,1,1)[12]
p, d, q = (1, 1, 1)
P, D, Q, period = (1, 1, 1, 12)
# Time series data: us_monthly_birth.shape = (372,)
observed = us_monthly_birth["birth_in_thousands"].values
# Integrated to seasonal order $D$
for _ in range(D):
observed = observed[period:] - observed[:-period]
# Integrated to order $d$
observed = tf.constant(np.diff(observed, n=d), tf.float32)

r = max(p, q, P * period, Q * period)
def likelihood(mu0, sigma, phi, theta, sphi, stheta):
batch_shape = tf.shape(mu0)
y_extended = tf.concat(
[tf.zeros(tf.concat([[r], batch_shape], axis=0), dtype=mu0.dtype),
tf.einsum('...,j->j...',
tf.ones_like(mu0, dtype=observed.dtype),
observed)],
axis=0)
eps_t = tf.zeros_like(y_extended, dtype=observed.dtype)

def arma_onestep(t, eps_t):
t_shift = t + r
# AR
y_past = tf.gather(y_extended, t_shift - (np.arange(p) + 1))
ar = tf.einsum("...p,p...->...", phi, y_past)
# MA
eps_past = tf.gather(eps_t, t_shift - (np.arange(q) + 1))
ma = tf.einsum("...q,q...->...", theta, eps_past)
# Seasonal AR
sy_past = tf.gather(y_extended, t_shift - (np.arange(P) + 1) * period)
sar = tf.einsum("...p,p...->...", sphi, sy_past)
# Seasonal MA
seps_past = tf.gather(eps_t, t_shift - (np.arange(Q) + 1) * period)
sma = tf.einsum("...q,q...->...", stheta, seps_past)

mu_at_t = ar + ma + sar + sma + mu0
eps_update = tf.gather(y_extended, t_shift) - mu_at_t
epsilon_t_next = tf.tensor_scatter_nd_update(
eps_t, [[t_shift]], eps_update[None, ...])
return t+1, epsilon_t_next

t, eps_output_ = tf.while_loop(
lambda t, *_: t < observed.shape[-1],
arma_onestep,
loop_vars=(0, eps_t),
maximum_iterations=observed.shape[-1])
eps_output = eps_output_[r:]
return tf.reduce_sum(
tfd.Normal(0, sigma[None, ...]).log_prob(eps_output), axis=0)

### Code 6.18#

@tfd.JointDistributionCoroutine
def sarima_priors():
mu0 = yield root(tfd.StudentT(df=6, loc=0, scale=2.5, name='mu0'))
sigma = yield root(tfd.HalfStudentT(df=7, loc=0, scale=1., name='sigma'))

# phi = yield root(tfd.Sample(tfd.Normal(loc=0, scale=0.5), p, name='phi'))
phi = yield root(
tfd.Sample(
tfd.TransformedDistribution(
tfd.Beta(concentration1=2., concentration0=2.),
tfb.Shift(-1.)(tfb.Scale(2.))),
p, name='phi')
)
theta = yield root(tfd.Sample(tfd.Normal(loc=0, scale=0.5), q, name='theta'))
sphi = yield root(tfd.Sample(tfd.Normal(loc=0, scale=0.5), P, name='sphi'))
stheta = yield root(tfd.Sample(tfd.Normal(loc=0, scale=0.5), Q, name='stheta'))
# a NUTS sampling routine with simple tuning
from tensorflow_probability.python.internal import unnest
from tensorflow_probability.python.internal import samplers

def run_mcmc_simple(
n_draws,
joint_dist,
n_chains=4,
return_compiled_function=False,
target_log_prob_fn=None,
bijector=None,
init_state=None,
seed=None,
**pins):
joint_dist_pinned = joint_dist.experimental_pin(**pins) if pins else joint_dist
if bijector is None:
bijector = joint_dist_pinned.experimental_default_event_space_bijector()
if target_log_prob_fn is None:
target_log_prob_fn = joint_dist_pinned.unnormalized_log_prob

if seed is None:
seed = 26401
run_mcmc_seed = samplers.sanitize_seed(seed, salt='run_mcmc_seed')

if init_state is None:
if pins:
init_state_ = joint_dist_pinned.sample_unpinned(n_chains)
else:
init_state_ = joint_dist_pinned.sample(n_chains)
ini_state_unbound = bijector.inverse(init_state_)
run_mcmc_seed, *init_seed = samplers.split_seed(
run_mcmc_seed, n=len(ini_state_unbound)+1)
init_state = bijector.forward(
tf.nest.map_structure(
lambda x, seed: tfd.Uniform(-1., tf.constant(1., x.dtype)).sample(
x.shape, seed=seed),
ini_state_unbound,
tf.nest.pack_sequence_as(ini_state_unbound, init_seed)))

@tf.function(autograph=False, jit_compile=True)
def run_inference_nuts(init_state, draws, tune, seed):
seed, tuning_seed, sample_seed = samplers.split_seed(seed, n=3)

def gen_kernel(step_size):
hmc = tfp.mcmc.NoUTurnSampler(
target_log_prob_fn=target_log_prob_fn, step_size=step_size)
hmc = tfp.mcmc.TransformedTransitionKernel(
hmc, bijector=bijector)
hmc, tune // 2, target_accept_prob=0.85)
return tuning_hmc

def tuning_trace_fn(_, pkr):
return pkr.inner_results.transformed_state, pkr.new_step_size

def get_tuned_stepsize(samples, step_size):
return tf.math.reduce_std(samples, axis=0) * step_size[-1]

step_size = tf.nest.map_structure(
tf.ones_like, bijector.inverse(init_state))
tuning_hmc = gen_kernel(step_size)
init_samples, (sample_unbounded, tuning_step_size) = tfp.mcmc.sample_chain(
num_results=200,
num_burnin_steps=tune // 2 - 200,
current_state=init_state,
kernel=tuning_hmc,
trace_fn=tuning_trace_fn,
seed=tuning_seed)

tuning_step_size = tf.nest.pack_sequence_as(
sample_unbounded, tuning_step_size)
step_size_new = tf.nest.map_structure(get_tuned_stepsize,
sample_unbounded,
tuning_step_size)
sample_hmc = gen_kernel(step_size_new)

def sample_trace_fn(_, pkr):
energy_diff = unnest.get_innermost(pkr, 'log_accept_ratio')
return {
'target_log_prob': unnest.get_innermost(pkr, 'target_log_prob'),
'n_steps': unnest.get_innermost(pkr, 'leapfrogs_taken'),
'diverging': unnest.get_innermost(pkr, 'has_divergence'),
'energy': unnest.get_innermost(pkr, 'energy'),
'accept_ratio': tf.minimum(1., tf.exp(energy_diff)),
'reach_max_depth': unnest.get_innermost(pkr, 'reach_max_depth'),
}
current_state = tf.nest.map_structure(lambda x: x[-1], init_samples)
return tfp.mcmc.sample_chain(
num_results=draws,
num_burnin_steps=tune // 2,
current_state=current_state,
kernel=sample_hmc,
trace_fn=sample_trace_fn,
seed=sample_seed)

mcmc_samples, mcmc_diagnostic = run_inference_nuts(

if return_compiled_function:
return mcmc_samples, mcmc_diagnostic, run_inference_nuts
else:
return mcmc_samples, mcmc_diagnostic
%%time
target_log_prob_fn = lambda *x: sarima_priors.log_prob(*x) + likelihood(*x)

mcmc_samples, sampler_stats = run_mcmc_simple(
target_log_prob_fn=target_log_prob_fn,
seed=tf.constant([623453, 456345], dtype=tf.int32),
)
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
CPU times: user 28.2 s, sys: 513 ms, total: 28.7 s
Wall time: 28.8 s
nuts_trace_arima = az.from_dict(
posterior={
k:np.swapaxes(v.numpy(), 1, 0)
for k, v in mcmc_samples._asdict().items()},
sample_stats={
k:np.swapaxes(sampler_stats[k], 1, 0)
for k in ["target_log_prob", "diverging", "accept_ratio", "n_steps"]}
)

axes = az.plot_trace(nuts_trace_arima);

#### SARIMAX Class#

import warnings
from statsmodels.tsa.statespace.tools import diff as tsa_diff
from tensorflow_probability.python.internal import distribution_util
from tensorflow_probability.python.internal import prefer_static as ps

class SARIMAX:
# Similar to https://www.statsmodels.org/stable/generated/statsmodels.tsa.statespace.sarimax.SARIMAX.html
# with simple differencing
def __init__(self,
observed,
design_matrix=None,
order=(1, 0, 0),
seasonal_order=(0, 0, 0, 0),
dtype=np.float32
):
"""
SARIMAX Likelihood for univariate time series

order - (p,d,q)
seasonal_order - (P,D,Q,s)
"""
# Static parameters
self.p, self.d, self.q = order
self.P, self.D, self.Q, self.period = seasonal_order

# Dynamic regression
if design_matrix is not None:
assert ps.rank(design_matrix) >= 2
assert ps.shape(design_matrix)[-2] == observed.shape[-1]

if self.period > 0:
warnings.warn("""
Can not model seasonal difference with Dynamic regressions,
Setting D to 0 ...
""")
self.period = 0

self.design_matrix = tf.convert_to_tensor(
np.diff(design_matrix, n=self.d, axis=0), dtype=dtype)
else:
self.design_matrix = None

if self.period <= 1:
self.P, self.D, self.Q = 0, 0, 0

self.r = max(self.p, self.q, self.P *
self.period, self.Q * self.period)

# Simple differencing
observed_diff = tsa_diff(
observed,
k_diff=self.d,
k_seasonal_diff=self.D,
seasonal_periods=self.period)
self.observed = tf.convert_to_tensor(observed_diff, dtype=dtype)

self.dtype = dtype

def _log_prob(self, *args):
mu0 = args[0]  # Location parameter
sigma = args[1]  # Scale parameter
i = 2
if self.design_matrix is not None:  # regression coefficients
reg_coeff = args[i]
mu_t = mu0[None, ...] + \
tf.einsum('...i,ji->j...', reg_coeff, self.design_matrix)
i += 1
else:
mu_t = tf.einsum('...,j->j...', mu0, ps.ones_like(self.observed))

# Parse args and prepare states
if self.p > 0:  # AR coefficients
phi = args[i]
i += 1
if self.q > 0:  # MA coefficients
theta = args[i]
i += 1
if self.P > 0:  # SAR coefficients
sphi = args[i]
i += 1
if self.Q > 0:  # SMA coefficients
stheta = args[i]
i += 1

batch_shape = ps.shape(mu0)
y_extended = ps.concat(
[ps.zeros(tf.concat([[self.r], batch_shape], axis=0), dtype=mu0.dtype),
tf.einsum('...,j->j...',
ps.ones_like(mu0, dtype=self.observed.dtype),
self.observed)],
axis=0
)
eps_t = ps.zeros_like(y_extended, dtype=self.observed.dtype)

def body_fn(t, mu_t, eps_t):
mu_temp = []
t_switch = t + self.r
if self.p > 0:  # AR
y_past = tf.gather(
y_extended, t_switch - (np.arange(self.p) + 1))
ar = tf.einsum("...p,p...->...", phi, y_past)
mu_temp.append(ar)
if self.q > 0:  # MA
eps_past = tf.gather(
eps_t, t_switch - (np.arange(self.q) + 1))
ma = tf.einsum("...q,q...->...", theta, eps_past)
mu_temp.append(ma)
if self.P > 0:  # SAR
y_past = tf.gather(
y_extended, t_switch - (np.arange(self.P) + 1) * self.period)
sar = tf.einsum("...p,p...->...", sphi, y_past)
mu_temp.append(sar)
if self.Q > 0:  # SMA
eps_past = tf.gather(
eps_t, t_switch - (np.arange(self.Q) + 1) * self.period)
sma = tf.einsum("...q,q...->...", stheta, eps_past)
mu_temp.append(sma)
mu_update = sum(mu_temp) + tf.gather(mu_t, t)
mu_t_next = tf.tensor_scatter_nd_update(
mu_t, [[t]], mu_update[None, ...])

eps_update = tf.gather(y_extended, t_switch) - mu_update
epsilon_t_next = tf.tensor_scatter_nd_update(
eps_t, [[t_switch]], eps_update[None, ...])
return t+1, mu_t_next, epsilon_t_next

t, mu_output, eps_output_ = tf.while_loop(
lambda t, *_: t < self.observed.shape[-1],
body_fn,
loop_vars=(0, mu_t, eps_t),
maximum_iterations=self.observed.shape[-1])
eps_output = eps_output_[self.r:]
return tfd.Normal(0, sigma[None, ...]).log_prob(eps_output), mu_output

def log_prob(self, *args):
log_prob_val, _ = self._log_prob(*args)
return ps.reduce_sum(log_prob_val, axis=0)

def log_prob_elementwise(self, *args):
sigma = args[1]  # Scale parameter
_, mu_output = self._log_prob(*args)
mu = distribution_util.move_dimension(mu_output, 0, -1)
return tfd.Normal(mu, sigma[..., None]).log_prob(self.observed)

#### y ~ Sarima(1,1,1)(1,1,1)[12]#

p, q, P, Q = 1, 1, 1, 1

@tfd.JointDistributionCoroutine
def sarima_priors():
mu0 = yield root(tfd.StudentT(df=6, loc=0, scale=2.5, name='mu0'))
sigma = yield root(tfd.HalfStudentT(df=7, loc=0, scale=1., name='sigma'))

phi = yield root(tfd.Sample(tfd.Normal(loc=0, scale=0.5), p, name='phi'))
theta = yield root(tfd.Sample(tfd.Normal(loc=0, scale=0.5), q, name='theta'))
sphi = yield root(tfd.Sample(tfd.Normal(loc=0, scale=0.5), P, name='sphi'))
stheta = yield root(tfd.Sample(tfd.Normal(loc=0, scale=0.5), Q, name='stheta'))

sarima_1 = SARIMAX(
us_monthly_birth["birth_in_thousands"],
order=(1, 1, 1),
seasonal_order=(1, 1, 1, 12))
%%time
target_log_prob_fn = lambda *x: sarima_priors.log_prob(
*x) + sarima_1.log_prob(*x)

mcmc_samples, sampler_stats = run_mcmc_simple(
target_log_prob_fn=target_log_prob_fn,
seed=tf.constant([623453, 456345], dtype=tf.int32),
)
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
CPU times: user 30.5 s, sys: 479 ms, total: 31 s
Wall time: 31.1 s
data_likelihood = np.swapaxes(sarima_1.log_prob_elementwise(*mcmc_samples), 1, 0)
sarima_0_idata = az.from_dict(
posterior={
k:np.swapaxes(v.numpy(), 1, 0)
for k, v in mcmc_samples._asdict().items()},
sample_stats={
k:np.swapaxes(sampler_stats[k], 1, 0)
for k in ["target_log_prob", "diverging", "accept_ratio", "n_steps"]},
observed_data={'observed': sarima_1.observed},
log_likelihood={'observed': data_likelihood},
)

axes = az.plot_trace(sarima_0_idata);

### Code 6.8#

def gen_fourier_basis(t, p=365.25, n=3):
x = 2 * np.pi * (np.arange(n) + 1) * t[:, None] / p
return np.concatenate((np.cos(x), np.sin(x)), axis=1)

p = 12
t_monthly = np.asarray([i % p for i in range(len(us_monthly_birth))]) + 1
monthly_X = gen_fourier_basis(t_monthly, p=p, n=2)

@tfd.JointDistributionCoroutine
def sarima_priors():
mu0 = yield root(tfd.StudentT(df=6, loc=0, scale=2.5, name='mu0'))
sigma = yield root(tfd.HalfStudentT(df=7, loc=0, scale=1., name='sigma'))
beta = yield root(tfd.Sample(
tfd.StudentT(df=6, loc=0, scale=2.5),
monthly_X.shape[-1],
name='beta'))
phi = yield root(tfd.Sample(tfd.Normal(loc=0, scale=0.5), 1, name='phi'))
theta = yield root(tfd.Sample(tfd.Normal(loc=0, scale=0.5), 1, name='theta'))

sarima_2 = SARIMAX(
us_monthly_birth["birth_in_thousands"],
design_matrix=monthly_X,
order=(1, 1, 1))
%%time
target_log_prob_fn = lambda *x: sarima_priors.log_prob(
*x) + sarima_2.log_prob(*x)

mcmc_samples, sampler_stats = run_mcmc_simple(
target_log_prob_fn=target_log_prob_fn,
seed=tf.constant([623453, 456345], dtype=tf.int32),
)
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
CPU times: user 33.9 s, sys: 598 ms, total: 34.5 s
Wall time: 34.6 s
data_likelihood = np.swapaxes(sarima_2.log_prob_elementwise(*mcmc_samples), 1, 0)
arimax_idata = az.from_dict(
posterior={
k:np.swapaxes(v.numpy(), 1, 0)
for k, v in mcmc_samples._asdict().items()},
sample_stats={
k:np.swapaxes(sampler_stats[k], 1, 0)
for k in ["target_log_prob", "diverging", "accept_ratio", "n_steps"]},
observed_data={'observed': sarima_2.observed},
log_likelihood={'observed': data_likelihood},
)

axes = az.plot_trace(arimax_idata);

#### y ~ Sarima(0,1,2)(1,1,1)[12]#

@tfd.JointDistributionCoroutine
def sarima_priors():
mu0 = yield root(tfd.StudentT(df=6, loc=0, scale=2.5, name='mu0'))
sigma = yield root(tfd.HalfStudentT(df=7, loc=0, scale=1., name='sigma'))

theta = yield root(tfd.Sample(tfd.Normal(loc=0, scale=0.5), 2, name='theta'))
sphi = yield root(tfd.Sample(tfd.Normal(loc=0, scale=0.5), 1, name='sphi'))
stheta = yield root(tfd.Sample(tfd.Normal(loc=0, scale=0.5), 1, name='stheta'))

sarima_3 = SARIMAX(
us_monthly_birth["birth_in_thousands"],
order=(0, 1, 2),
seasonal_order=(1, 1, 1, 12))
%%time
target_log_prob_fn = lambda *x: sarima_priors.log_prob(
*x) + sarima_3.log_prob(*x)

mcmc_samples, sampler_stats = run_mcmc_simple(
target_log_prob_fn=target_log_prob_fn,
seed=tf.constant([934563, 12356], dtype=tf.int32),
)
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
CPU times: user 1min 2s, sys: 812 ms, total: 1min 3s
Wall time: 1min 3s
data_likelihood = np.swapaxes(sarima_3.log_prob_elementwise(*mcmc_samples), 1, 0)
sarima_1_idata = az.from_dict(
posterior={
k:np.swapaxes(v.numpy(), 1, 0)
for k, v in mcmc_samples._asdict().items()},
sample_stats={
k:np.swapaxes(sampler_stats[k], 1, 0)
for k in ["target_log_prob", "diverging", "accept_ratio", "n_steps"]},
observed_data={'observed': sarima_2.observed},
log_likelihood={'observed': data_likelihood},
)

axes = az.plot_trace(sarima_1_idata);
az.summary(sarima_1_idata)
mean sd hdi_3% hdi_97% mcse_mean mcse_sd ess_bulk ess_tail r_hat
mu0 0.010 0.084 -0.148 0.179 0.002 0.001 2648.0 2030.0 1.0
sigma 7.354 0.277 6.820 7.863 0.005 0.004 3004.0 2521.0 1.0
theta[0] -0.373 0.060 -0.489 -0.265 0.001 0.001 2359.0 2522.0 1.0
theta[1] -0.137 0.048 -0.226 -0.046 0.001 0.001 2644.0 2533.0 1.0
sphi[0] -0.261 0.075 -0.404 -0.123 0.002 0.001 1731.0 2276.0 1.0
stheta[0] -0.297 0.082 -0.454 -0.148 0.002 0.002 1561.0 2073.0 1.0

### Table 6.1#

compare_dict = {"SARIMA(1,1,1)(1,1,1)[12]": sarima_0_idata,
#                 "ARIMAX(1,1,1)X[4]": arimax_idata,
"SARIMA(0,1,2)(1,1,1)[12]": sarima_1_idata}
cmp = az.compare(compare_dict, ic='loo')
cmp.round(2)
/opt/miniconda3/envs/bmcp/lib/python3.9/site-packages/arviz/stats/stats.py:145: UserWarning: The default method used to estimate the weights for each model,has changed from BB-pseudo-BMA to stacking
warnings.warn(
rank loo p_loo d_loo weight se dse warning loo_scale
SARIMA(1,1,1)(1,1,1)[12] 0 -1235.61 7.73 0.00 0.51 15.46 0.00 False log
SARIMA(0,1,2)(1,1,1)[12] 1 -1235.62 7.49 0.01 0.49 15.39 6.35 False log
# print(cmp.round(2).to_latex())

## State Space Models#

### Code 6.19#

theta0, theta1 = 1.2, 2.6
sigma = 0.4
num_timesteps = 100

time_stamp = tf.linspace(0., 1., num_timesteps)[..., None]
yhat = theta0 + theta1 * time_stamp
y = tfd.Normal(yhat, sigma).sample()

### Code 6.20#

# X_0
initial_state_prior = tfd.MultivariateNormalDiag(
loc=[0., 0.], scale_diag=[5., 5.])
# F_t
transition_matrix = lambda _: tf.linalg.LinearOperatorIdentity(2)
# eta_t ~ Normal(0, Q_t)
transition_noise = lambda _: tfd.MultivariateNormalDiag(
loc=[0., 0.], scale_diag=[0., 0.])
# H_t
H = tf.concat([tf.ones_like(time_stamp), time_stamp], axis=-1)
observation_matrix = lambda t: tf.linalg.LinearOperatorFullMatrix(
[tf.gather(H, t)])
# epsilon_t ~ Normal(0, R_t)
observation_noise = lambda _: tfd.MultivariateNormalDiag(
loc=[0.], scale_diag=[sigma])

linear_growth_model = tfd.LinearGaussianStateSpaceModel(
num_timesteps=num_timesteps,
transition_matrix=transition_matrix,
transition_noise=transition_noise,
observation_matrix=observation_matrix,
observation_noise=observation_noise,
initial_state_prior=initial_state_prior)

### Code 6.21#

# Run the Kalman filter
(
log_likelihoods,
mt_filtered, Pt_filtered,
mt_predicted, Pt_predicted,
observation_means, observation_cov  # observation_cov is S_t
) = linear_growth_model.forward_filter(y)

### Figure 6.16#

m0 = initial_state_prior.mean()
P0 = initial_state_prior.covariance()

P0_inv = tf.linalg.inv(P0)
P_t = tf.linalg.inv(P0_inv +
1 / sigma **2 * (tf.matmul(H, H, transpose_a=True)))
m_t = tf.matmul(P_t, (1 / sigma **2 * (tf.matmul(H, y, transpose_a=True)) +
tf.matmul(P0_inv, m0[..., None])))
filtered_vars = tf.linalg.diag_part(Pt_filtered)

_, ax = plt.subplots(1, 3, figsize=(10, 4))
ax[0].plot(time_stamp, y, '--o', alpha=.5);
ax[0].plot(time_stamp, observation_means, lw=1.5, color='k')
ax[0].set_title('Observed time series')
ax[0].legend(['Observed', 'Predicted'])
ax[0].set_xlabel('time');

color = ['C4', 'C1']
for i in range(2):
ax[1].plot(time_stamp, tf.transpose(mt_filtered[..., i]), color=color[i]);
ax[2].semilogy(time_stamp, tf.transpose(filtered_vars[..., i]), color=color[i]);
for i in range(2):
ax[i+1].set_xlabel('time')
ax[i+1].legend(['theta0', 'theta1'])

ax[1].hlines(m_t, time_stamp[0], time_stamp[-1], ls='--');
ax[1].set_title(r'$m_{t \mid t}$')
ax[2].hlines(tf.linalg.diag_part(P_t), time_stamp[0], time_stamp[-1], ls='--')
ax[2].set_title(r'diag($P_{t \mid t}$)')
ax[2].grid()
plt.tight_layout();

plt.savefig("img/chp06/fig16_linear_growth_lgssm.png");
/var/folders/7p/srk5qjp563l5f9mrjtp44bh800jqsw/T/ipykernel_14469/3981273553.py:31: UserWarning: This figure was using constrained_layout, but that is incompatible with subplots_adjust and/or tight_layout; disabling constrained_layout.
plt.tight_layout();

### Code 6.22#

num_timesteps = 300
phi1 = -.1
phi2 = .5
theta1 = -.25
sigma = 1.25

# X_0
initial_state_prior = tfd.MultivariateNormalDiag(
scale_diag=[sigma, sigma])
# F_t
transition_matrix = lambda _: tf.linalg.LinearOperatorFullMatrix(
[[phi1, 1], [phi2, 0]])
# eta_t ~ Normal(0, Q_t)
R_t = tf.constant([[sigma], [sigma*theta1]])
Q_t_tril = tf.concat([R_t, tf.zeros_like(R_t)], axis=-1)
transition_noise = lambda _: tfd.MultivariateNormalTriL(
scale_tril=Q_t_tril)
# H_t
observation_matrix = lambda t: tf.linalg.LinearOperatorFullMatrix(
[[1., 0.]])
# epsilon_t ~ Normal(0, 0)
observation_noise = lambda _: tfd.MultivariateNormalDiag(
loc=[0.], scale_diag=[0.])

arma = tfd.LinearGaussianStateSpaceModel(
num_timesteps=num_timesteps,
transition_matrix=transition_matrix,
transition_noise=transition_noise,
observation_matrix=observation_matrix,
observation_noise=observation_noise,
initial_state_prior=initial_state_prior
)

# Simulate from the model
sim_ts = arma.sample()
np.linalg.eigvals(Q_t_tril @ tf.transpose(Q_t_tril)) >= 0
array([ True,  True])
plt.plot(sim_ts);
arma.log_prob(sim_ts)
<tf.Tensor: shape=(), dtype=float32, numpy=-492.00345>

### Code 6.23#

@tfd.JointDistributionCoroutine
def arma_lgssm():
sigma = yield root(tfd.HalfStudentT(df=7, loc=0, scale=1., name='sigma'))
phi = yield root(tfd.Sample(tfd.Normal(loc=0, scale=0.5), 2, name='phi'))
theta = yield root(tfd.Sample(tfd.Normal(loc=0, scale=0.5), 1, name='theta'))

# X0
init_scale_diag = tf.concat([sigma[..., None], sigma[..., None]], axis=-1)
initial_state_prior = tfd.MultivariateNormalDiag(
scale_diag=init_scale_diag)

F_t = tf.concat([phi[..., None],
tf.concat([tf.ones_like(phi[..., 0, None]),
tf.zeros_like(phi[..., 0, None])],
axis=-1)[..., None]],
axis=-1)

def transition_matrix(_): return tf.linalg.LinearOperatorFullMatrix(F_t)

transition_scale_tril = tf.concat(
[sigma[..., None], theta * sigma[..., None]], axis=-1)[..., None]
scale_tril = tf.concat(
[transition_scale_tril,
tf.zeros_like(transition_scale_tril)],
axis=-1)

def transition_noise(_): return tfd.MultivariateNormalTriL(
scale_tril=scale_tril)

def observation_matrix(
t): return tf.linalg.LinearOperatorFullMatrix([[1., 0.]])

def observation_noise(t): return tfd.MultivariateNormalDiag(
loc=[0], scale_diag=[0.])

arma = yield tfd.LinearGaussianStateSpaceModel(
num_timesteps=num_timesteps,
transition_matrix=transition_matrix,
transition_noise=transition_noise,
observation_matrix=observation_matrix,
observation_noise=observation_noise,
initial_state_prior=initial_state_prior,
name='arma')
%%time
mcmc_samples, sampler_stats = run_mcmc_simple(
seed=tf.constant([23453, 94567], dtype=tf.int32),
arma=sim_ts)
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:5 out of the last 5 calls to <function run_mcmc_simple.<locals>.run_inference_nuts at 0x7fad6303a700> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has experimental_relax_shapes=True option that relaxes argument shapes that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.
CPU times: user 29.2 s, sys: 476 ms, total: 29.7 s
Wall time: 29.8 s

### Figure 6.17#

test_trace = az.from_dict(
posterior={
k:np.swapaxes(v.numpy(), 1, 0)
for k, v in mcmc_samples._asdict().items()},
sample_stats={
k:np.swapaxes(sampler_stats[k], 1, 0)
for k in ["target_log_prob", "diverging", "accept_ratio", "n_steps"]}
)

lines = (('sigma', {}, sigma), ('phi', {}, [phi1, phi2]), ('theta', {}, theta1),)
axes = az.plot_trace(test_trace, lines=lines);

plt.savefig("img/chp06/fig17_arma_lgssm_inference_result.png");

#### Bayesian Structural Time Series Models on Monthly Live Birth Data#

def generate_bsts_model(observed=None):
"""
Args:
observed: Observed time series, tfp.sts use it to generate data informed prior.
"""
# Trend
trend = tfp.sts.LocalLinearTrend(observed_time_series=observed)
# Seasonal
seasonal = tfp.sts.Seasonal(num_seasons=12, observed_time_series=observed)
# Full model
return tfp.sts.Sum([trend, seasonal], observed_time_series=observed)

observed = tf.constant(us_monthly_birth["birth_in_thousands"], dtype=tf.float32)
birth_model = generate_bsts_model(observed=observed)

# Generate the posterior distribution conditioned on the observed
# target_log_prob_fn = birth_model.joint_log_prob(observed_time_series=observed)

birth_model_jd = birth_model.joint_distribution(observed_time_series=observed)
WARNING:tensorflow:From /opt/miniconda3/envs/bmcp/lib/python3.9/site-packages/tensorflow/python/ops/linalg/linear_operator_composition.py:182: LinearOperator.graph_parents (from tensorflow.python.ops.linalg.linear_operator) is deprecated and will be removed in a future version.
Instructions for updating:
Do not call graph_parents.
WARNING:tensorflow:From /opt/miniconda3/envs/bmcp/lib/python3.9/site-packages/tensorflow_probability/python/distributions/distribution.py:342: MultivariateNormalFullCovariance.__init__ (from tensorflow_probability.python.distributions.mvn_full_covariance) is deprecated and will be removed after 2019-12-01.
Instructions for updating:
MultivariateNormalFullCovariance is deprecated, use MultivariateNormalTriL(loc=loc, scale_tril=tf.linalg.cholesky(covariance_matrix)) instead.

### Code 6.25#

birth_model.components
[<tensorflow_probability.python.sts.components.local_linear_trend.LocalLinearTrend at 0x7fae436e6cd0>,
<tensorflow_probability.python.sts.components.seasonal.Seasonal at 0x7fae76b00340>]

### Code 6.26#

birth_model.components[1].parameters
[Parameter(name='drift_scale', prior=<tfp.distributions.LogNormal 'LogNormal' batch_shape=[] event_shape=[] dtype=float32>, bijector=<tfp.bijectors.Chain 'chain_of_scale_of_softplus' batch_shape=[] forward_min_event_ndims=0 inverse_min_event_ndims=0 dtype_x=float32 dtype_y=float32 bijectors=[<tfp.bijectors.Scale 'scale' batch_shape=[] forward_min_event_ndims=0 inverse_min_event_ndims=0 dtype_x=float32 dtype_y=float32>, <tfp.bijectors.Softplus 'softplus' batch_shape=[] forward_min_event_ndims=0 inverse_min_event_ndims=0 dtype_x=float32 dtype_y=float32>]>)]
%%time
mcmc_samples, sampler_stats = run_mcmc(
seed=tf.constant([745678, 562345], dtype=tf.int32))
2021-11-19 18:01:06.021471: W tensorflow/compiler/tf2xla/kernels/assert_op.cc:38] Ignoring Assert operator mcmc_retry_init/assert_equal_1/Assert/AssertGuard/Assert
WARNING:tensorflow:6 out of the last 10 calls to <function windowed_adaptive_nuts at 0x7fae934af550> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has experimental_relax_shapes=True option that relaxes argument shapes that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.
CPU times: user 11min 21s, sys: 7.45 s, total: 11min 29s
Wall time: 11min 22s
birth_model_idata = az.from_dict(
posterior={
k:np.swapaxes(v.numpy(), 1, 0)
for k, v in mcmc_samples.items()},
sample_stats={
k:np.swapaxes(sampler_stats[k], 1, 0)
for k in ["target_log_prob", "diverging", "accept_ratio", "n_steps"]}
)

axes = az.plot_trace(birth_model_idata);

### Code 6.27#

# Using a subset of posterior samples.
parameter_samples = tf.nest.map_structure(lambda x: x[-100:, 0, ...], mcmc_samples)

# Get structual compoenent.
component_dists = tfp.sts.decompose_by_component(
birth_model,
observed_time_series=observed,
parameter_samples=parameter_samples)

# Get forecast for n_steps.
n_steps = 36
forecast_dist = tfp.sts.forecast(
birth_model,
observed_time_series=observed,
parameter_samples=parameter_samples,
num_steps_forecast=n_steps)

## Other Time Series Models#

### Figure 6.19#

birth_dates = us_monthly_birth.index
forecast_date = pd.date_range(
start=birth_dates[-1] + np.timedelta64(1, "M"),
end=birth_dates[-1] + np.timedelta64(1 + n_steps, "M"),
freq='M')

fig, axes = plt.subplots(
1 + len(component_dists.keys()), 1, figsize=(10, 9), sharex=True)

ax = axes[0]
ax.plot(us_monthly_birth, lw=1.5, label='observed')

forecast_mean = np.squeeze(forecast_dist.mean())
line = ax.plot(forecast_date, forecast_mean, lw=1.5,
label='forecast mean', color='C4')

forecast_std = np.squeeze(forecast_dist.stddev())
ax.fill_between(forecast_date,
forecast_mean - 2 * forecast_std,
forecast_mean + 2 * forecast_std,
color=line[0].get_color(), alpha=0.2)

for ax_, (key, dist) in zip(axes[1:], component_dists.items()):
comp_mean, comp_std = np.squeeze(dist.mean()), np.squeeze(dist.stddev())
line = ax_.plot(birth_dates, dist.mean(), lw=2.)
ax_.fill_between(birth_dates,
comp_mean - 2 * comp_std,
comp_mean + 2 * comp_std,
alpha=0.2)
ax_.set_title(key.name[:-1])

ax.legend()

ax.set_ylabel("Birth (thousands)")
ax_.set_xlabel("Year")
ax.set_title("Monthly live birth U.S.A", fontsize=15)
ax.text(0.99, .02,
"Source: Stoffer D (2019). “astsa: Applied Statistical Time Series Analysis.”",
transform=ax.transAxes,
horizontalalignment="right",
alpha=0.5)

fig.autofmt_xdate()
plt.tight_layout()

plt.savefig("img/chp06/fig19_bsts_lgssm_result.png")
/var/folders/7p/srk5qjp563l5f9mrjtp44bh800jqsw/T/ipykernel_14469/2909654341.py:43: UserWarning: This figure was using constrained_layout, but that is incompatible with subplots_adjust and/or tight_layout; disabling constrained_layout.
fig.autofmt_xdate()