1+ import numpy as np
2+ import pytest
3+
4+ from yaqs .core .data_structures .networks import MPO , MPS
5+ from yaqs .core .data_structures .simulation_parameters import Observable , PhysicsSimParams
6+
7+ from yaqs .core .methods .TDVP import (
8+ _split_mps_tensor ,
9+ _merge_mps_tensor_pair ,
10+ _merge_mpo_tensor_pair ,
11+ _contraction_operator_step_right ,
12+ _contraction_operator_step_left ,
13+ _compute_right_operator_blocks ,
14+ _apply_local_hamiltonian ,
15+ _apply_local_bond_contraction ,
16+ _local_hamiltonian_step ,
17+ _local_bond_step ,
18+ single_site_TDVP ,
19+ two_site_TDVP ,
20+ )
21+
22+ def test_split_mps_tensor_left_right_sqrt ():
23+ # Create an input tensor A with shape (d0*d1, D0, D2).
24+ # Let d0 = d1 = 2 so that A.shape[0]=4, and choose D0=3, D2=5.
25+ A = np .random .randn (4 , 3 , 5 )
26+ # For each svd_distr option, run the splitting and then reconstruct A.
27+ for distr in ['left' , 'right' , 'sqrt' ]:
28+ A0 , A1 = _split_mps_tensor (A , svd_distr = distr , threshold = 1e-8 )
29+ # A0 should have shape (2, 3, r) and A1 should have shape (2, r, 5),
30+ # where r is the effective rank.
31+ assert A0 .ndim == 3
32+ assert A1 .ndim == 3
33+ r = A0 .shape [2 ]
34+ assert A1 .shape [1 ] == r
35+ # Reconstruct A: undo the transpose on A1 then contract A0 and A1.
36+ A1_recon = A1 .transpose ((1 , 0 , 2 )) # now shape (r, 2, 5)
37+ # Contract along the rank index:
38+ # A0 has indices: (d0, D0, r), A1_recon has indices: (r, d1, D2).
39+ # Form a tensor of shape (d0, d1, D0, D2) then reshape back to (4, 3, 5)
40+ A_recon = np .tensordot (A0 , A1_recon , axes = (2 , 0 )) # shape (2, 3, 2, 5)
41+ A_recon = A_recon .transpose ((0 ,2 ,1 ,3 )).reshape (4 , 3 , 5 )
42+ # Up to SVD sign and ordering ambiguities, the reconstruction should be close.
43+ np .testing .assert_allclose (A , A_recon , atol = 1e-6 )
44+
45+ def test_split_mps_tensor_invalid_shape ():
46+ # If A.shape[0] is not divisible by 2, the function should raise a ValueError.
47+ A = np .random .randn (3 , 3 , 5 )
48+ with pytest .raises (ValueError ):
49+ _split_mps_tensor (A , svd_distr = 'left' )
50+
51+ def test_merge_mps_tensor_pair ():
52+ # Let A0 have shape (2, 3, 4) and A1 have shape (5, 4, 7).
53+ # In _merge_mps_tensor_pair the arrays are interpreted with label tuples
54+ # (0,2,3) for A0 and (1,3,4) for A1, so contraction is over the third axis of A0 and
55+ # the second axis of A1.
56+ A0 = np .random .randn (2 , 3 , 4 )
57+ A1 = np .random .randn (5 , 4 , 7 )
58+ merged = _merge_mps_tensor_pair (A0 , A1 )
59+ # Expected shape: first two axes merged from A0[0] and A1[0]:
60+ # output shape = ((2*5), 3, 7) i.e. (10, 3, 7)
61+ assert merged .shape == (10 , 3 , 7 )
62+
63+ def test_merge_mpo_tensor_pair ():
64+ # Let A0 be a 4D array with shape (2, 3, 4, 5) and
65+ # A1 be a 4D array with shape (7, 8, 5, 9).
66+ # The contract call uses label tuples (0,2,4,6) and (1,3,6,5) and contracts label 6.
67+ A0 = np .random .randn (2 , 3 , 4 , 5 )
68+ A1 = np .random .randn (7 , 8 , 5 , 9 )
69+ merged = _merge_mpo_tensor_pair (A0 , A1 )
70+ # Expected shape: merge first two axes of the result, where result (before reshape)
71+ # should have shape (2,7,3,8,4,9), then merged to (2*7, 3*8, 4,9) = (14,24,4,9).
72+ assert merged .shape == (14 , 24 , 4 , 9 )
73+
74+ def test_contraction_operator_step_right ():
75+ # Choose shapes as described in the function.
76+ A = np .random .randn (2 , 3 , 4 )
77+ # R: contract A's last axis (size 4) with R's first axis.
78+ R = np .random .randn (4 , 5 , 6 )
79+ # W must have shape with axes (1,3) matching T from tensordot(A, R, 1) of shape (2,3,5,6):
80+ # We require W.shape[1]=2 and W.shape[3]=5. Let W = (7,2,8,5)
81+ W = np .random .randn (7 , 2 , 8 , 5 )
82+ # B: choose shape so that B.conj() has axes (0,2) matching T from later step.
83+ # After the previous steps, T becomes shape (3,8,7,6); we contract with B.conj() axes ((2,3),(0,2)).
84+ # So let B be of shape (7, 9, 6).
85+ B = np .random .randn (7 , 9 , 6 )
86+ Rnext = _contraction_operator_step_right (A , B , W , R )
87+ # Expected shape: (3,8,9) (from the discussion above).
88+ assert Rnext .shape == (3 , 8 , 9 )
89+
90+ def test_contraction_operator_step_left ():
91+ # Set up dummy arrays with shapes so that the contraction is valid.
92+ # Let A be shape (3,4,10)
93+ A = np .random .randn (3 , 4 , 10 )
94+ # Let B be shape (7, 6, 8) so that B.conj() has shape (7,6,8).
95+ B = np .random .randn (7 , 6 , 8 )
96+ # Let L be shape (4,5,6) and require that L.shape[2] (6) matches B.shape[1] (6).
97+ L_arr = np .random .randn (4 , 5 , 6 )
98+ # Let W be shape (7,3,5,9) so that we contract axes ((0,2),(2,1)) with T later.
99+ W = np .random .randn (7 , 3 , 5 , 9 )
100+ Rnext = _contraction_operator_step_left (A , B , W , L_arr )
101+ # The expected shape from our reasoning is (10,9,8) (A's remaining axis 2 becomes output along with leftover dims from T).
102+ # We check that the result is 3-dimensional.
103+ assert Rnext .ndim == 3
104+
105+ def test_apply_local_hamiltonian ():
106+ # Let A: shape (2,3,4); R: shape (4,5,6) as before.
107+ A = np .random .randn (2 , 3 , 4 )
108+ R = np .random .randn (4 , 5 , 6 )
109+ # Let W be shape (7,2,8,5) as in test above.
110+ W = np .random .randn (7 , 2 , 8 , 5 )
111+ # Let L be shape (3,8,9) so that tensordot works.
112+ L_arr = np .random .randn (3 , 8 , 9 )
113+ out = _apply_local_hamiltonian (L_arr , R , W , A )
114+ # The function transposes the final result so we expect a 3D array.
115+ assert out .ndim == 3
116+
117+ def test_apply_local_bond_contraction ():
118+ # Let C: shape (2,3)
119+ C = np .random .randn (2 , 3 )
120+ # Let R: shape (3,4,5)
121+ R = np .random .randn (3 , 4 , 5 )
122+ # Let L: shape (2,4,6)
123+ L_arr = np .random .randn (2 , 4 , 6 )
124+ out = _apply_local_bond_contraction (L_arr , R , C )
125+ # Expected output shape: contraction gives shape (6,5)
126+ assert out .shape == (6 , 5 )
127+
128+ def test_local_hamiltonian_step ():
129+ # We choose an MPS tensor A with shape (d0, d0, d1) where d0=2, d1=4.
130+ # (The requirement here is that the first two dimensions are equal,
131+ # so that after the contraction chain the operator is square.)
132+ A = np .random .randn (2 , 2 , 4 ) # total elements: 2*2*4 = 16
133+ # Choose R of shape (d1, X, d1) with d1=4 and X=1.
134+ R = np .random .randn (4 , 1 , 4 ) # shape: (4,1,4)
135+ # Choose W of shape (d0, d0, X, X) with d0=2 and X=1.
136+ W = np .random .randn (2 , 2 , 1 , 1 ) # shape: (2,2,1,1)
137+ # Choose L so that the contraction makes sense.
138+ # In our contraction chain, after:
139+ # T1 = tensordot(A, R, axes=1) → shape (2,2,1,4)
140+ # T2 = tensordot(W, T1, axes=((1,3),(0,2))) → shape (2,1,2,4)
141+ # T3 = T2.transpose((2,1,0,3)) → shape (2,1,2,4) (here the permutation reorders axes)
142+ # Then we contract T3 with L along axes ((2,1),(0,1)).
143+ # To contract T3’s axes (axis2, axis1) = (2,1) we need L with shape (2,1,r).
144+ # Then T4 will have shape (remaining T3 axes: (axis0, axis3)) plus L’s remaining axis, i.e. (2,4,r).
145+ # Finally, a transpose (here, (0,2,1)) gives shape (2, r, 4).
146+ # We want the final shape to equal A’s shape (2,2,4), so we set r=2.
147+ L_arr = np .random .randn (2 , 1 , 2 ) # shape: (2,1,2)
148+ dt = 0.05
149+ numiter = 10
150+ out = _local_hamiltonian_step (L_arr , R , W , A , dt , numiter )
151+ # The operator should be square, so out.shape should equal A.shape.
152+ assert out .shape == A .shape , f"Expected shape { A .shape } { out .shape }
153+
154+ def test_local_bond_step ():
155+ # For the bond step we want the operator to be square.
156+ # Let C be a matrix of shape (p, q). We choose C to be square.
157+ C = np .random .randn (2 , 2 ) # total elements: 4
158+ # Choose R of shape (q, r, s). To contract C and get back the same number of elements,
159+ # we can choose R such that q=2, r=2, s=2.
160+ R = np .random .randn (2 , 2 , 2 ) # shape: (2,2,2)
161+ # Choose L of shape (p, r, t) with p=2 and r=2.
162+ # We want the final result to have shape (p, q) = (2,2); so t must equal q=2.
163+ L_arr = np .random .randn (2 , 2 , 2 ) # shape: (2,2,2)
164+ dt = 0.05
165+ numiter = 10
166+ out = _local_bond_step (L_arr , R , C , dt , numiter )
167+ # The output shape should equal the input shape.
168+ assert out .shape == C .shape , f"Expected shape { C .shape } { out .shape }
169+
170+ def test_single_site_TDVP ():
171+ L = 5
172+ d = 2
173+ J = 1
174+ g = 0.5
175+ H = MPO ()
176+ H .init_Ising (L , d , J , g )
177+ state = MPS (L )
178+ measurements = [Observable ('z' , site ) for site in range (L )]
179+ sim_params = PhysicsSimParams (measurements , T = 0.2 , dt = 0.1 , sample_timesteps = True , N = 1 , max_bond_dim = 4 , threshold = 1e-6 , order = 1 )
180+ single_site_TDVP (state , H , sim_params , numiter_lanczos = 5 )
181+ # Check that state still has L tensors and that each tensor is a numpy array.
182+ assert state .length == L
183+ for tensor in state .tensors :
184+ assert isinstance (tensor , np .ndarray )
185+
186+ canonical_site = state .check_canonical_form ()[0 ]
187+ assert canonical_site == 0 , (
188+ f"MPS should be site-canonical at site 0 after single-site TDVP, "
189+ f"but got canonical site: { canonical_site }
190+ )
191+
192+ def test_two_site_TDVP ():
193+ L = 5
194+ d = 2
195+ J = 1
196+ g = 0.5
197+ H = MPO ()
198+ H .init_Ising (L , d , J , g )
199+ state = MPS (L )
200+ measurements = [Observable ('z' , site ) for site in range (L )]
201+ sim_params = PhysicsSimParams (measurements , T = 0.2 , dt = 0.1 , sample_timesteps = True , N = 1 , max_bond_dim = 4 , threshold = 1e-6 , order = 1 )
202+ two_site_TDVP (state , H , sim_params , numiter_lanczos = 5 )
203+ # Check that state still has L tensors and that each tensor is a numpy array.
204+ assert state .length == L
205+ for tensor in state .tensors :
206+ assert isinstance (tensor , np .ndarray )
207+
208+ canonical_site = state .check_canonical_form ()[0 ]
209+ assert canonical_site == 0 , (
210+ f"MPS should be site-canonical at site 0 after two-site TDVP, "
211+ f"but got canonical site: { canonical_site }
212+ )
0 commit comments