#@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 tensorflow as tf
tf.compat.v2.enable_v2_behavior()

import tensorflow_probability as tfp
tfd = tfp.distributions
tfb = tfp.bijectors

"""
Explanation: TensorFlow Distributions の形状を理解する
<table class="tfo-notebook-buttons" align="left">
  <td><a target="_blank" href="https://www.tensorflow.org/probability/examples/Understanding_TensorFlow_Distributions_Shapes"><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/ja/probability/examples/Understanding_TensorFlow_Distributions_Shapes.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/ja/probability/examples/Understanding_TensorFlow_Distributions_Shapes.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/ja/probability/examples/Understanding_TensorFlow_Distributions_Shapes.ipynb"><img src="https://www.tensorflow.org/images/download_logo_32px.png">ノートブックをダウンロード</a></td>
</table>
End of explanation
"""


def describe_distributions(distributions):
  print('\n'.join([str(d) for d in distributions]))

"""
Explanation: 基礎
TensorFlow Distributions の形状には関連する 3 つの重要な概念があります。

イベントの形状は、分布からの 1 つの抽出の形状を表します。抽出は次元間で依存する場合があります。スカラー分布の場合、イベントの形状は [] です。5 次元の MultivariateNormal の場合、イベントの形状は [5] です。
バッチの形状は、独立した、同一に分布されていない抽出である「バッチ」の分布を表します。
サンプルの形状は、 分布ファミリからの独立した、同一に分布されたバッチの抽出を表します。

イベントの形状とバッチの形状は Distribution オブジェクトのプロパティですが、サンプルの形状は sample または log_prob への特定の呼び出しに関連付けられています。
このノートブックでは、例を使ってこれらの概念を説明していくので、すぐに分からなくても、心配する必要はありません。
また、これらの概念の概要については、このブログ記事を参照してください。
TensorFlow Eager に関する注意
このノートブックは、すべて TensorFlow Eager を使用して記述されています。提示された概念は Eager に依存していませんが、Eager では、Distribution オブジェクトが Python で作成されるときに、分布バッチとイベントの形状が評価されます（したがって既知です）。一方、グラフ（非 Eager モード）では、グラフが実行されるまでイベントとバッチの形状が決定されていない分布を定義することができます。
スカラー分布
上記のように、Distribution オブジェクトではイベントとバッチの形状が定義されています。まず、分布を説明するユーティリティから始めます。
End of explanation
"""


poisson_distributions = [
    tfd.Poisson(rate=1., name='One Poisson Scalar Batch'),
    tfd.Poisson(rate=[1., 10., 100.], name='Three Poissons'),
    tfd.Poisson(rate=[[1., 10., 100.,], [2., 20., 200.]],
                name='Two-by-Three Poissons'),
    tfd.Poisson(rate=[1.], name='One Poisson Vector Batch'),
    tfd.Poisson(rate=[[1.]], name='One Poisson Expanded Batch')
]

describe_distributions(poisson_distributions)

"""
Explanation: このセクションでは、スカラー分布（イベントの形状が [] の分布）について説明します。典型的な例は、rate で指定されたポアソン分布です。
End of explanation
"""


normal_distributions = [
    tfd.Normal(loc=0., scale=1., name='Standard'),
    tfd.Normal(loc=[0.], scale=1., name='Standard Vector Batch'),
    tfd.Normal(loc=[0., 1., 2., 3.], scale=1., name='Different Locs'),
    tfd.Normal(loc=[0., 1., 2., 3.], scale=[[1.], [5.]],
               name='Broadcasting Scale')
]

describe_distributions(normal_distributions)

"""
Explanation: ポアソン分布はスカラー分布であるため、そのイベントの形状は常に [] です。より多くのレートを指定すると、これらはバッチ形式で表示されます。例の最後のペアは興味深いものです。レートは 1 つだけですが、そのレートは空でない形状の numpy 配列に埋め込まれているため、その形状がバッチ形状になります。
標準の正規分布もスカラーです。イベントの形状は、ポアソンの場合と同じように [] ですが、ブロードキャストの最初の例で見ていきます。正規分布は、loc および scale パラメーターを使用して指定されます。
End of explanation
"""


describe_distributions(
    [tfd.Normal(loc=[[0., 1., 2., 3], [0., 1., 2., 3.]],
                scale=[[1., 1., 1., 1.], [5., 5., 5., 5.]])])

"""
Explanation: 上記の Broadcasting Scale 分布は興味深い例です。loc パラメーターは [4] の形状、scale パラメーターは [2, 1] の形状をもちます。Numpy ブロードキャストルールを使用すると、バッチ形状は [2, 4] になります。 "Broadcasting Scale" 分布を定義するための同等の（ただし、あまりエレガントではなく、推奨されない）方法は次のとおりです。
End of explanation
"""


def describe_sample_tensor_shape(sample_shape, distribution):
    print('Sample shape:', sample_shape)
    print('Returned sample tensor shape:',
          distribution.sample(sample_shape).shape)

def describe_sample_tensor_shapes(distributions, sample_shapes):
    started = False
    for distribution in distributions:
      print(distribution)
      for sample_shape in sample_shapes:
        describe_sample_tensor_shape(sample_shape, distribution)
      print()

sample_shapes = [1, 2, [1, 5], [3, 4, 5]]
describe_sample_tensor_shapes(poisson_distributions, sample_shapes)

describe_sample_tensor_shapes(normal_distributions, sample_shapes)

"""
Explanation: 以上のようにブロードキャストの表記は頭痛やバグの原因にもなりますが便利です。
スカラー分布のサンプリング
分布で実行できる主なことは sample と log_prob の 2 つです。まず、サンプリングについて見ていきましょう。基本的なルールは、分布からサンプリングする場合、結果のテンソルは形状 [sample_shape, batch_shape, event_shape] になります。batch_shape と event_shape は Distribution オブジェクトにより提供され、sample_shape は、sample の呼び出しにより提供されます。スカラー分布の場合、event_shape = [] であるため、サンプルから返されるテンソルの形状は [sample_shape, batch_shape] になります。では、試してみましょう。
End of explanation
"""


three_poissons = tfd.Poisson(rate=[1., 10., 100.], name='Three Poissons')
three_poissons

three_poissons.log_prob([[1., 10., 100.], [100., 10., 1]])  # sample_shape is [2].

three_poissons.log_prob([[[[1., 10., 100.], [100., 10., 1.]]]])  # sample_shape is [1, 1, 2].

"""
Explanation: sample についての説明は以上です。返されたサンプルテンソルの形状は [sample_shape, batch_shape, event_shape] です。
スカラー分布の log_prob の計算
次に、log_prob を見てみましょう。これは少し注意する必要があります。log_prob は、分布の log_prob を計算する場所を表す（空でない）テンソルを入力として受け取ります。最も単純なケースでは、このテンソルは [sample_shape, batch_shape, event_shape] の形式になります。batch_shape と event_shape は 分布のバッチおよびイベントの形状に一致します。スカラー分布の場合は、event_shape = [] なので、入力テンソルの形状は [sample_shape, batch_shape] です。この場合、[sample_shape, batch_shape] 形状のテンソルが返されます。
End of explanation
"""


three_poissons.log_prob([10.])

"""
Explanation: 最初の例では、入力と出力の形状が [2, 3] であり、2 番目の例では形状が [1, 1, 2, 3] であることに注意してください。
ブロードキャストがない場合はそれだけです。ブロードキャストを考慮する場合のルールは次のとおりです。これは一般的な説明であり、スカラー分布は簡略化されていることに注意してください。

n = len(batch_shape) + len(event_shape) を定義します。（スカラー分布の場合は、len(event_shape)=0。）
入力テンソル t の次元が n 未満の場合、正確に n 次元になるまで、左側にサイズ 1 の次元を追加して形状をパッディングします。
t' の右端の次元 n を log_prob 計算している分布の [batch_shape, event_shape] に対してブロードキャストします。詳しく説明すると、t' がすでに分布と一致している次元の場合は何もせず、t' の次元がシングルトンの場合は、そのシングルトンを適切な数で複製します。その他の場合はエラーです。（スカラー分布の場合、event_shape = [] であるため、 batch_shape に対してのみブロードキャストします。）
これで、log_prob を計算できるようになりました。結果のテンソルの形状は、[sample_shape, batch_shape] です。sample_shape は、右端の次元 n の左側にある t または t' の任意の次元として定義されます（sample_shape = shape(t)[:-n]）。

これが何を意味するのかわからないと混乱するかもしれないので、いくつかの例を見てみましょう。
End of explanation
"""


three_poissons.log_prob([[[1.], [10.]], [[100.], [1000.]]])

"""
Explanation: テンソル [10.] （形状 [1]）は 3 つのbatch_shape でブロードキャストされるため、値 10 での 3 つのポワソンの対数確率をすべて評価します。
End of explanation
"""


poisson_2_by_3 = tfd.Poisson(
    rate=[[1., 10., 100.,], [2., 20., 200.]],
    name='Two-by-Three Poissons')

poisson_2_by_3.log_prob(1.)

poisson_2_by_3.log_prob([1.])  # Exactly equivalent to above, demonstrating the scalar special case.

poisson_2_by_3.log_prob([[1., 1., 1.], [1., 1., 1.]])  # Another way to write the same thing. No broadcasting.

poisson_2_by_3.log_prob([[1., 10., 100.]])  # Input is [1, 3] broadcast to [2, 3].

poisson_2_by_3.log_prob([[1., 10., 100.], [1., 10., 100.]])  # Equivalent to above. No broadcasting.

poisson_2_by_3.log_prob([[1., 1., 1.], [2., 2., 2.]])  # No broadcasting.

poisson_2_by_3.log_prob([[1.], [2.]])  # Equivalent to above. Input shape [2, 1] broadcast to [2, 3].

"""
Explanation: 上記の例では、入力テンソルの形状は [2, 2, 1] ですが、分布オブジェクトの形状は 3 です。したがって、[2, 2] サンプル次元のそれぞれについて、提供された単一の値は、3 つのポワソンのそれぞれにブロードキャストします。
これは役に立つ考え方です。three_poissons には batch_shape = [2, 3] があるため、log_prob の呼び出しには最後の次元が 1 または 3 のテンソルが必要です。それ以外はエラーです。（numpy ブロードキャストルールは、スカラーの特殊なケースを、形状 [1] のテンソルと完全に同等であるものとして扱います。）
では、batch_shape = [2, 3] を使用して、より複雑なポアソン分布を使用して試してみましょう。
End of explanation
"""


poisson_2_by_3.log_prob([[[1., 1., 1.], [1., 1., 1.]], [[2., 2., 2.], [2., 2., 2.]]])  # Input shape [2, 2, 3].

"""
Explanation: 上記の例では、バッチを介したブロードキャストを見ていきましたが、サンプルの形状は空でした。値のコレクションがあり、バッチの各ポイントで各値の対数確率を取得する場合は、以下のように手動で実行できます。
End of explanation
"""


poisson_2_by_3.log_prob([[[1.], [1.]], [[2.], [2.]]])  # Input shape [2, 2, 1].

"""
Explanation: または、ブロードキャストに最後のバッチ次元を処理させることもできます。
End of explanation
"""


poisson_2_by_3.log_prob([[[1., 1., 1.]], [[2., 2., 2.]]])  # Input shape [2, 1, 3].

"""
Explanation: また、やや不自然ですがブロードキャストに最初のバッチ次元のみを処理させることもできます。
End of explanation
"""


poisson_2_by_3.log_prob([[[1.]], [[2.]]])  # Input shape [2, 1, 1].

"""
Explanation: または、ブロードキャストに両方のバッチ次元を処理させることもできます。
End of explanation
"""


poisson_2_by_3.log_prob(tf.constant([1., 2.])[..., tf.newaxis, tf.newaxis])

"""
Explanation: 上記は、必要な値が 2 つしかない場合は問題ありませんでした。しかし、すべてのバッチポイントで評価する値のリストが長い場合は、次の表記を使用します。形状の右側にサイズ 1 の余分な次元を追加すると、非常に便利です。
End of explanation
"""


three_poissons.log_prob([[1.], [10.], [50.], [100.]])

three_poissons.log_prob(tf.constant([1., 10., 50., 100.])[..., tf.newaxis])  # Equivalent to above.

"""
Explanation: これはストライドスライス表記のインスタンスであり、知っておく価値があります。
完全を期すために three_poissons に戻ると、同じ例は次のようになります。
End of explanation
"""


multinomial_distributions = [
    # Multinomial is a vector-valued distribution: if we have k classes,
    # an individual sample from the distribution has k values in it, so the
    # event_shape is `[k]`.
    tfd.Multinomial(total_count=100., probs=[.5, .4, .1],
                    name='One Multinomial'),
    tfd.Multinomial(total_count=[100., 1000.], probs=[.5, .4, .1],
                    name='Two Multinomials Same Probs'),
    tfd.Multinomial(total_count=100., probs=[[.5, .4, .1], [.1, .2, .7]],
                    name='Two Multinomials Same Counts'),
    tfd.Multinomial(total_count=[100., 1000.],
                    probs=[[.5, .4, .1], [.1, .2, .7]],
                    name='Two Multinomials Different Everything')

]

describe_distributions(multinomial_distributions)

"""
Explanation: 多変量分布
ここでは、空でないイベント形状を持つ多変量分布を見ていきます。まず、多項分布を見てみましょう。
End of explanation
"""


describe_sample_tensor_shapes(multinomial_distributions, sample_shapes)

"""
Explanation: 最後の 3 つの例では、batch_shape は常に [2] でしたが、ブロードキャストを使用して、共有する total_count または共有する probs 使用できます（または、使用しないこともできます）。内部では同じ形状になるようにブロードキャストされるためです。
既知の事柄を考慮すると、サンプリングは簡単です。
End of explanation
"""


two_multivariate_normals = tfd.MultivariateNormalDiag(loc=[1., 2., 3.], scale_identity_multiplier=[1., 2.])
two_multivariate_normals

"""
Explanation: 対数確率の計算も同様に簡単です。対角多変量正規分布の例を見てみましょう。（カウントと確率の制約により、ブロードキャストは許容できない値を生成することが多いため、多項分布はブロードキャストにあまり適していません。）平均は同じですがスケール（標準偏差）が異なる 2 つの 3 次元分布のバッチを使用します。
End of explanation
"""


two_multivariate_normals.log_prob([[[1., 2., 3.]], [[3., 4., 5.]]])  # Input has shape [2,1,3].

"""
Explanation: （スケールが ID の倍数である分布を使用したが、これは制限ではないことに注意してください。scale_identity_multiplier の代わりに scale を渡すことができます。）
次に、各バッチポイントの平均とシフトされた平均での対数確率を評価します。
End of explanation
"""


two_multivariate_normals.log_prob(
    tf.constant([[1., 2., 3.], [3., 4., 5.]])[:, tf.newaxis, :])  # Equivalent to above.

"""
Explanation: まったく同じように、[https://www.tensorflow.org/api_docs/cc/class/tensorflow/ops/strided-slice]（ストライドスライス表記）を使用して、定数の中央に追加の形状 = 1 次元を挿入できます。
End of explanation
"""


two_multivariate_normals.log_prob(tf.constant([[1., 2., 3.], [3., 4., 5.]]))

"""
Explanation: 一方、余分な次元を追加しない場合は、[1., 2., 3.] を最初のバッチポイントに渡し、[3., 4., 5.] を 2 番目のバッチポイントに渡します。
End of explanation
"""


six_way_multinomial = tfd.Multinomial(total_count=1000., probs=[.3, .25, .2, .15, .08, .02])
six_way_multinomial

"""
Explanation: 形状変換テクニック
Reshape Bijector
Reshape Bijector を使用すると、分布の event_shape の形状を変換できます。以下に例を示します。
End of explanation
"""


transformed_multinomial = tfd.TransformedDistribution(
    distribution=six_way_multinomial,
    bijector=tfb.Reshape(event_shape_out=[2, 3]))
transformed_multinomial

six_way_multinomial.log_prob([500., 100., 100., 150., 100., 50.])

transformed_multinomial.log_prob([[500., 100., 100.], [150., 100., 50.]])

"""
Explanation: [6] のイベント形状を持つ多項分布を作成しました。Reshape Bijector を使用すると、これを [2, 3] のイベント形状を持つ分布として扱うことができます。
Bijector は、${\mathbb R}^n$ の開集合上の微分可能な 1 対 1 の関数を表します。Bijectors は、TransformedDistribution と組み合わせて使用されます。これは、基本分布 $p(x)$ および$Y = g(X)$ を表す Bijector に関して分布 $p(y)$ をモデル化します。では、実際に見てみましょう。
End of explanation
"""


two_by_five_bernoulli = tfd.Bernoulli(
    probs=[[.05, .1, .15, .2, .25], [.3, .35, .4, .45, .5]],
    name="Two By Five Bernoulli")
two_by_five_bernoulli

"""
Explanation: これは、Reshape Bijector が実行できる唯一のことです。イベント次元をバッチ次元に、またはバッチ次元をイベント次元に変換することはできません。
Independent 分布
Independent 分布は、独立した、必ずしも同一ではない分布（バッチ）のコレクションを単一の分布として扱うために使用されます。より簡潔に言えば、Independent を使用すると、batch_shape の次元を event_shape の次元に変換できます。次に例を示します。
End of explanation
"""


pattern = [[1., 0., 0., 1., 0.], [0., 0., 1., 1., 1.]]
two_by_five_bernoulli.log_prob(pattern)

"""
Explanation: これは、表の確率が関連付けられた 2x5 のコインの配列として考えることができます。特定の任意の 1 と 0 のセットの確率を評価します。
End of explanation
"""


two_sets_of_five = tfd.Independent(
    distribution=two_by_five_bernoulli,
    reinterpreted_batch_ndims=1,
    name="Two Sets Of Five")
two_sets_of_five

"""
Explanation: Independent を使用すると、これを 2 つの異なる「5 つのベルヌーイのセット」に変換できます。これは、特定のパターンで出現するコイントスの「行」を単一の結果と見なす場合に役立ちます。
End of explanation
"""


two_sets_of_five.log_prob(pattern)

"""
Explanation: 数学的には、5 つの「セット」ごとの対数確率を計算しています。セット内の 5 つの「独立した」コイントスの対数確率を合計するため、分布は「independent」と呼ばれます。
End of explanation
"""


one_set_of_two_by_five = tfd.Independent(
    distribution=two_by_five_bernoulli, reinterpreted_batch_ndims=2,
    name="One Set Of Two By Five")
one_set_of_two_by_five.log_prob(pattern)

"""
Explanation: さらに、Independent を使用して、個々のイベントが 2x5 のベルヌーイのセットである分布を作成できます。
End of explanation
"""


describe_sample_tensor_shapes(
    [two_by_five_bernoulli,
     two_sets_of_five,
     one_set_of_two_by_five],
    [[3, 5]])

"""
Explanation: sample の観点では、Independent を使用しても何も変更されないことに注意してください。
End of explanation
"""


