surprise-minimizing-contextualizing-control/sampler.py at main · CognitiveModeling/surprise-minimizing-contextualizing-control · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
import numpy as np
from scipy.stats import entropy
from pgmpy.inference import BeliefPropagation


class Sampler():

    def __init__(self, model, threshold=100):
        self.model = model
        self.threshold = threshold

    def query(self, stim, task, outcome, max_samples=None, return_it=False, include_instruction=False, calculate_it_metrics=True):

        prior_c_new = np.array([.95, .05])

        inference = BeliefPropagation(self.model)
        prior = inference.query(['a', 'c'], evidence={'s0': stim[0], 's1': stim[1]}).values
        prior_c = prior.sum(axis=0)

        self.counts = np.ones(prior.shape)
        n_accept_samples = 0
        n_reject_samples = 0

        trace = [self.counts / self.counts.sum()]
        sample_cost = []
        sample_meta_cost = []
        sample_control_cost = []
        sample_error = []

        # Calculate IT metrics for first iteration
        sample_cost.append(0)
        sample_meta_cost.append(0)
        sample_control_cost.append(0)
        likelihood_o = self.model.get_cpds('o').values[outcome, :, stim[0], stim[1], :].T  # transpose to [a, c]
        goal_error = -1 * (prior.flatten() * np.log(likelihood_o.flatten())).sum()
        sample_error.append(goal_error)

        while n_accept_samples < self.threshold:

            # Sampling
            c = np.random.choice(np.arange(len(self.model.states['c'])), p=prior_c)
            a = np.random.choice([0, 1], p=self.model.get_cpds('a').values[:, c, stim[0], stim[1]])
            o = np.random.choice([0, 1], p=self.model.get_cpds('o').values[:, c, stim[0], stim[1], a])

            if include_instruction:
                i = np.random.choice([0, 1], p=self.model.get_cpds('i').values[:, c])
            else:
                i = task  # making sure that condition below is always true

            # Rejection
            if o == outcome and i == task:
                self.counts[a, c] += 1
                n_accept_samples += 1
            else:
                n_reject_samples += 1

            posterior = self.counts / self.counts.sum()
            trace.append(posterior)

            # Calculate IT metrics
            if calculate_it_metrics:
                sample_cost.append(entropy(posterior.flatten(), prior.flatten()))
                sample_meta_cost.append(entropy(posterior.sum(axis=0), prior_c))
                sample_control_cost.append(sample_cost[-1] - sample_meta_cost[-1])
                likelihood_o = self.model.get_cpds('o').values[outcome, :, stim[0], stim[1], :].T  # transpose to [a, c]
                goal_error = -1 * (posterior.flatten() * np.log(likelihood_o.flatten())).sum()
                sample_error.append(goal_error)

            if max_samples and len(self.trace) >= max_samples:  # add option to stop sampling after certain number of iterations
                break

        if not return_it:
            return posterior, len(trace)
        else:
            a_rate = n_accept_samples / (n_accept_samples + n_reject_samples)
            return posterior, len(trace), np.array(trace), np.array(sample_cost), np.array(sample_error), a_rate, n_reject_samples, n_accept_samples