#@title Licensed under the Apache License, Version 2.0 (the "License"); { display-mode: "form" }
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Explanation: Copyright 2018 The TensorFlow Probability Authors.
Licensed under the Apache License, Version 2.0 (the "License");
End of explanation
"""


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

"""
Explanation: TensorFlow 확률 사례 연구: 공분산 추정
<table class="tfo-notebook-buttons" align="left">
  <td>     <a target="_blank" href="https://www.tensorflow.org/probability/examples/TensorFlow_Probability_Case_Study_Covariance_Estimation"><img src="https://www.tensorflow.org/images/tf_logo_32px.png">TensorFlow.org에서 보기</a>   </td>
  <td><a target="_blank" href="https://colab.research.google.com/github/tensorflow/docs-l10n/blob/master/site/ko/probability/examples/TensorFlow_Probability_Case_Study_Covariance_Estimation.ipynb"><img src="https://www.tensorflow.org/images/colab_logo_32px.png">Google Colab에서 실행</a></td>
  <td><a target="_blank" href="https://github.com/tensorflow/docs-l10n/blob/master/site/ko/probability/examples/TensorFlow_Probability_Case_Study_Covariance_Estimation.ipynb"><img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png">GitHub에서 소스 보기</a></td>
  <td><a href="https://storage.googleapis.com/tensorflow_docs/docs-l10n/site/ko/probability/examples/TensorFlow_Probability_Case_Study_Covariance_Estimation.ipynb"><img src="https://www.tensorflow.org/images/download_logo_32px.png">노트북 다운로드</a></td>
</table>

저는 TensorFlow 확률을 배우기 위한 사례 연구를 목적으로 이 노트북을 작성했습니다. 해결할 문제는 2차원 평균 0 가우스 확률 변수의 샘플에 대한 공분산 행렬을 추정하는 것입니다. 이 문제에는 몇 가지 멋진 특징이 있습니다.

공분산에 역 Wishart 사전 분포(일반적인 접근법)를 사용하는 경우 이 문제는 해석적 솔루션을 갖고 있어 결과를 확인할 수 있습니다.
이 문제에는 제약된 매개변수를 샘플링하는 과정이 포함되어 있어 흥미로운 복잡성을 더해줍니다.
가장 간단한 해법은 가장 빠른 해법이 아니므로, 몇 가지 최적화 작업을 수행해야 합니다.

저의 경험을 순차적으로 나열하는 방식으로 해결 과정을 작성하기로 결정했습니다. TFP의 세부 사항을 파악하기 위한 시간을 잠시 가졌으며, 이러한 고민의 결과로 이 노트북은 매우 간단한 기능부터 시작하여 점점 복잡한 TFP 기능으로 나아갑니다. 그러한 과정 가운데 저는 많은 문제에 부딪쳤었고, 문제를 식별하는 데 도움이 되는 프로세스와 결국에 제가 찾은 해결책을 모두 담아내고자 했습니다. 또한, 개별 단계가 올바른지 확인하는 다수의 테스트를 포함하여 많은 세부 정보를 포함하려고 노력했습니다.
TensorFlow 확률을 알아야 하는 이유는?
저는 TensorFlow 확률이 제 프로젝트에 도움이 될 만한 부분을 몇 가지 발견했습니다.

TensorFlow 확률을 이용하면 노트북에서 대화형으로 복잡한 모델을 개발할 수 있는 프로토타입을 만들 수 있습니다. 코드를 작은 조각으로 나누어 단위 테스트를 도입하고 대화식으로 테스트할 수 있습니다.
확장할 준비가 되었으면 여러 기기의 최적화된 멀티 프로세서에서 TensorFlow를 실행할 수 있도록 준비한 모든 인프라를 활용할 수 있습니다.
마지막으로, 저는 Stan이라는 말을 참 좋아하는데요, 디버깅하는 일은 상당히 어렵다는 것을 알게 되었습니다. 코드를 찔러보고 중간 상태를 검사하는 등의 작업을 수행할 수 있는 도구가 거의 없는 독립 실행형 언어로 모든 모델링 코드를 작성해야 합니다.

안타까운 점은 TensorFlow 확률이 Stan과 PyMC3 보다 훨씬 새로운 것이기 때문에 설명서 작업이 진행 중이고 아직 구축되지 않은 기능이 많다는 것입니다. 다행스럽게도 저는 TFP의 기반이 탄탄하다는 것을 알게 되었으며, 모듈식으로 설계되었기 때문에 기능을 상당히 간단하게 확장할 수 있습니다. 이 노트북에서는 사례 연구를 해결하는 외에도 TFP를 확장하는 몇 가지 방법을 보여 드리겠습니다.
이 문서의 대상
저는 이 노트북의 독자들이 몇 가지 중요한 전제 조건을 갖추었으리라 가정합니다. 여러분은...

베이지안 추론의 기본 내용을 알고 있습니다(그렇지 않다면 Statistical Rethinking(통계적으로 다시 생각하기))라는 책을 가장 먼저 추천합니다.
Stan / PyMC3 / BUGS와 같은 MCMC 샘플링 라이브러리에 익숙합니다.
NumPy에 대해 확실히 알고 있습니다(좋은 입문서로는 Python for Data Analysis(데이터 분석을 위한 파이썬)가 있음).
TensorFlow에 대해 전문적인 수준은 아니더라도 약간은 알고 있습니다(Learning TensorFlow도 좋지만 TensorFlow는 급속한 발전을 이루고 있기에 대부분의 책은 다소 구식의 내용을 담고 있습니다. Stanford의 CS20 과정도 좋습니다).

첫 번째 시도
이 문제에 대한 저의 첫 번째 시도는 이러합니다. 스포일러: 제 솔루션은 효과가 없으며 문제를 해결하기 위해 여러 번 시도해야 합니다! 프로세스를 거치는 데 다소 시간이 걸리지만 아래의 각 시도는 TFP의 새로운 부분을 배우는 데 유용합니다.
한 가지 참고할 사항: TFP는 현재 역 Wishart 분포를 구현하지 않으며(마지막 부분에서 우리의 역 Wishart를 실행하는 방법을 볼 예정임) 대신, Wishart 사전 분포를 사용하여 정밀도 행렬을 추정하는 것으로 문제를 바꿔볼 것입니다.
End of explanation
"""


# 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))

"""
Explanation: 1단계: 관측 결과 수집
여기에 있는 모든 데이터는 합성 데이터이기 때문에 실제 사례보다 약간 정돈되어 보일 것입니다. 물론, 여러분도 자신만의 합성 데이터를 생성할 수 있습니다.
팁: 모델 형식을 결정하였다면 일부 매개변수 값을 선택하고 선택한 모델을 사용하여 일부 합성 데이터를 생성할 수 있습니다. 구현에 대한 온전성 검사의 일환으로 선택한 매개변수의 실제 값이 여러분의 추정에 포함되어 있는지 확인할 수 있습니다. 디버깅/테스트 주기를 더 빠르게 하기 위해 단순화된 모델 버전을 고려해야 할 수도 있습니다(예: 더 낮은 차원 또는 더 적은 수의 샘플 사용).
팁: 가장 쉽게 작업하려면 관측 작업 시 NumPy 배열을 사용하는 것입니다. 단, NumPy는 기본적으로 float64를 사용하지만 TensorFlow는 기본적으로 float32를 사용한다는 점을 알고 있어야 합니다.
일반적으로, TensorFlow 연산은 모든 인수가 동일한 형식을 갖기를 바라며 형식을 변경하려면 명시적 데이터 캐스팅을 수행해야 합니다. float64 관측을 사용하는 경우, 다수의 캐스트 연산을 추가해야 합니다. 이와 대조적으로 NumPy는 자동으로 캐스팅을 처리합니다. 따라서, TensorFlow에서 float64를 사용하도록 강제하는 것보다 Numpy 데이터를 float32로 변환하는 것이 훨씬 쉽습니다. 
일부 매개변수 값 선택
End of explanation
"""


# 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

"""
Explanation: 일부 합성 관측 생성
TensorFlow 확률은 데이터의 초기 차원이 샘플 인덱스를 나타내며 데이터의 최종 차원이 여러분의 샘플의 차원수를 나타낸다는 규칙을 사용합니다.
여기에서 우리는 각각 길이 2의 벡터인 샘플 100개가 필요합니다. 우리는 형상이 (100,2)인 배열 my_data를 생성합니다. my_data[i, :]는 $i$번째 샘플이며 길이 2의 벡터입니다.
(my_data는 float32 형식이어야 함을 기억하세요!)
End of explanation
"""


# 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()

print('mean of observations:', np.mean(my_data, axis=0))
print('true mean:', true_mean)

print('covariance of observations:\n', np.cov(my_data, rowvar=False))
print('true covariance:\n', true_cov)

"""
Explanation: 관측 결과 온전성 검사
잠재적인 버그 요인 중 하나는 합성 데이터를 엉망으로 만드는 것입니다! 간단한 검사를 해보도록 하겠습니다.
End of explanation
"""


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)

"""
Explanation: 네, 샘플이 합당한 것으로 보입니다. 다음 단계로 넘어가겠습니다.
2단계: NumPy에서 우도 함수 구현
TF 확률에서 MCMC 샘플링을 수행하기 위해 작성해야 할 주요 항목은 로그 우도 함수입니다. 일반적으로 NumPy보다 TF를 작성하는 것이 좀 더 까다롭기 때문에 NumPy에서 초기 구현을 수행하는 것이 도움이 된다고 생각합니다. 저는 우도 함수를 2개, 즉 $P(data | parameters)$에 해당하는 데이터 우도 함수와 $P(parameters)$에 해당하는 사전 우도 함수로 나눌 것입니다.
이러한 NumPy 함수는 극도로 최적화/벡터화할 필요는 없는데, 테스트를 위한 일부 값을 생성하는 것이 목표이기 때문입니다. 정확성이 핵심 고려 사항입니다!
먼저 우리는 데이터 로그 우도 부분을 구현할 것입니다. 이 작업은 매우 간단합니다. 한 가지 기억해야 할 점은 정밀도 행렬로 작업할 것이므로 그에 맞게 매개변수화할 것입니다.
End of explanation
"""


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)

"""
Explanation: 사후에 대한 해석적 해가 있기 때문에 정밀 행렬에 대해 Wishart 사전 분포를 사용할 것입니다(Wikipedia의 편리한 켤레 사전 분포 테이블 참조).
Wishart 분포는 2개의 매개변수를 갖습니다.

자유도의 수(Wikipedia에서 $\nu$로 레이블 표시됨)
스케일 행렬(Wikipedia에서 $V$로 레이블 표시됨)

매개변수 $\nu를 갖는 Wishart 분포의 평균 V$는 $E[W] = \nu V$, 분산은 $\text{Var}(W_{ij}) = \nu(v_{ij}^2+v_{ii}v_{jj})$입니다.
유용한 직관: 평균이 0이고 공분산이 $V$인 다변량 정규 확률 변수에서 $\nu$ 독립 추첨 $x_1 \ldots x_{\nu}$를 생성한 다음 합계 $W = \sum_{i=1}^{\nu} x_i x_i^T$를 형성하여 Wishart 샘플을 생성할 수 있습니다.
Wishart 샘플을 $\nu$로 나누어 스케일 조정하면 $x_i$의 샘플 공분산 행렬을 얻습니다. 이 샘플 공분산 행렬은 $\nu$가 증가함에 따라 $V$로 향하는 경향을 나타냅니다. $\nu$가 작으면 샘플 공분산 행렬에 많은 변동이 있으므로 작은 값의 $\nu$는 더 약한 사전 분포, 큰 값의 $\nu$는 더 강한 사전 분포에 해당하게 됩니다. $\nu$는 최소한 샘플링 공간의 차원만큼 커야 하며 그렇지 않은 경우 특이 행렬을 생성하게 됩니다.
우리는 $\nu = 3$을 사용하므로 약한 사전 분포를 갖게 되며 $V = \frac{1}{\nu}를 사용하여 공분산 추정치를 항등식으로 가져올 것입니다(평균은 $\nu V$임을 기억).
End of explanation
"""


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)))

"""
Explanation: Wishart 분포는 알려진 평균 $\mu$를 갖는 다변량 정규 분포의 정밀도 행렬을 추정하기 위한 켤레 사전 분포입니다.
Wishart 사전 분포 매개변수가 $\nu, V$ 및 다변량 정규 분포 $x_1, \ldots의 x_n$ 관측값으로 $n$를 갖는 경우를 가정합니다. 사후 확률 매개변수는 $n + \nu, \left(V^{-1} + \sum_{i=1}^n (x_i-\mu)(x_i-\mu)^T \right)^{-1}$입니다.
End of explanation
"""


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()

"""
Explanation: 사후 분포 및 실제 값을 빠르게 플롯한 것입니다. 사후 분포는 샘플 사후 분포에 근접하지만 항등식 방향으로 약간 축소됩니다. 또한, 실제 값은 사후 확률 모드로부터 상당히 멀어집니다. 그 이유는 사전 분포가 우리 데이터와 매우 잘 일치하지는 않기 때문일 것으로 생각합니다. 스케일이 조정된 역 Wishart 사전 분포와 같은 실제 문제에는 더 잘 대응할 수 있을 거라 생각하지만(이 문제와 관련해서는 Andrew Gelman의 설명 참조) 좋은 해석적 사후 분포는 얻을 수 없을 것입니다.
End of explanation
"""


# 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())

"""
Explanation: 3단계: TensorFlow에서 우도 함수 구현
스포일러: 우리의 첫 번째 시도는 작동하지 않을 것입니다. 그 이유는 아래에서 이야기를 나누겠습니다.
팁: 우도 함수를 개발할 때 TensorFlow 즉시 실행 모드를 사용하세요. 즉시 실행 모드를 사용하면 TF가 좀 더 NumPy와 같이 작동합니다. 모든 항목이 즉시 실행되므로 Session.run()를 사용하는 대신 대화식으로 디버그할 수 있습니다. 여기에서 메모를 참조하세요.
예비 단계: 분포 클래스
TFP에는 우리가 로그 확률 생성에 사용할 분포 클래스 모음이 있습니다. 한 가지 주목해야 할 점은 이러한 클래스가 단일 샘플이 아닌 샘플의 텐서와 함께 작동한다는 것입니다. 이를 통해 벡터화 및 관련 속도 향상이 가능합니다.
분포는 2가지 방식으로 샘플 텐서와 함께 작동할 수 있습니다. 단일 스칼라 매개변수가 있는 분포를 포함하는 구체적인 예를 통해 이러한 2가지 방식을 설명하는 것이 가장 간단합니다. rate 매개변수를 갖는 푸아송 분포를 사용하겠습니다.

rate 매개변수에 대해 단일 값으로 푸아송을 생성하는 경우 sample() 메서드 호출로 단일 값이 반환됩니다. 이 값을 event라고 부르며 이 경우 이벤트는 모두 스칼라입니다.
rate 매개변수에 대한 값의 텐서를 사용하여 푸아송을 생성하는 경우 sample() 메서드 호출로 rate 텐서의 각 값에 대해 하나씩, 여러 값이 반환됩니다. 이 객체는 각각 고유의 rate를 갖는 독립적인 푸아송 모음의 역할을 하며 sample()의 호출로 인해 반환된 각각의 값은 이러한 푸아송 중 하나에 해당합니다. 이렇게 독립적이지만 동일하게 분포되지 않은 이벤트 모음을 batch라고 부릅니다.
sample() 메서드는 기본값이 빈 튜플인 sample_shape 매개변수를 사용합니다. sample_shape에 비어 있지 않은 값을 전달하면 샘플이 여러 배치를 반환합니다. 이 배치 모음을 sample이라고 합니다.

분포의 log_prob() 메서드는 sample()이 데이터를 생성하는 방식과 유사한 방식으로 데이터를 소비합니다. log_prob()는 여러 개의 독립적인 이벤트 배치와 같은 샘플의 확률을 반환합니다.

스칼라 rate으로 생성한 푸아송 객체를 가진 경우 각 배치는 스칼라이며, 샘플의 텐서를 전달하면 동일한 크기의 로그 확률 텐서를 얻습니다.
rate 값의 T 형상 텐서로 생성한 푸아송 객체를 가진 경우 각 배치는 T 형상의 텐서입니다. D, T 형상의 샘플 텐서를 전달하면 D, T 형상의 로그 확률 텐서를 얻습니다.

아래는 이러한 사례를 보여주는 몇 가지 예입니다. 이벤트, 배치, 형상에 대한 자세한 튜토리얼은 이 노트북을 참조하세요.
End of explanation
"""


VALIDATE_ARGS = True
ALLOW_NAN_STATS = False

"""
Explanation: 데이터 로그 우도
먼저 데이터 로그 우도 함수를 구현합니다.
참고: 분포가 입력을 검증할 수는 있지만 기본적으로 그렇게 하지는 않습니다. 물론 여기서는 디버깅하는 동안 검증을 사용할 것입니다! 모두 문제가 없으면 속도가 정말 중요한 경우 검증을 해제할 수 있습니다.
End of explanation
"""


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)

"""
Explanation: NumPy 사례와의 한 가지 주요 차이점은 TensorFlow 우도 함수가 단일 행렬이 아닌 정밀도 행렬의 벡터를 처리해야 한다는 것입니다. 우리는 여러 체인으로부터 샘플링할 때 매개변수 벡터를 사용합니다.
하나의 정밀도 행렬 배치(예: 체인당 하나의 행렬)와 함께 작동하는 분포 객체를 생성합니다.
데이터 로그 확률을 계산할 때 배치 변수당 하나의 복사본을 갖도록 매개변수와 동일한 방식으로 데이터를 복제해야 합니다. 복제된 데이터의 형상은 다음과 같아야 합니다.
[sample shape, batch shape, event shape]
2차원 가우시안으로 작업하고 있기 때문에 여기서의 경우 이벤트 형상은 2입니다. 우리에게 100 샘플이 있으므로 샘플 형상은 100입니다. 배치 형상은 단순히 우리가 작업하는 정밀도 행렬의 수에 지나지 않습니다. 우도 함수를 호출할 때마다 데이터를 복제하는 것은 낭비이므로 미리 데이터를 복제하고 복제한 버전을 전달합니다.
이것은 비효율적인 구현이라는 점을 알아야 합니다. MultivariateNormalFullCovariance는 마지막 최적화 섹션에서 이야기를 나눌 몇 가지 대안에 비해 낭비가 심합니다.
End of explanation
"""


# 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])

"""
Explanation: 팁: 제가 발견한 아주 유용한 점 한 가지는 TensorFlow 함수의 온전성 검사를 거의 작성하지 않는 것입니다. TF에서 벡터화를 엉망으로 만들기가 정말 쉽기 때문에 더 간단한 NumPy 함수를 사용하면 훌륭하게 TF 출력을 확인할 수 있습니다. 작은 단위 테스트로 생각하면 이해가 쉽습니다.
End of explanation
"""


@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])

"""
Explanation: 사전 로그 우도
데이터 복제에 대해 걱정할 필요가 없기 때문에 사전 분포 계산이 더 쉽습니다.
End of explanation
"""


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

"""
Explanation: 결합 로그 우도 함수 작성
위의 데이터 로그 우도 함수는 우리의 관측에 따라 달라지지만 샘플러에는 그렇지 않습니다. [Closer](https://en.wikipedia.org/wiki/Closure_(computer_programming)를 사용하면 전역 변수를 사용하지 않고 종속성을 제거할 수 있습니다. Closer는 내부 함수에 필요한 변수를 포함하는 환경을 구축하는 외부 함수를 포함합니다.
End of explanation
"""


@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:]))

"""
Explanation: 4단계: 샘플
이제 샘플링할 시간입니다! 단순하게 작업할 수 있도록 1개의 체인만 사용하고 항등 행렬을 시작점으로 사용합니다. 나중에 더 세심하게 작업할 것입니다.
다시 말하지만 이것은 작동하지 않을 것이며 예외가 발생하게 됩니다.
End of explanation
"""


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:]))

"""
Explanation: 문제 확인
InvalidArgumentError (see above for traceback): Cholesky decomposition was not successful. The input might not be valid. 이는 큰 도움이 되지 않습니다. 무슨 일이 있었는지 자세히 알아보겠습니다.

실패 항목의 값을 볼 수 있도록 각 단계의 매개변수를 출력합니다.
특정 문제를 방지하기 위해 몇 가지 어설션을 추가합니다.

어설션은 TensorFlow 연산이기 때문에 까다로우며 어설션이 실행되고 최적화되어 그래프에서 빠지지 않도록 주의해야 합니다. TF 어설션에 익숙하지 않다면 TensorFlow 디버깅에 대한 이 개요를 읽어보는 것이 좋습니다. tf.control_dependencies를 사용하여 명시적으로 어설션을 강제 실행할 수 있습니다(아래 코드의 주석 참조).
TensorFlow의 기본 Print 함수는 어설션과 동일한 동작을 합니다. 이는 연산이므로 실행에 주의를 기울여야 합니다. Print는 노트북에서 작업할 때 부가적인 말썽을 일으킵니다. 즉, 출력이 stderr로 전송되며 stderr는 셀에 표시되지 않습니다. 여기서 트릭을 사용합니다. tf.Print를 사용하는 대신 tf.pyfunc를 통해 자체 TensorFlow 출력 연산을 생성합니다. 어설션과 마찬가지로 우리의 메서드가 실행되는지 확인해야 합니다.
End of explanation
"""


# 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())

"""
Explanation: 이것이 실패하는 이유
샘플러로 가장 처음에 시도하는 첫 번째 매개변수 값은 비대칭 행렬입니다. 이는 대칭(및 양정치) 행렬에 대해서만 정의되기 때문에 Cholesky 분해를 실패하게 만듭니다.
여기서 문제는 우리의 관심 매개변수가 정밀도 행렬이고, 정밀도 행렬은 실수, 대칭 및 양정치여야 한다는 것입니다. 샘플러는 이 제약 조건을 전혀 고려하지 않기 때문에(경사도법을 통한 경우는 예외일 수 있음) 특히 단계 크기가 큰 경우에 예외를 발생시키는 등 샘플러가 잘못된 값을 제안할 가능성이 대단히 높습니다.
Hamiltonian Monte Carlo 샘플러를 사용하는 경우, 매우 작은 단계 크기를 사용하여 문제를 해결할 수도 있지만(그래디언트가 유효하지 않은 영역으로부터 매개변수를 멀리 유지시킴) 작은 단계 크기를 사용하면 수렴이 느리게 됩니다. 경사도법과 관련이 없는 Metropolis-Hastings 샘플러를 사용하면 우리는 망하게 됩니다.
버전 2: 제약이 없는 매개변수로의 재매개변수화
위의 문제에 대한 간단한 해결책이 있습니다. 즉, 새 매개변수가 더 이상 이러한 제약 조건을 갖지 않도록 모델을 재매개변수화할 수 있습니다. TFP는 이러한 작업에 유용한 도구 세트인 Bijectors를 제공합니다.
Bijectors를 사용한 재매개변수화
우리의 정밀도 행렬은 실수이고 대칭이어야 합니다. 우리는 이러한 제약이 없는 대체 매개변수화를 원합니다. 시작점은 정밀도 행렬의 Cholesky 인수분해입니다. Cholesky 인수는 여전히 제약이 있는데, 이들은 하삼각이며 대각선 요소는 양수여야 합니다. 다만 우리가 Cholesky 인수의 대각선 로그를 취하면 로그는 더 이상 양수의 제약을 받지 않으며, 하삼각 부분을 1차원 벡터로 평면화하면 더 이상 하삼각 제약을 받지 않습니다. 여기서의 경우 결과는 제약이 없는 길이 3 벡터가 됩니다.
Stan 설명서에는 변환을 사용하여 매개변수에 대한 다양한 유형의 제약 조건을 제거하는 방법에 대한 훌륭한 내용이 들어 있습니다.
이 재매개변수화는 데이터 로그 우도 함수에 거의 영향을 미치지 않습니다. 변환을 반전시키기만 하면 정밀도 행렬을 얻을 수 있지만 사전 분포에 대한 영향은 더 복잡해집니다. 우리는 주어진 정밀도 행렬의 확률이 Wishart 분포에 의해 정해지도록 지정했습니다. 그러면 변환된 행렬의 확률은 얼마나 될까요?
다시 상기해 보면, 단조 함수 $g$를 1차원 확률 변수 $X$, $Y = g(X)$에 적용하면 $Y$의 밀도는 다음과 같이 주어집니다.
$$ f_Y(y) = | \frac{d}{dy}(g^{-1}(y)) | f_X(g^{-1}(y)) $$
$g^{-1}$ 항의 미분계수는 $g$가 로컬 볼륨을 변경하는 방식을 설명합니다. 더 높은 차원의 무작위 변수의 경우 수정 인자는 $g^{-1}$의 야고비안 행렬식의 절대값입니다(여기 참조).
역변환의 야고비안을 우리의 로그 사전 우도 함수에 추가해야 합니다. 다행히도 TFP의 Bijector 클래스가 이를 처리해줄 수 있습니다.
Bijector 클래스는 확률 밀도 함수에서 변수를 변경하는 데 사용하는 가역함수, 평활 함수를 나타낼 때 사용합니다. Bijectors에는 변환을 수행하는 forward() 메서드, 이를 반전시키는 inverse() 메서드, pdf를 재매개변수화할 때 필요한 야고비안 수정을 제공하는 forward_log_det_jacobian()와 inverse_log_det_jacobian() 메서드가 모두 있습니다.
TFP는 Chain 연산자를 통한 구성으로 결합하여 상당히 복잡한 변환을 형성할 수 있는 유용한 Bijectors 모음을 제공합니다. 우리의 경우 다음 3개의 Bijectors를 구성합니다(체인의 연산은 오른쪽에서 왼쪽으로 수행됨).

우리 변환의 첫 번째 단계는 정밀도 행렬에 대한 Cholesky 인수분해를 수행하는 것입니다. 이를 위한 Bijector 클래스는 없지만 CholeskyOuterProduct Bijector는 2개의 Cholesky 인수의 곱을 취합니다. 우리는 Invert 연산자를 사용하는 이 연산의 역을 사용할 수 있습니다.
다음 단계로 Cholesky 인수의 대각선 원소에 대한 로그를 취합니다. TransformDiagonal Bijector와 Exp Bijector의 역을 통해 이를 수행합니다.
마지막으로 FillTriangular Bijector의 역을 사용하여 행렬의 하삼각 부분을 벡터로 평면화합니다.
End of explanation
"""


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])

"""
Explanation: TransformedDistribution 클래스는 Bijector를 분포에 적용하고 log_prob()에 필요한 야고비안 수정을 수행하는 프로세스를 자동화합니다. 새로운 사전 분포는 다음과 같습니다.
End of explanation
"""


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])

"""
Explanation: 데이터 로그 우도의 변환을 역전시키기만 하면 됩니다.
precision = precision_to_unconstrained.inverse(transformed_precision)
우리는 실제로 정밀도 행렬의 Cholesky 인수분해를 원하기 때문에 여기에서 부분 역을 수행하는 것이 더 효율적일 수 있습니다. 그러나 최적화는 나중으로 미루고 부분적인 역은 독자를 위한 연습 문제로 남겨 두겠습니다.
End of explanation
"""


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)

"""
Explanation: 다시, 새로운 함수를 클로저로 래핑합니다.
End of explanation
"""


# 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()

"""
Explanation: 샘플링
이제 잘못된 매개변수 값으로 인해 우리의 샘플러가 날아갈 것에 대해 걱정할 필요가 없으므로 실제 샘플을 생성해 보겠습니다.
샘플러는 우리 매개변수의 제약 없는 버전으로 작동하므로 초기 값을 제약 없는 버전으로 변환해야 합니다. 우리가 생성하는 샘플도 모두 제약이 없는 형태이므로 이들을 다시 변환해야 합니다. Bijectors는 벡터화되어 있으므로 쉽게 작업을 수행할 수 있습니다.
End of explanation
"""


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

"""
Explanation: 샘플러 결과의 평균을 해석적 사후 확률 평균과 비교하겠습니다!
End of explanation
"""


np.reshape(precision_samples, [-1, 2, 2])

"""
Explanation: 우리가 완전히 틀린 것 같네요! 이유를 알아보겠습니다. 먼저 우리의 샘플을 살펴보겠습니다.
End of explanation
"""


# 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)))

"""
Explanation: 아, 이런! 모두 같은 값을 갖는 것 같습니다. 이유를 알아보겠습니다.
kernel_results_ 변수는 각 상태에서 샘플러에 대한 정보를 제공하는 명명된 튜플입니다. is_accepted필드가 여기에서 핵심적입니다.
End of explanation
"""


# 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() 

"""
Explanation: 우리의 모든 샘플이 거부되었습니다! 아마도 단계 크기가 너무 컸던 것 같습니다. 임의로 stepsize=0.1를 선택했습니다.
버전 3: 적응형 단계 크기로 샘플링
임의로 단계 크기를 선택하여 샘플링에 실패했기 때문에 우리는 몇 가지 논의 사항을 갖게 되었습니다.

적응형 단계 크기의 구현
몇 가지 수렴 검사의 수행

tensorflow_probability/python/mcmc/hmc.py에 적응형 단계 크기를 구현하기 위한 몇 가지 멋진 샘플 코드가 있습니다. 저는 아래와 같이 적용했습니다.
각 단계에 대해 개별 sess.run() 문이 있습니다. 필요한 경우 일부 단계별 진단을 쉽게 추가할 수 있으며 이는 디버깅에 매우 유용합니다. 예를 들어 증분 진행률, 각 단계의 시간 등을 나타낼 수 있습니다.
팁: 샘플링을 엉망으로 만드는 가장 일반적인 방법 중 하나는 루프에서 그래프를 늘리는 것입니다. (세션이 실행되기 전에 그래프를 종료하는 이유는 바로 이러한 문제를 방지하기 위함입니다.) 다만 finalize()을 사용하지 않은 경우 코드가 기어가는 속도로 느려지는지 확인하는 유용한 디버깅 방법은 len(mygraph.get_operations())를 통해 각 단계의 그래프 크기를 출력하는 것입니다. 길이가 증가하면 무언가 잘못하고 있는 것입니다.
여기에서는 3개의 독립 체인을 실행할 것입니다. 체인 간 비교를 수행하면 수렴을 확인하는 데 도움이 됩니다.
End of explanation
"""


print(np.mean(is_accepted))

"""
Explanation: 빠른 확인: 샘플링 중 허용률은 우리의 목표인 0.651에 근접합니다.
End of explanation
"""


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))

print('True posterior standard deviation:\n', posterior_sd)
print('Standard deviation of samples:\n', np.std(precision_samples_reshaped, axis=0))

"""
Explanation: 더 좋은 점은 샘플 평균과 표준 편차가 해석 해에 대해 기대되는 값에 가깝다는 것입니다.
End of explanation
"""


r_hat = tfp.mcmc.potential_scale_reduction(precision_samples).numpy()
print(r_hat)

"""
Explanation: 수렴 확인
일반적으로 우리는 확인할 해석 해가 없기 때문에 샘플러가 수렴되었는지 확인할 필요가 있습니다. 한 가지 표준 검사로 Gelman-Rubin $\hat{R}$ 통계가 있으며, 이를 위해 여러 샘플링 체인이 필요합니다. $\hat{R}$는 체인 간의 분산(평균의 분산)이 체인이 동일하게 분포된 경우 예측할 수 있는 수치를 초과하는 정도를 측정합니다. 1에 가까운 $\hat{R}$ 값이 근사 수렴을 나타내는 데 사용됩니다. 자세한 내용은 소스를 참조하세요.
End of explanation
"""


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()

"""
Explanation: 모델 비평
해석 해를 가지지 못했다면 지금이 실질적인 모델 비판을 해야 할 때일 수 있습니다.
다음은 검증 자료(빨간색)과 관련된 샘플 구성 요소의 몇 가지 간단한 히스토그램입니다. 샘플은 샘플 정밀도 행렬 값에서 단위 행렬 사전 값으로 축소되었습니다.
End of explanation
"""


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()

"""
Explanation: 정밀도 구성 요소 쌍에 대한 일부 산점도는 사후 확률의 상관 구조로 인해 실제 사후 확률 값이 위의 주변부에서 나타나는 것만큼 나타날 가능성이 적지 않음을 보여줍니다.
End of explanation
"""


# 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)

"""
Explanation: 버전 4: 제약이 있는 매개변수의 더 단순화된 샘플링
Bijectors는 정밀도 행렬 샘플링을 간단하게 만들었지만, 제약이 없는 표현으로 수동 변환하는 과정에서 많은 수작업이 필요했습니다. 더 쉬운 방법이 있습니다!
TransformedTransitionKernel
TransformedTransitionKernel는 이 프로세스를 단순화합니다. 샘플러를 래핑하고 모든 전환을 처리합니다. 또한 제약이 없는 매개변수 값을 제약이 있는 매개변수 값으로 매핑하는 Bijectors 목록을 인수로 취합니다. 따라서 여기에서는 위에서 사용한 precision_to_unconstrained Bijectors의 역이 필요합니다. 우리는 단순히 tfb.Invert(precision_to_unconstrained)를 사용할 수도 있지만 이는 역의 역을 취하는 작업(TensorFlow는 tf.Invert(tf.Invert())를 tf.Identity())로 단순화할만큼 스마트하지 않음)을 수반할 수 있기 때문에 대신 새 Bijector를 작성하겠습니다.
Bijector 제약
End of explanation
"""


@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()

"""
Explanation: TransformedTransitionKernel을 사용한 샘플링
TransformedTransitionKernel를 사용하면 더 이상 매개변수를 수동으로 변환할 필요가 없습니다. 초기 값과 우리의 샘플은 모두 정밀도 행렬입니다. 제약이 없는 Bijector를 커널로 전달하기만 하면 됩니다. 그러면 커널에서 모든 변환을 처리합니다.
End of explanation
"""


r_hat = tfp.mcmc.potential_scale_reduction(precision_samples).numpy()
print(r_hat)

"""
Explanation: 수렴 확인
$\hat{R}$ 수렴 확인이 잘 되는 것으로 보입니다!
End of explanation
"""


# 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))

print('True posterior standard deviation:\n', posterior_sd)
print('Standard deviation of samples:\n', np.std(precision_samples_reshaped, axis=0))

"""
Explanation: 해석적 사후 확률과의 비교
다시 해석적 사후 확률을 검사해보겠습니다.
End of explanation
"""


# 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)

"""
Explanation: 최적화
이제 전체 실행이 완료되었으니 더 최적화된 버전을 만들어 보겠습니다. 이 예시에서 속도는 그다지 중요하지 않지만, 일단 행렬이 커지면 몇 가지 최적화 작업으로 큰 차이가 만들어집니다.
우리가 할 수 있는 한 가지 큰 속도 향상 작업은 Cholesky 분해의 관점에서 재매개변수화를 수행하는 것입니다. 그 이유는 우리의 데이터 우도 함수에 공분산과 정밀도 행렬이 모두 필요하기 때문입니다. 행렬 역전은 비효율적이고($n \times n$ 행렬에 대해 $O(n^3)$) 공분산 또는 정밀도 행렬의 관점으로 매개변수화를 수행하면 다른 한쪽의 결과를 얻기 위해 역전을 수행해야 합니다.
다시 말하지만, 실수이며 양정치인 대칭 행렬 $M$는 $M = L L^T$ 형식의 곱으로 분해될 수 있습니다(여기서 행렬 $L$는 하삼각이며 양의 대각선을 가짐). $M$의 Cholesky 분해가 주어지면 $M$(하삼각 및 상삼각 행렬의 곱) 및 $M^{-1}$(후방 대입을 통함)를 더 효율적으로 얻을 수 있습니다. Cholesky 인수분해 자체는 계산에 자원이 적지 않게 들어가지만, Cholesky 인수의 측면에서 매개변수화하면 초기 매개변수 값의 Choleksy 인수분해만 계산하면 됩니다.
공분산 행렬의 Cholesky 분해 사용
TFP에는 공분산 행렬의 Cholesky 인수 측면에서 매개변수화된 다변량 정규 분포 버전의 MultivariateNormalTriL이 있습니다. 따라서 공분산 행렬의 Cholesky 인수 측면에서 매개변수화하면 데이터 로그 우도를 효율적으로 계산할 수 있습니다. 이때 어려운 점은 유사한 효율성으로 사전 로그 우도를 계산하는 것입니다.
샘플의 Cholesky 인수와 함께 작동하는 역 Wishart 분포 버전이 있다면 모든 것이 완전하게 갖추어지겠지만 안타깝게도 현실은 그렇지 못합니다(하지만 팀이 코드 제출을 환영합니다!). 대안으로 Bijectors 체인과 함께 샘플의 Cholesky 인수와 작동하는 Wishart 분포 버전을 사용할 수 있습니다.
이제 몇 개의 스톡 Bijectors만 있으면 정말 효율적으로 만들 수 있지만 TFP Bijectors의 강력함을 보여주는 유익한 예시로서, 그리고 연습으로서 이 프로세스를 보여드리려고 합니다.
Cholesky 인수에 작동하는 Wishart 분포
Wishart 분포에는 입력 및 출력 행렬이 Cholesky 인수가 되어야만 한다고 지정하는 유용한 플래그인 input_output_cholesky가 있습니다. 전체 행렬로 작업하는 것보다 Cholesky 인수로 작업하는 것이 더 효율적이고 수치적으로 유리하므로 그렇게 하는 것이 바람직합니다. 플래그의 의미 체계에 대한 중요한 참고 사항: 이는 분포에 대한 입력 및 출력의 표현이 변경되어야 한다는 표시일 뿐이며, 분포의 전체 재매개변수화를 나타내지는 않습니다. 이렇게 하려면 log_prob() 함수에 대한 야고비안 수정이 필요합니다. 우리는 실제로 이 전체 재매개변수화를 수행하기 원하므로 자체 분포를 구축할 것입니다.
End of explanation
"""


# 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)

"""
Explanation: 역 Wishart 분포 구축하기
우리의 공분산 행렬 $C$는 $C = L L^T$로 분해되며, 여기서 $L$는 하삼각이고 양의 대각선을 갖습니다. 우리는 $C \sim W^{-1}(\nu, V)$가 주어졌을 때 $L$의 확률을 알고 싶으며, 여기서 $W^{-1}$는 역 Wishart 분포입니다.
역 Wishart 분포는 $C \sim W^{-1}(\nu, V)$인 경우 정밀도 행렬이 $C^{-1} \sim W(\nu, V^{-1})$인 특성을 갖습니다. 따라서 우리는 매개변수로 Wishart 분포를 사용하는 TransformedDistribution와 정밀도 행렬의 Cholesky 인수를 그 역의 Cholesky 인수에 매핑하는 Bijector를 통해 $L$의 확률을 얻을 수 있습니다.
$C^{-1}의 Cholesky 인수에서 $L$로 가는 간단한(그러나 매우 효율적이지는 않은) 방법은 역해법으로 Cholesky 인수를 역전시킨 다음, 이러한 역전된 인수로부터 공분산 행렬을 형성한 다음, Cholesky 인수분해를 수행하는 것입니다.
$C^{-1} = M M^T$의 Cholesky 분해를 해보겠습니다. $M$는 하삼각이므로 MatrixInverseTriL Bijector를 사용하여 이를 역전시킬 수 있습니다.
$C = (M M^T)^{-1} = M^{-T}M^{-1} = M^{-T} (M^{-T})^T$와 같이 $M^{-1}$에서 $C$를 형성하는 것은 다소 까다롭습니다. $M$는 하삼각이므로 $M^{-1}$도 하삼각이 되고 $M^{-T}$는 상삼각이 됩니다. CholeskyOuterProduct() Bijector는 하삼각 행렬에서만 작동하므로 $M^{-T}$에서 $C$를 형성하는 데는 사용할 수 없습니다. 해결책은 행렬의 행과 열을 순열로 배치하는 Bijectors 체인입니다.
다행히 이 논리는 CholeskyToInvCholesky Bijector에 캡슐화되어 있습니다!
모든 조각을 결합
End of explanation
"""


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

"""
Explanation: 우리의 최종 분포
Cholesky 인수에서 작동하는 역 Wishart는 다음과 같습니다.
End of explanation
"""


# 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])

"""
Explanation: 우리는 역 Wishart를 갖게 되었지만 Bijector에서 Cholesky 분해를 수행해야 하기 때문에 속도가 좀 느립니다. 정밀도 행렬 매개변수화로 돌아가서 최적화를 위해 무엇을 할 수 있는지 알아보겠습니다.
최종(!) 버전: 정밀도 행렬의 Cholesky 분해 사용
또 다른 접근 방식은 정밀도 행렬의 Cholesky 인수로 작업하는 것입니다. 여기서 사전 우도 함수는 계산하기 쉽지만 TFP에는 정밀도로 매개변수화된 다변량 정규 버전이 없기 때문에 데이터 로그 우도 함수에 더 많은 작업이 필요합니다.
최적화된 사전 로그 우도
위에서 구축한 CholeskyWishart 분포를 사용하여 사전 분포를 구성합니다.
End of explanation
"""


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])

"""
Explanation: 최적화된 데이터 로그 우도
TFP Bijector를 사용하여 다변량 정규의 자체 버전을 구축할 수 있습니다. 핵심 아이디어는 다음과 같습니다.
요소가 $N(0, 1)$의 iid 샘플인 열 벡터 $X$가 있다고 가정합니다. $\text{mean}(X) = 0$이고 $\text{cov}(X) = I$인 조건을 갖습니다.
이제 $Y = A X + b$라고 하겠습니다. $\text{mean}(Y) = b$이고 $\text{cov}(Y) = A A^T$인 조건을 갖습니다.
따라서 $A A^T = C$라는 가정 하에 affine 변환 $Ax+b$를 사용하여 평균 $b$ 및 공분산 $C$로 이루어진 벡터를 iid 표준 정규 샘플의 벡터로 만들 수 있습니다. $C$의 Cholesky 분해에는 원하는 속성이 있습니다. 그러나 다른 해법도 있습니다.
$P = C^{-1}$라고 하고 $P$의 Cholesky 분해를 $B$, 즉 $B B^T = P$라고 하겠습니다. 이제 다음과 같습니다.
$P^{-1} = (B B^T)^{-1} = B^{-T} B^{-1} = B^{-T} (B^{-T})^T$
따라서, 원하는 평균과 공분산을 얻는 또 다른 방법은 affine 변환 $Y=B^{-T}X + b$를 사용하는 것입니다.
우리의 접근 방식(이 노트북에서 제공):

tfd.Independent()를 사용하여 1차원 Normal 무작위 변수 배치를 단일 다차원 무작위 변수로 결합합니다. Independent()의 reinterpreted_batch_ndims 매개변수는 이벤트 차원으로 재해석되어야 하는 배치 차원의 수를 지정합니다. 우리의 경우에는 길이가 2인 1차원 배치를 생성하여 길이가 2인 1차원 이벤트로 변환하므로 reinterpreted_batch_ndims=1이 됩니다.
Bijector를 적용하여 tfb.Invert(tfb.ScaleMatvecTriL(scale_tril=precision_cholesky, adjoint=True))와 같이 원하는 공분산을 추가합니다. 위에서 우리는 iid 정규 무작위 변수와 정밀도 행렬 $(B^{-T}X)$의 Cholesky 인수 역의 행/열 바꾸기를 곱합니다. tfb.Invert는 $B$ 역전을 처리하며, adjoint=True 플래그는 행/열 바꾸기를 수행합니다.
Bijector를 적용하여 tfb.Shift(shift=shift)와 같이 원하는 오프셋을 추가합니다. 초기 역전된 affine 변환과 별도의 단계로 이동을 수행해야 하는데, 그렇지 않으면 역전된 스케일이 이 이동에 적용되기 때문입니다($y=Ax+b$의 역전은 $x=A^{-1}y - A^{-1}b$이므로).
End of explanation
"""


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

"""
Explanation: 결합된 로그 우도 함수
이제 클로저에서 사전 및 데이터 로그 우도 함수를 결합합니다.
End of explanation
"""


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)

"""
Explanation: Bijector 제약
우리의 샘플은 유효한 Cholesky 인수인 것으로 제약됩니다. 즉, 샘플이 양의 대각선을 가진 하삼각 행렬이어야 합니다. TransformedTransitionKernel는 제약이 없는 텐서를 원하는 제약 조건의 텐서 사이에서 매핑하는 Bijector를 필요로 합니다. 우리는 Bijector의 역으로부터 Cholesky 분해를 제거하여 속도를 높였습니다.
End of explanation
"""


# 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)

"""
Explanation: 초기 값
초기 값의 텐서를 생성합니다. Cholesky 인수로 작업하고 있으므로 일부 Cholesky 인수 초기 값을 생성합니다.
End of explanation
"""


@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()

"""
Explanation: 샘플링
TransformedTransitionKernel를 사용하여 N_CHAINS 체인을 샘플링합니다.
End of explanation
"""


r_hat = tfp.mcmc.potential_scale_reduction(precision_samples).numpy()
print(r_hat)

"""
Explanation: 수렴 검사
빠른 수렴 검사의 결과가 좋아 보입니다.
End of explanation
"""


# 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])

"""
Explanation: 해석적 사후 확률과 결과 비교
End of explanation
"""


print('True posterior mean:\n', posterior_mean)
print('Mean of samples:\n', np.mean(precision_samples_reshaped, axis=0))

print('True posterior standard deviation:\n', posterior_sd)
print('Standard deviation of samples:\n', np.std(precision_samples_reshaped, axis=0))

"""
Explanation: 그리고 다시, 샘플 평균과 표준 편차가 해석적 사후 확률의 값과 일치합니다.
End of explanation
"""


