Project: s41467-024-50505-8 | Function: simulate_modelBased (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 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.

**KEY TERMS:** {
  "modelRatio": "belief ratio between pattern A and B",
  "actionValue4patternA": "fixed values for pattern A (1,2,3,4,5)",
  "actionValue4patternB": "fixed values for pattern B (5,4,3,2,1)",
  "alpha": "belief updating rate α",
  "rewardPredictionError": "prediction error based on weighted patterns"
}

**CONTEXT CODE:**   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)
    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)]
    #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)
    modelRatio = modelRatio + alpha[currentTrialGroup]*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

def logLikelihood_modelFree(stimSeq,rewardRulePattern,choiceSeq,alpha,beta):
  actionValue = np.zeros(5)
  stimuli = ['S1','S2','S3','S4','S5']
  logLikelihood = 0
  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)
    if stimSeq[index_trial][0]==choiceSeq[index_trial]:
      logLikelihood = logLikelihood + np.log(prob[0])
    elif stimSeq[index_trial][1]==choiceSeq[index_trial]:
      logLikelihood = logLikelihood + np.log(prob[1])
    else:
      print('Error: choiceSeq[n] must be stimSeq[n][0] or stimSeq[n][1].')
      sys.exit(1)
    currentChoice = choiceSeq[index_trial]
    givenReward = returnReward(currentChoice, rewardRulePattern[index_trial])
    rewardPredictionError = givenReward - actionValue[stimuli.index(currentChoice)]
    actionValue[stimuli.index(currentChoice)] = actionValue[stimuli.index(currentChoice)] + alpha*rewardPredictionError
  if (alpha < 0) or (beta < 0):
    logLikelihood = np.log(0.1)*len(stimSeq)
  AIC = -2.0*logLikelihood + 2.0*2.0
  BIC = -2.0*logLikelihood + 2.0*np.log(len(stimSeq))
  return logLikelihood, AIC, BIC

def logLikelihood_modelFree_asymmetric(stimSeq,rewardRulePattern,choiceSeq,alpha_plus,alpha_minus,beta):
  actionValue = np.zeros(5)
  stimuli = ['S1','S2','S3','S4','S5']
  logLikelihood = 0
  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)


Complete the function implementation replacing "raise NotImplementedError".