32 lines
1012 B
Python
32 lines
1012 B
Python
# 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() |