changed how mutual information was estimated

Author: sjgershm

analyze_collins.m

@@ -6,7 +6,7 @@
         data = load_data('collins18');
     end
 
-    beta = linspace(0.1,15,30);
+    beta = linspace(0.1,15,50);
 
     for s = 1:length(data)
         B = unique(data(s).learningblock);
@@ -15,19 +15,17 @@
         V_data =zeros(length(B),1);
         for b = 1:length(B)
             ix = data(s).learningblock==B(b) & data(s).phase==0;
-            stim = data(s).stim(ix)';
+            state = data(s).state(ix);
             c = data(s).corchoice(ix);
-            choice = data(s).choice(ix)';
-            lowerx = 1; upperx = max(stim); lowery = 1; uppery = max(stim);
-            descriptor = [lowerx,upperx,upperx-lowerx;lowery,uppery,uppery-lowery];
-            R_data(b) = information(stim,choice,descriptor);
-            V_data(b) = mean(data(s).cor(ix));
+            action = data(s).action(ix);
+            R_data(b) = mutual_information(state,action,0.7);
+            V_data(b) = mean(data(s).reward(ix));
 
-            S = unique(stim);
+            S = unique(state);
             Q = zeros(length(S),3);
             Ps = zeros(1,length(S));
             for i = 1:length(S)
-                ii = stim==S(i);
+                ii = state==S(i);
                 Ps(i) = mean(ii);
                 a = c(ii); a = a(1);
                 Q(i,a) = 1;
@@ -41,27 +39,27 @@
                 cond(b) = 2;
             end
 
-            ix = data(s).learningblock==B(b) & data(s).phase==1;
-            stim = data(s).stim(ix)';
-            choice = data(s).choice(ix)';
-            try
-                R_test(b) = information(stim,choice);
-                V_test(b) = mean(data(s).cor(ix));
-            catch
-                R_test(b) = nan;
-                V_test(b) = nan;
-            end
         end
 
         for c = 1:2
             results.R(s,:,c) = nanmean(R(cond==c,:));
             results.V(s,:,c) = nanmean(V(cond==c,:));
             results.R_data(s,c) = nanmean(R_data(cond==c));
             results.V_data(s,c) = nanmean(V_data(cond==c));
-            results.V_test(s,c) = nanmean(V_test(cond==c));
-            results.R_test(s,c) = nanmean(R_test(cond==c));
         end
 
         clear R V
 
-    end
+    end
+    
+    p = signrank(results.R_data(:,1),results.R_data(:,2))
+    
+    R = squeeze(nanmean(results.R));
+    V = squeeze(nanmean(results.V));
+    for c = 1:2
+        Vd2(:,c) =  interp1(R(:,c),V(:,c),results.R_data(:,c));
+        results.bias(:,c) = results.V_data(:,c) - Vd2(:,c);
+    end
+    
+    [r,p] = corr([results.V_data(:,1); results.V_data(:,2)],[Vd2(:,1); Vd2(:,2)])
+    [r,p] = corr([results.R_data(:,1); results.R_data(:,2)],[results.bias(:,1); results.bias(:,2)])

analyze_steyvers.m

@@ -7,12 +7,10 @@
     end
 
     beta = linspace(1.5,5,30);
-    
-    lowerx = 1; upperx = 32; lowery = 1; uppery = 4;
-    descriptor = [lowerx,upperx,upperx-lowerx;lowery,uppery,uppery-lowery];
-    
+        
     for s = 1:length(data)
-        results.R_data(s,1) = information(data(s).state',data(s).action',descriptor);
+        results.R_data(s,1) = mutual_information(data(s).state,data(s).action);
+        if isnan(results.R_data(s)); keyboard; end
         results.V_data(s,1) = mean(data(s).reward);
         for state = 1:32
             Ps(s,state) = mean(data(s).state==state);
@@ -31,4 +29,7 @@
         Q(i,a) = 1;
     end
     [results.R,results.V] = blahut_arimoto(Ps,Q,beta);
-    results.Q = Q; results.Ps = Ps;
+    results.Q = Q; results.Ps = Ps;
+    
+    Vd = interp1(results.R,results.V,results.R_data,'cubic');   % for some reason linear interpoloation doesn't work here
+    [r,p] = corr(Vd,results.V_data)

blahut_arimoto.m

@@ -20,11 +20,11 @@
     nIter = 50;
     if nargin < 3; b = linspace(0.1,15,30); end
     R = zeros(length(b),1); V = zeros(length(b),1); Pa = zeros(length(b),A);
-    q = ones(1,A)./A;
 
     for j = 1:length(b)
         F = b(j).*Q;
         v0 = mean(Q(:));
+        q = ones(1,A)./A;
         for i = 1:nIter
             logP = log(q) + F;
             Z = logsumexp(logP,2);

fit_models_collins.m

@@ -1,8 +1,30 @@
-function [results, bms_results] = fit_models_collins(data)
+function [results, bms_results] = fit_models_collins(data,models,results)
 
     % Fit Collins (2018) data. Requires mfit package.
 
-    for m = 1:2
+    if nargin < 2; models = 1:2; end
+    
+    for s = 1:length(data)
+        trials = find(data(s).phase==0);
+        data(s).N = length(trials);
+        data(s).C = 3;
+        A = zeros(data(s).N,3);
+        for t = 1:data(s).N
+            if data(s).action(t)>0
+                A(t,data(s).action(t)) = 1;
+            end
+        end
+        
+        for b = 1:max(data(s).learningblock)
+            ix = data(s).learningblock==b;
+            for i = 1:size(A,2)
+                A(ix,i) = eps + smooth(A(ix,i));
+            end
+        end
+        data(s).logPa = log(A./sum(A,2));
+    end
+    
+    for m = models
         disp(['... fitting model ',num2str(m)]);
 
         switch m
@@ -11,13 +33,15 @@
 
                 param(1) = struct('name','b1','logpdf',@(x) 0);
                 param(2) = struct('name','b2','logpdf',@(x) 0);
-                param(3) = struct('name','sticky','logpdf',@(x) 0);
+                param(3) = struct('name','lr','logpdf',@(x) 0);
+                param(4) = struct('name','tau','logpdf',@(x) 0);
                 fun = @lik_collins;
 
             case 2
 
                 param(1) = struct('name','b1','logpdf',@(x) 0);
                 param(2) = struct('name','b2','logpdf',@(x) 0);
+                param(3) = struct('name','lr','logpdf',@(x) 0);
                 fun = @lik_collins;
 
         end
@@ -36,24 +60,33 @@
 function lik = lik_collins(x,data)
 
     B = x(1:2);
-    if length(x) > 2
-        sticky = x(3);
+    lr = 1./(1+exp(-x(3)));
+    if length(x) > 3
+        tau = x(end);
     else
-        sticky = 1;
+        tau = 1;
     end
 
     lik = 0;
 
-    for t = 1:size(data.Q,1)
-        a = data.choice(t);
+    for t = 1:data.N
+        
+        if t==1 || data.learningblock(t)~=data.learningblock(t-1)
+           Q = zeros(data.ns(t),3);
+        end
+        
+        a = data.action(t);
+        s = data.state(t);
+        r = data.reward(t);
         if a > 0
             if data.ns(t)==3
                 b = B(1);
             else
                 b = B(2);
             end
-            d = b*data.Q(t,:) + sticky*data.logPa(t,:);
+            d = b*Q(data.state(t),:) + tau*data.logPa(t,:);
             lik = lik + d(a) - logsumexp(d,2);
+            Q(s,a) = Q(s,a) + lr*(r-Q(s,a));
         end
     end
 end