all strategy models drafted

Author: mdhumphries

Functions/__pycache__/plotSessionStructure.cpython-39.pyc

@@ -0,0 +1,32 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Created on Wed Aug 24 11:48:26 2022
+
+@author: Lowri Powell & Mark Humphries
+"""
+import matplotlib.pyplot as plt
+import numpy as np
+
+
+def plotSessionStructure(TestData):
+    no_Trials = np.size(TestData.TrialIndex)
+    
+    sessionLines = TestData[TestData['NewSessionTrials'] == 1].index  # indices list  when new session was started
+    plt.vlines(sessionLines, 0, 1, colors='lightgray', linestyles='--', linewidth=0.75,
+               label="New Sessions")  # vertical lines indicate the new session trials
+    ruleLines = np.array(TestData[TestData['RuleChangeTrials'] == 1].index)  # indices list  when new session was started
+    ruleLines = np.insert(ruleLines, 0, 0)  # sets array for x values of rule change
+
+    for x in ruleLines:  # to get shade change for rule change and labels
+        minx = x / no_Trials
+        plt.axhspan(1, 1.25, xmin=minx, alpha=0.3, edgecolor='k')  # change transparency for separation
+        plt.axvline(x, 0.8, 1)  # dividing lines
+
+    #  creating labels
+    plt.text(1, 1.125, "Right Arm", label='Go to the Right')
+    plt.text(120, 1.125, "Lit Arm", label='Go to the Lit Arm')
+    plt.text(225, 1.125, "Left Arm", label='Go to the Left')
+    plt.text(330, 1.125, "Unlit Arm", label='Go to the Dark Arm')
+    plt.text(150, 1.3, "Rule for Reward")
+    plt.legend()  # add legend

Functions/plotSessionStructure.py

@@ -4,12 +4,10 @@
 
 import strategymodels
 
-#from strategy_models.go_left import go_left
-#from strategy_models.go_right import go_right
-#from strategy_models.alternate import alternate
 from Functions.set_Beta_prior import set_priors
 from Functions.update_strategy_posterior_probability import update_strategy_posterior_probability
 from Functions.Summaries_of_Beta_distribution import summaries_of_Beta_Distribution
+from Functions.plotSessionStructure import plotSessionStructure
 
 # initiate TestData variable so that rat 2 testdata can be loaded
 TestData = pd.read_csv('data.csv')
@@ -23,8 +21,6 @@
 
 [alpha0, beta0] = set_priors(prior_type)  # define priors
 
-no_Trials = np.size(TestData.TrialIndex)
-
 
 #%% initialise storage
 Output_collection = {} # empty dict in which to store dataframes
@@ -41,7 +37,6 @@
 
     rows_of_data = TestData.iloc[0:trial+1]     # select all rows of data up to the curren trial; is trial+1 as dataframe includes colu   
 
-    print(rows_of_data)
    #%% 
     for index_strategy in range(len(strategies)):
        # run current strategy model on data up to current trial 
@@ -61,31 +56,18 @@
        Output_collection[strategies[index_strategy]] = pd.concat([Output_collection[strategies[index_strategy]], new_df], ignore_index=True)       # add new row to dataframe
 
 #%% plot results
+no_Trials = np.size(TestData.TrialIndex)
+
 # plotting time series of MAPprobability
 plt.figure(figsize=(10, 5))
-plt.plot(Output_collection['go_right'].MAPprobability, linewidth=0.75)  # plots the time series
+plt.plot(Output_collection['go_left'].MAPprobability, linewidth=0.75)  # plots the time series
 plt.axis([0, no_Trials, 0, 1.25])  # establishes axis limits
 plt.xlabel('Trials'), plt.ylabel('P(Strategy)')  # labelling the axis
 plt.axhline(y=0.5, color='firebrick', linestyle='--', linewidth=0.75,
             label="Chance")  # shows the line at which Chance is exceeded
-sessionLines = TestData[TestData['NewSessionTrials'] == 1].index  # indices list  when new session was started
-plt.vlines(sessionLines, 0, 1, colors='lightgray', linestyles='--', linewidth=0.75,
-           label="New Sessions")  # vertical lines indicate the new session trials
-ruleLines = np.array(TestData[TestData['RuleChangeTrials'] == 1].index)  # indices list  when new session was started
-ruleLines = np.insert(ruleLines, 0, 0)  # sets array for x values of rule change
 
-for x in ruleLines:  # to get shade change for rule change and labels
-    minx = x / no_Trials
-    plt.axhspan(1, 1.25, xmin=minx, alpha=0.3, edgecolor='k')  # change transparency for separation
-    plt.axvline(x, 0.8, 1)  # dividing lines
+plotSessionStructure(TestData)
 
-#  creating labels
-plt.text(1, 1.125, "Right Arm", label='Go to the Right')
-plt.text(120, 1.125, "Lit Arm", label='Go to the Lit Arm')
-plt.text(225, 1.125, "Left Arm", label='Go to the Left')
-plt.text(330, 1.125, "Unlit Arm", label='Go to the Dark Arm')
-plt.text(150, 1.3, "Rule for Reward")
-plt.legend()  # add legend
 plt.show()