Spaces:
Sleeping
Sleeping
| """Tests for multi-spinor (tensor-product) states.""" | |
| import math | |
| import pytest | |
| import torch | |
| from spinor_lib import ( | |
| Spinor, | |
| bloch_vector, | |
| concurrence, | |
| density_matrix, | |
| embed_operator, | |
| entanglement_entropy, | |
| expectation, | |
| fidelity, | |
| generate_su2_rotation, | |
| identity, | |
| inner_product, | |
| kron, | |
| partial_trace, | |
| pauli, | |
| purity, | |
| reduced_density_matrix, | |
| rotate, | |
| tensor_product, | |
| von_neumann_entropy, | |
| ) | |
| C128 = torch.complex128 | |
| F64 = torch.float64 | |
| def cs(*vals, dim=2): | |
| return Spinor(torch.tensor(vals, dtype=C128), dim=dim) | |
| def random_spinor(*batch, dim=2, generator=None): | |
| return Spinor(torch.randn(*batch, dim, dtype=C128, generator=generator), dim=dim) | |
| UP, DOWN = cs(1, 0), cs(0, 1) | |
| BELL = cs(1, 0, 0, 1, dim=4) * (1 / math.sqrt(2)) # (|00> + |11>) / sqrt 2 | |
| class TestTensorProduct: | |
| def test_basis_states(self): | |
| assert torch.allclose(tensor_product(UP, DOWN).data, torch.tensor([0, 1, 0, 0], dtype=C128)) | |
| assert torch.allclose(tensor_product(DOWN, UP).data, torch.tensor([0, 0, 1, 0], dtype=C128)) | |
| def test_matches_torch_kron(self): | |
| g = torch.Generator().manual_seed(0) | |
| a, b, c = (random_spinor(generator=g) for _ in range(3)) | |
| out = tensor_product(a, b, c) | |
| assert out.dim == 8 and out.shape == (8,) | |
| assert torch.allclose(out.data, torch.kron(torch.kron(a.data, b.data), c.data)) | |
| def test_batch_broadcast(self): | |
| g = torch.Generator().manual_seed(1) | |
| a = random_spinor(4, 1, generator=g) | |
| b = random_spinor(3, generator=g) | |
| out = tensor_product(a, b) | |
| assert out.shape == (4, 3, 4) | |
| assert torch.allclose(out.data[2, 1], torch.kron(a.data[2, 0], b.data[1])) | |
| def test_norm_is_product_of_norms(self): | |
| g = torch.Generator().manual_seed(2) | |
| a, b = random_spinor(5, generator=g), random_spinor(5, generator=g) | |
| assert torch.allclose(tensor_product(a, b).norm(), a.norm() * b.norm()) | |
| def test_mixed_dims_and_dtype_promotion(self): | |
| a = Spinor(torch.tensor([1.0, 2.0, 3.0]), dim=3) # real float32 | |
| b = cs(1j, 1) | |
| out = tensor_product(a, b) | |
| assert out.dim == 6 and out.dtype == C128 | |
| def test_gradient(self): | |
| data = torch.tensor([1.0, 0.5j], dtype=C128, requires_grad=True) | |
| out = tensor_product(Spinor(data), DOWN) | |
| torch.real(inner_product(out, out)).backward() | |
| assert torch.allclose(data.grad, 2 * data.detach()) | |
| class TestKronAndEmbed: | |
| def test_kron_matches_torch(self): | |
| g = torch.Generator().manual_seed(3) | |
| A = torch.randn(2, 2, dtype=C128, generator=g) | |
| B = torch.randn(3, 3, dtype=C128, generator=g) | |
| assert torch.allclose(kron(A, B), torch.kron(A, B)) | |
| def test_kron_batched(self): | |
| g = torch.Generator().manual_seed(4) | |
| A = torch.randn(5, 2, 2, dtype=C128, generator=g) | |
| B = torch.randn(5, 2, 2, dtype=C128, generator=g) | |
| out = kron(A, B) | |
| assert out.shape == (5, 4, 4) | |
| assert torch.allclose(out[3], torch.kron(A[3], B[3])) | |
| def test_embed_operator_positions(self): | |
| sz = pauli(dtype=C128)[2] | |
| eye = identity(2, dtype=C128) | |
| assert torch.allclose(embed_operator(sz, (2, 2), 0), torch.kron(sz, eye)) | |
| assert torch.allclose(embed_operator(sz, (2, 2), 1), torch.kron(eye, sz)) | |
| assert torch.allclose(embed_operator(sz, (2, 2, 2), 1), torch.kron(torch.kron(eye, sz), eye)) | |
| def test_embed_rejects_wrong_sizes(self): | |
| with pytest.raises(ValueError): | |
| embed_operator(torch.eye(3, dtype=C128), (2, 2), 0) | |
| with pytest.raises(IndexError): | |
| embed_operator(torch.eye(2, dtype=C128), (2, 2), 2) | |
| def test_rotating_one_factor_of_a_product_state(self): | |
| g = torch.Generator().manual_seed(5) | |
| a, b = random_spinor(generator=g), random_spinor(generator=g) | |
| R = generate_su2_rotation([1, 2, 0], torch.tensor(0.9, dtype=F64)) | |
| lhs = rotate(tensor_product(a, b), embed_operator(R, (2, 2), 1)) | |
| rhs = tensor_product(a, rotate(b, R)) | |
| assert torch.allclose(lhs.data, rhs.data) | |
| def test_local_expectation_matches_single_spinor(self): | |
| g = torch.Generator().manual_seed(6) | |
| a, b = random_spinor(generator=g), random_spinor(generator=g) | |
| sx = pauli(dtype=C128)[0] | |
| ab = tensor_product(a, b).normalize() | |
| got = expectation(ab, embed_operator(sx, (2, 2), 0)) | |
| assert got.item() == pytest.approx(expectation(a.normalize(), sx).item()) | |
| class TestPartialTrace: | |
| def test_product_state_reduces_to_factors(self): | |
| g = torch.Generator().manual_seed(7) | |
| a, b = random_spinor(generator=g), random_spinor(generator=g) | |
| ab = tensor_product(a, b) | |
| assert torch.allclose(reduced_density_matrix(ab, (2, 2), 0), density_matrix(a)) | |
| assert torch.allclose(reduced_density_matrix(ab, (2, 2), 1), density_matrix(b)) | |
| def test_bell_state_reduces_to_maximally_mixed(self): | |
| rho_a = reduced_density_matrix(BELL, (2, 2), 0) | |
| assert torch.allclose(rho_a, torch.eye(2, dtype=C128) / 2) | |
| assert purity(rho_a).item() == pytest.approx(0.5) | |
| assert torch.allclose(bloch_vector(rho_a), torch.zeros(3, dtype=F64), atol=1e-12) | |
| def test_keep_everything_is_identity(self): | |
| rho = density_matrix(BELL) | |
| assert torch.allclose(partial_trace(rho, (2, 2), [0, 1]), rho) | |
| def test_three_subsystems_keep_two(self): | |
| g = torch.Generator().manual_seed(8) | |
| a, b, c = (random_spinor(generator=g) for _ in range(3)) | |
| abc = tensor_product(a, b, c) | |
| rho_ac = partial_trace(density_matrix(abc), (2, 2, 2), [0, 2]) | |
| assert rho_ac.shape == (4, 4) | |
| assert torch.allclose(rho_ac, density_matrix(tensor_product(a, c))) | |
| def test_uneven_dims_and_batch(self): | |
| g = torch.Generator().manual_seed(9) | |
| a = random_spinor(6, dim=2, generator=g) | |
| b = random_spinor(6, dim=3, generator=g) | |
| rho = density_matrix(tensor_product(a, b)) | |
| rho_b = partial_trace(rho, (2, 3), 1) | |
| assert rho_b.shape == (6, 3, 3) | |
| assert torch.allclose(rho_b, density_matrix(b)) | |
| def test_trace_is_preserved(self): | |
| g = torch.Generator().manual_seed(10) | |
| rho = density_matrix(random_spinor(4, dim=8, generator=g)) | |
| reduced = partial_trace(rho, (2, 2, 2), 1) | |
| assert torch.allclose(torch.einsum("bii->b", reduced), torch.ones(4, dtype=C128)) | |
| def test_rejects_bad_inputs(self): | |
| rho = density_matrix(BELL) | |
| with pytest.raises(ValueError): | |
| partial_trace(rho, (2, 3), 0) | |
| with pytest.raises(ValueError): | |
| partial_trace(rho, (2, 2), [0, 0]) | |
| with pytest.raises(IndexError): | |
| partial_trace(rho, (2, 2), 5) | |
| class TestEntanglement: | |
| def test_von_neumann_entropy_values(self): | |
| eye = torch.eye(2, dtype=C128) | |
| assert von_neumann_entropy(eye / 2).item() == pytest.approx(math.log(2)) | |
| assert von_neumann_entropy(eye / 2, base=2).item() == pytest.approx(1.0) | |
| assert von_neumann_entropy(density_matrix(UP)).item() == pytest.approx(0.0, abs=1e-12) | |
| def test_entanglement_entropy_bell_vs_product(self): | |
| assert entanglement_entropy(BELL, (2, 2), 0, base=2).item() == pytest.approx(1.0) | |
| product = tensor_product(cs(1, 1j), cs(2, 0.5)) | |
| assert entanglement_entropy(product, (2, 2), 0).item() == pytest.approx(0.0, abs=1e-12) | |
| def test_entanglement_entropy_symmetric_in_bipartition(self): | |
| g = torch.Generator().manual_seed(11) | |
| s = random_spinor(7, dim=4, generator=g) | |
| assert torch.allclose( | |
| entanglement_entropy(s, (2, 2), 0), entanglement_entropy(s, (2, 2), 1) | |
| ) | |
| def test_concurrence_values(self): | |
| assert concurrence(BELL).item() == pytest.approx(1.0) | |
| assert concurrence(BELL * 3).item() == pytest.approx(1.0) # normalisation-free | |
| assert concurrence(tensor_product(cs(1, 2j), cs(0.3, 1))).item() == pytest.approx(0.0, abs=1e-12) | |
| def test_concurrence_agrees_with_entropy(self): | |
| # For two-qubit pure states S = h((1 + sqrt(1 - C^2)) / 2) in bits. | |
| g = torch.Generator().manual_seed(12) | |
| s = random_spinor(20, dim=4, generator=g) | |
| C = concurrence(s) | |
| lam = (1 + torch.sqrt(1 - C**2)) / 2 | |
| h = -(lam * torch.log2(lam) + (1 - lam) * torch.log2(1 - lam)) | |
| assert torch.allclose(entanglement_entropy(s, (2, 2), 0, base=2), h, atol=1e-8) | |
| def test_concurrence_rejects_non_two_qubit(self): | |
| with pytest.raises(ValueError): | |
| concurrence(UP) | |
| def test_gradients_through_entanglement_measures(self): | |
| data = torch.tensor([1.0, 0.2, 0.3j, 0.7], dtype=C128, requires_grad=True) | |
| s = Spinor(data, dim=4) | |
| (concurrence(s) + entanglement_entropy(s, (2, 2), 0)).backward() | |
| assert torch.isfinite(data.grad).all() | |
| def test_entanglement_is_invariant_under_local_rotations(self): | |
| g = torch.Generator().manual_seed(13) | |
| s = random_spinor(5, dim=4, generator=g) | |
| R = generate_su2_rotation(torch.randn(5, 3, dtype=F64, generator=g), torch.rand(5, dtype=F64, generator=g)) | |
| local = rotate(s, embed_operator(R, (2, 2), 0)) | |
| assert torch.allclose(concurrence(local), concurrence(s)) | |
| assert torch.allclose(entanglement_entropy(local, (2, 2), 1), entanglement_entropy(s, (2, 2), 1)) | |
| def test_bell_state_via_fidelity_is_far_from_all_products(self): | |
| g = torch.Generator().manual_seed(14) | |
| products = tensor_product(random_spinor(200, generator=g), random_spinor(200, generator=g)) | |
| assert (fidelity(BELL, products) <= 0.5 + 1e-12).all() | |