#@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/qcnn"><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/qcnn.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/qcnn.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/qcnn.ipynb"><img src="https://www.tensorflow.org/images/download_logo_32px.png">ノートブックをダウンロード</a></td>
</table>

このチュートリアルでは、単純な<a href="https://www.nature.com/articles/s41567-019-0648-8" class="external">量子畳み込みニューラルネットワーク</a>（QCNN）を実装します。QCNN は、並進的に不変でもある古典的な畳み込みニューラルネットワークに提案された量子アナログです。
この例では、デバイスの量子センサまたは複雑なシミュレーションなど、量子データソースの特定のプロパティを検出する方法を実演します。量子データソースは、励起の有無にかかわらず<a href="https://arxiv.org/pdf/quant-ph/0504097.pdf" class="external">クラスタ状態</a>です。QCNN はこの検出を学習します（論文で使用されたデータセットは SPT フェーズ分類です）。
セットアップ
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)

# Define some circuits.
circuit1 = cirq.Circuit(cirq.X(qubit))
circuit2 = cirq.Circuit(cirq.H(qubit))

# Convert to a tensor.
input_circuit_tensor = tfq.convert_to_tensor([circuit1, circuit2])

# Define a circuit that we want to append
y_circuit = cirq.Circuit(cirq.Y(qubit))

# Instantiate our layer
y_appender = tfq.layers.AddCircuit()

# Run our circuit tensor through the layer and save the output.
output_circuit_tensor = y_appender(input_circuit_tensor, append=y_circuit)

"""
Explanation: 1. QCNN を構築する
1.1 TensorFlow グラフで回路を組み立てる
TensorFlow Quantum（TFQ）には、グラフ内で回路を構築するために設計されたレイヤークラスがあります。たとえば tfq.layers.AddCircuit レイヤーがあり、tf.keras.Layer を継承しています。このレイヤーは、次の図で示すように、回路の入力バッチの前後いずれかに追加できます。
<img src="./images/qcnn_1.png" width="700">
次のスニペットには、このレイヤーが使用されています。
End of explanation
"""


print(tfq.from_tensor(input_circuit_tensor))

"""
Explanation: 入力テンソルを調べます。
End of explanation
"""


print(tfq.from_tensor(output_circuit_tensor))

"""
Explanation: 次に、出力テンソルを調べます。
End of explanation
"""


def generate_data(qubits):
    """Generate training and testing data."""
    n_rounds = 20  # Produces n_rounds * n_qubits datapoints.
    excitations = []
    labels = []
    for n in range(n_rounds):
        for bit in qubits:
            rng = np.random.uniform(-np.pi, np.pi)
            excitations.append(cirq.Circuit(cirq.rx(rng)(bit)))
            labels.append(1 if (-np.pi / 2) <= rng <= (np.pi / 2) else -1)

    split_ind = int(len(excitations) * 0.7)
    train_excitations = excitations[:split_ind]
    test_excitations = excitations[split_ind:]

    train_labels = labels[:split_ind]
    test_labels = labels[split_ind:]

    return tfq.convert_to_tensor(train_excitations), np.array(train_labels), \
        tfq.convert_to_tensor(test_excitations), np.array(test_labels)

"""
Explanation: 以下の例は tfq.layers.AddCircuit を使用せずに実行できますが、TensorFlow 計算グラフに複雑な機能を埋め込む方法を理解する上で役立ちます。
1.2 問題の概要
クラスター状態を準備し、「励起」があるかどうかを検出する量子分類器をトレーニングします。クラスター状態は極めてこじれていますが、古典的コンピュータにおいては必ずしも困難ではありません。わかりやすく言えば、これは論文で使用されているデータセットよりも単純です。
この分類タスクでは、次の理由により、ディープ <a href="https://arxiv.org/pdf/quant-ph/0610099.pdf" class="external">MERA</a> のような QCNN アーキテクチャを実装します。

QCNN と同様に、リングのクラスター状態は並進的に不変である
クラスター状態は非常にもつれている

このアーキテクチャはエンタングルメントを軽減し、単一のキュービットを読み出すことで分類を取得する上で効果があります。
<img src="./images/qcnn_2.png" width="1000">
「励起」のあるクラスター状態は、cirq.rx ゲートがすべてのキュービットに適用されたクラスター状態として定義されます。Qconv と QPool については、このチュートリアルの後の方で説明しています。
1.3 TensorFlow のビルディングブロック
<img src="./images/qcnn_3.png" width="1000">
TensorFlow Quantum を使ってこの問題を解決する方法として、次を実装することが挙げられます。

モデルへの入力は回路で、空の回路か励起を示す特定のキュー人における X ゲートです。
モデルの残りの量子コンポーネントは、tfq.layers.AddCircuit レイヤーで作成されます。
たとえば tfq.layers.PQC レイヤーが使用されているとした場合、$\langle \hat{Z} \rangle$ を読み取って、励起のある状態には 1 のラベルと、励起のない状態には -1 のラベルと比較します。

1.4 データ
モデルを構築する前に、データを生成することができます。この場合には、クラスター状態に励起がは一斉思案す（元の論文では、より複雑なデータセットが使用されています）。励起は、cirq.rx ゲートで表されます。十分に大きい回転は励起と見なされ、1 とラベル付けされ、十分に大きくない回転は -1 とラベル付けされ、励起ではないと見なされます。
End of explanation
"""


sample_points, sample_labels, _, __ = generate_data(cirq.GridQubit.rect(1, 4))
print('Input:', tfq.from_tensor(sample_points)[0], 'Output:', sample_labels[0])
print('Input:', tfq.from_tensor(sample_points)[1], 'Output:', sample_labels[1])

"""
Explanation: 通常の機械学習と同じように、モデルのベンチマークに使用するトレーニングとテストのセットを作成していることがわかります。次のようにすると、データポイントを素早く確認できます。
End of explanation
"""


def cluster_state_circuit(bits):
    """Return a cluster state on the qubits in `bits`."""
    circuit = cirq.Circuit()
    circuit.append(cirq.H.on_each(bits))
    for this_bit, next_bit in zip(bits, bits[1:] + [bits[0]]):
        circuit.append(cirq.CZ(this_bit, next_bit))
    return circuit

"""
Explanation: 1.5 レイヤーを定義する
上記の図で示すレイヤーを TensorFlow で定義しましょう。
1.5.1 クラスター状態
まず始めに、<a href="https://arxiv.org/pdf/quant-ph/0504097.pdf" class="external">クラスター状態</a>を定義しますが、これには Google が量子回路のプログラミング用に提供している <a href="https://github.com/quantumlib/Cirq" class="external">Cirq</a> フレームワークを使用します。モデルの静的な部分であるため、tfq.layers.AddCircuit 機能を使用して埋め込みます。
End of explanation
"""


SVGCircuit(cluster_state_circuit(cirq.GridQubit.rect(1, 4)))

"""
Explanation: 矩形の <a href="https://cirq.readthedocs.io/en/stable/generated/cirq.GridQubit.html" class="external"><code>cirq.GridQubit</code></a> のクラスター状態回路を表示します。
End of explanation
"""


def one_qubit_unitary(bit, symbols):
    """Make a Cirq circuit enacting a rotation of the bloch sphere about the X,
    Y and Z axis, that depends on the values in `symbols`.
    """
    return cirq.Circuit(
        cirq.X(bit)**symbols[0],
        cirq.Y(bit)**symbols[1],
        cirq.Z(bit)**symbols[2])


def two_qubit_unitary(bits, symbols):
    """Make a Cirq circuit that creates an arbitrary two qubit unitary."""
    circuit = cirq.Circuit()
    circuit += one_qubit_unitary(bits[0], symbols[0:3])
    circuit += one_qubit_unitary(bits[1], symbols[3:6])
    circuit += [cirq.ZZ(*bits)**symbols[6]]
    circuit += [cirq.YY(*bits)**symbols[7]]
    circuit += [cirq.XX(*bits)**symbols[8]]
    circuit += one_qubit_unitary(bits[0], symbols[9:12])
    circuit += one_qubit_unitary(bits[1], symbols[12:])
    return circuit


def two_qubit_pool(source_qubit, sink_qubit, symbols):
    """Make a Cirq circuit to do a parameterized 'pooling' operation, which
    attempts to reduce entanglement down from two qubits to just one."""
    pool_circuit = cirq.Circuit()
    sink_basis_selector = one_qubit_unitary(sink_qubit, symbols[0:3])
    source_basis_selector = one_qubit_unitary(source_qubit, symbols[3:6])
    pool_circuit.append(sink_basis_selector)
    pool_circuit.append(source_basis_selector)
    pool_circuit.append(cirq.CNOT(control=source_qubit, target=sink_qubit))
    pool_circuit.append(sink_basis_selector**-1)
    return pool_circuit

"""
Explanation: 1.5.2 QCNN レイヤー
<a href="https://arxiv.org/abs/1810.03787" class="external">Cong and Lukin の QCNN に関する論文</a>を使用して、モデルを構成するレイヤーを定義します。これには次の前提条件があります。

<a href="https://arxiv.org/abs/quant-ph/0507171" class="external">Tucci の論文</a>にある 1 キュービットと 2 キュービットのパラメータ化されたユニタリ―行列
一般的なパラメータ化された 2 キュービットプーリング演算
End of explanation
"""


SVGCircuit(one_qubit_unitary(cirq.GridQubit(0, 0), sympy.symbols('x0:3')))

"""
Explanation: 作成したものを確認するために、1 キュービットのユニタリー回路を出力しましょう。
End of explanation
"""


SVGCircuit(two_qubit_unitary(cirq.GridQubit.rect(1, 2), sympy.symbols('x0:15')))

"""
Explanation: 次に、2 キュービットのユニタリー回路を出力します。
End of explanation
"""


SVGCircuit(two_qubit_pool(*cirq.GridQubit.rect(1, 2), sympy.symbols('x0:6')))

"""
Explanation: そして 2 キュービットのプーリング回路を出力します。
End of explanation
"""


def quantum_conv_circuit(bits, symbols):
    """Quantum Convolution Layer following the above diagram.
    Return a Cirq circuit with the cascade of `two_qubit_unitary` applied
    to all pairs of qubits in `bits` as in the diagram above.
    """
    circuit = cirq.Circuit()
    for first, second in zip(bits[0::2], bits[1::2]):
        circuit += two_qubit_unitary([first, second], symbols)
    for first, second in zip(bits[1::2], bits[2::2] + [bits[0]]):
        circuit += two_qubit_unitary([first, second], symbols)
    return circuit

"""
Explanation: 1.5.2.1 量子畳み込み
<a href="https://arxiv.org/abs/1810.03787" class="external">Cong と Lukin</a> の論文にあるとおり、1 次元量子畳み込みを、ストライド 1 の隣接するすべてのキュービットペアに 2 キュービットのパラメーター化されたユニタリの適用として定義します。
End of explanation
"""


SVGCircuit(
    quantum_conv_circuit(cirq.GridQubit.rect(1, 8), sympy.symbols('x0:15')))

"""
Explanation: （非常に水平な）回路を表示します。
End of explanation
"""


def quantum_pool_circuit(source_bits, sink_bits, symbols):
    """A layer that specifies a quantum pooling operation.
    A Quantum pool tries to learn to pool the relevant information from two
    qubits onto 1.
    """
    circuit = cirq.Circuit()
    for source, sink in zip(source_bits, sink_bits):
        circuit += two_qubit_pool(source, sink, symbols)
    return circuit

"""
Explanation: 1.5.2.2 量子プーリング
量子プーリングレイヤーは、上記で定義された 2 キュービットプールを使用して、$N$ キュービットから $\frac{N}{2}$ キュービットまでをプーリングします。
End of explanation
"""


test_bits = cirq.GridQubit.rect(1, 8)

SVGCircuit(
    quantum_pool_circuit(test_bits[:4], test_bits[4:], sympy.symbols('x0:6')))

"""
Explanation: プーリングコンポーネント回路を調べます。
End of explanation
"""


def create_model_circuit(qubits):
    """Create sequence of alternating convolution and pooling operators 
    which gradually shrink over time."""
    model_circuit = cirq.Circuit()
    symbols = sympy.symbols('qconv0:63')
    # Cirq uses sympy.Symbols to map learnable variables. TensorFlow Quantum
    # scans incoming circuits and replaces these with TensorFlow variables.
    model_circuit += quantum_conv_circuit(qubits, symbols[0:15])
    model_circuit += quantum_pool_circuit(qubits[:4], qubits[4:],
                                          symbols[15:21])
    model_circuit += quantum_conv_circuit(qubits[4:], symbols[21:36])
    model_circuit += quantum_pool_circuit(qubits[4:6], qubits[6:],
                                          symbols[36:42])
    model_circuit += quantum_conv_circuit(qubits[6:], symbols[42:57])
    model_circuit += quantum_pool_circuit([qubits[6]], [qubits[7]],
                                          symbols[57:63])
    return model_circuit


# Create our qubits and readout operators in Cirq.
cluster_state_bits = cirq.GridQubit.rect(1, 8)
readout_operators = cirq.Z(cluster_state_bits[-1])

# Build a sequential model enacting the logic in 1.3 of this notebook.
# Here you are making the static cluster state prep as a part of the AddCircuit and the
# "quantum datapoints" are coming in the form of excitation
excitation_input = tf.keras.Input(shape=(), dtype=tf.dtypes.string)
cluster_state = tfq.layers.AddCircuit()(
    excitation_input, prepend=cluster_state_circuit(cluster_state_bits))

quantum_model = tfq.layers.PQC(create_model_circuit(cluster_state_bits),
                               readout_operators)(cluster_state)

qcnn_model = tf.keras.Model(inputs=[excitation_input], outputs=[quantum_model])

# Show the keras plot of the model
tf.keras.utils.plot_model(qcnn_model,
                          show_shapes=True,
                          show_layer_names=False,
                          dpi=70)

"""
Explanation: 1.6 モデルの定義
定義したレイヤーを使用して純粋な量子 CNN を構築します。8 キュービットで開始し、1 キュービットまでプールダウンしてから、$\langle \hat{Z} \rangle$ を測定します。
End of explanation
"""


# Generate some training data.
train_excitations, train_labels, test_excitations, test_labels = generate_data(
    cluster_state_bits)


# Custom accuracy metric.
@tf.function
def custom_accuracy(y_true, y_pred):
    y_true = tf.squeeze(y_true)
    y_pred = tf.map_fn(lambda x: 1.0 if x >= 0 else -1.0, y_pred)
    return tf.keras.backend.mean(tf.keras.backend.equal(y_true, y_pred))


qcnn_model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.02),
                   loss=tf.losses.mse,
                   metrics=[custom_accuracy])

history = qcnn_model.fit(x=train_excitations,
                         y=train_labels,
                         batch_size=16,
                         epochs=25,
                         verbose=1,
                         validation_data=(test_excitations, test_labels))

plt.plot(history.history['loss'][1:], label='Training')
plt.plot(history.history['val_loss'][1:], label='Validation')
plt.title('Training a Quantum CNN to Detect Excited Cluster States')
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.legend()
plt.show()

"""
Explanation: 1.7 モデルをトレーニングする
この例を単純化するために、完全なバッチでモデルをトレーニングします。
End of explanation
"""


# 1-local operators to read out
readouts = [cirq.Z(bit) for bit in cluster_state_bits[4:]]


def multi_readout_model_circuit(qubits):
    """Make a model circuit with less quantum pool and conv operations."""
    model_circuit = cirq.Circuit()
    symbols = sympy.symbols('qconv0:21')
    model_circuit += quantum_conv_circuit(qubits, symbols[0:15])
    model_circuit += quantum_pool_circuit(qubits[:4], qubits[4:],
                                          symbols[15:21])
    return model_circuit


# Build a model enacting the logic in 2.1 of this notebook.
excitation_input_dual = tf.keras.Input(shape=(), dtype=tf.dtypes.string)

cluster_state_dual = tfq.layers.AddCircuit()(
    excitation_input_dual, prepend=cluster_state_circuit(cluster_state_bits))

quantum_model_dual = tfq.layers.PQC(
    multi_readout_model_circuit(cluster_state_bits),
    readouts)(cluster_state_dual)

d1_dual = tf.keras.layers.Dense(8)(quantum_model_dual)

d2_dual = tf.keras.layers.Dense(1)(d1_dual)

hybrid_model = tf.keras.Model(inputs=[excitation_input_dual], outputs=[d2_dual])

# Display the model architecture
tf.keras.utils.plot_model(hybrid_model,
                          show_shapes=True,
                          show_layer_names=False,
                          dpi=70)

"""
Explanation: 2. ハイブリッドモデル
量子畳み込みを使用して 8 キュービットから 1 キュービットにする必要はありません。量子畳み込みの 1～2 ラウンドを実行し、結果を従来のニューラルネットワークにフィードすることも可能です。このセクションでは、量子と従来のハイブリッドモデルを説明します。
2.1 単一量子フィルタを備えたハイブリッドモデル
量子畳み込みのレイヤーを 1 つ適用し、すべてのビットの $\langle \hat{Z}_n \rangle$ を読み取り、続いて密に接続されたニューラルネットワークを読み取ります。
<img src="./images/qcnn_5.png" width="1000">  
2.1.1 モデルの定義
End of explanation
"""


hybrid_model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.02),
                     loss=tf.losses.mse,
                     metrics=[custom_accuracy])

hybrid_history = hybrid_model.fit(x=train_excitations,
                                  y=train_labels,
                                  batch_size=16,
                                  epochs=25,
                                  verbose=1,
                                  validation_data=(test_excitations,
                                                   test_labels))

plt.plot(history.history['val_custom_accuracy'], label='QCNN')
plt.plot(hybrid_history.history['val_custom_accuracy'], label='Hybrid CNN')
plt.title('Quantum vs Hybrid CNN performance')
plt.xlabel('Epochs')
plt.legend()
plt.ylabel('Validation Accuracy')
plt.show()

"""
Explanation: 2.1.2 モデルをトレーニングする
End of explanation
"""


excitation_input_multi = tf.keras.Input(shape=(), dtype=tf.dtypes.string)

cluster_state_multi = tfq.layers.AddCircuit()(
    excitation_input_multi, prepend=cluster_state_circuit(cluster_state_bits))

# apply 3 different filters and measure expectation values

quantum_model_multi1 = tfq.layers.PQC(
    multi_readout_model_circuit(cluster_state_bits),
    readouts)(cluster_state_multi)

quantum_model_multi2 = tfq.layers.PQC(
    multi_readout_model_circuit(cluster_state_bits),
    readouts)(cluster_state_multi)

quantum_model_multi3 = tfq.layers.PQC(
    multi_readout_model_circuit(cluster_state_bits),
    readouts)(cluster_state_multi)

# concatenate outputs and feed into a small classical NN
concat_out = tf.keras.layers.concatenate(
    [quantum_model_multi1, quantum_model_multi2, quantum_model_multi3])

dense_1 = tf.keras.layers.Dense(8)(concat_out)

dense_2 = tf.keras.layers.Dense(1)(dense_1)

multi_qconv_model = tf.keras.Model(inputs=[excitation_input_multi],
                                   outputs=[dense_2])

# Display the model architecture
tf.keras.utils.plot_model(multi_qconv_model,
                          show_shapes=True,
                          show_layer_names=True,
                          dpi=70)

"""
Explanation: ご覧のとおり、非常に控えめな古典的支援により、ハイブリッドモデルは通常、純粋な量子バージョンよりも速く収束します。
2.2 多重量子フィルタを備えたハイブリッド畳み込み
多重量子畳み込みと従来のニューラルネットワークを使用してそれらを組み合わせるアーキテクチャを試してみましょう。
<img src="./images/qcnn_6.png" width="1000">  
2.2.1 モデルの定義
End of explanation
"""


multi_qconv_model.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=0.02),
    loss=tf.losses.mse,
    metrics=[custom_accuracy])

multi_qconv_history = multi_qconv_model.fit(x=train_excitations,
                                            y=train_labels,
                                            batch_size=16,
                                            epochs=25,
                                            verbose=1,
                                            validation_data=(test_excitations,
                                                             test_labels))

plt.plot(history.history['val_custom_accuracy'][:25], label='QCNN')
plt.plot(hybrid_history.history['val_custom_accuracy'][:25], label='Hybrid CNN')
plt.plot(multi_qconv_history.history['val_custom_accuracy'][:25],
         label='Hybrid CNN \n Multiple Quantum Filters')
plt.title('Quantum vs Hybrid CNN performance')
plt.xlabel('Epochs')
plt.legend()
plt.ylabel('Validation Accuracy')
plt.show()

"""
Explanation: 2.2.2 モデルをトレーニングする
End of explanation
"""


