In [2]:
#library
import matplotlib.pyplot as plt
import random

#Situation Parameters
Ts = 500            #Time Steps
VI = [5, 30]       #Variable Interval 
Sigma = .3          #Sigma Variable Interval 
Gain = [5, 5]       #Amount Reinforcement
Distance = .25      #Distance Alternatives
View = False        

#Agent Parameters
Energy = 50         #Initial Energy
Cost = .25          #Cost TimeStep
X = .8              #Memory

#Variables
Position = 0
TAValue = 0
Choice = 0
Reward = 0
TTime = [0,0]
AValue = [.5, .5]   

#Data
Data = [[0]*9 for i in range(Ts)]
Res = [0, 0]
TRes = [0, 0]
Ref = [0, 0]

#Funtions
def GetCF(AList,Column):
    return [float(i[Column]) for i in AList]

TVI = [round(random.gauss(VI[0], VI[0]*Sigma)), round(random.gauss(VI[1], VI[1]*Sigma))]
for t in range(Ts):
    
    TTime[0] += 1
    TTime[1] += 1
    TAValue = round(AValue[0] / (AValue[0] + AValue[1]),3)
    
    if random.random() <= TAValue:
        Choice = -1
        if Position != -1:
            Position -= Distance
        if (Position == -1) & (TTime[0] >= TVI[0]):
            Energy += Gain[0]-Cost
            TTime[0] = 0
            Reward = -1
            Ref[0] += Gain[0]
            TVI[0] = round(random.gauss(VI[0], VI[0]*Sigma))
        else:
            Energy += -Cost
            Reward = -100
        AValue[0] = round((X*AValue[0])+((1-X)*Energy),3) 
        if AValue[0] < 0:
            AValue[0] = .001
        elif AValue[0] > 1: 
            AValue[0] = 1   
    else:
        Choice = 1
        if Position != 1:
            Position += Distance
        if (Position == 1) & (TTime[1] >= VI[1]):
            Energy += Gain[1]-Cost
            TTime[1] = 0
            Reward = 1
            Ref[1] += Gain[1]
            TVI[1] = round(random.gauss(VI[1], VI[1]*Sigma))
        else:
            Energy += -Cost
            Reward = -100
        AValue[1] = round((X*AValue[1])+((1-X)*Energy),3) 
        if AValue[1] < 0:
            AValue[1] = .001
        elif AValue[1] > 1: 
            AValue[1] = 1   
            
    if Position < 0:
        TRes[0] += 1
    elif Position > 0:
        TRes[1] += 1
    if Position == -1:
        Res[0] += 1
    elif Position == 1:
        Res[1] += 1
        
    Data[t][0] = t
    Data[t][1] = Position 
    Data[t][2] = AValue[0]
    Data[t][3] = AValue[1]
    Data[t][4] = TAValue
    Data[t][5] = Reward 
    Data[t][6] = TTime[0]
    Data[t][7] = TTime[1]
    Data[t][8] = Energy
    
    if View == True:
        plt.figure()    
        if Reward == -1:
            plt.gca().add_patch(plt.Rectangle((-.8,.3), -.4, .4, fc='k',ec='w'))
        elif Reward == 1:
            plt.gca().add_patch(plt.Rectangle((.8,.3), .4, .4, fc='k',ec='w'))
        plt.gca().add_patch(plt.Rectangle((-.9,.4), -.2, .2, fc='w',ec='k'))
        plt.gca().add_patch(plt.Rectangle((.9,.4), .2, .2, fc='w',ec='k'))
        if Position == -1:
            plt.gca().add_patch(plt.Rectangle((-.9,.4), -.2, .2, fc='w',ec='r'))
        elif Position == 1:
            plt.gca().add_patch(plt.Rectangle((.9,.4), .2, .2, fc='w',ec='r'))
        plt.plot(Position,.5, 'k-o')
        plt.axis([-1.5, 1.5, 0, 1])
    

plt.figure(1)    
plt.plot(GetCF(Data,1),GetCF(Data,0), 'k-')
plt.plot(GetCF(Data,5),GetCF(Data,0), 'ro')
plt.xlabel('Time Step')
plt.ylabel('Allocation')
plt.axis([-1.5, 1.5, 0, len(Data)+1])

print('Responses    :',Res,'',round(Res[0]/(Res[0]+Res[1]),2) )
print('Allocation   :',TRes,'',round(TRes[0]/(TRes[0]+TRes[1]),2))
print('Reinforcement:',Ref,'',round(Ref[0]/(Ref[0]+Ref[1]),2))
Responses    : [90, 42]  0.68
Allocation   : [291, 165]  0.64
Reinforcement: [160, 25]  0.86