函数内部的 For 循环无法迭代变量

Question

I'm willing to compute all values of all of variables combinations. However, the code always produces the same values for all iterations.

#chsh
import numpy as np
from bell_state import bell_state
from pauli import pauli

#setup of measurement

N_A = 2  #measurement outcomes
N_B = 2 

#Bell_state used
rho = bell_state(1)

#Alice 

A_1 =pauli('Z')
A_2 =pauli('X')

#Bob

B_1 =(1/np.sqrt(2))*(pauli('Z')+pauli('X'))
B_2 =(1/np.sqrt(2))*(pauli('Z')-pauli('X'))

def AB(a,b):
    if a==1 and b==1:
        aabb = np.kron(A_1,B_1)
    elif a==2 and b==1:
        aabb = np.kron(A_2,B_1)
    elif a==1 and b==2:
        aabb = np.kron(A_1,B_2)
    elif a==2 and b==2:
        aabb = np.kron(A_2,B_2)
    else:
        print('error')
    return aabb

def Obs(a,b):
    for a in range(1,N_A):
        for b in range(1,N_B):
            ob =np.trace(np.dot(AB(a,b),rho))
            return ob

IObs = Obs(1,1) + Obs(1,2) + Obs(2,1) - Obs(2,2)

The code showed this results, those not what I want

Obs(1,1) : 0.7071067811865475 
 Obs(2,1) : 0.7071067811865475 
 Obs(1,2) : 0.7071067811865475 
 Obs(2,2) : 0.7071067811865475 

IObs : 
 1.414213562373095

The results what I expected actually as same as the results that (I computed manually without loops by this below code

Ob11= np.trace(np.dot(AB(1,1),rho))
Ob21= np.trace(np.dot(AB(2,1),rho))
Ob12= np.trace(np.dot(AB(1,2),rho))
Ob22= np.trace(np.dot(AB(2,2),rho))

results of what I want

Obs(1,1)  : 0.7071067811865475 
Obs(2,1)  : 0.7071067811865475 
Obs(1,2)  : 0.7071067811865475 
Obs(2,2)  : -0.7071067811865475 

IObs : 
 2.82842712474619

Perhaps somebody could help me find what I've done wrong with my for-loop function. Thank you!

Answer 1

I think the issue is in Obs function here

def Obs(a,b):
    for a in range(1,N_A):
        for b in range(1,N_B):
            ob =np.trace(np.dot(AB(a,b),rho))
            return ob

As you might see, the return statement is placed inside the loop block. This makes your code for calculating ob is never executed more than 1, in other word, it is executed exactly once for whatever a and b you pass to the function.

Based on my assumption (cause I am not too sure about the logic you wanted to implement), you may want to try to use this instead:

def Obs(a,b):
    for a in range(1,N_A):
        for b in range(1,N_B):
            ob =np.trace(np.dot(AB(a,b),rho))
    return ob