add utils.plot_density_graph v0.2

PaddlePaddle · Nov 17, 2021 · 16e7d3b · 16e7d3b
1 parent 136cb1f
commit 16e7d3b
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 5 deletions.
diff --git a/paddle_quantum/utils.py b/paddle_quantum/utils.py
@@ -1758,8 +1758,8 @@ def plot_density_graph(density_matrix: Union[paddle.Tensor, np.ndarray],
     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)
+    ax_real.bar3d(xx, yy, np.zeros_like(real), size, size, np.abs(real))
+    ax_imag.bar3d(xx, yy, np.zeros_like(imag), size, size, np.abs(imag))
 
     return figure
 

diff --git a/test_and_documents/test.py b/test_and_documents/test.py
@@ -1,6 +1,7 @@
 from paddle_quantum.circuit import UAnsatz
-from paddle import kron
-from paddle_quantum.state import vec,density_op
+import matplotlib.pyplot as plt
+from paddle_quantum.utils import plot_density_graph
+import numpy as np
 import paddle
 import unittest
 
@@ -35,7 +36,7 @@ def 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_np = (np.random.rand(8, 8) + np.random.rand(8, 8) * 1j)-0.5-0.5j
         self.x_tensor = paddle.to_tensor(self.x_np)
 
     def test_input_type(self):