From d0b1b1598bbc54ec7cbd289ed6c19d11f1263669 Mon Sep 17 00:00:00 2001 From: Aaron Sander <61705296+aaronleesander@users.noreply.github.com> Date: Wed, 12 Feb 2025 17:24:31 +0100 Subject: [PATCH] added new tests --- tests/test_PhysicsTJM.py | 108 +++++++++++++++++ tests/test_TDVP.py | 212 ++++++++++++++++++++++++++++++++ tests/test_gate_library.py | 241 +++++++++++++++++++++++++++++++++++++ 3 files changed, 561 insertions(+) create mode 100644 tests/test_PhysicsTJM.py create mode 100644 tests/test_TDVP.py create mode 100644 tests/test_gate_library.py diff --git a/tests/test_PhysicsTJM.py b/tests/test_PhysicsTJM.py new file mode 100644 index 0000000..fb6304c --- /dev/null +++ b/tests/test_PhysicsTJM.py @@ -0,0 +1,108 @@ +import pytest +from unittest.mock import patch + +from yaqs.core.data_structures.networks import MPO, MPS +from yaqs.core.data_structures.noise_model import NoiseModel +from yaqs.core.data_structures.simulation_parameters import Observable, PhysicsSimParams +from yaqs.physics.PhysicsTJM import initialize, step_through, sample, PhysicsTJM_2, PhysicsTJM_1 + + +def test_initialize(): + L = 5 + d = 2 + J = 1 + g = 0.5 + H = MPO() + H.init_Ising(L, d, J, g) + state = MPS(L) + noise_model = NoiseModel(['relaxation'], [0.1]) + measurements = [Observable('x', site) for site in range(L)] + sim_params = PhysicsSimParams(measurements, T=0.2, dt=0.2, sample_timesteps=False, N=1, max_bond_dim=2, threshold=1e-6, order=1) + with patch('yaqs.physics.PhysicsTJM.apply_dissipation') as mock_dissipation, \ + patch('yaqs.physics.PhysicsTJM.stochastic_process') as mock_stochastic_process: + initialize(state, noise_model, sim_params) + mock_dissipation.assert_called_once_with(state, noise_model, sim_params.dt/2) + mock_stochastic_process.assert_called_once_with(state, noise_model, sim_params.dt) + +def test_step_through(): + L = 5 + d = 2 + J = 1 + g = 0.5 + H = MPO() + H.init_Ising(L, d, J, g) + state = MPS(L) + noise_model = NoiseModel(['relaxation'], [0.1]) + measurements = [Observable('x', site) for site in range(L)] + sim_params = PhysicsSimParams(measurements, T=0.2, dt=0.2, sample_timesteps=False, N=1, max_bond_dim=2, threshold=1e-6, order=1) + with patch('yaqs.physics.PhysicsTJM.dynamic_TDVP') as mock_dynamic_TDVP, \ + patch('yaqs.physics.PhysicsTJM.apply_dissipation') as mock_dissipation, \ + patch('yaqs.physics.PhysicsTJM.stochastic_process') as mock_stochastic_process: + step_through(state, H, noise_model, sim_params) + mock_dynamic_TDVP(state, H, sim_params) + mock_dissipation.assert_called_once_with(state, noise_model, sim_params.dt) + mock_stochastic_process.assert_called_once_with(state, noise_model, sim_params.dt) + +def test_PhysicsTJM_2(): + L = 5 + d = 2 + J = 1 + g = 0.5 + H = MPO() + H.init_Ising(L, d, J, g) + state = MPS(L) + noise_model = None + measurements = [Observable('z', site) for site in range(L)] + sim_params = PhysicsSimParams(measurements, T=0.2, dt=0.1, sample_timesteps=False, N=1, max_bond_dim=4, threshold=1e-6, order=2) + args = (0, state, noise_model, sim_params, H) + results = PhysicsTJM_2(args) + # When sample_timesteps is True, results should have shape (num_observables, len(times)) + assert results.shape == (len(measurements), 1), "Results incorrect shape" + +def test_PhysicsTJM_2_sample_timesteps(): + L = 5 + d = 2 + J = 1 + g = 0.5 + H = MPO() + H.init_Ising(L, d, J, g) + state = MPS(L) + noise_model = None + measurements = [Observable('z', site) for site in range(L)] + sim_params = PhysicsSimParams(measurements, T=0.2, dt=0.1, sample_timesteps=True, N=1, max_bond_dim=4, threshold=1e-6, order=2) + args = (0, state, noise_model, sim_params, H) + results = PhysicsTJM_2(args) + # When sample_timesteps is True, results should have shape (num_observables, len(times)) + assert results.shape == (len(measurements), len(sim_params.times)), "Results incorrect shape" + +def test_PhysicsTJM_1(): + L = 5 + d = 2 + J = 1 + g = 0.5 + H = MPO() + H.init_Ising(L, d, J, g) + state = MPS(L) + noise_model = None + measurements = [Observable('z', site) for site in range(L)] + sim_params = PhysicsSimParams(measurements, T=0.2, dt=0.1, sample_timesteps=False, N=1, max_bond_dim=4, threshold=1e-6, order=1) + args = (0, state, noise_model, sim_params, H) + results = PhysicsTJM_1(args) + # When sample_timesteps is True, results should have shape (num_observables, len(times)) + assert results.shape == (len(measurements), 1), "Results incorrect shape" + +def test_PhysicsTJM_1_sample_timesteps(): + L = 5 + d = 2 + J = 1 + g = 0.5 + H = MPO() + H.init_Ising(L, d, J, g) + state = MPS(L) + noise_model = None + measurements = [Observable('z', site) for site in range(L)] + sim_params = PhysicsSimParams(measurements, T=0.2, dt=0.1, sample_timesteps=True, N=1, max_bond_dim=4, threshold=1e-6, order=1) + args = (0, state, noise_model, sim_params, H) + results = PhysicsTJM_1(args) + # When sample_timesteps is True, results should have shape (num_observables, len(times)) + assert results.shape == (len(measurements), len(sim_params.times)), "Results incorrect shape" diff --git a/tests/test_TDVP.py b/tests/test_TDVP.py new file mode 100644 index 0000000..921ea2b --- /dev/null +++ b/tests/test_TDVP.py @@ -0,0 +1,212 @@ +import numpy as np +import pytest + +from yaqs.core.data_structures.networks import MPO, MPS +from yaqs.core.data_structures.simulation_parameters import Observable, PhysicsSimParams + +from yaqs.core.methods.TDVP import ( + _split_mps_tensor, + _merge_mps_tensor_pair, + _merge_mpo_tensor_pair, + _contraction_operator_step_right, + _contraction_operator_step_left, + _compute_right_operator_blocks, + _apply_local_hamiltonian, + _apply_local_bond_contraction, + _local_hamiltonian_step, + _local_bond_step, + single_site_TDVP, + two_site_TDVP, +) + +def test_split_mps_tensor_left_right_sqrt(): + # Create an input tensor A with shape (d0*d1, D0, D2). + # Let d0 = d1 = 2 so that A.shape[0]=4, and choose D0=3, D2=5. + A = np.random.randn(4, 3, 5) + # For each svd_distr option, run the splitting and then reconstruct A. + for distr in ['left', 'right', 'sqrt']: + A0, A1 = _split_mps_tensor(A, svd_distr=distr, threshold=1e-8) + # A0 should have shape (2, 3, r) and A1 should have shape (2, r, 5), + # where r is the effective rank. + assert A0.ndim == 3 + assert A1.ndim == 3 + r = A0.shape[2] + assert A1.shape[1] == r + # Reconstruct A: undo the transpose on A1 then contract A0 and A1. + A1_recon = A1.transpose((1, 0, 2)) # now shape (r, 2, 5) + # Contract along the rank index: + # A0 has indices: (d0, D0, r), A1_recon has indices: (r, d1, D2). + # Form a tensor of shape (d0, d1, D0, D2) then reshape back to (4, 3, 5) + A_recon = np.tensordot(A0, A1_recon, axes=(2, 0)) # shape (2, 3, 2, 5) + A_recon = A_recon.transpose((0,2,1,3)).reshape(4, 3, 5) + # Up to SVD sign and ordering ambiguities, the reconstruction should be close. + np.testing.assert_allclose(A, A_recon, atol=1e-6) + +def test_split_mps_tensor_invalid_shape(): + # If A.shape[0] is not divisible by 2, the function should raise a ValueError. + A = np.random.randn(3, 3, 5) + with pytest.raises(ValueError): + _split_mps_tensor(A, svd_distr='left') + +def test_merge_mps_tensor_pair(): + # Let A0 have shape (2, 3, 4) and A1 have shape (5, 4, 7). + # In _merge_mps_tensor_pair the arrays are interpreted with label tuples + # (0,2,3) for A0 and (1,3,4) for A1, so contraction is over the third axis of A0 and + # the second axis of A1. + A0 = np.random.randn(2, 3, 4) + A1 = np.random.randn(5, 4, 7) + merged = _merge_mps_tensor_pair(A0, A1) + # Expected shape: first two axes merged from A0[0] and A1[0]: + # output shape = ((2*5), 3, 7) i.e. (10, 3, 7) + assert merged.shape == (10, 3, 7) + +def test_merge_mpo_tensor_pair(): + # Let A0 be a 4D array with shape (2, 3, 4, 5) and + # A1 be a 4D array with shape (7, 8, 5, 9). + # The contract call uses label tuples (0,2,4,6) and (1,3,6,5) and contracts label 6. + A0 = np.random.randn(2, 3, 4, 5) + A1 = np.random.randn(7, 8, 5, 9) + merged = _merge_mpo_tensor_pair(A0, A1) + # Expected shape: merge first two axes of the result, where result (before reshape) + # should have shape (2,7,3,8,4,9), then merged to (2*7, 3*8, 4,9) = (14,24,4,9). + assert merged.shape == (14, 24, 4, 9) + +def test_contraction_operator_step_right(): + # Choose shapes as described in the function. + A = np.random.randn(2, 3, 4) + # R: contract A's last axis (size 4) with R's first axis. + R = np.random.randn(4, 5, 6) + # W must have shape with axes (1,3) matching T from tensordot(A, R, 1) of shape (2,3,5,6): + # We require W.shape[1]=2 and W.shape[3]=5. Let W = (7,2,8,5) + W = np.random.randn(7, 2, 8, 5) + # B: choose shape so that B.conj() has axes (0,2) matching T from later step. + # After the previous steps, T becomes shape (3,8,7,6); we contract with B.conj() axes ((2,3),(0,2)). + # So let B be of shape (7, 9, 6). + B = np.random.randn(7, 9, 6) + Rnext = _contraction_operator_step_right(A, B, W, R) + # Expected shape: (3,8,9) (from the discussion above). + assert Rnext.shape == (3, 8, 9) + +def test_contraction_operator_step_left(): + # Set up dummy arrays with shapes so that the contraction is valid. + # Let A be shape (3,4,10) + A = np.random.randn(3, 4, 10) + # Let B be shape (7, 6, 8) so that B.conj() has shape (7,6,8). + B = np.random.randn(7, 6, 8) + # Let L be shape (4,5,6) and require that L.shape[2] (6) matches B.shape[1] (6). + L_arr = np.random.randn(4, 5, 6) + # Let W be shape (7,3,5,9) so that we contract axes ((0,2),(2,1)) with T later. + W = np.random.randn(7, 3, 5, 9) + Rnext = _contraction_operator_step_left(A, B, W, L_arr) + # The expected shape from our reasoning is (10,9,8) (A's remaining axis 2 becomes output along with leftover dims from T). + # We check that the result is 3-dimensional. + assert Rnext.ndim == 3 + +def test_apply_local_hamiltonian(): + # Let A: shape (2,3,4); R: shape (4,5,6) as before. + A = np.random.randn(2, 3, 4) + R = np.random.randn(4, 5, 6) + # Let W be shape (7,2,8,5) as in test above. + W = np.random.randn(7, 2, 8, 5) + # Let L be shape (3,8,9) so that tensordot works. + L_arr = np.random.randn(3, 8, 9) + out = _apply_local_hamiltonian(L_arr, R, W, A) + # The function transposes the final result so we expect a 3D array. + assert out.ndim == 3 + +def test_apply_local_bond_contraction(): + # Let C: shape (2,3) + C = np.random.randn(2, 3) + # Let R: shape (3,4,5) + R = np.random.randn(3, 4, 5) + # Let L: shape (2,4,6) + L_arr = np.random.randn(2, 4, 6) + out = _apply_local_bond_contraction(L_arr, R, C) + # Expected output shape: contraction gives shape (6,5) + assert out.shape == (6, 5) + +def test_local_hamiltonian_step(): + # We choose an MPS tensor A with shape (d0, d0, d1) where d0=2, d1=4. + # (The requirement here is that the first two dimensions are equal, + # so that after the contraction chain the operator is square.) + A = np.random.randn(2, 2, 4) # total elements: 2*2*4 = 16 + # Choose R of shape (d1, X, d1) with d1=4 and X=1. + R = np.random.randn(4, 1, 4) # shape: (4,1,4) + # Choose W of shape (d0, d0, X, X) with d0=2 and X=1. + W = np.random.randn(2, 2, 1, 1) # shape: (2,2,1,1) + # Choose L so that the contraction makes sense. + # In our contraction chain, after: + # T1 = tensordot(A, R, axes=1) → shape (2,2,1,4) + # T2 = tensordot(W, T1, axes=((1,3),(0,2))) → shape (2,1,2,4) + # T3 = T2.transpose((2,1,0,3)) → shape (2,1,2,4) (here the permutation reorders axes) + # Then we contract T3 with L along axes ((2,1),(0,1)). + # To contract T3’s axes (axis2, axis1) = (2,1) we need L with shape (2,1,r). + # Then T4 will have shape (remaining T3 axes: (axis0, axis3)) plus L’s remaining axis, i.e. (2,4,r). + # Finally, a transpose (here, (0,2,1)) gives shape (2, r, 4). + # We want the final shape to equal A’s shape (2,2,4), so we set r=2. + L_arr = np.random.randn(2, 1, 2) # shape: (2,1,2) + dt = 0.05 + numiter = 10 + out = _local_hamiltonian_step(L_arr, R, W, A, dt, numiter) + # The operator should be square, so out.shape should equal A.shape. + assert out.shape == A.shape, f"Expected shape {A.shape}, got {out.shape}" + +def test_local_bond_step(): + # For the bond step we want the operator to be square. + # Let C be a matrix of shape (p, q). We choose C to be square. + C = np.random.randn(2, 2) # total elements: 4 + # Choose R of shape (q, r, s). To contract C and get back the same number of elements, + # we can choose R such that q=2, r=2, s=2. + R = np.random.randn(2, 2, 2) # shape: (2,2,2) + # Choose L of shape (p, r, t) with p=2 and r=2. + # We want the final result to have shape (p, q) = (2,2); so t must equal q=2. + L_arr = np.random.randn(2, 2, 2) # shape: (2,2,2) + dt = 0.05 + numiter = 10 + out = _local_bond_step(L_arr, R, C, dt, numiter) + # The output shape should equal the input shape. + assert out.shape == C.shape, f"Expected shape {C.shape}, got {out.shape}" + +def test_single_site_TDVP(): + L = 5 + d = 2 + J = 1 + g = 0.5 + H = MPO() + H.init_Ising(L, d, J, g) + state = MPS(L) + measurements = [Observable('z', site) for site in range(L)] + sim_params = PhysicsSimParams(measurements, T=0.2, dt=0.1, sample_timesteps=True, N=1, max_bond_dim=4, threshold=1e-6, order=1) + single_site_TDVP(state, H, sim_params, numiter_lanczos=5) + # Check that state still has L tensors and that each tensor is a numpy array. + assert state.length == L + for tensor in state.tensors: + assert isinstance(tensor, np.ndarray) + + canonical_site = state.check_canonical_form()[0] + assert canonical_site == 0, ( + f"MPS should be site-canonical at site 0 after single-site TDVP, " + f"but got canonical site: {canonical_site}" + ) + +def test_two_site_TDVP(): + L = 5 + d = 2 + J = 1 + g = 0.5 + H = MPO() + H.init_Ising(L, d, J, g) + state = MPS(L) + measurements = [Observable('z', site) for site in range(L)] + sim_params = PhysicsSimParams(measurements, T=0.2, dt=0.1, sample_timesteps=True, N=1, max_bond_dim=4, threshold=1e-6, order=1) + two_site_TDVP(state, H, sim_params, numiter_lanczos=5) + # Check that state still has L tensors and that each tensor is a numpy array. + assert state.length == L + for tensor in state.tensors: + assert isinstance(tensor, np.ndarray) + + canonical_site = state.check_canonical_form()[0] + assert canonical_site == 0, ( + f"MPS should be site-canonical at site 0 after two-site TDVP, " + f"but got canonical site: {canonical_site}" + ) \ No newline at end of file diff --git a/tests/test_gate_library.py b/tests/test_gate_library.py new file mode 100644 index 0000000..610521a --- /dev/null +++ b/tests/test_gate_library.py @@ -0,0 +1,241 @@ +import numpy as np +import pytest +from numpy.testing import assert_allclose, assert_array_equal + +from yaqs.core.libraries.gate_library import _split_tensor, _extend_gate, GateLibrary + +from yaqs.core.data_structures.networks import MPO + +def test_split_tensor_valid_shape(): + # Create a simple tensor of shape (2,2,2,2). + # Here we use a tensor with values 0..15. + tensor = np.arange(16, dtype=float).reshape(2, 2, 2, 2) + tensors = _split_tensor(tensor) + # Expect a list of two tensors. + assert isinstance(tensors, list) + assert len(tensors) == 2 + + t1, t2 = tensors + # t1 is obtained by reshaping U then expanding a dummy dimension. + # Its shape should be (2, 2, 1, r) for some r. + assert t1.ndim == 4 + assert t1.shape[0] == 2 + assert t1.shape[1] == 2 + + # t2 is reshaped from V, then transposed and expanded. + # Its shape should be (2, 2, r, 1). + assert t2.ndim == 4 + assert t2.shape[0] == 2 + assert t2.shape[1] == 2 + +def test_split_tensor_invalid_shape(): + # A tensor that does not have shape (2,2,2,2) should trigger the assertion. + tensor = np.zeros((2, 2, 2)) + with pytest.raises(AssertionError): + _split_tensor(tensor) + +def test_extend_gate_no_identity(): + # Use a simple tensor. For example, use the 4x4 identity reshaped to (2,2,2,2). + tensor = np.eye(4).reshape(2, 2, 2, 2) + # Choose sites such that |site0 - site1| == 1 so no identity tensor is added. + sites = [0, 1] + mpo = _extend_gate(tensor, sites) + # Check that we got an MPO instance. + assert isinstance(mpo, MPO) + # With no gap, _split_tensor returns 2 tensors and no identity is inserted. + assert len(mpo.tensors) == 2 + +def test_extend_gate_with_identity(): + # Use a simple tensor. + tensor = np.eye(4).reshape(2, 2, 2, 2) + # Use sites with a gap (difference 2 → one identity tensor inserted). + sites = [0, 2] + mpo = _extend_gate(tensor, sites) + assert isinstance(mpo, MPO) + # Expected list: + # [first tensor from _split_tensor, + # one identity tensor, + # second tensor from _split_tensor] + assert len(mpo.tensors) == 3 + + # Verify that the inserted identity tensor has the expected structure. + identity_tensor = mpo.tensors[1] + # The identity tensor shape is (2,2,prev_bond,prev_bond) + prev_bond = mpo.tensors[0].shape[3] + assert identity_tensor.shape == (2, 2, prev_bond, prev_bond) + # Check that each diagonal slice is a 2x2 identity. + for i in range(prev_bond): + assert_array_equal(identity_tensor[:, :, i, i], np.eye(2)) + +def test_extend_gate_reverse_order(): + # Check that if sites are provided in reverse order, the MPO tensors are reversed + # and each tensor is transposed on its last two indices. + tensor = np.eye(4).reshape(2, 2, 2, 2) + mpo_forward = _extend_gate(tensor, [0, 1]) + mpo_reverse = _extend_gate(tensor, [1, 0]) + mpo_reverse.tensors.reverse() + # For each pair of corresponding tensors, the reverse version should be the + # transpose of the forward one on axes (0, 1, 3, 2). + for t_f, t_r in zip(mpo_forward.tensors, mpo_reverse.tensors): + assert_allclose(t_r, np.transpose(t_f, (0, 1, 3, 2))) + +def test_gate_x(): + gate = GateLibrary.x() + gate.set_sites(0) + assert gate.sites == [0] + # For X, the tensor should equal the matrix. + assert_array_equal(gate.tensor, gate.matrix) + +def test_gate_y(): + gate = GateLibrary.y() + gate.set_sites(0) + assert gate.sites == [0] + # For Y, the tensor should equal the matrix. + assert_array_equal(gate.tensor, gate.matrix) + +def test_gate_z(): + gate = GateLibrary.z() + gate.set_sites(0) + assert gate.sites == [0] + # For Z, the tensor should equal the matrix. + assert_array_equal(gate.tensor, gate.matrix) + +def test_gate_id(): + gate = GateLibrary.id() + gate.set_sites(0) + assert gate.sites == [0] + assert_array_equal(gate.tensor, gate.matrix) + +def test_gate_sx(): + gate = GateLibrary.sx() + gate.set_sites(0) + assert gate.sites == [0] + assert_array_equal(gate.tensor, gate.matrix) + +def test_gate_h(): + gate = GateLibrary.h() + gate.set_sites(0) + assert gate.sites == [0] + assert_array_equal(gate.tensor, gate.matrix) + +def test_gate_phase(): + gate = GateLibrary.p() + theta = np.pi / 3 + gate.set_params([theta]) + gate.set_sites(4) + assert gate.sites == [4] + expected_gen = (theta / 2) * np.array([[1, 0], [0, -1]]) + assert_allclose(gate.generator, expected_gen) + # For Phase, tensor equals matrix. + assert_array_equal(gate.tensor, gate.matrix) + +def test_gate_rx(): + gate = GateLibrary.rx() + theta = np.pi / 2 + gate.set_params([theta]) + gate.set_sites(1) + expected = np.array([ + [np.cos(theta / 2), -1j * np.sin(theta / 2)], + [-1j * np.sin(theta / 2), np.cos(theta / 2)] + ]) + assert_allclose(gate.tensor, expected) + +def test_gate_ry(): + gate = GateLibrary.ry() + theta = np.pi / 3 + gate.set_params([theta]) + gate.set_sites(1) + expected = np.array([ + [np.cos(theta / 2), -np.sin(theta / 2)], + [np.sin(theta / 2), np.cos(theta / 2)] + ]) + assert_allclose(gate.matrix, expected) + # For Ry, tensor equals matrix. + assert_allclose(gate.tensor, expected) + +def test_gate_rz(): + gate = GateLibrary.rz() + theta = np.pi / 4 + gate.set_params([theta]) + gate.set_sites(2) + expected = np.array([ + [np.exp(-1j * theta / 2), 0], + [0, np.exp(1j * theta / 2)] + ]) + assert_allclose(gate.matrix, expected) + # For Rz, tensor equals matrix. + assert_allclose(gate.tensor, expected) + +def test_gate_cx(): + gate = GateLibrary.cx() + gate.set_sites(0, 1) + assert gate.sites == [0, 1] + # Check that the tensor is reshaped to (2,2,2,2) + assert gate.tensor.shape == (2, 2, 2, 2) + # The gate should also have an MPO constructed. + assert hasattr(gate, "mpo") + assert isinstance(gate.mpo, MPO) + # Basic check: there should be at least 2 tensors in the MPO. + assert len(gate.mpo.tensors) >= 2 + +def test_gate_cz(): + # Forward order + gate = GateLibrary.cz() + gate.set_sites(0, 1) + tensor_forward = gate.tensor.copy() + assert gate.sites == [0, 1] + # Reverse order: tensor should be transposed on axes (1,0,3,2) + gate_rev = GateLibrary.cz() + gate_rev.set_sites(1, 0) + expected = np.transpose(tensor_forward, (1, 0, 3, 2)) + np.testing.assert_allclose(gate_rev.tensor, expected) + +def test_gate_swap(): + gate = GateLibrary.swap() + gate.set_sites(2, 3) + assert gate.sites == [2, 3] + # Check that tensor is equal to the matrix reshaped to (2,2,2,2) + expected = gate.matrix.reshape(2, 2, 2, 2) + assert_array_equal(gate.tensor, expected) + + +def test_gate_rxx(): + gate = GateLibrary.rxx() + theta = np.pi / 3 + gate.set_params([theta]) + gate.set_sites(0, 1) + expected = gate.matrix.reshape(2, 2, 2, 2) + assert_allclose(gate.tensor, expected) + +def test_gate_ryy(): + gate = GateLibrary.ryy() + theta = np.pi / 4 + gate.set_params([theta]) + gate.set_sites(1, 2) + expected = gate.matrix.reshape(2, 2, 2, 2) + assert_allclose(gate.tensor, expected) + +def test_gate_rzz(): + gate = GateLibrary.rzz() + theta = np.pi / 6 + gate.set_params([theta]) + gate.set_sites(0, 1) + expected = gate.matrix.reshape(2, 2, 2, 2) + assert_allclose(gate.tensor, expected) + +def test_gate_cphase_forward(): + gate = GateLibrary.cp() + theta = np.pi / 2 + gate.set_params([theta]) + gate.set_sites(0, 1) # Forward order (site0 < site1) so no transpose + expected = np.reshape(gate.matrix, (2, 2, 2, 2)) + assert_array_equal(gate.tensor, expected) + +def test_gate_cphase_reverse(): + gate = GateLibrary.cp() + theta = np.pi / 2 + gate.set_params([theta]) + gate.set_sites(1, 0) # Reverse order; tensor should be transposed on (1,0,3,2) + expected = np.reshape(gate.matrix, (2, 2, 2, 2)) + expected = np.transpose(expected, (1, 0, 3, 2)) + assert_allclose(gate.tensor, expected)