#!/usr/bin/env python """ Synthesize data for Kanen reversals paper (controls/OCD/substance abuse disorder; placebo/amisulpride/pramipexole. Rudolf Cardinal, 13 May 2019 onwards. Possible languages -- R, Python, C++ all easy; Python probably the quickest to write (and less intrinsically prone to namespacing errors than R). Also offers the opportunity to structure this well. Last update: 20 May 2019. """ # ============================================================================= # Imports # ============================================================================= import argparse import copy import csv import os import logging import pprint import random import statistics import sys from typing import Any, Dict, Generator, List, Optional, Set, TextIO, Tuple from cardinal_pythonlib.dicts import HashableDict from cardinal_pythonlib.fileops import mkdir_p from cardinal_pythonlib.logs import main_only_quicksetup_rootlogger from cardinal_pythonlib.maths_numpy import softmax from cardinal_pythonlib.randomness import coin from cardinal_pythonlib.typing_helpers import CSVWriterType import cardinal_pythonlib.version import numpy as np from openpyxl.workbook.workbook import Workbook as XLWorkbook from openpyxl.worksheet.worksheet import Worksheet as XLWorksheet if sys.version_info < (3, 6): raise AssertionError("Need Python 3.6 or higher") cardinal_pythonlib.version.assert_version_ge("1.0.54") log = logging.getLogger(__name__) # ============================================================================= # Constants # ============================================================================= # Groups CONTROL = "control" OCD = "OCD" SUD = "SUD" GROUPS = [CONTROL, OCD, SUD] # Drugs PLACEBO = "placebo" AMISULPRIDE = "amisulpride" PRAMIPEXOLE = "pramipexole" DRUGS = [PLACEBO, AMISULPRIDE, PRAMIPEXOLE] # Parameter names ALPHA_REWARD = "alpha_reward" ALPHA_PUNISHMENT = "alpha_punishment" TAU_REINFORCEMENT_SENSITIVITY = "tau_reinforcement_sensitivity" TAU_LOCATION_STICKINESS = "tau_location_stickiness" TAU_STIMULUS_STICKINESS = "tau_stimulus_stickiness" PARAMETER_NAMES = [ALPHA_REWARD, ALPHA_PUNISHMENT, TAU_REINFORCEMENT_SENSITIVITY, TAU_LOCATION_STICKINESS, TAU_STIMULUS_STICKINESS] PARAMS_TYPE = Dict[str, Dict[Tuple[str, str], float]] # ============================================================================= # Data from the winning model # ============================================================================= """ Below are group mean values, from the Bayesian analysis. To retrieve them, from within R: .. code-block:: R source("bayesian_reversals_jk.R") run_models("2") s <- summarize_for_group_level_params(fit2) cols <- c("parameter", "mean") # use the mean s[grep("^reward_rate_by_group_drug.*", parameter), ..cols] s[grep("^punish_rate_by_group_drug.*", parameter), ..cols] s[grep("^reinf_sensitivity_by_group_drug.*", parameter), ..cols] s[grep("^side_stickiness_by_group_drug.*", parameter), ..cols] s[grep("^stimulus_stickiness_by_group_drug.*", parameter), ..cols] Note that group 1 = control, 2 = SUD, 3 = OCD. Note that drug 1 = placebo, 2 = amisulpride, 3 = pramipexole. So ``reward_rate_by_group_drug[1,1]`` means "control, placebo"; etc. """ MODEL_2_FITTED_PARAMETERS = { ALPHA_REWARD: { (CONTROL, PLACEBO): 0.8474892, (CONTROL, AMISULPRIDE): 0.8865761, (CONTROL, PRAMIPEXOLE): 0.8891623, (SUD, PLACEBO): 0.2684286, (SUD, AMISULPRIDE): 0.9652802, (SUD, PRAMIPEXOLE): 0.9017787, (OCD, PLACEBO): 0.7751281, (OCD, AMISULPRIDE): 0.8971247, (OCD, PRAMIPEXOLE): 0.6166932, }, ALPHA_PUNISHMENT: { (CONTROL, PLACEBO): 0.5594376, (CONTROL, AMISULPRIDE): 0.6072152, (CONTROL, PRAMIPEXOLE): 0.6208992, (SUD, PLACEBO): 0.7226903, (SUD, AMISULPRIDE): 0.6940049, (SUD, PRAMIPEXOLE): 0.6551279, (OCD, PLACEBO): 0.6732506, (OCD, AMISULPRIDE): 0.7725250, (OCD, PRAMIPEXOLE): 0.7667571, }, TAU_REINFORCEMENT_SENSITIVITY: { (CONTROL, PLACEBO): 4.844720, (CONTROL, AMISULPRIDE): 4.553197, (CONTROL, PRAMIPEXOLE): 4.562372, (SUD, PLACEBO): 4.525492, (SUD, AMISULPRIDE): 3.132476, (SUD, PRAMIPEXOLE): 2.542957, (OCD, PLACEBO): 5.572169, (OCD, AMISULPRIDE): 4.978025, (OCD, PRAMIPEXOLE): 5.393172, }, TAU_LOCATION_STICKINESS: { (CONTROL, PLACEBO): 0.082065452, (CONTROL, AMISULPRIDE): 0.056370747, (CONTROL, PRAMIPEXOLE): 0.039643831, (SUD, PLACEBO): 0.156305868, (SUD, AMISULPRIDE): -0.051518697, (SUD, PRAMIPEXOLE): 0.038808212, (OCD, PLACEBO): -0.065349797, (OCD, AMISULPRIDE): 0.032245186, (OCD, PRAMIPEXOLE): -0.007937715, }, TAU_STIMULUS_STICKINESS: { (CONTROL, PLACEBO): 0.01990708, (CONTROL, AMISULPRIDE): 0.11881618, (CONTROL, PRAMIPEXOLE): -0.00927243, (SUD, PLACEBO): 0.31427687, (SUD, AMISULPRIDE): 0.17301462, (SUD, PRAMIPEXOLE): 0.32213736, (OCD, PLACEBO): -0.36730167, (OCD, AMISULPRIDE): -0.31840116, (OCD, PRAMIPEXOLE): -0.25849867, } } """ Now a version with the location stickiness parameter fixed to the overall mean (meaning: the mean of the 3x3 group means). Obtain this as follows: .. code-block:: R source("bayesian_reversals_jk.R") run_models("2") """ _OVERALL_MEAN_LOC_STICKINESS = statistics.mean( MODEL_2_FITTED_PARAMETERS[TAU_LOCATION_STICKINESS].values() ) MODEL_2_FITTED_PARAMETERS_FIX_LOC_STICKINESS = copy.deepcopy( MODEL_2_FITTED_PARAMETERS) MODEL_2_FITTED_PARAMETERS_FIX_LOC_STICKINESS[TAU_LOCATION_STICKINESS] = { (CONTROL, PLACEBO): _OVERALL_MEAN_LOC_STICKINESS, (CONTROL, AMISULPRIDE): _OVERALL_MEAN_LOC_STICKINESS, (CONTROL, PRAMIPEXOLE): _OVERALL_MEAN_LOC_STICKINESS, (SUD, PLACEBO): _OVERALL_MEAN_LOC_STICKINESS, (SUD, AMISULPRIDE): _OVERALL_MEAN_LOC_STICKINESS, (SUD, PRAMIPEXOLE): _OVERALL_MEAN_LOC_STICKINESS, (OCD, PLACEBO): _OVERALL_MEAN_LOC_STICKINESS, (OCD, AMISULPRIDE): _OVERALL_MEAN_LOC_STICKINESS, (OCD, PRAMIPEXOLE): _OVERALL_MEAN_LOC_STICKINESS, } """ Now a version with the stimulus stickiness parameter fixed to the overall mean. """ _OVERALL_MEAN_STIM_STICKINESS = statistics.mean( MODEL_2_FITTED_PARAMETERS[TAU_STIMULUS_STICKINESS].values() ) MODEL_2_FITTED_PARAMETERS_FIX_STIM_STICKINESS = copy.deepcopy( MODEL_2_FITTED_PARAMETERS) MODEL_2_FITTED_PARAMETERS_FIX_STIM_STICKINESS[TAU_STIMULUS_STICKINESS] = { (CONTROL, PLACEBO): _OVERALL_MEAN_STIM_STICKINESS, (CONTROL, AMISULPRIDE): _OVERALL_MEAN_STIM_STICKINESS, (CONTROL, PRAMIPEXOLE): _OVERALL_MEAN_STIM_STICKINESS, (SUD, PLACEBO): _OVERALL_MEAN_STIM_STICKINESS, (SUD, AMISULPRIDE): _OVERALL_MEAN_STIM_STICKINESS, (SUD, PRAMIPEXOLE): _OVERALL_MEAN_STIM_STICKINESS, (OCD, PLACEBO): _OVERALL_MEAN_STIM_STICKINESS, (OCD, AMISULPRIDE): _OVERALL_MEAN_STIM_STICKINESS, (OCD, PRAMIPEXOLE): _OVERALL_MEAN_STIM_STICKINESS, } # ============================================================================= # Parameters for (limited) parameter recovery demonstration # ============================================================================= DEMO_PARAM_CENTRE = { ALPHA_REWARD: 0.5, ALPHA_PUNISHMENT: 0.5, TAU_REINFORCEMENT_SENSITIVITY: 4.5, TAU_LOCATION_STICKINESS: 0, TAU_STIMULUS_STICKINESS: 0, } DEMO_PARAM_ALL = { # should include the central values from DEMO_PARAM_CENTRE ALPHA_REWARD: [0.1, 0.3, 0.5, 0.7, 0.9], ALPHA_PUNISHMENT: [0.1, 0.3, 0.5, 0.7, 0.9], TAU_REINFORCEMENT_SENSITIVITY: [1, 1.5, 3, 4.5, 6, 7.5], TAU_LOCATION_STICKINESS: [-0.2, -0.1, 0, 0.1, 0.2], TAU_STIMULUS_STICKINESS: [-0.5, -0.25, 0, 0.25, 0.5], } # ============================================================================= # Infrastructure classes # ============================================================================= class GroupDrugConfig(object): """ Represents a simulation's configuration for a group/drug combination. """ def __init__(self, group_name: str, drug_name: str, alpha_reward: float, alpha_punishment: float, tau_reinforcement_sensitivity: float, tau_location_stickiness: float, tau_stimulus_stickiness: float, n_subjects: int) -> None: self.group_name = group_name self.drugname = drug_name self.alpha_reward = alpha_reward self.alpha_punishment = alpha_punishment self.tau_reinforcement_sensitivity = tau_reinforcement_sensitivity self.tau_location_stickiness = tau_location_stickiness self.tau_stimulus_stickiness = tau_stimulus_stickiness self.n_subjects = n_subjects def group_drug_name(self) -> str: return f"{self.group_name}_{self.drugname}" class SubjectConfig(object): """ Represents a simulation's configuration for a particular subject. """ def __init__(self, subject_name: str, subject_num: int, group_drug_config: GroupDrugConfig) -> None: self.subject_name = subject_name self.subject_num = subject_num self.group_drug_config = group_drug_config def group_and_subject_name(self) -> str: return f"{self.group_drug_config.group_drug_name()}_{self.subject_name}" # noqa class TrialRecord(object): """ Represents the results of a single trial. """ def __init__(self, subject_config: SubjectConfig, part: int, trial_number: int, block: int, trial_in_block: int, left_stimulus: int, right_stimulus: int, correct_stimulus: int, incorrect_stimulus: int, v_left: float, v_right: float, p_choose_left: float, p_choose_right: float, chosen_stimulus: int, chose_correctly: bool, responded_left: bool, responded_right: bool, outcome: int, correct_responses_in_block: int, reversal_criterion: int) -> None: self.subject_config = subject_config self.part = part self.trial_number = trial_number self.block = block self.trial_in_block = trial_in_block self.left_stimulus = left_stimulus self.right_stimulus = right_stimulus self.correct_stimulus = correct_stimulus self.incorrect_stimulus = incorrect_stimulus self.v_left = v_left self.v_right = v_right self.p_choose_left = p_choose_left self.p_choose_right = p_choose_right self.chosen_stimulus = chosen_stimulus self.chose_correctly = chose_correctly self.responded_left = responded_left self.responded_right = responded_right self.outcome = outcome self.correct_responses_in_block = correct_responses_in_block self.reversal_criterion = reversal_criterion @staticmethod def header() -> List[str]: """ Returns headings for a CSV file. Order must match :meth:`csv_row`. """ return [ "group_name", "subject_name", "subject_num", "drug_name", "alpha_reward", "alpha_punishment", "tau_reinforcement_sensitivity", "tau_location_stickiness", "tau_stimulus_stickiness", "part", "trial_number", "block", "trial_in_block", "left_stimulus", "right_stimulus", "correct_stimulus", "incorrect_stimulus", "v_left", "v_right", "p_choose_left", "p_choose_right", "chosen_stimulus", "chose_correctly", "responded_left", "responded_right", "outcome", "correct_responses_in_block", "reversal_criterion", ] def row(self) -> List[str]: """ Returns values for a CSV or Excel row. Order must match :meth:`csv_header`. """ sc = self.subject_config gc = sc.group_drug_config return [ gc.group_name, sc.subject_name, sc.subject_num, gc.drugname, gc.alpha_reward, gc.alpha_punishment, gc.tau_reinforcement_sensitivity, gc.tau_location_stickiness, gc.tau_stimulus_stickiness, self.part, self.trial_number, self.block, self.trial_in_block, self.left_stimulus, self.right_stimulus, self.correct_stimulus, self.incorrect_stimulus, self.v_left, self.v_right, self.p_choose_left, self.p_choose_right, self.chosen_stimulus, int(self.chose_correctly), int(self.responded_left), int(self.responded_right), self.outcome, self.correct_responses_in_block, self.reversal_criterion, ] class Writer(object): """ Class to write to one or more CSV or XLSX files. """ def __init__(self, output_filename: str, output_dir: str, excel: bool) -> None: """ Args: output_filename: use single-file mode and store all data in this file output_dir: use multi-file mode and store files (one per subject) within this directory excel: use XLSX rather than CSV format? """ assert bool(output_filename) != bool(output_dir) self.output_filename = output_filename self.output_dir = output_dir self.separate_files = bool(output_dir) self.current_filename = "" self.excel = excel self.outfile = None # type: Optional[TextIO] self.csv_writer = None # type: Optional[CSVWriterType] self.wb = None # type: Optional[XLWorkbook] self.ws = None # type: Optional[XLWorksheet] def new_subject(self, subject_name: str) -> None: """ Start a new subject. Args: subject_name: Base name of the filename (for multi-file mode). """ self._close() if self.separate_files: extension = ".xlsx" if self.excel else ".csv" filestem = subject_name + extension mkdir_p(self.output_dir) self.current_filename = os.path.join(self.output_dir, filestem) else: self.current_filename = self.output_filename if os.path.isfile(self.current_filename): raise RuntimeError( f"File {self.current_filename!r} already exists!") if self.excel: self.wb = XLWorkbook() self.wb.remove(self.wb.active) # remove autocreated sheet self.ws = self.wb.create_sheet(title="sim") else: log.info(f"Writing to {self.current_filename!r}") self.outfile = open(self.current_filename, "w") self.csv_writer = csv.writer(self.outfile) # type: CSVWriterType headings = TrialRecord.header() self.writerow(headings) def writerow(self, row: List[Any]) -> None: """ Write one row of data. """ if self.excel: self.ws.append(row) else: row_str = [str(x) for x in row] self.csv_writer.writerow(row_str) def _close(self) -> None: """ Close any open file. """ if self.wb: log.info(f"Writing to {self.current_filename!r}") f = open(self.current_filename, "wb") self.wb.save(f) f.close() self.wb = None elif self.outfile: self.outfile.close() self.outfile = None def __del__(self) -> None: """ Destructor. Ensure files are safely saved and closed. """ self._close() # ============================================================================= # Task ("world") structure # ============================================================================= class ExperimentalWorldProbabilisticReversal(object): """ Represents the task -- the experimental "world" -- instantiated for a particular session. """ def __init__(self, writer: Writer, subject_config: SubjectConfig) -> None: self.writer = writer self.subject_config = subject_config # Fixed features self.p_rewarded_if_correct = 0.85 # ... checked empirically from first 10k rows of JK data, plus Ersche # 2011 (exact figure not given, p755); 0.85 is right self.p_rewarded_if_incorrect = 0.0 # ... checked empirically from JK data, plus Ersche 2011 (no mention of # positive feedback for incorrect response); 0.0 is right self.n_stimuli = 4 self._n_blocks_per_part = 10 self._n_parts = 2 # Dynamic features self.part = 1 # "part" = original notation for "pair index" self.trial_number = 1 self.block = 1 self.trial_in_block = 1 self.reversed = False # currently reversed from starting situation? self.finished = False self.reversal_criterion = self._get_reversal_criterion() self.correct_responses_in_block = 0 # Per trial: # - Task self.left_correct = None # type: Optional[bool] self.left_stimulus = None # type: Optional[int] self.right_stimulus = None # type: Optional[int] self.correct_stimulus = None # type: Optional[int] self.incorrect_stimulus = None # type: Optional[int] # - Internals self.v_left = None # type: Optional[float] self.v_right = None # type: Optional[float] self.p_choose_left = None # type: Optional[float] self.p_choose_right = None # type: Optional[float] # - Response self.chosen_stimulus_num = None # type: Optional[int] self.chose_correctly = None # type: Optional[bool] self.responded_left = None # type: Optional[bool] self.responded_right = None # type: Optional[bool] # - Outcome self.outcome = None # type: Optional[int] # Start writing self.writer.new_subject(self.subject_config.group_and_subject_name()) @staticmethod def _get_reversal_criterion(old_criterion: int = None) -> int: """ Returns a reversal criterion -- always different from the last. """ while True: criterion = random.randint(10, 15) # inclusive if old_criterion is None or criterion != old_criterion: return criterion def stimuli_offered(self) -> Tuple[int, int]: """ Determines and returns ``(left_stim, right_stim)`` for a given trial. Uses 1-based stimulus numbering. """ first_stimulus = 2 * (self.part - 1) + 1 # 1, 3, ... second_stimulus = first_stimulus + 1 # 2, 4, ... if self.reversed: # Reversed situation, e.g. 2 correct, 1 incorrect self.correct_stimulus = second_stimulus self.incorrect_stimulus = first_stimulus else: # Starting situation, e.g. 1 correct, 2 incorrect self.correct_stimulus = first_stimulus self.incorrect_stimulus = second_stimulus # Side assignment is random self.left_correct = coin(0.5) if self.left_correct: self.left_stimulus = self.correct_stimulus self.right_stimulus = self.incorrect_stimulus else: # right correct self.left_stimulus = self.incorrect_stimulus self.right_stimulus = self.correct_stimulus return self.left_stimulus, self.right_stimulus def choose(self, chosen: int, p_chose_chosen: float, v_left: float, v_right: float) -> int: """ Records a choice. Returns the reinforcement outcome. Args: chosen: the stimulus number chosen by the subject p_chose_chosen: the probability with which this was chosen v_left: value of left stimulus v_right: value of right stimulus Returns: tuple: ``(chose_correctly, outcome)`` where ``outcome`` is 1 for reward, 0 for nonreward """ # Record self.chosen_stimulus_num = chosen self.chose_correctly = chosen == self.correct_stimulus self.responded_left = chosen == self.left_stimulus self.responded_right = chosen == self.right_stimulus if self.responded_left: self.p_choose_left = p_chose_chosen self.p_choose_right = 1 - p_chose_chosen else: # responded right self.p_choose_left = 1 - p_chose_chosen self.p_choose_right = p_chose_chosen self.v_left = v_left self.v_right = v_right # Decide outcome p_reward = ( self.p_rewarded_if_correct if self.chose_correctly else self.p_rewarded_if_incorrect ) self.outcome = int(coin(p_reward)) if self.chose_correctly: self.correct_responses_in_block += 1 return self.outcome def next_trial(self) -> None: """ Saves data from this trial and moves to the next trial. """ # Write data self._write_trial_row() # Update our state self.trial_number += 1 self.trial_in_block += 1 # Finished the block? if self.correct_responses_in_block >= self.reversal_criterion: # Reverse stimuli; start new block self.reversed = not self.reversed self.block += 1 self.trial_in_block = 1 self.correct_responses_in_block = 0 # Pick new reversal criterion self.reversal_criterion = self._get_reversal_criterion() # Finished the part? if self.block > self._n_blocks_per_part: self.block = 1 self.part += 1 # Finished the task? if self.part > self._n_parts: self.finished = True # Clear trial recording variables # - Task self.left_correct = None self.left_stimulus = None self.right_stimulus = None self.correct_stimulus = None self.incorrect_stimulus = None # - Internals self.v_left = None self.v_right = None self.p_choose_left = None self.p_choose_right = None # - Response self.chosen_stimulus_num = None self.chose_correctly = None self.responded_left = None self.responded_right = None # - Outcome self.outcome = None def _write_trial_row(self) -> None: """ Record the results. """ tr = TrialRecord( subject_config=self.subject_config, part=self.part, trial_number=self.trial_number, block=self.block, trial_in_block=self.trial_in_block, left_stimulus=self.left_stimulus, right_stimulus=self.right_stimulus, correct_stimulus=self.correct_stimulus, incorrect_stimulus=self.incorrect_stimulus, v_left=self.v_left, v_right=self.v_right, p_choose_left=self.p_choose_left, p_choose_right=self.p_choose_right, chosen_stimulus=self.chosen_stimulus_num, chose_correctly=self.chose_correctly, responded_left=self.responded_left, responded_right=self.responded_right, outcome=self.outcome, correct_responses_in_block=self.correct_responses_in_block, reversal_criterion=self.reversal_criterion, ) self.writer.writerow(tr.row()) # ============================================================================= # Core RL model # ============================================================================= LEFT = 0 RIGHT = 1 def run_subject_model_2(writer: Writer, subject_config: SubjectConfig) -> None: """ Args: writer: write data to this :class:`Writer` subject_config: :class:`SubjectConfig` for this subject's session """ log.debug(f"Running subject {subject_config.subject_name}") gc = subject_config.group_drug_config # Parameters, fixed for this session alpha_reward = gc.alpha_reward alpha_punish = gc.alpha_punishment tau_reinf = gc.tau_reinforcement_sensitivity tau_loc = gc.tau_location_stickiness tau_stim = gc.tau_stimulus_stickiness # Set up task ew = ExperimentalWorldProbabilisticReversal(writer=writer, subject_config=subject_config) # Starting state v = [0] * ew.n_stimuli # stimulus value side_stickiness = [0, 0] # left, right previously_chosen_stim_idx = None # type: Optional[int] # stimulus index chosen on previous trial # noqa # Run trials while not ew.finished: # --------------------------------------------------------------------- # Set up the trial # --------------------------------------------------------------------- t = ew.trial_number log.debug(f"Subject {subject_config.subject_name}, trial {t}") left_stimulus_num, right_stimulus_num = ew.stimuli_offered() # Zero-based indexes: left_stimulus_idx = left_stimulus_num - 1 right_stimulus_idx = right_stimulus_num - 1 # --------------------------------------------------------------------- # Implement behaviour # --------------------------------------------------------------------- # 1. Obtain values and choose left_stim_stickiness = int(previously_chosen_stim_idx == left_stimulus_idx) # noqa # ... 0 or 1 q_left = [ v[left_stimulus_idx] * tau_reinf + left_stim_stickiness * tau_stim + side_stickiness[LEFT] * tau_loc ] right_stim_stickiness = int(previously_chosen_stim_idx == right_stimulus_idx) # noqa # ... 0 or 1 q_right = [ v[right_stimulus_idx] * tau_reinf + right_stim_stickiness * tau_stim + side_stickiness[RIGHT] * tau_loc ] q = np.array([q_left, q_right]) p = softmax(q) # fixed beta = 1; the tau components are effectively beta # noqa p_choose_left = float(p[LEFT]) chose_left = coin(p_choose_left) if chose_left: chosen_stimulus_num = left_stimulus_num p_chose_chosen = p_choose_left else: # chose right chosen_stimulus_num = right_stimulus_num p_chose_chosen = 1 - p_choose_left chosen_stimulus_idx = chosen_stimulus_num - 1 # zero-based # For testing: how to do "random choice": # chosen_stimulus = left_stimulus if coin(0.5) else right_stimulus # 2. Interact with the world; obtain reinforcement outcome = ew.choose( chosen=chosen_stimulus_num, p_chose_chosen=p_chose_chosen, v_left=v[left_stimulus_idx], v_right=v[right_stimulus_idx], ) # 3. Update our internal state. # -- Value (RL) prediction_error = outcome - v[chosen_stimulus_idx] alpha = alpha_reward if prediction_error > 0 else alpha_punish v[chosen_stimulus_idx] += alpha * prediction_error # -- Side stickiness if chose_left: side_stickiness[LEFT] = 1 side_stickiness[RIGHT] = 0 else: # chose right side_stickiness[LEFT] = 0 side_stickiness[RIGHT] = 1 # -- Stimulus stickiness previously_chosen_stim_idx = chosen_stimulus_idx # --------------------------------------------------------------------- # Move to the next trial (or finish), and record results # --------------------------------------------------------------------- ew.next_trial() log.debug(f"... subject {subject_config.subject_name} took " f"{ew.trial_number - 1} trials") # ============================================================================= # Group/experiment-level functions # ============================================================================= def run_group_model_2(writer: Writer, group_drug_config: GroupDrugConfig) -> None: """ Args: writer: write data to this :class:`Writer` object group_drug_config: :class:`GroupDrugConfig` for this group/drug combination Note: ``csv.writer()`` seems to get confused by ``multiprocessing``, and spits out multiple CSV headers. Skip it; serial is fast enough. Also, parallel processing introduces variability to the random number seed system. """ gc = group_drug_config log.info(f"Running group: {gc.group_name}/{gc.drugname} " f"for {gc.n_subjects} subjects") for subjectnum in range(1, gc.n_subjects + 1): subjectname = f"s{subjectnum}" sc = SubjectConfig(subject_name=subjectname, subject_num=subjectnum, group_drug_config=gc) run_subject_model_2(writer, sc) log.info("Finished group") def run_synthetic_experiment_model_2(writer: Writer, n_subjects_per_group: int, params: PARAMS_TYPE) -> None: """ Runs an entire synthetic experiment. """ for groupname in GROUPS: for drugname in DRUGS: lookup = (groupname, drugname) gc = GroupDrugConfig( group_name=groupname, drug_name=drugname, alpha_reward=params[ALPHA_REWARD][lookup], alpha_punishment=params[ALPHA_PUNISHMENT][lookup], tau_reinforcement_sensitivity=params[TAU_REINFORCEMENT_SENSITIVITY][lookup], # noqa tau_location_stickiness=params[TAU_LOCATION_STICKINESS][lookup], # noqa tau_stimulus_stickiness=params[TAU_STIMULUS_STICKINESS][lookup], # noqa n_subjects=n_subjects_per_group, ) run_group_model_2(writer, gc) # ============================================================================= # Limited parameter recovery demonstration # ============================================================================= def gen_recovery_parameters() -> Generator[Dict[str, float], None, None]: done = set() # type: Set[HashableDict] for varying_param in PARAMETER_NAMES: varying_param_values = DEMO_PARAM_ALL[varying_param] for vpv in varying_param_values: alpha_reward = ( vpv if varying_param == ALPHA_REWARD else DEMO_PARAM_CENTRE[ALPHA_REWARD] ) alpha_punishment = ( vpv if varying_param == ALPHA_PUNISHMENT else DEMO_PARAM_CENTRE[ALPHA_PUNISHMENT] ) tau_reinf = ( vpv if varying_param == TAU_REINFORCEMENT_SENSITIVITY else DEMO_PARAM_CENTRE[TAU_REINFORCEMENT_SENSITIVITY] ) tau_loc = ( vpv if varying_param == TAU_LOCATION_STICKINESS else DEMO_PARAM_CENTRE[TAU_LOCATION_STICKINESS] ) tau_stim = ( vpv if varying_param == TAU_STIMULUS_STICKINESS else DEMO_PARAM_CENTRE[TAU_STIMULUS_STICKINESS] ) params = HashableDict({ ALPHA_REWARD: alpha_reward, ALPHA_PUNISHMENT: alpha_punishment, TAU_REINFORCEMENT_SENSITIVITY: tau_reinf, TAU_LOCATION_STICKINESS: tau_loc, TAU_STIMULUS_STICKINESS: tau_stim, }) if params in done: continue done.add(params) yield params def demo_parameter_recovery(writer: Writer, n_subjects: int) -> None: for i, params in enumerate(gen_recovery_parameters(), start=1): gc = GroupDrugConfig( group_name=f"group_{i}", drug_name="thedrug", alpha_reward=params[ALPHA_REWARD], alpha_punishment=params[ALPHA_PUNISHMENT], tau_reinforcement_sensitivity=params[TAU_REINFORCEMENT_SENSITIVITY], # noqa tau_location_stickiness=params[TAU_LOCATION_STICKINESS], tau_stimulus_stickiness=params[TAU_STIMULUS_STICKINESS], n_subjects=n_subjects, ) log.info(f"Generating data for: {params}") run_group_model_2(writer, gc) # ============================================================================= # Tests # ============================================================================= def test_gen_recovery_parameters() -> None: """ Tests the generation of parameters to be recovered. """ n = 0 for params in gen_recovery_parameters(): ar = params[ALPHA_REWARD] ap = params[ALPHA_PUNISHMENT] tr = params[TAU_REINFORCEMENT_SENSITIVITY] tl = params[TAU_LOCATION_STICKINESS] ts = params[TAU_STIMULUS_STICKINESS] print(f"alpha_rew={ar}, alpha_pun={ap}, tau_reinf={tr}, " f"tau_loc={tl}, tau_stim={ts}") n += 1 print(f"NUMBER OF PARAMETER SETS: {n}") # ============================================================================= # Main # ============================================================================= def main() -> None: """ Command-line entry point. """ # ------------------------------------------------------------------------- # Arguments # ------------------------------------------------------------------------- params_standard = "standard" params_fix_stim_stickiness = "fix_stimulus_stickiness" params_fix_loc_stickiness = "fix_location_stickiness" params_parameter_recovery = "parameter_recovery_demo" parser = argparse.ArgumentParser( """ Generate synthetic data: Kanen et al., 2019. - All "subjects" within a group have identical parameters, as the purpose of this simulation is to establish whether the detected group-level parameter changes are sufficient to explain the behaviour seen via a conventional analysis; inter-subject variability is not required for this, and we would like arbitrarily high power, so we use a large number of identical "subjects" from each group. There is, therefore, no need to represent the within-subjects structure/intersubject variability. - Examples: ./sim_from_model.py --parameters standard --output_filename synthetic_m2_standard.csv ./sim_from_model.py --parameters fix_stimulus_stickiness --output_filename synthetic_m2_single_tau_stim.csv ./sim_from_model.py --parameters fix_location_stickiness --output_filename synthetic_m2_single_tau_loc.csv ./sim_from_model.py --parameters standard --output_dir synthetic_m2_standard_output --excel ./sim_from_model.py --parameters fix_stimulus_stickiness --output_dir synthetic_m2_single_tau_stim_output --excel ./sim_from_model.py --parameters fix_location_stickiness --output_dir synthetic_m2_single_tau_loc_output --excel ./sim_from_model.py --parameters parameter_recovery_demo --output_filename synth_parameter_recovery.csv """, # noqa formatter_class=argparse.ArgumentDefaultsHelpFormatter ) parser.add_argument( "--output_filename", type=str, help="Name of output file (e.g. 'synthetic_m2.csv', " "'something.xlsx'), to use one file for all data" ) parser.add_argument( "--output_dir", type=str, help="Name of output directory, to use one file per subject" ) parser.add_argument( "--verbose", action="store_true", help="Verbose" ) parser.add_argument( "--excel", action="store_true", help="Use .XLSX rather than .CSV" ) parser.add_argument( "--n_subjects_per_group", type=int, default=100, help="Number of subjects to simulate, per group/drug combination (or," "for the parameter recovery demo, in total)" ) parser.add_argument( "--seed", type=int, default=1234, help="Random number seed" ) parser.add_argument( "--parameters", type=str, choices=[params_standard, params_fix_stim_stickiness, params_fix_loc_stickiness, params_parameter_recovery], required=True, help=( f"Simulation type. " f"{params_standard!r}: per-group parameters extracted from " f"Bayesian model. " f"{params_fix_stim_stickiness!r}: stimulus stickiness parameter " f"is fixed to its overall (not per-group/per-drug) mean." f"{params_fix_loc_stickiness!r}: location (side) stickiness " f"parameter is fixed to its overall (not per-group/per-drug) mean." ) ) args = parser.parse_args() assert bool(args.output_filename) != bool(args.output_dir), ( "Specify either --output_filename or --output_dir") # ------------------------------------------------------------------------- # Logging # ------------------------------------------------------------------------- main_only_quicksetup_rootlogger( logging.DEBUG if args.verbose else logging.INFO) # ------------------------------------------------------------------------- # Choose simulation type # ------------------------------------------------------------------------- param_recovery = False params = None # type: Optional[PARAMS_TYPE] if args.parameters == params_standard: params = MODEL_2_FITTED_PARAMETERS elif args.parameters == params_fix_loc_stickiness: params = MODEL_2_FITTED_PARAMETERS_FIX_LOC_STICKINESS elif args.parameters == params_fix_stim_stickiness: params = MODEL_2_FITTED_PARAMETERS_FIX_STIM_STICKINESS elif args.parameters == params_parameter_recovery: param_recovery = True else: raise AssertionError("bug") # ------------------------------------------------------------------------- # Random number generator # ------------------------------------------------------------------------- log.info(f"Using RNG seed: {args.seed}") random.seed(args.seed) # ------------------------------------------------------------------------- # Open file # ------------------------------------------------------------------------- writer = Writer(output_filename=args.output_filename, output_dir=args.output_dir, excel=args.excel) # ------------------------------------------------------------------------- # Run sim # ------------------------------------------------------------------------- if param_recovery: demo_parameter_recovery( writer=writer, n_subjects=args.n_subjects_per_group, ) else: log.info(f"Using parameters:\n{pprint.pformat(params)}") run_synthetic_experiment_model_2( writer=writer, n_subjects_per_group=args.n_subjects_per_group, params=params, ) # ============================================================================= # Command-line entry point # ============================================================================= if __name__ == "__main__": # test_gen_recovery_parameters() main()