text: # AUTOGENERATED! DO NOT EDIT! File to edit: 02_Data_Preparation.ipynb (unless otherwise specified).
  
  __all__ = ['prepare_murddata', 'prepare_repdata', 'simulate_data']
  
  # Cell
  # hide
  import scipy.io as sio
  import numpy as np
  import pandas as pd
  from psifr import fr
  
  def prepare_murddata(path, dataset_index):
      """
      Prepares data formatted like `data/MurdData_clean.mat` for fitting.
  
      Loads data from `path` with same format as `data/MurdData_clean.mat` and
      returns a selected dataset as an array of unique recall trials and a
      dataframe of unique study and recall events organized according to `psifr`
      specifications.
  
      **Arguments**:
      - path: source of data file
      - dataset_index: index of the dataset to be extracted from the file
  
      **Returns**:
      - trials: int64-array where rows identify a unique trial of responses and
          columns corresponds to a unique recall index.
      - merged: as a long format table where each row describes one study or
          recall event.
      - list_length: length of lists studied in the considered dataset
      """
      # load all the data
      matfile = sio.loadmat(path, squeeze_me=True)
      murd_data = [matfile['data'].item()[0][i].item() for i in range(3)]
  
      # encode dataset into psifr format
      trials, list_length, subjects = murd_data[dataset_index][:3]
      trials = trials.astype('int64')
  
      data = []
      for trial_index, trial in enumerate(trials):
  
          # every time the subject changes, reset list_index
          if not data or data[-1][0] != subjects[trial_index]:
              list_index = 0
          list_index += 1
  
          # add study events
          for i in range(list_length):
              data += [[subjects[trial_index],
                        list_index, 'study', i+1, i+1]]
  
          # add recall events
          for recall_index, recall_event in enumerate(trial):
              if recall_event != 0:
                  data += [[subjects[trial_index], list_index,
                            'recall', recall_index+1, recall_event]]
  
      data = pd.DataFrame(data, columns=[
          'subject', 'list', 'trial_type', 'position', 'item'])
      merged = fr.merge_free_recall(data)
      return trials, merged, list_length
  
  # Cell
  
  import scipy.io as sio
  import numpy as np
  import pandas as pd
  from psifr import fr
  
  def prepare_repdata(path):
      """
      Prepares data formatted like `data/repFR.mat` for fitting.
      """
  
      # load all the data
      matfile = sio.loadmat(path, squeeze_me=True)['data'].item()
      subjects = matfile[0]
      pres_itemnos = matfile[4]
      recalls = matfile[6]
      list_types = matfile[7]
      list_length = matfile[12]
  
      # convert pres_itemnos into rows of unique indices for easier model encoding
      presentations = []
      for i in range(len(pres_itemnos)):
          seen = []
          presentations.append([])
          for p in pres_itemnos[i]:
              if p not in seen:
                  seen.append(p)
              presentations[-1].append(seen.index(p))
      presentations = np.array(presentations)
  
      # discard intrusions from recalls
      trials = []
      for i in range(len(recalls)):
          trials.append([])
  
          trial = list(recalls[i])
          for t in trial:
              if (t > 0) and (t not in trials[-1]):
                  trials[-1].append(t)
  
          while len(trials[-1]) < list_length:
              trials[-1].append(0)
  
      trials = np.array(trials)
  
      # encode dataset into psifr format
      data = []
      for trial_index, trial in enumerate(trials):
          presentation = presentations[trial_index]
  
          # every time the subject changes, reset list_index
          if not data or data[-1][0] != subjects[trial_index]:
              list_index = 0
          list_index += 1
  
          # add study events
          for presentation_index, presentation_event in enumerate(presentation):
              data += [[subjects[trial_index],
                        list_index, 'study', presentation_index+1, presentation_event,  list_types[trial_index]
                       ]]
  
          # add recall events
          for recall_index, recall_event in enumerate(trial):
              if recall_event != 0:
                  data += [[subjects[trial_index], list_index,
                            'recall', recall_index+1, presentation[recall_event-1], list_types[trial_index]
                           ]]
  
      data = pd.DataFrame(data, columns=[
          'subject', 'list', 'trial_type', 'position', 'item', 'condition'])
      merged = fr.merge_free_recall(data, list_keys=['condition'])
  
      return trials, merged, list_length, presentations, list_types, data, subjects
  
  # Cell
  
  def simulate_data(model, experiment_count, first_recall_item=None):
      """
      Initialize a model with specified parameters and experience sequences and
      then populate a psifr-formatted dataframe with the outcomes of performing `free recall`.
  
      **Required model attributes**:
      - item_count: specifies number of items encoded into memory
      - context: vector representing an internal contextual state
      - experience: adding a new trace to the memory model
      - free_recall: function that freely recalls a given number of items or until recall stops
      """
  
      # encode items
      try:
          model.experience(np.eye(model.item_count, model.item_count + 1, 1))
      except ValueError:
          # so we can apply to CMR
          model.experience(np.eye(model.item_count, model.item_count))
  
      # simulate retrieval for the specified number of times, tracking results in df
      data = []
      for experiment in range(experiment_count):
          data += [[experiment, 0, 'study', i + 1, i] for i in range(model.item_count)]
      for experiment in range(experiment_count):
          if first_recall_item is not None:
              model.force_recall(first_recall_item)
          data += [[experiment, 0, 'recall', i + 1, o] for i, o in enumerate(model.free_recall())]
      data = pd.DataFrame(data, columns=['subject', 'list', 'trial_type', 'position', 'item'])
      merged = fr.merge_free_recall(data)
  
      return merged