def test_update_pA_single_factor_one_modality(self):
"""
Test for updating prior Dirichlet parameters over sensory likelihood (pA)
in the case that ONE observation modalities is updated and the generative model
has a single hidden state factor
"""
n_states = [3]
qs = Categorical(values=construct_init_qs(n_states))
l_rate = 1.0
# multiple observation modalities
num_obs = [3, 4]
modality_to_update = [np.random.randint(len(num_obs))]
A = Categorical(values=construct_generic_A(num_obs, n_states))
pA = Dirichlet(values=construct_pA(num_obs, n_states))
observation = A.dot(qs, return_numpy=False).sample()
pA_updated = learning.update_likelihood_dirichlet(
pA,
A,
observation,
qs,
lr=l_rate,
modalities=modality_to_update,
return_numpy=True)
for modality, no in enumerate(num_obs):
if modality in modality_to_update:
update = maths.spm_cross(
np.eye(no)[observation[modality]], qs.values)
validation_pA = pA[modality] + l_rate * update
else:
validation_pA = pA[modality]
self.assertTrue(
np.all(pA_updated[modality] == validation_pA.values))
def test_cross_function_d(self):
""" Test case d: outer-producting a vector and a sequence of vectors:
Options:
- First vector is a Categorical, second sequence of vectors is a numpy ndarray
(dtype = object)
- First vector is a Categorical, second sequence of vectors is a Categorical
(where self.IS_AOA = True))
"""
array_path = os.path.join(os.getcwd(), DATA_PATH + "cross_d.mat")
mat_contents = loadmat(file_name=array_path)
result_1 = mat_contents["result1"]
random_vec = Categorical(values=mat_contents["random_vec"])
states = mat_contents["s"]
for i in range(len(states)):
states[i] = states[i].squeeze()
# First way, where first array is a Categorical, second array is a numpy ndarray
# (dtype = object)
result_1a_py = random_vec.cross(states, return_numpy=True)
self.assertTrue(np.isclose(result_1, result_1a_py).all())
# Second way, where first array is a Categorical, second array is a Categorical
# (where self.IS_AOA = True)
states = Categorical(values=states[0])
result_1b_py = random_vec.cross(states, return_numpy=True)
self.assertTrue(np.isclose(result_1, result_1b_py).all())
def test_update_pB_single_dactor_with_actions(self):
"""
Test for updating prior Dirichlet parameters over transition likelihood (pB)
in the case that the one and only hidden state factor is updated, and there
are actions.
"""
n_states = [3]
n_control = [3]
qs_prev = Categorical(values=construct_init_qs(n_states))
qs = Categorical(values=construct_init_qs(n_states))
l_rate = 1.0
B = Categorical(values=construct_generic_B(n_states, n_control))
pB = Dirichlet(values=np.ones_like(B.values))
action = np.array([np.random.randint(nc) for nc in n_control])
pB_updated = core.update_transition_dirichlet(pB,
B,
action,
qs,
qs_prev,
lr=l_rate,
factors="all",
return_numpy=True)
validation_pB = pB.copy()
validation_pB[:, :, action[0]] += (
l_rate * core.spm_cross(qs.values, qs_prev.values) *
(B[:, :, action[0]].values > 0))
self.assertTrue(np.all(pB_updated == validation_pB.values))
def test_update_pA_multi_factor_some_modalities(self):
"""
Test for updating prior Dirichlet parameters over sensory likelihood (pA)
in the case that SOME observation modalities are updated and the generative model
has multiple hidden state factors
"""
n_states = [2, 6]
qs = Categorical(values=construct_init_qs(n_states))
l_rate = 1.0
# multiple observation modalities
num_obs = [3, 4, 5]
modalities_to_update = [0, 2]
A = Categorical(values=construct_generic_A(num_obs, n_states))
pA = Dirichlet(values=construct_pA(num_obs, n_states))
observation = A.dot(qs, return_numpy=False).sample()
pA_updated = core.update_likelihood_dirichlet(
pA,
A,
observation,
qs,
lr=l_rate,
modalities=modalities_to_update,
return_numpy=True)
for modality, no in enumerate(num_obs):
if modality in modalities_to_update:
update = core.spm_cross(
np.eye(no)[observation[modality]], qs.values)
validation_pA = pA[modality] + l_rate * update
else:
validation_pA = pA[modality]
self.assertTrue(
np.all(pA_updated[modality] == validation_pA.values))
def test_dot_function_b(self):
""" Continuous states and outcomes """
array_path = os.path.join(os.getcwd(), DATA_PATH + "dot_b.mat")
mat_contents = loadmat(file_name=array_path)
A = mat_contents["A"]
obs = mat_contents["o"]
states = mat_contents["s"]
states = np.array(states, dtype=object)
result_1 = mat_contents["result1"]
result_2 = mat_contents["result2"]
result_3 = mat_contents["result3"]
A = Categorical(values=A)
result_1_py = A.dot(obs, obs_mode=True, return_numpy=True)
self.assertTrue(np.isclose(result_1, result_1_py).all())
result_2_py = A.dot(states, return_numpy=True)
result_2_py = result_2_py.astype("float64")[:,
np.newaxis] # type: ignore
self.assertTrue(np.isclose(result_2, result_2_py).all())
result_3_py = A.dot(states, dims_to_omit=[0], return_numpy=True)
self.assertTrue(np.isclose(result_3, result_3_py).all())
def test_cross_function_b(self):
"""Test case b: outer-producting two vectors together:
Options:
- both vectors are stored in single Categorical (with two entries, where self.AoA == True)
- first vector is a Categorical (self.AoA = False) and second array is a numpy ndarray
(non-object array)
- first vector is a Categorical, second vector is also Categorical
"""
array_path = os.path.join(os.getcwd(), DATA_PATH + "cross_b.mat")
mat_contents = loadmat(file_name=array_path)
result_1 = mat_contents["result1"]
result_2 = mat_contents["result2"]
# first way, where both arrays as stored as two entries in a single AoA Categorical
states = Categorical(values=mat_contents["s"][0])
result_1_py = states.cross(return_numpy=True)
self.assertTrue(np.isclose(result_1, result_1_py).all())
# second way (type 1), where first array is a Categorical, second array is a
# straight numpy array
states_first_factor = Categorical(values=mat_contents["s"][0][0])
states_second_factor = mat_contents["s"][0][1]
result_2a_py = states_first_factor.cross(states_second_factor,
return_numpy=True)
self.assertTrue(np.isclose(result_2, result_2a_py).all())
# second way (type 2), where first array is a Categorical, second array
# is another Categorical
states_first_factor = Categorical(values=mat_contents["s"][0][0])
states_second_factor = Categorical(values=mat_contents["s"][0][1])
result_2b_py = states_first_factor.cross(states_second_factor,
return_numpy=True)
self.assertTrue(np.isclose(result_2, result_2b_py).all())
def test_normalize_multi_factor(self):
values_1 = np.random.rand(5)
values_2 = np.random.rand(4, 3)
values = np.array([values_1, values_2])
c = Categorical(values=values)
c.normalize()
self.assertTrue(c.is_normalized())
def test_pB_info_gain(self):
"""
Test the pB_info_gain function. Demonstrates operation
by manipulating shape of the Dirichlet priors over likelihood parameters
(pB), which affects information gain for different states
"""
n_states = [2]
n_control = [2]
qs = Categorical(values=np.eye(n_states[0])[0])
B = Categorical(values=construct_generic_B(n_states, n_control))
pB_matrix = construct_pB(n_states, n_control)
# create prior over dirichlets such that there is a skew
# in the parameters about the likelihood mapping from the
# hidden states to hidden states under the second action,
# such that hidden state 0 is considered to be more likely than the other,
# given the action in question
# Therefore taking that action would yield an expected state that afford
# high information gain about that part of the likelihood distribution.
#
pB_matrix[0, :, 1] = 2.0
pB = Dirichlet(values=pB_matrix)
# single timestep
n_step = 1
policies = core.construct_policies(n_states,
n_control,
policy_len=n_step)
pB_info_gains = np.zeros(len(policies))
for idx, policy in enumerate(policies):
qs_pi = core.get_expected_states(qs, B, policy)
pB_info_gains[idx] += core.calc_pB_info_gain(pB, qs_pi, qs, policy)
self.assertGreater(pB_info_gains[1], pB_info_gains[0])
def test_dot_function_e(self):
""" Continuous states and outcomes, but add a final (fourth) hidden state factor """
array_path = os.path.join(os.getcwd(), DATA_PATH + "dot_e.mat")
mat_contents = loadmat(file_name=array_path)
A = mat_contents["A"]
obs = mat_contents["o"]
states = mat_contents["s"]
states_array_version = np.empty(states.shape[1], dtype=object)
for i in range(states.shape[1]):
states_array_version[i] = states[0][i][0]
result_1 = mat_contents["result1"]
result_2 = mat_contents["result2"]
result_3 = mat_contents["result3"]
A = Categorical(values=A)
result_1_py = A.dot(obs, obs_mode=True, return_numpy=True)
self.assertTrue(np.isclose(result_1, result_1_py).all())
result_2_py = A.dot(states_array_version, return_numpy=True)
result_2_py = result_2_py.astype("float64")[:,
np.newaxis] # type: ignore
self.assertTrue(np.isclose(result_2, result_2_py).all())
result_3_py = A.dot(states_array_version,
dims_to_omit=[0],
return_numpy=True)
self.assertTrue(np.isclose(result_3, result_3_py).all())
def test_normalize_multi_factor(self):
values_1 = np.random.rand(5)
values_2 = np.random.rand(4, 3)
values = np.array([values_1, values_2], dtype=object)
d = Dirichlet(values=values)
normed = Categorical(values=d.mean(return_numpy=True))
self.assertTrue(normed.is_normalized())
def test_copy(self):
values = np.random.rand(3, 2)
c = Categorical(values=values)
c_copy = c.copy()
self.assertTrue(np.array_equal(c_copy.values, c.values))
c_copy.values = c_copy.values * 2
self.assertFalse(np.array_equal(c_copy.values, c.values))
def infer_action(qs, A, B, C, n_control, policies):
n_policies = len(policies)
# negative expected free energy
neg_G = np.zeros([n_policies, 1])
for i, policy in enumerate(policies):
neg_G[i] = evaluate_policy(policy, qs, A, B, C)
# get distribution over policies
q_pi = core.softmax(neg_G)
# probabilites of control states
qu = Categorical(dims=n_control)
# sum probabilites of controls
for i, policy in enumerate(policies):
# control state specified by policy
u = int(policy[0, :])
# add probability of policy
qu[u] += q_pi[i]
# normalize
qu.normalize()
# sample control
u = qu.sample()
return u
def reset(self, state=None):
if state is None:
loc_state = np.zeros(self.n_locations)
loc_state[0] = 1.0
scene_state = np.zeros(self.n_scenes)
self._true_scene = np.random.randint(self.n_scenes)
scene_state[self._true_scene] = 1.0
full_state = np.empty(self.n_factors, dtype=object)
full_state[LOCATION_ID] = loc_state
full_state[SCENE_ID] = scene_state
self._state = Categorical(values=full_state)
else:
self._state = Categorical(values=state)
return self._get_observation()
def reset(self, init_qs=None):
self.curr_timestep = 1
if init_qs is None:
if self.inference_algo == 'VANILLA':
self.qs = utils.obj_array_uniform(self.n_states)
else: # in the case you're doing MMP (i.e. you have an inference_horizon > 1), we have to account for policy- and timestep-conditioned posterior beliefs
self.qs = utils.obj_array(len(self.policies))
for p_i, _ in enumerate(self.policies):
self.qs[p_i] = utils.obj_array(
self.inference_horizon + self.policy_len +
1) # + 1 to include belief about current timestep
# self.qs[p_i][0] = copy.deepcopy(self.D) # initialize the very first belief of the inference_horizon as the prior over initial hidden states
self.qs[p_i][0] = utils.obj_array_uniform(self.n_states)
first_belief = utils.obj_array(len(self.policies))
for p_i, _ in enumerate(self.policies):
first_belief[p_i] = copy.deepcopy(self.D)
if self.edge_handling_params['policy_sep_prior']:
self.set_latest_beliefs(last_belief=first_belief)
else:
self.set_latest_beliefs(last_belief=self.D)
else:
if isinstance(init_qs, Categorical):
self.qs = init_qs
else:
self.qs = Categorical(values=init_qs)
return self.qs
def sample_action(q_pi, policies, n_control, sampling_type="marginal_action"):
"""
Samples action from posterior over policies, using one of two methods.
Parameters
----------
q_pi [1D numpy.ndarray or Categorical]:
Posterior beliefs about (possibly multi-step) policies.
policies [list of numpy ndarrays]:
List of arrays that indicate the policies under consideration. Each element
within the list is a matrix that stores the
the indices of the actions upon the separate hidden state factors, at
each timestep (n_step x n_control_factor)
n_control [list of integers]:
List of the dimensionalities of the different (controllable)) hidden state factors
sampling_type [string, 'marginal_action' or 'posterior_sample']:
Indicates whether the sampled action for a given hidden state factor is given by
the evidence for that action, marginalized across different policies ('marginal_action')
or simply the action entailed by a sample from the posterior over policies
Returns
----------
selectedPolicy [1D numpy ndarray]:
Numpy array containing the indices of the actions along each control factor
"""
n_factors = len(n_control)
if sampling_type == "marginal_action":
if utils.is_distribution(q_pi):
q_pi = utils.to_numpy(q_pi)
action_marginals = np.empty(n_factors, dtype=object)
for c_idx in range(n_factors):
action_marginals[c_idx] = np.zeros(n_control[c_idx])
# weight each action according to its integrated posterior probability over policies and timesteps
for pol_idx, policy in enumerate(policies):
for t in range(policy.shape[0]):
for factor_i, action_i in enumerate(policy[t, :]):
action_marginals[factor_i][action_i] += q_pi[pol_idx]
action_marginals = Categorical(values=action_marginals)
action_marginals.normalize()
selected_policy = np.array(action_marginals.sample())
elif sampling_type == "posterior_sample":
if utils.is_distribution(q_pi):
policy_index = q_pi.sample()
selected_policy = policies[policy_index]
else:
q_pi = Categorical(values=q_pi)
policy_index = q_pi.sample()
selected_policy = policies[policy_index]
else:
raise ValueError(f"{sampling_type} not supported")
return selected_policy
def test_multi_factor_init_values_expand(self):
values_1 = np.random.rand(5)
values_2 = np.random.rand(4)
values = np.array([values_1, values_2])
c = Categorical(values=values)
self.assertEqual(c.shape, (2, ))
self.assertEqual(c[0].shape, (5, 1))
self.assertEqual(c[1].shape, (4, 1))
def test_multi_factor_init_values(self):
values_1 = np.random.rand(5, 4)
values_2 = np.random.rand(4, 3)
values = np.array([values_1, values_2], dtype=object)
c = Categorical(values=values)
self.assertEqual(c.shape, (2, ))
self.assertEqual(c[0].shape, (5, 4))
self.assertEqual(c[1].shape, (4, 3))
def get_uniform_posterior(self):
values = np.array(
[
np.ones(self.n_states[f]) / self.n_states[f]
for f in range(self.n_factors)
]
)
return Categorical(values=values)
def test_dot_function_c_cat(self):
""" Test with vectors and matrices, discrete state / outcomes but with a
third hidden state factor. Now, when arguments themselves are
instances of Categorical
"""
array_path = os.path.join(os.getcwd(), DATA_PATH + "dot_c.mat")
mat_contents = loadmat(file_name=array_path)
A = mat_contents["A"]
obs = Categorical(values=mat_contents["o"])
states = mat_contents["s"]
states_array_version = np.empty(states.shape[1], dtype=object)
for i in range(states.shape[1]):
states_array_version[i] = states[0][i][0]
states_array_version = Categorical(values=states_array_version)
result_1 = mat_contents["result1"]
result_2 = mat_contents["result2"]
result_3 = mat_contents["result3"]
A = Categorical(values=A)
result_1_py = A.dot(obs, obs_mode=True, return_numpy=True)
self.assertTrue(np.isclose(result_1, result_1_py).all())
result_2_py = A.dot(states_array_version, return_numpy=True)
result_2_py = result_2_py.astype("float64")[:,
np.newaxis] # type: ignore
self.assertTrue(np.isclose(result_2, result_2_py).all())
result_3_py = A.dot(states_array_version,
dims_to_omit=[0],
return_numpy=True)
self.assertTrue(np.isclose(result_3, result_3_py).all())
def test_dot_function_a_cat(self):
""" Test with vectors and matrices, discrete state / outcomes
Now, when arguments themselves are instances of Categorical
"""
array_path = os.path.join(os.getcwd(), DATA_PATH + "dot_a.mat")
mat_contents = loadmat(file_name=array_path)
A = mat_contents["A"]
obs = Categorical(values=mat_contents["o"])
states = Categorical(values=mat_contents["s"][0])
result_1 = mat_contents["result1"]
result_2 = mat_contents["result2"]
result_3 = mat_contents["result3"]
A = Categorical(values=A)
result_1_py = A.dot(obs, obs_mode=True, return_numpy=True)
self.assertTrue(np.isclose(result_1, result_1_py).all())
result_2_py = A.dot(states, return_numpy=True)
result_2_py = result_2_py.astype("float64")[:,
np.newaxis] # type: ignore
self.assertTrue(np.isclose(result_2, result_2_py).all())
result_3_py = A.dot(states, dims_to_omit=[0], return_numpy=True)
self.assertTrue(np.isclose(result_3, result_3_py).all())
def _construct_transition_dist(self):
B_locs = np.eye(self.n_locations)
B_locs = B_locs.reshape(self.n_locations, self.n_locations, 1)
B_locs = np.tile(B_locs, (1, 1, self.n_locations))
B_locs = B_locs.transpose(1, 2, 0)
B = np.empty(self.n_factors, dtype=object)
B[LOCATION_ID] = B_locs
B[SCENE_ID] = np.eye(self.n_scenes).reshape(self.n_scenes, self.n_scenes, 1)
return Categorical(values=B)
def reset(self, init_qs=None):
if init_qs is None:
self.qs = self._construct_D_prior()
else:
if isinstance(init_qs, Categorical):
self.qs = init_qs
else:
self.qs = Categorical(values=init_qs)
return self.qs
def test_pA_info_gain(self):
"""
Test the pA_info_gain function. Demonstrates operation
by manipulating shape of the Dirichlet priors over likelihood parameters
(pA), which affects information gain for different expected observations
"""
n_states = [2]
n_control = [2]
qs = Categorical(values=np.eye(n_states[0])[0])
B = Categorical(values=construct_generic_B(n_states, n_control))
# single timestep
n_step = 1
policies = core.construct_policies(n_states,
n_control,
policy_len=n_step)
# single observation modality
num_obs = [2]
# create noiseless identity A matrix
A = Categorical(values=np.eye(num_obs[0]))
# create prior over dirichlets such that there is a skew
# in the parameters about the likelihood mapping from the
# second hidden state (index 1) to observations, such that one
# observation is considered to be more likely than the other conditioned on that state.
# Therefore sampling that observation would afford high info gain
# about parameters for that part of the likelhood distribution.
pA_matrix = construct_pA(num_obs, n_states)
pA_matrix[0, 1] = 2.0
pA = Dirichlet(values=pA_matrix)
pA_info_gains = np.zeros(len(policies))
for idx, policy in enumerate(policies):
qs_pi = core.get_expected_states(qs, B, policy)
qo_pi = core.get_expected_obs(qs_pi, A)
pA_info_gains[idx] += core.calc_pA_info_gain(pA, qo_pi, qs_pi)
self.assertGreater(pA_info_gains[1], pA_info_gains[0])
def step(self, actions):
prob_states = np.empty(self.n_factors, dtype=object)
for f in range(self.n_factors):
prob_states[f] = (
self._transition_dist[f][:, :, actions[f]]
.dot(self._state[f], return_numpy=True)
.flatten()
)
state = Categorical(values=prob_states).sample()
self._state = self._construct_state(state)
return self._get_observation()
def _construct_B_distribution(self):
if self.n_factors == 1:
B = np.eye(*self.n_states)[:, :, np.newaxis]
if 0 in self.control_fac_idx:
B = np.tile(B, (1, 1, self.n_controls[0]))
B = B.transpose(1, 2, 0)
else:
B = np.empty(self.n_factors, dtype=object)
for factor, ns in enumerate(self.n_states):
B_basic = np.eye(ns)[:, :, np.newaxis]
if factor in self.control_fac_idx:
B[factor] = np.tile(B_basic,
(1, 1, self.n_controls[factor]))
B[factor] = B[factor].transpose(1, 2, 0)
else:
B[factor] = B_basic
B = Categorical(values=B)
B.normalize()
return B
def test_cross_function_c(self):
"""Test case c: outer-producting a vector and a matrix together:
Options:
- First vector is a Categorical, and the matrix argument is a numpy ndarray
(non-object array)
- First vector is a Categorical, and the matrix argument is also a Categorical
"""
array_path = os.path.join(os.getcwd(), DATA_PATH + "cross_c.mat")
mat_contents = loadmat(file_name=array_path)
result_1 = mat_contents["result1"]
random_vec = Categorical(values=mat_contents["random_vec"])
# first way, where first array is a Categorical, second array is a numpy ndarray
random_matrix = mat_contents["random_matrix"]
result_1a_py = random_vec.cross(random_matrix, return_numpy=True)
self.assertTrue(np.isclose(result_1, result_1a_py).all())
# second way, where first array is a Categorical, second array is a Categorical
random_matrix = Categorical(values=mat_contents["random_matrix"])
result_1b_py = random_vec.cross(random_matrix, return_numpy=True)
self.assertTrue(np.isclose(result_1, result_1b_py).all())
def test_update_pB_multi_factor_with_actions_all_factors(self):
"""
Test for updating prior Dirichlet parameters over transition likelihood (pB)
in the case that there are mulitple hidden state factors, and there
are actions. All factors are updated
"""
n_states = [3, 4, 5]
n_control = [3, 4, 5]
qs_prev = Categorical(values=construct_init_qs(n_states))
qs = Categorical(values=construct_init_qs(n_states))
l_rate = 1.0
B = Categorical(values=construct_generic_B(n_states, n_control))
B.normalize()
pB = Dirichlet(values=construct_pB(n_states, n_control))
action = np.array([np.random.randint(nc) for nc in n_control])
pB_updated = core.update_transition_dirichlet(pB,
B,
action,
qs,
qs_prev,
lr=l_rate,
factors="all",
return_numpy=True)
validation_pB = pB.copy()
for factor, _ in enumerate(n_control):
validation_pB = pB[factor].copy()
validation_pB[:, :, action[factor]] += (
l_rate *
core.spm_cross(qs[factor].values, qs_prev[factor].values) *
(B[factor][:, :, action[factor]].values > 0))
self.assertTrue(np.all(pB_updated[factor] == validation_pB.values))
def test_update_pB_single_factor_no_actions(self):
"""
Test for updating prior Dirichlet parameters over transition likelihood (pB)
in the case that the one and only hidden state factor is updated, and there
are no actions.
"""
n_states = [3]
n_control = [
1
] # this is how we encode the fact that there aren't any actions
qs_prev = Categorical(values=construct_init_qs(n_states))
qs = Categorical(values=construct_init_qs(n_states))
l_rate = 1.0
B = Categorical(
values=np.random.rand(n_states[0], n_states[0], n_control[0]))
B.normalize()
pB = Dirichlet(values=np.ones_like(B.values))
action = np.array([np.random.randint(nc) for nc in n_control])
pB_updated = learning.update_transition_dirichlet(pB,
B,
action,
qs,
qs_prev,
lr=l_rate,
factors="all",
return_numpy=True)
validation_pB = pB.copy()
validation_pB[:, :, 0] += (l_rate *
maths.spm_cross(qs.values, qs_prev.values) *
(B[:, :, action[0]].values > 0))
self.assertTrue(np.all(pB_updated == validation_pB.values))
def test_contains_zeros(self):
values = np.array([[1.0, 0.0], [1.0, 1.0]])
c = Categorical(values=values)
self.assertTrue(c.contains_zeros())
values = np.array([[1.0, 1.0], [1.0, 1.0]])
c = Categorical(values=values)
self.assertFalse(c.contains_zeros())
def test_is_normalized(self):
values = np.array([[0.7, 0.5], [0.3, 0.5]])
c = Categorical(values=values)
self.assertTrue(c.is_normalized())
values = np.array([[0.2, 0.8], [0.3, 0.5]])
c = Categorical(values=values)
self.assertFalse(c.is_normalized())
