Пример использования вероятности TensorFlow: оценка ковариации

Посмотреть на TensorFlow.org Запускаем в Google Colab Посмотреть исходный код на GitHub Скачать блокнот

Я написал этот блокнот в качестве примера для изучения TensorFlow Probability. Проблема, которую я решил решить, - это оценка ковариационной матрицы для выборок двумерной гауссовской случайной величины со средним значением 0. У проблемы есть несколько приятных особенностей:

  • Если мы используем обратный априор Уишарта для ковариации (общий подход), проблема имеет аналитическое решение, поэтому мы можем проверить наши результаты.
  • Проблема включает выборку ограниченного параметра, что добавляет интересную сложность.
  • Самое простое решение - не самое быстрое, поэтому необходимо провести некоторую оптимизацию.

Я решил описать свой опыт по ходу дела. Мне потребовалось некоторое время, чтобы осмыслить тонкости TFP, так что этот ноутбук запускается довольно просто, а затем постепенно переходит к более сложным функциям TFP. По пути я столкнулся с множеством проблем и попытался уловить как процессы, которые помогли мне их идентифицировать, так и обходные пути, которые я в итоге нашел. Я попытался включить много деталей (включая множество тестов, чтобы убедиться, что отдельные шаги верны).

Зачем изучать вероятность TensorFlow?

Я нашел TensorFlow Probability привлекательным для моего проекта по нескольким причинам:

  • Вероятность TensorFlow позволяет создавать прототипы сложных моделей в интерактивном режиме в записной книжке. Вы можете разбить свой код на небольшие части, которые можно тестировать в интерактивном режиме и с помощью модульных тестов.
  • Когда вы будете готовы к масштабированию, вы сможете воспользоваться преимуществами всей инфраструктуры, которая у нас есть для запуска TensorFlow на нескольких оптимизированных процессорах на нескольких машинах.
  • Наконец, хотя мне очень нравится Стэн, мне довольно сложно отлаживать его. Вы должны написать весь свой код моделирования на отдельном языке, который имеет очень мало инструментов, позволяющих вам тыкать в свой код, проверять промежуточные состояния и так далее.

Обратной стороной является то, что TensorFlow Probability намного новее, чем Stan и PyMC3, поэтому документация находится в стадии разработки, и есть много функций, которые еще предстоит создать. К счастью, я обнаружил, что фундамент TFP прочный, и он спроектирован по модульному принципу, что позволяет довольно просто расширять его функциональные возможности. В этой записной книжке, помимо решения тематического исследования, я покажу несколько способов расширения TFP.

Для кого это

Я предполагаю, что читатели приходят к этой записной книжке с некоторыми важными предпосылками. Вам следует:

Первая попытка

Вот моя первая попытка решить проблему. Спойлер: мое решение не работает, и потребуется несколько попыток, чтобы все исправить! Хотя этот процесс занимает некоторое время, каждая попытка, описанная ниже, была полезна для изучения новой части TFP.

Одно замечание: в настоящее время TFP не реализует обратное распределение Уишарта (в конце мы увидим, как развернуть собственное обратное распределение Уишарта), поэтому вместо этого я заменю задачу на оценку матрицы точности с использованием предварительного распределения Уишарта.

import collections
import math
import os
import time

import numpy as np
import pandas as pd
import scipy
import scipy.stats
import matplotlib.pyplot as plt

import tensorflow.compat.v2 as tf
tf.enable_v2_behavior()
import tensorflow_probability as tfp
tfd = tfp.distributions
tfb = tfp.bijectors

Шаг 1: соберите наблюдения вместе

Все мои данные здесь синтетические, так что это покажется немного более аккуратным, чем реальный пример. Однако нет причин, по которым вы не можете генерировать собственные синтетические данные.

Совет : После того, как вы определились с формой вашей модели, вы можете выбрать некоторые значения параметров и использовать выбранную вами модель для генерации некоторых синтетических данных. В качестве проверки работоспособности вашей реализации вы можете затем убедиться, что ваши оценки включают истинные значения выбранных вами параметров. Чтобы ускорить цикл отладки / тестирования, вы можете рассмотреть упрощенную версию вашей модели (например, использовать меньшее количество измерений или меньшее количество образцов).

Совет: проще всего работать со своими наблюдениями как с массивами NumPy. Важно отметить, что NumPy по умолчанию использует float64, а TensorFlow по умолчанию использует float32.

Как правило, операции TensorFlow хотят, чтобы все аргументы имели один и тот же тип, и вам необходимо выполнить явное приведение данных для изменения типов. Если вы используете наблюдения float64, вам нужно будет добавить много операций приведения. NumPy, напротив, позаботится о преобразовании автоматически. Следовательно, гораздо проще преобразовать ваши данные Numpy в float32, чем заставить TensorFlow использовать float64.

Выберите некоторые значения параметров

# We're assuming 2-D data with a known true mean of (0, 0)
true_mean = np.zeros([2], dtype=np.float32)
# We'll make the 2 coordinates correlated
true_cor = np.array([[1.0, 0.9], [0.9, 1.0]], dtype=np.float32)
# And we'll give the 2 coordinates different variances
true_var = np.array([4.0, 1.0], dtype=np.float32)
# Combine the variances and correlations into a covariance matrix
true_cov = np.expand_dims(np.sqrt(true_var), axis=1).dot(
    np.expand_dims(np.sqrt(true_var), axis=1).T) * true_cor
# We'll be working with precision matrices, so we'll go ahead and compute the
# true precision matrix here
true_precision = np.linalg.inv(true_cov)
# Here's our resulting covariance matrix
print(true_cov)
# Verify that it's positive definite, since np.random.multivariate_normal
# complains about it not being positive definite for some reason.
# (Note that I'll be including a lot of sanity checking code in this notebook -
# it's a *huge* help for debugging)
print('eigenvalues: ', np.linalg.eigvals(true_cov))
[[4.  1.8]
 [1.8 1. ]]
eigenvalues:  [4.843075   0.15692513]

Произведите синтетические наблюдения

Обратите внимание, что TensorFlow Probability использует соглашение о том, что начальные измерения ваших данных представляют собой индексы выборки, а окончательные измерения ваших данных представляют размерность ваших выборок.

Здесь нам нужно 100 выборок, каждая из которых является вектором длины 2. Мы сгенерируем массив my_data с shape (100, 2). my_data[i, :] - это $ i $ -я выборка, и это вектор длины 2.

(Не забудьте сделать my_data типом float32!)

# Set the seed so the results are reproducible.
np.random.seed(123)

# Now generate some observations of our random variable.
# (Note that I'm suppressing a bunch of spurious about the covariance matrix
# not being positive semidefinite via check_valid='ignore' because it really is
# positive definite!)
my_data = np.random.multivariate_normal(
    mean=true_mean, cov=true_cov, size=100,
    check_valid='ignore').astype(np.float32)
my_data.shape
(100, 2)

Вменяемость проверьте наблюдения

Один из потенциальных источников ошибок - испортить ваши синтетические данные! Сделаем несколько простых проверок.

# Do a scatter plot of the observations to make sure they look like what we
# expect (higher variance on the x-axis, y values strongly correlated with x)
plt.scatter(my_data[:, 0], my_data[:, 1], alpha=0.75)
plt.show()

PNG

print('mean of observations:', np.mean(my_data, axis=0))
print('true mean:', true_mean)
mean of observations: [-0.24009615 -0.16638893]
true mean: [0. 0.]
print('covariance of observations:\n', np.cov(my_data, rowvar=False))
print('true covariance:\n', true_cov)
covariance of observations:
 [[3.95307734 1.68718486]
 [1.68718486 0.94910269]]
true covariance:
 [[4.  1.8]
 [1.8 1. ]]

Хорошо, наши образцы выглядят разумно. Следующий шаг.

Шаг 2. Реализуйте функцию правдоподобия в NumPy

Главное, что нам нужно будет написать для выполнения выборки MCMC в TF Probability, - это функция правдоподобия журнала. В общем, написать TF немного сложнее, чем NumPy, поэтому я считаю полезным выполнить первоначальную реализацию в NumPy. Я собираюсь разделить функцию правдоподобия на 2 части: функцию правдоподобия данных, которая соответствует $ P (data | parameters) $, и априорную функцию правдоподобия, которая соответствует $ P (parameters) $.

Обратите внимание, что эти функции NumPy не должны быть супероптимизированы / векторизованы, поскольку цель состоит в том, чтобы просто сгенерировать некоторые значения для тестирования. Правильность - ключевое соображение!

Сначала мы реализуем часть вероятности журнала данных. Это довольно просто. Единственное, что нужно помнить, это то, что мы собираемся работать с матрицами точности, поэтому мы параметризуем соответственно.

def log_lik_data_numpy(precision, data):
  # np.linalg.inv is a really inefficient way to get the covariance matrix, but
  # remember we don't care about speed here
  cov = np.linalg.inv(precision)
  rv = scipy.stats.multivariate_normal(true_mean, cov)
  return np.sum(rv.logpdf(data))

# test case: compute the log likelihood of the data given the true parameters
log_lik_data_numpy(true_precision, my_data)
-280.81822950593767

Мы собираемся использовать априор Уишарта для матрицы точности, поскольку есть аналитическое решение для апостериорного (см . Удобную таблицу сопряженных априорных значений в Википедии ).

Распределение Уишарта имеет 2 параметра:

  • количество степеней свободы (обозначено в Википедии $ \ nu $)
  • масштабная матрица (в Википедии обозначена как $ V $)

Среднее значение для распределения Уишарта с параметрами $ \ nu, V $ равно $ E [W] = \ nu V $, а дисперсия равна $ \ text {Var} (W_ {ij}) = \ nu (v_ {ij} ^ 2 + v_ {ii} v_ {jj}) $

Некоторая полезная интуиция: вы можете сгенерировать образец Уишарта, сгенерировав $ \ nu $ независимых отрисовок $ x_1 \ ldots x _ {\ nu} $ из многомерной нормальной случайной величины со средним 0 и ковариацией $ V $, а затем сформировав сумму $ W = \ sum_ {i = 1} ^ {\ nu} x_i x_i ^ T $.

Если вы измените масштаб выборок Wishart, разделив их на $ \ nu $, вы получите образец ковариационной матрицы $ x_i $. Эта примерная ковариационная матрица должна стремиться к $ V $ при увеличении $ \ nu $. Когда $ \ nu $ мало, есть много вариаций в выборочной матрице ковариации, поэтому малые значения $ \ nu $ соответствуют более слабым априорным признакам, а большие значения $ \ nu $ соответствуют более сильным априорным значениям. Обратите внимание, что размер $ \ nu $ должен быть не меньше размера пространства, в котором вы производите выборку, иначе вы создадите особые матрицы.

Мы будем использовать $ \ nu = 3 $, чтобы иметь слабую априорность, и мы возьмем $ V = \ frac {1} {\ nu} I $, который подтянет нашу оценку ковариации к тождеству (напомним, что среднее равно $ \ nu V $).

PRIOR_DF = 3
PRIOR_SCALE = np.eye(2, dtype=np.float32) / PRIOR_DF

def log_lik_prior_numpy(precision):
  rv = scipy.stats.wishart(df=PRIOR_DF, scale=PRIOR_SCALE)
  return rv.logpdf(precision)

# test case: compute the prior for the true parameters
log_lik_prior_numpy(true_precision)
-9.103606346649766

Распределение Уишарта является сопряженным априорным методом для оценки матрицы точности многомерной нормали с известным средним значением $ \ mu $.

Предположим, что априорными параметрами Уишарта являются $ \ nu, V $ и что у нас есть $ n $ наблюдений многомерной нормали, $ x_1, \ ldots, x_n $. Апостериорные параметры: $ n + \ nu, \ left (V ^ {- 1} + \ sum_ {i = 1} ^ n (x_i- \ mu) (x_i- \ mu) ^ T \ right) ^ {- 1 } $.

n = my_data.shape[0]
nu_prior = PRIOR_DF
v_prior = PRIOR_SCALE
nu_posterior = nu_prior + n
v_posterior = np.linalg.inv(np.linalg.inv(v_prior) + my_data.T.dot(my_data))
posterior_mean = nu_posterior * v_posterior
v_post_diag = np.expand_dims(np.diag(v_posterior), axis=1)
posterior_sd = np.sqrt(nu_posterior *
                       (v_posterior ** 2.0 + v_post_diag.dot(v_post_diag.T)))

Быстрый сюжет апостериоров и истинных значений. Обратите внимание, что задние части близки к задним частям образца, но немного сужены в сторону идентичности. Также обратите внимание, что истинные значения довольно далеки от режима апостериорного распределения - по-видимому, это связано с тем, что предыдущие значения не очень хорошо подходят для наших данных. В реальной задаче мы, вероятно, лучше справимся с чем-то вроде масштабированного обратного априорного значения Уишарта для ковариации (см., Например, комментарий Эндрю Гельмана по этому поводу), но тогда у нас не было бы хорошей аналитической апостериорной оценки.

sample_precision = np.linalg.inv(np.cov(my_data, rowvar=False, bias=False))
fig, axes = plt.subplots(2, 2)
fig.set_size_inches(10, 10)
for i in range(2):
  for j in range(2):
    ax = axes[i, j]
    loc = posterior_mean[i, j]
    scale = posterior_sd[i, j]
    xmin = loc - 3.0 * scale
    xmax = loc + 3.0 * scale
    x = np.linspace(xmin, xmax, 1000)
    y = scipy.stats.norm.pdf(x, loc=loc, scale=scale)
    ax.plot(x, y)
    ax.axvline(true_precision[i, j], color='red', label='True precision')
    ax.axvline(sample_precision[i, j], color='red', linestyle=':', label='Sample precision')
    ax.set_title('precision[%d, %d]' % (i, j))
plt.legend()
plt.show()

PNG

Шаг 3. Реализуйте функцию правдоподобия в TensorFlow

Спойлер: Наша первая попытка не сработает; О том, почему, поговорим ниже.

Совет : используйте активный режим TensorFlow при разработке функций вероятности. Активный режим заставляет TF вести себя больше как NumPy - все выполняется немедленно, поэтому вы можете выполнять отладку в интерактивном режиме, вместо того, чтобы использовать Session.run() . См. Примечания здесь .

Предварительно: классы распределения

В TFP есть набор классов распределения, которые мы будем использовать для генерации вероятностей журнала. Следует отметить, что эти классы работают с тензорами выборок, а не только с отдельными выборками - это позволяет осуществлять векторизацию и связанные с этим ускорения.

Распределение может работать с тензором выборок двумя разными способами. Проще всего проиллюстрировать эти два способа на конкретном примере, включающем распределение с одним скалярным параметром. Я буду использовать распределение Пуассона , у которого есть параметр rate .

  • Если мы создадим Пуассон с одним значением для параметра rate , вызов его метода sample() вернет одно значение. Это значение называется event , и в этом случае все события являются скалярами.
  • Если мы создадим пуассоновский тензор значений для параметра rate , вызов его метода sample() теперь возвращает несколько значений, по одному для каждого значения в тензоре скорости. Объект действует как набор независимых Пуассонов, каждый со своей скоростью, и каждое из значений, возвращаемых вызовом sample() соответствует одному из этих Пуассонов. Этот набор независимых, но не одинаково распределенных событий называется batch .
  • Метод sample() принимает параметр sample_shape который по умолчанию sample_shape пустому кортежу. Передача непустого значения для sample_shape приводит к тому, что образец возвращает несколько пакетов. Этот набор партий называется sample .

Метод распределения log_prob() потребляет данные таким же образом, как и метод sample() генерирует. log_prob() возвращает вероятности для выборок, то есть для нескольких независимых пакетов событий.

  • Если у нас есть объект Пуассона, который был создан со скалярной rate , каждый пакет является скаляром, и если мы передадим тензор выборок, мы получим тензор того же размера логарифмических вероятностей.
  • Если у нас есть объект Пуассона, который был создан с тензором формы T значений rate , каждая партия является тензором формы T Если мы передадим тензор выборок формы D, T, мы получим тензор логарифмических вероятностей формы D, T.

Ниже приведены несколько примеров, иллюстрирующих эти случаи. См. Этот блокнот для более подробного руководства по событиям, пакетам и фигурам.

# case 1: get log probabilities for a vector of iid draws from a single
# normal distribution
norm1 = tfd.Normal(loc=0., scale=1.)
probs1 = norm1.log_prob(tf.constant([1., 0.5, 0.]))

# case 2: get log probabilities for a vector of independent draws from
# multiple normal distributions with different parameters.  Note the vector
# values for loc and scale in the Normal constructor.
norm2 = tfd.Normal(loc=[0., 2., 4.], scale=[1., 1., 1.])
probs2 = norm2.log_prob(tf.constant([1., 0.5, 0.]))

print('iid draws from a single normal:', probs1.numpy())
print('draws from a batch of normals:', probs2.numpy())
iid draws from a single normal: [-1.4189385 -1.0439385 -0.9189385]
draws from a batch of normals: [-1.4189385 -2.0439386 -8.918939 ]

Вероятность регистрации данных

Сначала мы реализуем функцию правдоподобия журнала данных.

VALIDATE_ARGS = True
ALLOW_NAN_STATS = False

Одним из ключевых отличий от случая NumPy является то, что наша функция правдоподобия TensorFlow должна будет обрабатывать векторы матриц точности, а не просто отдельные матрицы. Векторы параметров будут использоваться при выборке из нескольких цепочек.

Мы создадим объект распределения, который работает с пакетом матриц точности (то есть по одной матрице на цепочку).

При вычислении вероятностей журнала наших данных нам нужно, чтобы наши данные реплицировались таким же образом, как и наши параметры, чтобы на каждую переменную пакета была одна копия. Форма наших реплицированных данных должна быть следующей:

[sample shape, batch shape, event shape]

В нашем случае форма события равна 2 (поскольку мы работаем с двумерными гауссианами). Форма образца - 100, так как у нас 100 образцов. Форма пакета будет просто количеством матриц точности, с которыми мы работаем. Реплицировать данные каждый раз, когда мы вызываем функцию правдоподобия, расточительно, поэтому мы заранее реплицируем данные и передадим реплицированную версию.

Обратите внимание, что это неэффективная реализация: MultivariateNormalFullCovariance стоит дорого по сравнению с некоторыми альтернативами, о которых мы поговорим в разделе оптимизации в конце.

def log_lik_data(precisions, replicated_data):
  n = tf.shape(precisions)[0]  # number of precision matrices
  # We're estimating a precision matrix; we have to invert to get log
  # probabilities.  Cholesky inversion should be relatively efficient,
  # but as we'll see later, it's even better if we can avoid doing the Cholesky
  # decomposition altogether.
  precisions_cholesky = tf.linalg.cholesky(precisions)
  covariances = tf.linalg.cholesky_solve(
      precisions_cholesky, tf.linalg.eye(2, batch_shape=[n]))
  rv_data = tfd.MultivariateNormalFullCovariance(
      loc=tf.zeros([n, 2]),
      covariance_matrix=covariances,
      validate_args=VALIDATE_ARGS,
      allow_nan_stats=ALLOW_NAN_STATS)

  return tf.reduce_sum(rv_data.log_prob(replicated_data), axis=0)
# For our test, we'll use a tensor of 2 precision matrices.
# We'll need to replicate our data for the likelihood function.
# Remember, TFP wants the data to be structured so that the sample dimensions
# are first (100 here), then the batch dimensions (2 here because we have 2
# precision matrices), then the event dimensions (2 because we have 2-D
# Gaussian data).  We'll need to add a middle dimension for the batch using
# expand_dims, and then we'll need to create 2 replicates in this new dimension
# using tile.
n = 2
replicated_data = np.tile(np.expand_dims(my_data, axis=1), reps=[1, 2, 1])
print(replicated_data.shape)
(100, 2, 2)

Совет: Одна вещь, которую я нашел чрезвычайно полезной, - это написание небольших проверок работоспособности моих функций TensorFlow. Очень легко испортить векторизацию в TF, поэтому наличие более простых функций NumPy - отличный способ проверить вывод TF. Думайте об этом как о небольших модульных тестах.

# check against the numpy implementation
precisions = np.stack([np.eye(2, dtype=np.float32), true_precision])
n = precisions.shape[0]
lik_tf = log_lik_data(precisions, replicated_data=replicated_data).numpy()

for i in range(n):
  print(i)
  print('numpy:', log_lik_data_numpy(precisions[i], my_data))
  print('tensorflow:', lik_tf[i])
0
numpy: -430.71218815801365
tensorflow: -430.71207
1
numpy: -280.81822950593767
tensorflow: -280.8182

Вероятность предыдущего журнала

Предыдущее проще, поскольку нам не нужно беспокоиться о репликации данных.

@tf.function(autograph=False)
def log_lik_prior(precisions):
  rv_precision = tfd.WishartTriL(
      df=PRIOR_DF,
      scale_tril=tf.linalg.cholesky(PRIOR_SCALE),
      validate_args=VALIDATE_ARGS,
      allow_nan_stats=ALLOW_NAN_STATS)
  return rv_precision.log_prob(precisions)
# check against the numpy implementation
precisions = np.stack([np.eye(2, dtype=np.float32), true_precision])
n = precisions.shape[0]
lik_tf = log_lik_prior(precisions).numpy()

for i in range(n):
  print(i)
  print('numpy:', log_lik_prior_numpy(precisions[i]))
  print('tensorflow:', lik_tf[i])
0
numpy: -2.2351873809649625
tensorflow: -2.2351875
1
numpy: -9.103606346649766
tensorflow: -9.103608

Постройте функцию правдоподобия совместного журнала

Приведенная выше функция правдоподобия журнала данных зависит от наших наблюдений, но в сэмплере их не будет. Мы можем избавиться от зависимости без использования глобальной переменной с помощью [закрытия] (https://en.wikipedia.org/wiki/Closure_ (computer_programming)). Замыкания включают внешнюю функцию, которая создает среду, содержащую переменные, необходимые для внутренняя функция.

def get_log_lik(data, n_chains=1):
  # The data argument that is passed in will be available to the inner function
  # below so it doesn't have to be passed in as a parameter.
  replicated_data = np.tile(np.expand_dims(data, axis=1), reps=[1, n_chains, 1])

  @tf.function(autograph=False)
  def _log_lik(precision):
    return log_lik_data(precision, replicated_data) + log_lik_prior(precision)

  return _log_lik

Шаг 4: образец

Хорошо, пора пробовать! Чтобы не усложнять задачу, мы будем использовать одну цепочку, а в качестве отправной точки будем использовать единичную матрицу. Позже мы сделаем все более тщательно.

Опять же, это не сработает - мы получим исключение.

@tf.function(autograph=False)
def sample():
  tf.random.set_seed(123)
  init_precision = tf.expand_dims(tf.eye(2), axis=0)

  # Use expand_dims because we want to pass in a tensor of starting values
  log_lik_fn = get_log_lik(my_data, n_chains=1)

  # we'll just do a few steps here
  num_results = 10
  num_burnin_steps = 10
  states = tfp.mcmc.sample_chain(
     num_results=num_results,
     num_burnin_steps=num_burnin_steps,
     current_state=[
         init_precision,
     ],
     kernel=tfp.mcmc.HamiltonianMonteCarlo(
         target_log_prob_fn=log_lik_fn,
         step_size=0.1,
         num_leapfrog_steps=3),
     trace_fn=None,
     seed=123)
  return states

try:
  states = sample()
except Exception as e:
  # shorten the giant stack trace
  lines = str(e).split('\n')
  print('\n'.join(lines[:5]+['...']+lines[-3:]))
Cholesky decomposition was not successful. The input might not be valid.
     [[{ {node mcmc_sample_chain/trace_scan/while/body/_79/smart_for_loop/while/body/_371/mh_one_step/hmc_kernel_one_step/leapfrog_integrate/while/body/_537/leapfrog_integrate_one_step/maybe_call_fn_and_grads/value_and_gradients/StatefulPartitionedCall/Cholesky} }]] [Op:__inference_sample_2849]

Function call stack:
sample
...
Function call stack:
sample

Выявление проблемы

InvalidArgumentError (see above for traceback): Cholesky decomposition was not successful. The input might not be valid. Это не очень полезно. Посмотрим, сможем ли мы узнать больше о том, что произошло.

  • Мы распечатаем параметры для каждого шага, чтобы увидеть значение, при котором что-то не получается.
  • Мы добавим несколько утверждений для защиты от конкретных проблем.

Утверждения сложны, потому что они являются операциями TensorFlow, и мы должны позаботиться о том, чтобы они выполнялись и не выходили за пределы графика. Если вы не знакомы с утверждениями TF, стоит прочитать этот обзор отладки TensorFlow. Вы можете явно принудительно выполнить утверждения с помощью tf.control_dependencies (см. Комментарии в коде ниже).

Собственная функция Print TensorFlow имеет то же поведение, что и утверждения - это операция, и вам нужно позаботиться о ее выполнении. Print вызывает дополнительную головную боль, когда мы работаем в записной книжке: ее вывод отправляется на stderr , а stderr не отображается в ячейке. Здесь мы воспользуемся уловкой: вместо использования tf.Print мы создадим нашу собственную операцию печати tf.pyfunc через tf.pyfunc . Как и в случае с утверждениями, мы должны убедиться, что наш метод выполняется.

def get_log_lik_verbose(data, n_chains=1):
  # The data argument that is passed in will be available to the inner function
  # below so it doesn't have to be passed in as a parameter.
  replicated_data = np.tile(np.expand_dims(data, axis=1), reps=[1, n_chains, 1])

  def _log_lik(precisions):
    # An internal method we'll make into a TensorFlow operation via tf.py_func
    def _print_precisions(precisions):
      print('precisions:\n', precisions)
      return False  # operations must return something!
    # Turn our method into a TensorFlow operation
    print_op = tf.compat.v1.py_func(_print_precisions, [precisions], tf.bool)

    # Assertions are also operations, and some care needs to be taken to ensure
    # that they're executed
    assert_op = tf.assert_equal(
        precisions, tf.linalg.matrix_transpose(precisions),
        message='not symmetrical', summarize=4, name='symmetry_check')

    # The control_dependencies statement forces its arguments to be executed
    # before subsequent operations
    with tf.control_dependencies([print_op, assert_op]):
      return (log_lik_data(precisions, replicated_data) +
              log_lik_prior(precisions))

  return _log_lik
@tf.function(autograph=False)
def sample():
  tf.random.set_seed(123)
  init_precision = tf.eye(2)[tf.newaxis, ...]
  log_lik_fn = get_log_lik_verbose(my_data)
  # we'll just do a few steps here
  num_results = 10
  num_burnin_steps = 10
  states = tfp.mcmc.sample_chain(
      num_results=num_results,
      num_burnin_steps=num_burnin_steps,
      current_state=[
          init_precision,
      ],
      kernel=tfp.mcmc.HamiltonianMonteCarlo(
          target_log_prob_fn=log_lik_fn,
          step_size=0.1,
          num_leapfrog_steps=3),
      trace_fn=None,
      seed=123)

try:
  states = sample()
except Exception as e:
  # shorten the giant stack trace
  lines = str(e).split('\n')
  print('\n'.join(lines[:5]+['...']+lines[-3:]))
precisions:
 [[[1. 0.]
  [0. 1.]]]
precisions:
 [[[1. 0.]
  [0. 1.]]]
precisions:
 [[[ 0.24315196 -0.2761638 ]
  [-0.33882257  0.8622    ]]]
 assertion failed: [not symmetrical] [Condition x == y did not hold element-wise:] [x (leapfrog_integrate_one_step/add:0) = ] [[[0.243151963 -0.276163787][-0.338822573 0.8622]]] [y (leapfrog_integrate_one_step/maybe_call_fn_and_grads/value_and_gradients/matrix_transpose/transpose:0) = ] [[[0.243151963 -0.338822573][-0.276163787 0.8622]]]
     [[{ {node mcmc_sample_chain/trace_scan/while/body/_96/smart_for_loop/while/body/_381/mh_one_step/hmc_kernel_one_step/leapfrog_integrate/while/body/_503/leapfrog_integrate_one_step/maybe_call_fn_and_grads/value_and_gradients/symmetry_check_1/Assert/AssertGuard/else/_577/Assert} }]] [Op:__inference_sample_4837]

Function call stack:
sample
...
Function call stack:
sample

Почему это не удается

Самым первым новым значением параметра, которое пробует пробоотборник, является асимметричная матрица. Это приводит к сбою разложения Холецкого, поскольку оно определено только для симметричных (и положительно определенных) матриц.

Проблема здесь в том, что наш интересующий параметр представляет собой матрицу точности, а матрицы точности должны быть действительными, симметричными и положительно определенными. Сэмплер ничего не знает об этом ограничении (кроме, возможно, градиентов), поэтому вполне возможно, что сэмплер предложит недопустимое значение, что приведет к исключению, особенно если размер шага велик.

С гамильтоновым сэмплером Монте-Карло мы можем обойти проблему, используя очень маленький размер шага, поскольку градиент должен удерживать параметры вдали от недопустимых областей, но маленькие размеры шага означают медленную сходимость. С семплером Metropolis-Hastings, который ничего не знает о градиентах, мы обречены.

Версия 2: изменение параметров до неограниченных

Существует простое решение проблемы, описанной выше: мы можем изменить параметры нашей модели так, чтобы новые параметры больше не имели этих ограничений. TFP предоставляет полезный набор инструментов - бижекторов - для этого.

Репараметризация бижекторами

Наша матрица точности должна быть реальной и симметричной; нам нужна альтернативная параметризация, не имеющая этих ограничений. Отправной точкой является факторизация Холецкого матрицы точности. Факторы Холецкого по-прежнему ограничены - они имеют нижний треугольник, а их диагональные элементы должны быть положительными. Однако, если мы возьмем логарифм диагоналей фактора Холецкого, то бревна больше не будут ограничены положительными значениями, а затем, если мы сгладим нижнюю треугольную часть в одномерный вектор, у нас больше не будет ограничения нижнего треугольника. . Результатом в нашем случае будет вектор длины 3 без ограничений.

руководстве Stan есть отличная глава об использовании преобразований для снятия различных типов ограничений на параметры.)

Эта повторная параметризация мало влияет на нашу функцию правдоподобия журнала данных - нам просто нужно инвертировать наше преобразование, чтобы получить обратно матрицу точности, - но влияние на предыдущее сложнее. Мы указали, что вероятность данной матрицы точности задается распределением Уишарта; какова вероятность нашей преобразованной матрицы?

Напомним, что если мы применим монотонную функцию $ g $ к одномерной случайной величине $ X $, $ Y = g (X) $, плотность для $ Y $ будет равна

$$ f_Y(y) = | \frac{d}{dy}(g^{-1}(y)) | f_X(g^{-1}(y)) $$

Производная от $ g ^ {- 1} $ term учитывает способ, которым $ g $ изменяет локальные объемы. Для случайных величин более высокой размерности корректирующим фактором является абсолютное значение определителя якобиана $ g ^ {- 1} $ (см. Здесь ).

Нам придется добавить якобиан обратного преобразования в нашу логарифмическую функцию априорного правдоподобия. К счастью, класс Bijector TFP может позаботиться об этом за нас.

Класс Bijector используется для представления обратимых гладких функций, используемых для изменения переменных в функциях плотности вероятности. У всех бижекторов есть метод forward() который выполняет преобразование, метод inverse() который его инвертирует, и forward_log_det_jacobian() и inverse_log_det_jacobian() которые обеспечивают якобианские поправки, которые нам нужны при повторной параметризации PDF.

TFP предоставляет набор полезных бижекторов, которые мы можем комбинировать посредством композиции с помощью оператора Chain для формирования довольно сложных преобразований. В нашем случае составим следующие 3 бижектора (операции в цепочке выполняются справа налево):

  1. Первым шагом нашего преобразования является факторизация Холецкого матрицы точности. Для этого не существует класса Bijector; ОднакоCholeskyOuterProduct bijector принимает продукт 2 Холецких факторов. Мы можем использовать операцию, обратную этой операции, с помощью оператора Invert .
  2. Следующим шагом будет логарифм диагональных элементов фактора Холецкого. Мы достигаем этого через биектор TransformDiagonal и инверсию биектора Exp .
  3. Наконец, мы сглаживаем нижнюю треугольную часть матрицы до вектора, используя инверсию FillTriangular FillTriangular.
# Our transform has 3 stages that we chain together via composition:
precision_to_unconstrained = tfb.Chain([
    # step 3: flatten the lower triangular portion of the matrix
    tfb.Invert(tfb.FillTriangular(validate_args=VALIDATE_ARGS)),
    # step 2: take the log of the diagonals    
    tfb.TransformDiagonal(tfb.Invert(tfb.Exp(validate_args=VALIDATE_ARGS))),
    # step 1: decompose the precision matrix into its Cholesky factors
    tfb.Invert(tfb.CholeskyOuterProduct(validate_args=VALIDATE_ARGS)),
])
# sanity checks
m = tf.constant([[1., 2.], [2., 8.]])
m_fwd = precision_to_unconstrained.forward(m)
m_inv = precision_to_unconstrained.inverse(m_fwd)

# bijectors handle tensors of values, too!
m2 = tf.stack([m, tf.eye(2)])
m2_fwd = precision_to_unconstrained.forward(m2)
m2_inv = precision_to_unconstrained.inverse(m2_fwd)

print('single input:')
print('m:\n', m.numpy())
print('precision_to_unconstrained(m):\n', m_fwd.numpy())
print('inverse(precision_to_unconstrained(m)):\n', m_inv.numpy())
print()

print('tensor of inputs:')
print('m2:\n', m2.numpy())
print('precision_to_unconstrained(m2):\n', m2_fwd.numpy())
print('inverse(precision_to_unconstrained(m2)):\n', m2_inv.numpy())
single input:
m:
 [[1. 2.]
 [2. 8.]]
precision_to_unconstrained(m):
 [0.6931472 2.        0.       ]
inverse(precision_to_unconstrained(m)):
 [[1. 2.]
 [2. 8.]]

tensor of inputs:
m2:
 [[[1. 2.]
  [2. 8.]]

 [[1. 0.]
  [0. 1.]]]
precision_to_unconstrained(m2):
 [[0.6931472 2.        0.       ]
 [0.        0.        0.       ]]
inverse(precision_to_unconstrained(m2)):
 [[[1. 2.]
  [2. 8.]]

 [[1. 0.]
  [0. 1.]]]

Класс TransformedDistribution автоматизирует процесс применения биектора к распределению и внесения необходимой поправки Якоби в log_prob() . Наш новый приор становится:

def log_lik_prior_transformed(transformed_precisions):
  rv_precision = tfd.TransformedDistribution(
      tfd.WishartTriL(
          df=PRIOR_DF,
          scale_tril=tf.linalg.cholesky(PRIOR_SCALE),
          validate_args=VALIDATE_ARGS,
          allow_nan_stats=ALLOW_NAN_STATS),
      bijector=precision_to_unconstrained,
      validate_args=VALIDATE_ARGS)
  return rv_precision.log_prob(transformed_precisions)
# Check against the numpy implementation.  Note that when comparing, we need
# to add in the Jacobian correction.
precisions = np.stack([np.eye(2, dtype=np.float32), true_precision])
transformed_precisions = precision_to_unconstrained.forward(precisions)
lik_tf = log_lik_prior_transformed(transformed_precisions).numpy()
corrections = precision_to_unconstrained.inverse_log_det_jacobian(
    transformed_precisions, event_ndims=1).numpy()
n = precisions.shape[0]

for i in range(n):
  print(i)
  print('numpy:', log_lik_prior_numpy(precisions[i]) + corrections[i])
  print('tensorflow:', lik_tf[i])
0
numpy: -0.8488930160357633
tensorflow: -0.84889317
1
numpy: -7.305657151741624
tensorflow: -7.305659

Нам просто нужно инвертировать преобразование для вероятности нашего журнала данных:

precision = precision_to_unconstrained.inverse(transformed_precision)

Поскольку нам действительно нужна факторизация Холецкого матрицы точности, было бы более эффективно сделать здесь частичное обратное. Однако мы оставим оптимизацию на потом, а частичное обратное оставим читателю в качестве упражнения.

def log_lik_data_transformed(transformed_precisions, replicated_data):
  # We recover the precision matrix by inverting our bijector.  This is
  # inefficient since we really want the Cholesky decomposition of the
  # precision matrix, and the bijector has that in hand during the inversion,
  # but we'll worry about efficiency later.
  n = tf.shape(transformed_precisions)[0]
  precisions = precision_to_unconstrained.inverse(transformed_precisions)
  precisions_cholesky = tf.linalg.cholesky(precisions)
  covariances = tf.linalg.cholesky_solve(
      precisions_cholesky, tf.linalg.eye(2, batch_shape=[n]))
  rv_data = tfd.MultivariateNormalFullCovariance(
      loc=tf.zeros([n, 2]),
      covariance_matrix=covariances,
      validate_args=VALIDATE_ARGS,
      allow_nan_stats=ALLOW_NAN_STATS)

  return tf.reduce_sum(rv_data.log_prob(replicated_data), axis=0)
# sanity check
precisions = np.stack([np.eye(2, dtype=np.float32), true_precision])
transformed_precisions = precision_to_unconstrained.forward(precisions)
lik_tf = log_lik_data_transformed(
    transformed_precisions, replicated_data).numpy()

for i in range(precisions.shape[0]):
  print(i)
  print('numpy:', log_lik_data_numpy(precisions[i], my_data))
  print('tensorflow:', lik_tf[i])
0
numpy: -430.71218815801365
tensorflow: -430.71207
1
numpy: -280.81822950593767
tensorflow: -280.8182

Снова мы завершаем наши новые функции закрытием.

def get_log_lik_transformed(data, n_chains=1):
  # The data argument that is passed in will be available to the inner function
  # below so it doesn't have to be passed in as a parameter.
  replicated_data = np.tile(np.expand_dims(data, axis=1), reps=[1, n_chains, 1])

  @tf.function(autograph=False)
  def _log_lik_transformed(transformed_precisions):
    return (log_lik_data_transformed(transformed_precisions, replicated_data) +
            log_lik_prior_transformed(transformed_precisions))

  return _log_lik_transformed
# make sure everything runs
log_lik_fn = get_log_lik_transformed(my_data)
m = tf.eye(2)[tf.newaxis, ...]
lik = log_lik_fn(precision_to_unconstrained.forward(m)).numpy()
print(lik)
[-431.5611]

Отбор проб

Теперь, когда нам не нужно беспокоиться о взрыве нашего сэмплера из-за недопустимых значений параметров, давайте сгенерируем несколько реальных сэмплов.

Сэмплер работает с неограниченной версией наших параметров, поэтому нам нужно преобразовать наше начальное значение в его неограниченную версию. Сэмплы, которые мы генерируем, также будут в неограниченной форме, поэтому нам нужно преобразовать их обратно. Бижекторы векторизованы, поэтому это легко сделать.

# We'll choose a proper random initial value this time
np.random.seed(123)
initial_value_cholesky = np.array(
    [[0.5 + np.random.uniform(), 0.0],
     [-0.5 + np.random.uniform(), 0.5 + np.random.uniform()]],
    dtype=np.float32)
initial_value =  initial_value_cholesky.dot(
  initial_value_cholesky.T)[np.newaxis, ...]

# The sampler works with unconstrained values, so we'll transform our initial
# value
initial_value_transformed = precision_to_unconstrained.forward(
  initial_value).numpy()
# Sample!
@tf.function(autograph=False)
def sample():
  tf.random.set_seed(123)
  log_lik_fn = get_log_lik_transformed(my_data, n_chains=1)

  num_results = 1000
  num_burnin_steps = 1000

  states, is_accepted = tfp.mcmc.sample_chain(
      num_results=num_results,
      num_burnin_steps=num_burnin_steps,
      current_state=[
          initial_value_transformed,
      ],
      kernel=tfp.mcmc.HamiltonianMonteCarlo(
          target_log_prob_fn=log_lik_fn,
          step_size=0.1,
          num_leapfrog_steps=3),
      trace_fn=lambda _, pkr: pkr.is_accepted,
      seed=123)
  # transform samples back to their constrained form
  precision_samples = [precision_to_unconstrained.inverse(s) for s in states]
  return states, precision_samples, is_accepted

states, precision_samples, is_accepted = sample()

Давайте сравним среднее значение выходных данных нашего пробоотборника с аналитическим апостериорным средним!

print('True posterior mean:\n', posterior_mean)
print('Sample mean:\n', np.mean(np.reshape(precision_samples, [-1, 2, 2]), axis=0))
True posterior mean:
 [[ 0.9641779 -1.6534661]
 [-1.6534661  3.8683164]]
Sample mean:
 [[ 1.4315274  -0.25587553]
 [-0.25587553  0.5740424 ]]

Мы далеко! Разберемся почему. Сначала посмотрим на наши образцы.

np.reshape(precision_samples, [-1, 2, 2])
array([[[ 1.4315385, -0.2558777],
        [-0.2558777,  0.5740494]],

       [[ 1.4315385, -0.2558777],
        [-0.2558777,  0.5740494]],

       [[ 1.4315385, -0.2558777],
        [-0.2558777,  0.5740494]],

       ...,

       [[ 1.4315385, -0.2558777],
        [-0.2558777,  0.5740494]],

       [[ 1.4315385, -0.2558777],
        [-0.2558777,  0.5740494]],

       [[ 1.4315385, -0.2558777],
        [-0.2558777,  0.5740494]]], dtype=float32)

Ой, похоже, все они имеют одинаковую ценность. Разберемся почему.

Переменная kernel_results_ - это именованный кортеж, который дает информацию о сэмплере в каждом состоянии. Ключевым здесь является поле is_accepted .

# Look at the acceptance for the last 100 samples
print(np.squeeze(is_accepted)[-100:])
print('Fraction of samples accepted:', np.mean(np.squeeze(is_accepted)))
[False False False False False False False False False False False False
 False False False False False False False False False False False False
 False False False False False False False False False False False False
 False False False False False False False False False False False False
 False False False False False False False False False False False False
 False False False False False False False False False False False False
 False False False False False False False False False False False False
 False False False False False False False False False False False False
 False False False False]
Fraction of samples accepted: 0.0

Все наши образцы были отклонены! Предположительно наш размер шага был слишком большим. Я stepsize=0.1 выбрал stepsize=0.1 .

Версия 3: выборка с адаптивным размером шага

Поскольку выборка с произвольным выбором размера шага не удалась, у нас есть несколько пунктов повестки дня:

  1. реализовать адаптивный размер шага и
  2. выполнить некоторые проверки сходимости.

В tensorflow_probability/python/mcmc/hmc.py есть хороший пример кода для реализации адаптивных размеров шага. Я адаптировал его ниже.

Обратите внимание, что для каждого шага существует отдельный sess.run() . Это действительно полезно для отладки, так как позволяет при необходимости легко добавлять пошаговую диагностику. Например, мы можем показать постепенный прогресс, время каждого шага и т. Д.

Совет: Один из распространенных способов испортить вашу выборку - заставить ваш график расти в цикле. (Причина завершения графика перед запуском сеанса состоит в том, чтобы предотвратить именно такие проблемы.) Однако, если вы не использовали finalize (), полезная проверка отладки, если ваш код замедляется до сканирования, - это распечатать график size на каждом шаге через len(mygraph.get_operations()) - если длина увеличивается, вы, вероятно, делаете что-то плохое.

Мы собираемся запустить здесь 3 независимых цепочки. Проведение некоторых сравнений между цепочками поможет нам проверить сходимость.

# The number of chains is determined by the shape of the initial values.
# Here we'll generate 3 chains, so we'll need a tensor of 3 initial values.
N_CHAINS = 3

np.random.seed(123)

initial_values = []
for i in range(N_CHAINS):
  initial_value_cholesky = np.array(
      [[0.5 + np.random.uniform(), 0.0],
       [-0.5 + np.random.uniform(), 0.5 + np.random.uniform()]],
      dtype=np.float32)
  initial_values.append(initial_value_cholesky.dot(initial_value_cholesky.T))
initial_values = np.stack(initial_values)

initial_values_transformed = precision_to_unconstrained.forward(
  initial_values).numpy()
@tf.function(autograph=False)
def sample():
  tf.random.set_seed(123)
  log_lik_fn = get_log_lik_transformed(my_data)

  # Tuning acceptance rates:
  dtype = np.float32
  num_burnin_iter = 3000
  num_warmup_iter = int(0.8 * num_burnin_iter) 
  num_chain_iter = 2500

  # Set the target average acceptance ratio for the HMC as suggested by
  # Beskos et al. (2013):
  # https://projecteuclid.org/download/pdfview_1/euclid.bj/1383661192
  target_accept_rate = 0.651

  # Initialize the HMC sampler.
  hmc = tfp.mcmc.HamiltonianMonteCarlo(
      target_log_prob_fn=log_lik_fn,
      step_size=0.01,
      num_leapfrog_steps=3)

  # Adapt the step size using standard adaptive MCMC procedure. See Section 4.2
  # of Andrieu and Thoms (2008):
  # http://www4.ncsu.edu/~rsmith/MA797V_S12/Andrieu08_AdaptiveMCMC_Tutorial.pdf
  adapted_kernel = tfp.mcmc.SimpleStepSizeAdaptation(
      inner_kernel=hmc,
      num_adaptation_steps=num_warmup_iter,
      target_accept_prob=target_accept_rate)

  states, is_accepted = tfp.mcmc.sample_chain(
      num_results=num_chain_iter,
      num_burnin_steps=num_burnin_iter,
      current_state=initial_values_transformed,
      kernel=adapted_kernel,
      trace_fn=lambda _, pkr: pkr.inner_results.is_accepted,
      parallel_iterations=1)
  # transform samples back to their constrained form
  precision_samples = precision_to_unconstrained.inverse(states)
  return states, precision_samples, is_accepted

states, precision_samples, is_accepted = sample()

Быстрая проверка: наш показатель приемлемости во время выборки близок к нашему целевому значению 0,651.

print(np.mean(is_accepted))
0.6190666666666667

Более того, наше выборочное среднее и стандартное отклонение близки к тому, что мы ожидаем от аналитического решения.

precision_samples_reshaped = np.reshape(precision_samples, [-1, 2, 2])
print('True posterior mean:\n', posterior_mean)
print('Mean of samples:\n', np.mean(precision_samples_reshaped, axis=0))
True posterior mean:
 [[ 0.9641779 -1.6534661]
 [-1.6534661  3.8683164]]
Mean of samples:
 [[ 0.96426415 -1.6519215 ]
 [-1.6519215   3.8614824 ]]
print('True posterior standard deviation:\n', posterior_sd)
print('Standard deviation of samples:\n', np.std(precision_samples_reshaped, axis=0))
True posterior standard deviation:
 [[0.13435492 0.25050813]
 [0.25050813 0.53903675]]
Standard deviation of samples:
 [[0.13622096 0.25235635]
 [0.25235635 0.5394968 ]]

Проверка сходимости

В общем, у нас не будет аналитического решения для проверки, поэтому нам нужно убедиться, что сэмплер сходится. Одной из стандартных проверок является статистика $ \ hat {R} $ Гельмана-Рубина, которая требует нескольких цепочек выборки. $ \ hat {R} $ измеряет степень, в которой дисперсия (средних) между цепочками превышает то, что можно было бы ожидать, если бы цепи были распределены одинаково. Значения $ \ hat {R} $, близкие к 1, используются для обозначения приблизительной сходимости. Подробности см. В источнике .

r_hat = tfp.mcmc.potential_scale_reduction(precision_samples).numpy()
print(r_hat)
[[1.0038308 1.0005717]
 [1.0005717 1.0006068]]

Модельная критика

Если бы у нас не было аналитического решения, пришло бы время сделать реальную критику модели.

Вот несколько быстрых гистограмм компонентов выборки относительно нашей базовой истины (выделены красным). Обратите внимание, что выборки были сокращены от значений матрицы выборочной точности до единичной матрицы ранее.

fig, axes = plt.subplots(2, 2, sharey=True)
fig.set_size_inches(8, 8)
for i in range(2):
  for j in range(2):
    ax = axes[i, j]
    ax.hist(precision_samples_reshaped[:, i, j])
    ax.axvline(true_precision[i, j], color='red',
               label='True precision')
    ax.axvline(sample_precision[i, j], color='red', linestyle=':',
               label='Sample precision')
    ax.set_title('precision[%d, %d]' % (i, j))
plt.tight_layout()
plt.legend()
plt.show()

PNG

Некоторые диаграммы разброса пар компонентов точности показывают, что из-за корреляционной структуры апостериорных значений истинные апостериорные значения не так маловероятны, как они выглядят из маргинальных чисел выше.

fig, axes = plt.subplots(4, 4)
fig.set_size_inches(12, 12)
for i1 in range(2):
  for j1 in range(2):
    index1 = 2 * i1 + j1
    for i2 in range(2):
      for j2 in range(2):
        index2 = 2 * i2 + j2
        ax = axes[index1, index2]
        ax.scatter(precision_samples_reshaped[:, i1, j1],
                   precision_samples_reshaped[:, i2, j2], alpha=0.1)
        ax.axvline(true_precision[i1, j1], color='red')
        ax.axhline(true_precision[i2, j2], color='red')
        ax.axvline(sample_precision[i1, j1], color='red', linestyle=':')
        ax.axhline(sample_precision[i2, j2], color='red', linestyle=':')
        ax.set_title('(%d, %d) vs (%d, %d)' % (i1, j1, i2, j2))
plt.tight_layout()
plt.show()

PNG

Версия 4: более простая выборка ограниченных параметров

Bijectors упростили выборку прецизионной матрицы, но потребовалось изрядное количество ручных преобразований в неограниченное представление и обратно. Есть способ попроще!

Ядро TransformedTransition

TransformedTransitionKernel упрощает этот процесс. Он обертывает ваш сэмплер и обрабатывает все преобразования. Он принимает в качестве аргумента список бижекторов, которые сопоставляют неограниченные значения параметров с ограниченными. Итак, здесь нам понадобится инверсия бижектора precision_to_unconstrained который мы использовали выше. Мы могли бы просто использовать tfb.Invert(precision_to_unconstrained) , но это потребовало бы взятия инверсий обратных (TensorFlow недостаточно умен, чтобы упростить tf.Invert(tf.Invert()) до tf.Identity()) , поэтому вместо этого мы Просто напишу новый бижектор.

Сдерживающий биектор

# The bijector we need for the TransformedTransitionKernel is the inverse of
# the one we used above
unconstrained_to_precision = tfb.Chain([
    # step 3: take the product of Cholesky factors
    tfb.CholeskyOuterProduct(validate_args=VALIDATE_ARGS),
    # step 2: exponentiate the diagonals    
    tfb.TransformDiagonal(tfb.Exp(validate_args=VALIDATE_ARGS)),
    # step 1: map a vector to a lower triangular matrix
    tfb.FillTriangular(validate_args=VALIDATE_ARGS),
])
# quick sanity check
m = [[1., 2.], [2., 8.]]
m_inv = unconstrained_to_precision.inverse(m).numpy()
m_fwd = unconstrained_to_precision.forward(m_inv).numpy()

print('m:\n', m)
print('unconstrained_to_precision.inverse(m):\n', m_inv)
print('forward(unconstrained_to_precision.inverse(m)):\n', m_fwd)
m:
 [[1.0, 2.0], [2.0, 8.0]]
unconstrained_to_precision.inverse(m):
 [0.6931472 2.        0.       ]
forward(unconstrained_to_precision.inverse(m)):
 [[1. 2.]
 [2. 8.]]

Выборка с помощью TransformedTransitionKernel

С TransformedTransitionKernel нам больше не нужно вручную преобразовывать наши параметры. Наши начальные значения и наши образцы - все матрицы точности; нам просто нужно передать наш неограничивающий бижектор (ы) ядру, и оно позаботится обо всех преобразованиях.

@tf.function(autograph=False)
def sample():
  tf.random.set_seed(123)
  log_lik_fn = get_log_lik(my_data)

  # Tuning acceptance rates:
  dtype = np.float32
  num_burnin_iter = 3000
  num_warmup_iter = int(0.8 * num_burnin_iter) 
  num_chain_iter = 2500

  # Set the target average acceptance ratio for the HMC as suggested by
  # Beskos et al. (2013):
  # https://projecteuclid.org/download/pdfview_1/euclid.bj/1383661192
  target_accept_rate = 0.651

  # Initialize the HMC sampler.
  hmc = tfp.mcmc.HamiltonianMonteCarlo(
      target_log_prob_fn=log_lik_fn,
      step_size=0.01,
      num_leapfrog_steps=3)

  ttk = tfp.mcmc.TransformedTransitionKernel(
      inner_kernel=hmc, bijector=unconstrained_to_precision)

  # Adapt the step size using standard adaptive MCMC procedure. See Section 4.2
  # of Andrieu and Thoms (2008):
  # http://www4.ncsu.edu/~rsmith/MA797V_S12/Andrieu08_AdaptiveMCMC_Tutorial.pdf
  adapted_kernel = tfp.mcmc.SimpleStepSizeAdaptation(
      inner_kernel=ttk,
      num_adaptation_steps=num_warmup_iter,
      target_accept_prob=target_accept_rate)

  states = tfp.mcmc.sample_chain(
      num_results=num_chain_iter,
      num_burnin_steps=num_burnin_iter,
      current_state=initial_values,
      kernel=adapted_kernel,
      trace_fn=None,
      parallel_iterations=1)
  # transform samples back to their constrained form
  return states

precision_samples  = sample()

Проверка сходимости

Проверка сходимости $ \ hat {R} $ выглядит неплохо!

r_hat = tfp.mcmc.potential_scale_reduction(precision_samples).numpy()
print(r_hat)
[[1.0013582 1.0019467]
 [1.0019467 1.0011805]]

Сравнение с аналитическим апостериорным

Снова проверим с аналитическим апостериорным.

# The output samples have shape [n_steps, n_chains, 2, 2]
# Flatten them to [n_steps * n_chains, 2, 2] via reshape:
precision_samples_reshaped = np.reshape(precision_samples, [-1, 2, 2])
print('True posterior mean:\n', posterior_mean)
print('Mean of samples:\n', np.mean(precision_samples_reshaped, axis=0))
True posterior mean:
 [[ 0.9641779 -1.6534661]
 [-1.6534661  3.8683164]]
Mean of samples:
 [[ 0.96687526 -1.6552585 ]
 [-1.6552585   3.867676  ]]
print('True posterior standard deviation:\n', posterior_sd)
print('Standard deviation of samples:\n', np.std(precision_samples_reshaped, axis=0))
True posterior standard deviation:
 [[0.13435492 0.25050813]
 [0.25050813 0.53903675]]
Standard deviation of samples:
 [[0.13329624 0.24913791]
 [0.24913791 0.53983927]]

Оптимизация

Теперь, когда у нас есть сквозная работа, давайте сделаем более оптимизированную версию. Скорость не имеет большого значения для этого примера, но как только матрицы станут больше, несколько оптимизаций будут иметь большое значение.

Одно большое улучшение скорости, которое мы можем сделать, - это изменить параметры в терминах разложения Холецкого. Причина в том, что для нашей функции правдоподобия данных требуются как ковариационная, так и прецизионная матрицы. Инверсия матриц обходится дорого ($ O (n ^ 3) $ для матрицы $ n \ times n $), и если мы параметризуем либо ковариацию, либо матрицу точности, нам нужно сделать инверсию, чтобы получить другую.

Напомним, что вещественная положительно определенная симметричная матрица $ M $ может быть разложена в произведение вида $ M = LL ^ T $, где матрица $ L $ имеет нижнюю треугольную форму и положительные диагонали. Учитывая разложение Холецкого матрицы $ M $, мы можем более эффективно получить как $ M $ (произведение нижней и верхней треугольной матрицы), так и $ M ^ {- 1} $ (с помощью обратной подстановки). Сама по себе факторизация Холецкого обходится недешево, но если мы параметризуем с точки зрения факторов Холецкого, нам нужно только вычислить факторизацию Холексы для начальных значений параметров.

Использование разложения Холецкого ковариационной матрицы

TFP имеет версию многомерного нормального распределения MultivariateNormalTriL , которая параметризована в терминах фактора Холецкого ковариационной матрицы. Так что, если бы мы параметризовали ковариационную матрицу с помощью коэффициента Холецкого, мы могли бы эффективно вычислить логарифмическое правдоподобие. Проблема состоит в том, чтобы вычислить предыдущую логарифмическую вероятность с аналогичной эффективностью.

Если бы у нас была версия обратного распределения Вишарта, которая работала бы с факторами Холецкого выборок, все было бы готово. Увы, нет. (Однако команда приветствовала бы отправку кода!) В качестве альтернативы мы можем использовать версию дистрибутива Wishart, которая работает с факторами Холецкого образцов вместе с цепочкой бижекторов.

На данный момент нам не хватает нескольких стандартных бижекторов, чтобы сделать вещи действительно эффективными, но я хочу показать этот процесс как упражнение и полезную иллюстрацию силы бижекторов TFP.

Распределение Уишарта, работающее на факторах Холецкого

В распределении Wishart есть полезный флаг input_output_cholesky , который указывает, что матрицы ввода и вывода должны быть факторами Холецкого. Работать с факторами Холецкого более эффективно и численно выгоднее, чем с полными матрицами, поэтому это желательно. Важный момент о семантике флага: это всего лишь указание на то, что представление входных и выходных данных для распределения должно измениться - это не указывает на полную репараметрацию распределения, которая потребовала бы исправления Якоби для log_prob() функция. На самом деле мы хотим полностью изменить параметры, поэтому создадим собственный дистрибутив.

# An optimized Wishart distribution that has been transformed to operate on
# Cholesky factors instead of full matrices.  Note that we gain a modest
# additional speedup by specifying the Cholesky factor of the scale matrix
# (i.e. by passing in the scale_tril parameter instead of scale).

class CholeskyWishart(tfd.TransformedDistribution):
  """Wishart distribution reparameterized to use Cholesky factors."""
  def __init__(self,
      df,
      scale_tril,
      validate_args=False,
      allow_nan_stats=True,
      name='CholeskyWishart'):
    # Wishart has a bunch of methods that we want to support but not
    # implement.  We'll subclass TransformedDistribution here to take care of
    # those.  We'll override the few for which speed is critical and implement
    # them with a separate Wishart for which input_output_cholesky=True
    super(CholeskyWishart, self).__init__(
        distribution=tfd.WishartTriL(
            df=df,
            scale_tril=scale_tril,
            input_output_cholesky=False,
            validate_args=validate_args,
            allow_nan_stats=allow_nan_stats),
        bijector=tfb.Invert(tfb.CholeskyOuterProduct()),
        validate_args=validate_args,
        name=name
    )
    # Here's the Cholesky distribution we'll use for log_prob() and sample()
    self.cholesky = tfd.WishartTriL(
        df=df,
        scale_tril=scale_tril,
        input_output_cholesky=True,
        validate_args=validate_args,
        allow_nan_stats=allow_nan_stats)

  def _log_prob(self, x):
    return (self.cholesky.log_prob(x) +
            self.bijector.inverse_log_det_jacobian(x, event_ndims=2))

  def _sample_n(self, n, seed=None):
    return self.cholesky._sample_n(n, seed)
# some checks
PRIOR_SCALE_CHOLESKY = np.linalg.cholesky(PRIOR_SCALE)

@tf.function(autograph=False)
def compute_log_prob(m):
  w_transformed = tfd.TransformedDistribution(
      tfd.WishartTriL(df=PRIOR_DF, scale_tril=PRIOR_SCALE_CHOLESKY),
      bijector=tfb.Invert(tfb.CholeskyOuterProduct()))
  w_optimized = CholeskyWishart(
      df=PRIOR_DF, scale_tril=PRIOR_SCALE_CHOLESKY)
  log_prob_transformed = w_transformed.log_prob(m)
  log_prob_optimized = w_optimized.log_prob(m)
  return log_prob_transformed, log_prob_optimized

for matrix in [np.eye(2, dtype=np.float32),
               np.array([[1., 0.], [2., 8.]], dtype=np.float32)]:
  log_prob_transformed, log_prob_optimized = [
      t.numpy() for t in compute_log_prob(matrix)]
  print('Transformed Wishart:', log_prob_transformed)
  print('Optimized Wishart', log_prob_optimized)
Transformed Wishart: -0.84889317
Optimized Wishart -0.84889317
Transformed Wishart: -99.269455
Optimized Wishart -99.269455

Построение обратного распределения Вишарта

Наша ковариационная матрица $ C $ разложена на $ C = LL ^ T $, где $ L $ имеет нижнюю треугольную форму и положительную диагональ. Мы хотим знать вероятность $ L $ при условии, что $ C \ sim W ^ {- 1} (\ nu, V) $, где $ W ^ {- 1} $ - обратное распределение Уишарта.

Обратное распределение Уишарта обладает тем свойством, что если $ C \ sim W ^ {- 1} (\ nu, V) $, то матрица точности $ C ^ {- 1} \ sim W (\ nu, V ^ {- 1 }) $. Таким образом, мы можем получить вероятность $ L $ через TransformedDistribution которая принимает в качестве параметров распределение Уишарта и биектор, отображающий фактор Холецкого матрицы точности в фактор Холецкого его обратного.

Простой (но не суперэффективный) способ перейти от фактора Холецкого $ C ^ {- 1} $ к $ L $ - это инвертировать фактор Холецкого путем обратного решения, затем сформировать ковариационную матрицу из этих инвертированных факторов и затем делаем факторизацию Холецкого.

Пусть разложение Холецкого $ C ^ {- 1} = MM ^ T $. $ M $ является нижним треугольником, поэтому мы можем инвертировать его с помощью MatrixInverseTriL MatrixInverseTriL.

Сформировать $ C $ из $ M ^ {- 1} $ немного сложно: $ C = (MM ^ T) ^ {- 1} = M ^ {- T} M ^ {- 1} = M ^ {- T } (M ^ {- T}) ^ T $. $ M $ является нижним треугольником, поэтому $ M ^ {- 1} $ также будет нижним треугольником, а $ M ^ {- T} $ будет верхним треугольником. Биектор CholeskyOuterProduct CholeskyOuterProduct() работает только с нижнетреугольными матрицами, поэтому мы не можем использовать его для формирования $ C $ из $ M ^ {- T} $. Наше обходное решение - это цепочка бижекторов, которые переставляют строки и столбцы матрицы.

К счастью , эта логика заключена в CholeskyToInvCholesky bijector!

Объединение всех частей

# verify that the bijector works
m = np.array([[1., 0.], [2., 8.]], dtype=np.float32)
c_inv = m.dot(m.T)
c = np.linalg.inv(c_inv)
c_chol = np.linalg.cholesky(c)
wishart_cholesky_to_iw_cholesky = tfb.CholeskyToInvCholesky()
w_fwd = wishart_cholesky_to_iw_cholesky.forward(m).numpy()

print('numpy =\n', c_chol)
print('bijector =\n', w_fwd)
numpy =
 [[ 1.0307764   0.        ]
 [-0.03031695  0.12126781]]
bijector =
 [[ 1.0307764   0.        ]
 [-0.03031695  0.12126781]]

Наш финальный дистрибутив

Наша инверсия Уишарта, оперирующая факторами Холецкого, выглядит следующим образом:

inverse_wishart_cholesky = tfd.TransformedDistribution(
    distribution=CholeskyWishart(
        df=PRIOR_DF,
        scale_tril=np.linalg.cholesky(np.linalg.inv(PRIOR_SCALE))),
    bijector=tfb.CholeskyToInvCholesky())

У нас есть обратный Wishart, но он немного медленный, потому что нам нужно выполнить разложение Холецкого в биекторе. Вернемся к параметризации матрицы точности и посмотрим, что мы можем там сделать для оптимизации.

Финальная (!) Версия: использование разложения Холецкого матрицы точности

Альтернативный подход - работать с коэффициентами Холецкого матрицы точности. Здесь априорную функцию правдоподобия легко вычислить, но функция правдоподобия журнала данных требует больше работы, поскольку TFP не имеет версии многомерной нормали, параметризованной по точности.

Оптимизированная вероятность предыдущего журнала

Мы используем дистрибутив CholeskyWishart мы построили выше, чтобы построить предыдущий.

# Our new prior.
PRIOR_SCALE_CHOLESKY = np.linalg.cholesky(PRIOR_SCALE)

def log_lik_prior_cholesky(precisions_cholesky):
  rv_precision = CholeskyWishart(
      df=PRIOR_DF,
      scale_tril=PRIOR_SCALE_CHOLESKY,
      validate_args=VALIDATE_ARGS,
      allow_nan_stats=ALLOW_NAN_STATS)
  return rv_precision.log_prob(precisions_cholesky)
# Check against the slower TF implementation and the NumPy implementation.
# Note that when comparing to NumPy, we need to add in the Jacobian correction.
precisions = [np.eye(2, dtype=np.float32),
              true_precision]
precisions_cholesky = np.stack([np.linalg.cholesky(m) for m in precisions])
precisions = np.stack(precisions)
lik_tf = log_lik_prior_cholesky(precisions_cholesky).numpy()
lik_tf_slow = tfd.TransformedDistribution(
    distribution=tfd.WishartTriL(
        df=PRIOR_DF, scale_tril=tf.linalg.cholesky(PRIOR_SCALE)),
    bijector=tfb.Invert(tfb.CholeskyOuterProduct())).log_prob(
    precisions_cholesky).numpy()
corrections = tfb.Invert(tfb.CholeskyOuterProduct()).inverse_log_det_jacobian(
    precisions_cholesky, event_ndims=2).numpy()
n = precisions.shape[0]

for i in range(n):
  print(i)
  print('numpy:', log_lik_prior_numpy(precisions[i]) + corrections[i])
  print('tensorflow slow:', lik_tf_slow[i])
  print('tensorflow fast:', lik_tf[i])
0
numpy: -0.8488930160357633
tensorflow slow: -0.84889317
tensorflow fast: -0.84889317
1
numpy: -7.442875031036973
tensorflow slow: -7.442877
tensorflow fast: -7.442876

Оптимизированная вероятность регистрации данных

Мы можем использовать бижекторы TFP, чтобы построить нашу собственную версию многомерной нормали. Вот ключевая идея:

Предположим, у меня есть вектор-столбец $ X $, элементы которого являются iid-образцами $ N (0, 1) $. У нас есть $ \ text {mean} (X) = 0 $ и $ \ text {cov} (X) = I $.

Пусть теперь $ Y = AX + b $. У нас есть $ \ text {mean} (Y) = b $ и $ \ text {cov} (Y) = AA ^ T $.

Следовательно, мы можем создавать векторы со средним значением $ b $ и ковариацией $ C $, используя аффинное преобразование $ Ax + b $ в векторы стандартных нормальных выборок iid при условии $ AA ^ T = C $. Разложение Холецкого $ C $ обладает желаемым свойством. Однако есть и другие решения.

Пусть $ P = C ^ {- 1} $ и разложение Холецкого для $ P $ имеет вид $ B $, т.е. $ BB ^ T = P $. Сейчас

$ P ^ {- 1} = (BB ^ T) ^ {- 1} = B ^ {- T} B ^ {- 1} = B ^ {- T} (B ^ {- T}) ^ T $

Итак, еще один способ получить желаемое среднее значение и ковариацию - использовать аффинное преобразование $ Y = B ^ {- T} X + b $.

Наш подход (любезно предоставлен этой записной книжкой ):

  1. Используйте tfd.Independent() чтобы объединить пакет одномерных Normal случайных величин в одну многомерную случайную величину. Параметр reinterpreted_batch_ndims для Independent() указывает количество измерений пакета, которые следует переинтерпретировать как измерения событий. In our case we create a 1-D batch of length 2 that we transform into a 1-D event of length 2, so reinterpreted_batch_ndims=1 .
  2. Apply a bijector to add the desired covariance: tfb.Invert(tfb.ScaleMatvecTriL(scale_tril=precision_cholesky, adjoint=True)) . Note that above we're multiplying our iid normal random variables by the transpose of the inverse of the Cholesky factor of the precision matrix $(B^{-T}X)$. The tfb.Invert takes care of inverting $B$, and the adjoint=True flag performs the transpose.
  3. Apply a bijector to add the desired offset: tfb.Shift(shift=shift) Note that we have to do the shift as a separate step from the initial inverted affine transform because otherwise the inverted scale is applied to the shift (since the inverse of $y=Ax+b$ is $x=A^{-1}y - A^{-1}b$).
class MVNPrecisionCholesky(tfd.TransformedDistribution):
  """Multivariate normal parameterized by loc and Cholesky precision matrix."""

  def __init__(self, loc, precision_cholesky, name=None):
    super(MVNPrecisionCholesky, self).__init__(
        distribution=tfd.Independent(
            tfd.Normal(loc=tf.zeros_like(loc),
                       scale=tf.ones_like(loc)),
            reinterpreted_batch_ndims=1),
        bijector=tfb.Chain([
            tfb.Shift(shift=loc),
            tfb.Invert(tfb.ScaleMatvecTriL(scale_tril=precision_cholesky,
                                  adjoint=True)),
        ]),
        name=name)
@tf.function(autograph=False)
def log_lik_data_cholesky(precisions_cholesky, replicated_data):
  n = tf.shape(precisions_cholesky)[0]  # number of precision matrices
  rv_data = MVNPrecisionCholesky(
      loc=tf.zeros([n, 2]),
      precision_cholesky=precisions_cholesky)
  return tf.reduce_sum(rv_data.log_prob(replicated_data), axis=0)
# check against the numpy implementation
true_precision_cholesky = np.linalg.cholesky(true_precision)
precisions = [np.eye(2, dtype=np.float32), true_precision]
precisions_cholesky = np.stack([np.linalg.cholesky(m) for m in precisions])
precisions = np.stack(precisions)
n = precisions_cholesky.shape[0]
replicated_data = np.tile(np.expand_dims(my_data, axis=1), reps=[1, 2, 1])
lik_tf = log_lik_data_cholesky(precisions_cholesky, replicated_data).numpy()

for i in range(n):
  print(i)
  print('numpy:', log_lik_data_numpy(precisions[i], my_data))
  print('tensorflow:', lik_tf[i])
0
numpy: -430.71218815801365
tensorflow: -430.71207
1
numpy: -280.81822950593767
tensorflow: -280.81824

Combined log likelihood function

Now we combine our prior and data log likelihood functions in a closure.

def get_log_lik_cholesky(data, n_chains=1):
  # The data argument that is passed in will be available to the inner function
  # below so it doesn't have to be passed in as a parameter.
  replicated_data = np.tile(np.expand_dims(data, axis=1), reps=[1, n_chains, 1])

  @tf.function(autograph=False)
  def _log_lik_cholesky(precisions_cholesky):
    return (log_lik_data_cholesky(precisions_cholesky, replicated_data) +
            log_lik_prior_cholesky(precisions_cholesky))

  return _log_lik_cholesky

Constraining bijector

Our samples are constrained to be valid Cholesky factors, which means they must be lower triangular matrices with positive diagonals. The TransformedTransitionKernel needs a bijector that maps unconstrained tensors to/from tensors with our desired constraints. We've removed the Cholesky decomposition from the bijector's inverse, which speeds things up.

unconstrained_to_precision_cholesky = tfb.Chain([
    # step 2: exponentiate the diagonals    
    tfb.TransformDiagonal(tfb.Exp(validate_args=VALIDATE_ARGS)),
    # step 1: expand the vector to a lower triangular matrix
    tfb.FillTriangular(validate_args=VALIDATE_ARGS),
])
# some checks
inv = unconstrained_to_precision_cholesky.inverse(precisions_cholesky).numpy()
fwd = unconstrained_to_precision_cholesky.forward(inv).numpy()
print('precisions_cholesky:\n', precisions_cholesky)
print('\ninv:\n', inv)
print('\nfwd(inv):\n', fwd)
precisions_cholesky:
 [[[ 1.         0.       ]
  [ 0.         1.       ]]

 [[ 1.1470785  0.       ]
  [-2.0647411  1.0000004]]]

inv:
 [[ 0.0000000e+00  0.0000000e+00  0.0000000e+00]
 [ 3.5762781e-07 -2.0647411e+00  1.3721828e-01]]

fwd(inv):
 [[[ 1.         0.       ]
  [ 0.         1.       ]]

 [[ 1.1470785  0.       ]
  [-2.0647411  1.0000004]]]

Initial values

We generate a tensor of initial values. We're working with Cholesky factors, so we generate some Cholesky factor initial values.

# The number of chains is determined by the shape of the initial values.
# Here we'll generate 3 chains, so we'll need a tensor of 3 initial values.
N_CHAINS = 3

np.random.seed(123)

initial_values_cholesky = []
for i in range(N_CHAINS):
  initial_values_cholesky.append(np.array(
      [[0.5 + np.random.uniform(), 0.0],
       [-0.5 + np.random.uniform(), 0.5 + np.random.uniform()]],
      dtype=np.float32))
initial_values_cholesky = np.stack(initial_values_cholesky)

Sampling

We sample N_CHAINS chains using the TransformedTransitionKernel .

@tf.function(autograph=False)
def sample():
  tf.random.set_seed(123)
  log_lik_fn = get_log_lik_cholesky(my_data)

  # Tuning acceptance rates:
  dtype = np.float32
  num_burnin_iter = 3000
  num_warmup_iter = int(0.8 * num_burnin_iter) 
  num_chain_iter = 2500

  # Set the target average acceptance ratio for the HMC as suggested by
  # Beskos et al. (2013):
  # https://projecteuclid.org/download/pdfview_1/euclid.bj/1383661192
  target_accept_rate = 0.651

  # Initialize the HMC sampler.
  hmc = tfp.mcmc.HamiltonianMonteCarlo(
      target_log_prob_fn=log_lik_fn,
      step_size=0.01,
      num_leapfrog_steps=3)

  ttk = tfp.mcmc.TransformedTransitionKernel(
      inner_kernel=hmc, bijector=unconstrained_to_precision_cholesky)

  # Adapt the step size using standard adaptive MCMC procedure. See Section 4.2
  # of Andrieu and Thoms (2008):
  # http://www4.ncsu.edu/~rsmith/MA797V_S12/Andrieu08_AdaptiveMCMC_Tutorial.pdf
  adapted_kernel = tfp.mcmc.SimpleStepSizeAdaptation(
      inner_kernel=ttk,
      num_adaptation_steps=num_warmup_iter,
      target_accept_prob=target_accept_rate)

  states = tfp.mcmc.sample_chain(
      num_results=num_chain_iter,
      num_burnin_steps=num_burnin_iter,
      current_state=initial_values,
      kernel=adapted_kernel,
      trace_fn=None,
      parallel_iterations=1)
  # transform samples back to their constrained form
  samples = tf.linalg.matmul(states, states, transpose_b=True)
  return samples

precision_samples = sample()

Convergence check

A quick convergence check looks good:

r_hat = tfp.mcmc.potential_scale_reduction(precision_samples).numpy()
print(r_hat)
[[1.0013583 1.0019467]
 [1.0019467 1.0011804]]

Comparing results to the analytic posterior

# The output samples have shape [n_steps, n_chains, 2, 2]
# Flatten them to [n_steps * n_chains, 2, 2] via reshape:
precision_samples_reshaped = np.reshape(precision_samples, newshape=[-1, 2, 2])

And again, the sample means and standard deviations match those of the analytic posterior.

print('True posterior mean:\n', posterior_mean)
print('Mean of samples:\n', np.mean(precision_samples_reshaped, axis=0))
True posterior mean:
 [[ 0.9641779 -1.6534661]
 [-1.6534661  3.8683164]]
Mean of samples:
 [[ 0.9668749 -1.6552604]
 [-1.6552604  3.8676758]]
print('True posterior standard deviation:\n', posterior_sd)
print('Standard deviation of samples:\n', np.std(precision_samples_reshaped, axis=0))
True posterior standard deviation:
 [[0.13435492 0.25050813]
 [0.25050813 0.53903675]]
Standard deviation of samples:
 [[0.13329637 0.24913797]
 [0.24913797 0.53983945]]

Ok, all done! We've got our optimized sampler working.