#@title Licensed under the Apache License, Version 2.0 (the "License");
# 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 2020 The TensorFlow Authors.
End of explanation
"""


!pip install tensorflow==2.7.0

"""
Explanation: 勾配の計算
<table class="tfo-notebook-buttons" align="left">
  <td>     <a target="_blank" href="https://www.tensorflow.org/quantum/tutorials/gradients">     <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/quantum/tutorials/gradients.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/quantum/tutorials/gradients.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/quantum/tutorials/gradients.ipynb"><img src="https://www.tensorflow.org/images/download_logo_32px.png">ノートブックをダウンロード</a></td>
</table>

このチュートリアルでは、量子回路の期待値の勾配計算アルゴリズムについて説明します。
量子回路で特定のオブザーバブルの期待値の勾配を計算することは、複雑なプロセスです。行列の乗算やベクトルの加算などの従来の機械学習変換では簡単に使用できる解析的勾配式がありますが、オブザーバブルの期待値には、このような解析的勾配式は必ずしもありません。そのため、シナリオに適したさまざまな量子勾配計算方法を使用する必要があります。このチュートリアルでは、2 つの異なる微分スキームを比較対照します。
セットアップ
End of explanation
"""


!pip install tensorflow-quantum

# Update package resources to account for version changes.
import importlib, pkg_resources
importlib.reload(pkg_resources)

"""
Explanation: TensorFlow Quantum をインストールします。
End of explanation
"""


import tensorflow as tf
import tensorflow_quantum as tfq

import cirq
import sympy
import numpy as np

# visualization tools
%matplotlib inline
import matplotlib.pyplot as plt
from cirq.contrib.svg import SVGCircuit

"""
Explanation: 次に、TensorFlow とモジュールの依存関係をインポートします。
End of explanation
"""


qubit = cirq.GridQubit(0, 0)
my_circuit = cirq.Circuit(cirq.Y(qubit)**sympy.Symbol('alpha'))
SVGCircuit(my_circuit)

"""
Explanation: 1. 予備
量子回路の勾配計算の概念をもう少し具体的に見てみましょう。次のようなパラメータ化された回路があるとします。
End of explanation
"""


pauli_x = cirq.X(qubit)
pauli_x

"""
Explanation: オブザーバブルは以下のとおりです。
End of explanation
"""


def my_expectation(op, alpha):
    """Compute ⟨Y(alpha)| `op` | Y(alpha)⟩"""
    params = {'alpha': alpha}
    sim = cirq.Simulator()
    final_state_vector = sim.simulate(my_circuit, params).final_state_vector
    return op.expectation_from_state_vector(final_state_vector, {qubit: 0}).real


my_alpha = 0.3
print("Expectation=", my_expectation(pauli_x, my_alpha))
print("Sin Formula=", np.sin(np.pi * my_alpha))

"""
Explanation: この演算子を見ると、$⟨Y(\alpha)| X | Y(\alpha)⟩ = \sin(\pi \alpha)$ であることが分かります。
End of explanation
"""


def my_grad(obs, alpha, eps=0.01):
    grad = 0
    f_x = my_expectation(obs, alpha)
    f_x_prime = my_expectation(obs, alpha + eps)
    return ((f_x_prime - f_x) / eps).real


print('Finite difference:', my_grad(pauli_x, my_alpha))
print('Cosine formula:   ', np.pi * np.cos(np.pi * my_alpha))

"""
Explanation: $f_{1}(\alpha) = ⟨Y(\alpha)| X | Y(\alpha)⟩$ と定義すると $f_{1}^{'}(\alpha) = \pi \cos(\pi \alpha)$ になります。確認しましょう。
End of explanation
"""


expectation_calculation = tfq.layers.Expectation(
    differentiator=tfq.differentiators.ForwardDifference(grid_spacing=0.01))

expectation_calculation(my_circuit,
                        operators=pauli_x,
                        symbol_names=['alpha'],
                        symbol_values=[[my_alpha]])

"""
Explanation: 2. 微分器の必要性
より大きな回路では、与えられた量子回路の勾配を正確に計算する式は必ずしもありません。単純な式では勾配を計算できない場合、tfq.differentiators.Differentiatorクラスを使用すると、回路の勾配を計算するためのアルゴリズムを定義できます。たとえば、TensorFlow Quantum（TFQ）で上記の例を次のように再現できます。
End of explanation
"""


sampled_expectation_calculation = tfq.layers.SampledExpectation(
    differentiator=tfq.differentiators.ForwardDifference(grid_spacing=0.01))

sampled_expectation_calculation(my_circuit,
                                operators=pauli_x,
                                repetitions=500,
                                symbol_names=['alpha'],
                                symbol_values=[[my_alpha]])

"""
Explanation: ただし、サンプリングに基づいて期待値を推定するように切り替えると（実際のデバイスで何が起こるか）、値が少し変わる可能性があります。これは、期待値が不正確になることを意味します。
End of explanation
"""


# Make input_points = [batch_size, 1] array.
input_points = np.linspace(0, 5, 200)[:, np.newaxis].astype(np.float32)
exact_outputs = expectation_calculation(my_circuit,
                                        operators=pauli_x,
                                        symbol_names=['alpha'],
                                        symbol_values=input_points)
imperfect_outputs = sampled_expectation_calculation(my_circuit,
                                                    operators=pauli_x,
                                                    repetitions=500,
                                                    symbol_names=['alpha'],
                                                    symbol_values=input_points)
plt.title('Forward Pass Values')
plt.xlabel('$x$')
plt.ylabel('$f(x)$')
plt.plot(input_points, exact_outputs, label='Analytic')
plt.plot(input_points, imperfect_outputs, label='Sampled')
plt.legend()

# Gradients are a much different story.
values_tensor = tf.convert_to_tensor(input_points)

with tf.GradientTape() as g:
    g.watch(values_tensor)
    exact_outputs = expectation_calculation(my_circuit,
                                            operators=pauli_x,
                                            symbol_names=['alpha'],
                                            symbol_values=values_tensor)
analytic_finite_diff_gradients = g.gradient(exact_outputs, values_tensor)

with tf.GradientTape() as g:
    g.watch(values_tensor)
    imperfect_outputs = sampled_expectation_calculation(
        my_circuit,
        operators=pauli_x,
        repetitions=500,
        symbol_names=['alpha'],
        symbol_values=values_tensor)
sampled_finite_diff_gradients = g.gradient(imperfect_outputs, values_tensor)

plt.title('Gradient Values')
plt.xlabel('$x$')
plt.ylabel('$f^{\'}(x)$')
plt.plot(input_points, analytic_finite_diff_gradients, label='Analytic')
plt.plot(input_points, sampled_finite_diff_gradients, label='Sampled')
plt.legend()

"""
Explanation: これは、勾配における深刻な精度の問題につながる可能性があります。
End of explanation
"""


# A smarter differentiation scheme.
gradient_safe_sampled_expectation = tfq.layers.SampledExpectation(
    differentiator=tfq.differentiators.ParameterShift())

with tf.GradientTape() as g:
    g.watch(values_tensor)
    imperfect_outputs = gradient_safe_sampled_expectation(
        my_circuit,
        operators=pauli_x,
        repetitions=500,
        symbol_names=['alpha'],
        symbol_values=values_tensor)

sampled_param_shift_gradients = g.gradient(imperfect_outputs, values_tensor)

plt.title('Gradient Values')
plt.xlabel('$x$')
plt.ylabel('$f^{\'}(x)$')
plt.plot(input_points, analytic_finite_diff_gradients, label='Analytic')
plt.plot(input_points, sampled_param_shift_gradients, label='Sampled')
plt.legend()

"""
Explanation: ここでは、解析の場合は有限差分式は勾配自体を高速に計算できますが、サンプリングベースの方法の場合ではノイズが多すぎることが分かります。適切な勾配を計算するには、より注意深い手法を使用する必要があります。次に、解析的期待値の勾配計算にはあまり適していませんが、実際のサンプルベースの方法の場合ではより優れたパフォーマンスを発揮する、大幅に低速な手法を見ていきます。
End of explanation
"""


pauli_z = cirq.Z(qubit)
pauli_z

"""
Explanation: 上記から、特定の微分器が特定の研究シナリオに最適であることがわかります。一般に、デバイスノイズなどに対して堅牢な、低速のサンプルベースの方法は、より「現実的」設定でアルゴリズムをテストまたは実装する場合に適した微分器です。有限差分のようなより高速な方法はアルゴリズムのデバイスにおける実行可能性にはまだ関心がなく、解析的計算やより高いスループットが必要な場合に最適です。
3. 複数のオブザーバブル
2 番目のオブザーバブルを使用し、TensorFlow Quantum が 1 つの回路に対して複数のオブザーバブルをサポートする方法を見てみましょう。
End of explanation
"""


test_value = 0.

print('Finite difference:', my_grad(pauli_z, test_value))
print('Sin formula:      ', -np.pi * np.sin(np.pi * test_value))

"""
Explanation: このオブザーバブルが以前と同じ回路で使用されている場合、$f_{2}(\alpha) = ⟨Y(\alpha)| Z | Y(\alpha)⟩ = \cos(\pi \alpha)$ および$f_{2}^{'}(\alpha) = -\pi \sin(\pi \alpha)$になります。確認します。
End of explanation
"""


sum_of_outputs = tfq.layers.Expectation(
    differentiator=tfq.differentiators.ForwardDifference(grid_spacing=0.01))

sum_of_outputs(my_circuit,
               operators=[pauli_x, pauli_z],
               symbol_names=['alpha'],
               symbol_values=[[test_value]])

"""
Explanation: （ほぼ）一致します。
次に、$g(\alpha) = f_{1}(\alpha) + f_{2}(\alpha)$ を定義すると、$g'(\alpha) = f_{1}^{'}(\alpha) + f^{'}_{2}(\alpha)$になります。TensorFlow Quantum で複数のオブザーバブルを定義して回路と共に使用するには、$g$ にさらに項を追加します。
これは、回路内の特定のシンボルの勾配が、その回路に適用されたそのシンボルの各オブザーバブルに関する勾配の合計に等しいことを意味します。これは、TensorFlow の勾配取得およびバックプロパゲーションと互換性があります（特定のシンボルの勾配として、すべてのオブザーバブルの勾配の合計を指定します）。
End of explanation
"""


test_value_tensor = tf.convert_to_tensor([[test_value]])

with tf.GradientTape() as g:
    g.watch(test_value_tensor)
    outputs = sum_of_outputs(my_circuit,
                             operators=[pauli_x, pauli_z],
                             symbol_names=['alpha'],
                             symbol_values=test_value_tensor)

sum_of_gradients = g.gradient(outputs, test_value_tensor)

print(my_grad(pauli_x, test_value) + my_grad(pauli_z, test_value))
print(sum_of_gradients.numpy())

"""
Explanation: ここで、最初のエントリは期待値 w.r.t Pauli X であり、2 番目のエントリは期待値 w.r.t Pauli Z です。勾配は以下のとおりです。
End of explanation
"""


class MyDifferentiator(tfq.differentiators.Differentiator):
    """A Toy differentiator for <Y^alpha | X |Y^alpha>."""

    def __init__(self):
        pass

    def get_gradient_circuits(self, programs, symbol_names, symbol_values):
        """Return circuits to compute gradients for given forward pass circuits.
        
        Every gradient on a quantum computer can be computed via measurements
        of transformed quantum circuits.  Here, you implement a custom gradient
        for a specific circuit.  For a real differentiator, you will need to
        implement this function in a more general way.  See the differentiator
        implementations in the TFQ library for examples.
        """

        # The two terms in the derivative are the same circuit...
        batch_programs = tf.stack([programs, programs], axis=1)

        # ... with shifted parameter values.
        shift = tf.constant(1/2)
        forward = symbol_values + shift
        backward = symbol_values - shift
        batch_symbol_values = tf.stack([forward, backward], axis=1)
  
        # Weights are the coefficients of the terms in the derivative.
        num_program_copies = tf.shape(batch_programs)[0]
        batch_weights = tf.tile(tf.constant([[[np.pi/2, -np.pi/2]]]),
                                [num_program_copies, 1, 1])

        # The index map simply says which weights go with which circuits.
        batch_mapper = tf.tile(
            tf.constant([[[0, 1]]]), [num_program_copies, 1, 1])

        return (batch_programs, symbol_names, batch_symbol_values,
                batch_weights, batch_mapper)

"""
Explanation: ここで、各オブザーバブルの勾配の合計が実際に $\alpha$ の勾配であることを確認しました。この動作は、すべての TensorFlow Quantum 微分器によってサポートされており、TensorFlow の他の部分との互換性において重要な役割を果たします。
4. 高度な使用
TensorFlow Quantum 内にあるすべての微分器は tfq.differentiators.Differentiator をサブクラス化します。微分器を実装するには、ユーザーは 2 つのインターフェースのいずれかを実装する必要があります。標準的なのは get_gradient_circuits を実装することで、これは勾配の推定を取得するためにどの回路を測定するのかを基本クラスに指定します。または、differentiate_analytic と differentiate_sampled をオーバーロードすることもできます。クラス tfq.differentiators.Adjoint はこの経路を使用します。
以下は TensorFlow Quantum を使用して、サーキットの勾配を実装しています。パラメーターシフトの小さな例を使用します。
上記で定義した $|\alpha⟩ = Y^{\alpha}|0⟩$ という回路を思い出してみましょう。前と同様に、$X$ 観測可能に対するこの回路の期待値として関数を定義できます（$f(\alpha) = ⟨\alpha|X|\alpha⟩$）。パラメーターシフトの規則を使用すると、この回路では、導関数は $$\frac{\partial}{\partial \alpha} f(\alpha) = \frac{\pi}{2} f\left(\alpha + \frac{1}{2}\right) -  \frac{ \pi}{2} f\left(\alpha - \frac{1}{2}\right)$$ であることを見つけ出すことができます。get_gradient_circuits 関数は、この導関数の成分を返します。
End of explanation
"""


custom_dif = MyDifferentiator()
custom_grad_expectation = tfq.layers.Expectation(differentiator=custom_dif)

# Now let's get the gradients with finite diff.
with tf.GradientTape() as g:
    g.watch(values_tensor)
    exact_outputs = expectation_calculation(my_circuit,
                                            operators=[pauli_x],
                                            symbol_names=['alpha'],
                                            symbol_values=values_tensor)

analytic_finite_diff_gradients = g.gradient(exact_outputs, values_tensor)

# Now let's get the gradients with custom diff.
with tf.GradientTape() as g:
    g.watch(values_tensor)
    my_outputs = custom_grad_expectation(my_circuit,
                                         operators=[pauli_x],
                                         symbol_names=['alpha'],
                                         symbol_values=values_tensor)

my_gradients = g.gradient(my_outputs, values_tensor)

plt.subplot(1, 2, 1)
plt.title('Exact Gradient')
plt.plot(input_points, analytic_finite_diff_gradients.numpy())
plt.xlabel('x')
plt.ylabel('f(x)')
plt.subplot(1, 2, 2)
plt.title('My Gradient')
plt.plot(input_points, my_gradients.numpy())
plt.xlabel('x')

"""
Explanation: Differentiator の基本クラスは、上記で見たパラメーターシフトの式のように、get_gradient_circuits から返された成分を使用して導関数を計算します。この新しい微分器を tfq.layer オブジェクトで使用することができます。
End of explanation
"""


# Create a noisy sample based expectation op.
expectation_sampled = tfq.get_sampled_expectation_op(
    cirq.DensityMatrixSimulator(noise=cirq.depolarize(0.01)))

# Make it differentiable with your differentiator:
# Remember to refresh the differentiator before attaching the new op
custom_dif.refresh()
differentiable_op = custom_dif.generate_differentiable_op(
    sampled_op=expectation_sampled)

# Prep op inputs.
circuit_tensor = tfq.convert_to_tensor([my_circuit])
op_tensor = tfq.convert_to_tensor([[pauli_x]])
single_value = tf.convert_to_tensor([[my_alpha]])
num_samples_tensor = tf.convert_to_tensor([[5000]])

with tf.GradientTape() as g:
    g.watch(single_value)
    forward_output = differentiable_op(circuit_tensor, ['alpha'], single_value,
                                       op_tensor, num_samples_tensor)

my_gradients = g.gradient(forward_output, single_value)

print('---TFQ---')
print('Foward:  ', forward_output.numpy())
print('Gradient:', my_gradients.numpy())
print('---Original---')
print('Forward: ', my_expectation(pauli_x, my_alpha))
print('Gradient:', my_grad(pauli_x, my_alpha))

"""
Explanation: この新しい微分器を使用して、微分可能な演算を生成できるようになりました。
重要点：微分器は一度に 1 つの演算にしか接続できないため、以前に演算に接続されていた微分器は、新しい演算に接続する前に更新する必要があります。
End of explanation
"""


