Project: s41467-024-50505-8 | Function: simulate_modelFree (python)

system: You are an expert Python programmer implementing algorithms from scientific papers.
Complete the masked function based on the algorithm description.
Follow Python best practices and maintain consistent code style.
Return **only** the complete Python function code.

user: Complete this Python function:

**DESCRIPTION:** Implements the core model-free reinforcement learning algorithm (EXP model) used to simulate experience-based value updating in the NOVEL task. Uses Q-learning with reward prediction errors and softmax action selection to model direct learning from stimulus-reward associations.

**KEY TERMS:** {
  "actionValue": "stimulus values V(s)",
  "alpha": "learning rate α",
  "rewardPredictionError": "reward prediction error RPE",
  "softmax": "softmax choice policy with temperature β",
  "givenReward": "obtained reward r"
}

**CONTEXT CODE:**       reward = 2
    elif choice == 'S3':
      reward = 3
    elif choice == 'S4':
      reward = 4
    elif choice == 'S5':
      reward = 5
    else:
      print('Error: The chosen item must be S1, S2, S3, S4, or S5.')
      sys.exit(1)
  elif rewardRulePattern == 'B':
    if choice == 'S1':
      reward = 5
    elif choice == 'S2':
      reward = 4
    elif choice == 'S3':
      reward = 3
    elif choice == 'S4':
      reward = 2
    elif choice == 'S5':
      reward = 1
    else:
      print('Error: The chosen item must be S1, S2, S3, S4, or S5.')
      sys.exit(1) 
  else:
    print('Error: Reward rule pattern must be A or B.')
    sys.exit(1) 
  return reward

def optimalChoiceRate(stimSeq, rewardRulePattern,choiceSeq):
  stimuli = ['S1','S2','S3','S4','S5']
  optimalOrNot = np.zeros(len(stimSeq))
  for index_trial in range(len(optimalOrNot)):
    stimLeft = stimuli.index(stimSeq[index_trial][0])
    stimRight = stimuli.index(stimSeq[index_trial][1])
    currentChoice = stimuli.index(choiceSeq[index_trial])
    if rewardRulePattern[index_trial]=='A':
      if currentChoice==np.max([stimLeft,stimRight]):
        optimalOrNot[index_trial] = 1
    elif rewardRulePattern[index_trial]=='B':
      if currentChoice==np.min([stimLeft,stimRight]):
        optimalOrNot[index_trial] = 1
    else:
      print('Error: Reward rule pattern must be A or B.')
      sys.exit(1)
  rate = np.zeros(round(len(stimSeq)/10))
  for index_block in range(round(len(stimSeq)/10)):
    rate[index_block] = np.mean(optimalOrNot[(index_block*10):((index_block+1)*10)])
  return rate

def simulate_modelFree(stimSeq,rewardRulePattern,alpha,beta):
  # MASKED: Implements the core model-free reinforcement learning algorithm (EXP model) used to simulate experience-based value updating in the NOVEL task. Uses Q-learning with reward prediction errors and softmax action selection to model direct learning from stimulus-reward associations.
  raise NotImplementedError("Function has been masked for testing purposes")

#Added on 20210108.
#In this model, two learning rate parameters are used.
def simulate_modelFree_asymmetric(stimSeq,rewardRulePattern,alpha_plus,alpha_minus,beta):
  actionValue = np.zeros(5)
  stimuli = ['S1','S2','S3','S4','S5']
  choiceSeq = list()
  for index_trial in range(len(stimSeq)):
    stimLeft = stimuli.index(stimSeq[index_trial][0])
    stimRight = stimuli.index(stimSeq[index_trial][1])
    prob = softmax(np.array([actionValue[stimLeft],actionValue[stimRight]]), beta)
    currentChoice = np.random.choice([stimuli[stimLeft], stimuli[stimRight]], size=(1,), p=prob)
    currentChoice = currentChoice[0]
    choiceSeq.append(currentChoice)
    givenReward = returnReward(currentChoice, rewardRulePattern[index_trial])
    rewardPredictionError = givenReward - actionValue[stimuli.index(currentChoice)]
    if rewardPredictionError > 0:
      actionValue[stimuli.index(currentChoice)] = actionValue[stimuli.index(currentChoice)] + alpha_plus*rewardPredictionError
    else:
      actionValue[stimuli.index(currentChoice)] = actionValue[stimuli.index(currentChoice)] + alpha_minus*rewardPredictionError
  return choiceSeq

#Added on 20210221.
#In this model, trials are grouped by the user and alpha is estimated for each trial group.
#For example, if you group trials 0-69 trials, 70-89 trials, 90-159 trials, & 159-299 trials into the four groups (four periods).
#Four alpha parameters are estimated, which are cprresponding to 'learning phase', 'stable phase 1', 're-learning phase' and 'stable phase 2'.
def simulate_modelFree_temporallyChangingAlpha(stimSeq,rewardRulePattern,alpha,beta,trialGroup):
  actionValue = np.zeros(5)
  stimuli = ['S1','S2','S3','S4','S5']
  choiceSeq = list()
  for index_trial in range(len(stimSeq)):
    stimLeft = stimuli.index(stimSeq[index_trial][0])
    stimRight = stimuli.index(stimSeq[index_trial][1])
    prob = softmax(np.array([actionValue[stimLeft],actionValue[stimRight]]), beta)
    currentChoice = np.random.choice([stimuli[stimLeft], stimuli[stimRight]], size=(1,), p=prob)
    currentChoice = currentChoice[0]
    choiceSeq.append(currentChoice)
    givenReward = returnReward(currentChoice, rewardRulePattern[index_trial])
    rewardPredictionError = givenReward - actionValue[stimuli.index(currentChoice)]
    #choose used alpha depending on the trial group.
    currentTrialGroup = -1
    for index_trialGroup in range(len(trialGroup)):
      if np.any(trialGroup[index_trialGroup]==index_trial):
        if currentTrialGroup == -1:
          currentTrialGroup = index_trialGroup
        else:
          print('Error: the same trial is used in multiple groups you specified.')
          sys.exit(1)
    if currentTrialGroup == -1:
      print('Error: trial # ' + str(index_trial) + ' does not belong to any group you specified.')
      sys.exit(1)
    actionValue[stimuli.index(currentChoice)] = actionValue[stimuli.index(currentChoice)] + alpha[currentTrialGroup]*rewardPredictionError
  return choiceSeq

def simulate_modelBased(stimSeq,rewardRulePattern,alpha,beta):
  # MASKED: Implements the model-based reinforcement learning algorithm (INF model) used to simulate inference-based value updating in the FAMILIAR task. Uses prior knowledge of two possible reward patterns and updates belief ratios based on prediction errors to enable rapid adaptation through inference.
  raise NotImplementedError("Function has been masked for testing purposes")

def simulate_modelBased_asymmetric(stimSeq,rewardRulePattern,alpha_plus,alpha_minus,beta):
  modelRatio = 0.5
  stimuli = ['S1','S2','S3','S4','S5']
  choiceSeq = list()
  actionValue4patternA = np.array([1.0,2.0,3.0,4.0,5.0]) 
  actionValue4patternB = np.array([5.0,4.0,3.0,2.0,1.0]) 
  for index_trial in range(len(stimSeq)):
    stimLeft = stimuli.index(stimSeq[index_trial][0])
    stimRight = stimuli.index(stimSeq[index_trial][1])
    prob = modelRatio*softmax(np.array([actionValue4patternA[stimLeft],actionValue4patternA[stimRight]]), beta)+(1.0-modelRatio)*softmax(np.array([actionValue4patternB[stimLeft],actionValue4patternB[stimRight]]), beta)
    currentChoice = np.random.choice([stimuli[stimLeft], stimuli[stimRight]], size=(1,), p=prob)
    currentChoice = currentChoice[0]
    choiceSeq.append(currentChoice)
    givenReward = returnReward(currentChoice, rewardRulePattern[index_trial])
    rewardPredictionError = givenReward - modelRatio*actionValue4patternA[stimuli.index(currentChoice)] - (1.0-modelRatio)*actionValue4patternB[stimuli.index(currentChoice)]
    if rewardPredictionError > 0:
      modelRatio = modelRatio + alpha_plus*rewardPredictionError*(actionValue4patternA[stimuli.index(currentChoice)]-actionValue4patternB[stimuli.index(currentChoice)])
    else:
      modelRatio = modelRatio + alpha_minus*rewardPredictionError*(actionValue4patternA[stimuli.index(currentChoice)]-actionValue4patternB[stimuli.index(currentChoice)])
    if modelRatio > 1.0:
      modelRatio = 1.0
    elif modelRatio < 0.0:
      modelRatio = 0.0
  return choiceSeq

#Added on 20210221.
#In this model, trials are grouped by the user and alpha is estimated for each trial group.
#For example, if you group trials 0-69 trials, 70-89 trials, 90-159 trials, & 159-299 trials into the four groups (four periods).
#Four alpha parameters are estimated, which are cprresponding to 'learning phase', 'stable phase 1', 're-learning phase' and 'stable phase 2'.
def simulate_modelBased_temporallyChangingAlpha(stimSeq,rewardRulePattern,alpha,beta,trialGroup):
  modelRatio = 0.5
  stimuli = ['S1','S2','S3','S4','S5']
  choiceSeq = list()
  actionValue4patternA = np.array([1.0,2.0,3.0,4.0,5.0]) 
  actionValue4patternB = np.array([5.0,4.0,3.0,2.0,1.0]) 
  for index_trial in range(len(stimSeq)):
    stimLeft = stimuli.index(stimSeq[index_trial][0])
    stimRight = stimuli.index(stimSeq[index_trial][1])
    prob = modelRatio*softmax(np.array([actionValue4patternA[stimLeft],actionValue4patternA[stimRight]]), beta)+(1.0-modelRatio)*softmax(np.array([actionValue4patternB[stimLeft],actionValue4patternB[stimRight]]), beta)


Complete the function implementation replacing "raise NotImplementedError".