model_config.py

import numpy as np


def make_cpd_arrays(config):
    n_c = config['n_c']
    hs = config['hs']
    h_dim = config['h_dim']
    tss = config['tss']
    ts_dim = config['ts_dim']
    cts = config['cts']

    p_a__c_s = np.zeros((2, n_c, 2, 2))
    for c in range(n_c):
        for s0 in range(2):
            for s1 in range(2):
                if h_dim[c] == 0:
                    p_a__c_s[0, c, s0, s1] = hs[c] if s0 == 0 else 1 - hs[c]
                else:
                    p_a__c_s[0, c, s0, s1] = hs[c] if s1 == 0 else 1 - hs[c]
    p_a__c_s[1, :, :, :] = 1 - p_a__c_s[0, :, :, :]

    p_o__c_s_a = np.zeros((2, n_c, 2, 2, 2))
    for c in range(n_c):
        for s0 in range(2):
            for s1 in range(2):
                for a in range(2):
                    if ts_dim[c] == 0:
                        p_o__c_s_a[0, c, s0, s1, a] = tss[c] if (s0 == a) else 1 - tss[c]
                    else:
                        p_o__c_s_a[0, c, s0, s1, a] = tss[c] if (s1 == a) else 1 - tss[c]
    p_o__c_s_a[1, :, :, :, :] = 1 - p_o__c_s_a[0, :, :, :, :]

    if n_c > 1:
        p_c__c = np.ones((n_c, n_c)) * ((1 - cts) / (n_c - 1))
        np.fill_diagonal(p_c__c, cts)
        p_c__c_s = np.stack([p_c__c for _ in range(2)], axis=-1)
    else:
        p_c__c_s = np.ones((1, 1, 2))

    #                      s0   0  0    0  0      1  1    1  1
    #                      s1   0  0    1  1      0  0    1  1
    #                       a   0  1    0  1      0  1    0  1
    ground_truth = np.array([[[[0, 1], [0, 1]], [[1, 0], [1, 0]]],   # t=1
                             [[[0, 1], [1, 0]], [[0, 1], [1, 0]]]])  # t=0

    return p_a__c_s, p_o__c_s_a, p_c__c_s, ground_truth

stroop_config = {
    'n_c': 2,
    'hs': [.9, .5],
    'h_dim': [1, 1],
    'tss': [.99, .99],
    'ts_dim': [0, 0],
    'cts': .7,
}