optimize entropy

codes/test.py

@@ -0,0 +1,32 @@
+# unit test for the functions in quantum_states.py
+
+import unittest
+import numpy as np
+from quantum_states import random_pure_state, von_neumann_entropy_bipartite_pure_state
+
+class LearningCase(unittest.TestCase):
+    def test_random_pure_state_shape_and_norm(self):
+        dim_a = 2
+        dim_b = 2
+        state = random_pure_state(dim_a, dim_b)
+        self.assertEqual(state.shape, (dim_a * dim_b,))
+        self.assertAlmostEqual(np.linalg.norm(state), 1)
+
+    def test_partial_trace_entropy(self):
+        dim_a = 2
+        dim_b = 2
+        state = random_pure_state(dim_a, dim_b)
+        self.assertAlmostEqual(von_neumann_entropy_bipartite_pure_state(state, dim_a, dim_b), von_neumann_entropy_bipartite_pure_state(state, dim_b, dim_a))
+
+    def test_sample_uniformly(self):
+        # calculate the distribution of the random pure state
+        dim_a = 2
+        dim_b = 2
+        state = random_pure_state(dim_a, dim_b)
+        
+
+def main():
+    unittest.main()
+
+if __name__ == "__main__":
+    main()