From 136cb1faebcc5102df35de2e2e0c2a9f9ee06977 Mon Sep 17 00:00:00 2001
From: gsq7474741 <gsq7474741@icloud.com>
Date: Tue, 9 Nov 2021 02:53:53 +0800
Subject: [PATCH] add utils.plot_density_graph v0.1

---
 paddle_quantum/utils.py      | 37 ++++++++++++++++++++++++++++++++++++
 test_and_documents/readme.md | 12 +++++++++++-
 test_and_documents/test.py   | 31 ++++++++++++++++++++++++++++--
 3 files changed, 77 insertions(+), 3 deletions(-)

diff --git a/paddle_quantum/utils.py b/paddle_quantum/utils.py
index bbfd15b..66a0651 100644
--- a/paddle_quantum/utils.py
+++ b/paddle_quantum/utils.py
@@ -34,6 +34,7 @@ import matplotlib as mpl
 from paddle_quantum import simulator
 import matplotlib.animation as animation
 import matplotlib.image
+from typing import Union, Optional
 
 __all__ = [
     "partial_trace",
@@ -1726,6 +1727,42 @@ def decompose(matrix):
 
     return pauli_form
 
+def plot_density_graph(density_matrix: Union[paddle.Tensor, np.ndarray],
+                       size: Optional[float]=.3) -> plt.Figure:
+    r"""密度矩阵可视化工具。
+    Args:
+        density_matrix (numpy.ndarray or paddle.Tensor): 多量子比特的量子态的状态向量或者密度矩阵,要求量子数大于1
+        size (float): 条宽度，在0到1之间，默认0.3
+    Returns:
+        plt.Figure: 对应的密度矩阵可视化3D直方图
+    """
+    if not isinstance(density_matrix, (np.ndarray, paddle.Tensor)):
+        msg = f'Expected density_matrix to be np.ndarray or paddle.Tensor, but got {type(density_matrix)}'
+        raise TypeError(msg)
+    if isinstance(density_matrix, paddle.Tensor):
+        density_matrix = density_matrix.numpy()
+    if density_matrix.shape[0] != density_matrix.shape[1]:
+        msg = f'Expected density matrix dim0 equal to dim1, but got dim0={density_matrix.shape[0]}, dim1={density_matrix.shape[1]}'
+        raise ValueError(msg)
+
+    real = density_matrix.real
+    imag = density_matrix.imag
+
+    figure = plt.figure()
+    ax_real = figure.add_subplot(121, projection='3d', title="real")
+    ax_imag = figure.add_subplot(122, projection='3d', title="imag")
+
+    xx, yy = np.meshgrid(
+        list(range(real.shape[0])), list(range(real.shape[1])))
+    xx, yy = xx.ravel(), yy.ravel()
+    real = real.reshape(-1)
+    imag = imag.reshape(-1)
+
+    ax_real.bar3d(xx, yy, np.zeros_like(real), size, size, real)
+    ax_imag.bar3d(xx, yy, np.zeros_like(imag), size, size, imag)
+
+    return figure
+
 def img_to_density_matrix(img_file):
     r"""将图片编码为密度矩阵
     Args:
diff --git a/test_and_documents/readme.md b/test_and_documents/readme.md
index f8d685c..6853f8b 100644
--- a/test_and_documents/readme.md
+++ b/test_and_documents/readme.md
@@ -1 +1,11 @@
-通过在UAnsatz类中添加新的成员函数expand来实现扩展
+- 通过在UAnsatz类中添加新的成员函数expand来实现扩展
+
+
+- 增加utils.plot_density_graph密度矩阵可视化工具。 
+```
+Args:
+density_matrix (numpy.ndarray or paddle.Tensor): 多量子比特的量子态的状态向量或者密度矩阵,要求量子数大于1
+size (float): 条宽度，在0到1之间，默认0.3
+Returns:
+plt.Figure: 对应的密度矩阵可视化3D直方图
+```
\ No newline at end of file
diff --git a/test_and_documents/test.py b/test_and_documents/test.py
index 9d998d8..9423751 100644
--- a/test_and_documents/test.py
+++ b/test_and_documents/test.py
@@ -2,6 +2,7 @@ from paddle_quantum.circuit import UAnsatz
 from paddle import kron
 from paddle_quantum.state import vec,density_op
 import paddle
+import unittest
 
 
 #density_matrix
@@ -30,5 +31,31 @@ def test_state_vector():
     cir2.run_state_vector()
     print(cir2.get_state())
 
-test_density_matrix()
-test_state_vector()
+
+class TestPlotDensityGraph(unittest.TestCase):
+    def setUp(self):
+        self.func = plot_density_graph
+        self.x_np = np.random.rand(4, 4) + np.random.rand(4, 4) * 1j
+        self.x_tensor = paddle.to_tensor(self.x_np)
+
+    def test_input_type(self):
+        self.assertRaises(TypeError, self.func, 1)
+        self.assertRaises(TypeError, self.func, [1, 2, 3])
+
+    def test_input_shape(self):
+        x = np.zeros((2, 3))
+        self.assertRaises(ValueError, self.func, x)
+
+    def test_ndarray_input_inputs(self):
+        res = self.func(self.x_np)
+        res.show()
+
+    def test_tensor_input(self):
+        res = self.func(self.x_tensor)
+        res.show()
+
+
+if __name__ == '__main__':
+    test_density_matrix()
+    test_state_vector()
+    unittest.main()
\ No newline at end of file
-- 
GitLab