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 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_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_dot_function_e_cat(self): """ Continuous states and outcomes, but add a final (fourth) hidden state factor Now, when arguments themselves are instances of Categorical """ array_path = os.path.join(os.getcwd(), DATA_PATH + "dot_e.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_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_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_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_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_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_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_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_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())
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 _construct_D_prior(self): if self.n_factors == 1: D = Categorical(values=np.ones(*self.n_states)) else: D = Categorical( values=np.array([np.ones(Ns) for Ns in self.n_states])) D.normalize() return D
def test_sample_single(self): # values are already normalized values = np.array([1.0, 0.0]) c = Categorical(values=values) self.assertEqual(0, c.sample()) # values are not normalized values = np.array([0, 10.0]) c = Categorical(values=values) self.assertEqual(1, c.sample())
def _construct_A_distribution(self): if self.n_modalities == 1: A = Categorical(values=np.random.rand(*(self.n_observations[0] + self.n_states))) else: A = np.empty(self.n_modalities, dtype=object) for modality, no in enumerate(self.n_observations): A[modality] = np.random.rand(*([no] + self.n_states)) A = Categorical(values=A) A.normalize() return A
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 test_sample_AoA(self): # values are already normalized values_1 = np.array([1.0, 0.0]) values_2 = np.array([0.0, 1.0, 0.0]) values = np.array([values_1, values_2]) c = Categorical(values=values) self.assertTrue(np.isclose(np.array([0, 1]), c.sample()).all()) # values are not normalized values_1 = np.array([10.0, 0.0]) values_2 = np.array([0.0, 10.0, 0.0]) values = np.array([values_1, values_2]) c = Categorical(values=values) self.assertTrue(np.isclose(np.array([0, 1]), c.sample()).all())
def test_state_info_gain(self): """ Test the states_info_gain function. Demonstrates working by manipulating uncertainty in the likelihood matrices (A or B) in a ways that alternatively change the resolvability of uncertainty (via an imprecise expected state and a precise mapping, or high ambiguity and imprecise mapping). """ n_states = [2] n_control = [2] qs = Categorical(values=np.eye(n_states[0])[0]) # add some uncertainty into the consequences of the second policy, which # leads to increased epistemic value of observations, in case of pursuing # that policy -- in the case of a precise observation likelihood model B_matrix = construct_generic_B(n_states, n_control) B_matrix[:, :, 1] = core.softmax(B_matrix[:, :, 1]) B = Categorical(values=B_matrix) # 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])) state_info_gains = np.zeros(len(policies)) for idx, policy in enumerate(policies): qs_pi = core.get_expected_states(qs, B, policy) state_info_gains[idx] += core.calc_states_info_gain(A, qs_pi) self.assertGreater(state_info_gains[1], state_info_gains[0]) # we can 'undo' the epistemic bonus of the second policy by making the A matrix # totally ambiguous, thus observations cannot resolve uncertainty about hidden states # - in this case, uncertainty in the posterior beliefs doesn't matter A = Categorical(values=np.ones((num_obs[0], num_obs[0]))) A.normalize() state_info_gains = np.zeros(len(policies)) for idx, policy in enumerate(policies): qs_pi = core.get_expected_states(qs, B, policy) state_info_gains[idx] += core.calc_states_info_gain(A, qs_pi) self.assertEqual(state_info_gains[0], state_info_gains[1])
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_dot_function_c(self): """ Discrete states and outcomes, but also a third hidden state factor """ array_path = os.path.join(os.getcwd(), DATA_PATH + "dot_c.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) result_1_py = A.dot_likelihood(obs, 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]) c = Categorical(values=values) c.normalize() self.assertTrue(c.is_normalized())
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 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 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 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 test_update_pB_multi_factor_no_actions_one_factor(self): """ Test for updating prior Dirichlet parameters over transition likelihood (pB) in the case that there are mulitple hidden state factors, and there are no actions. One factor is updated """ n_states = [3, 4] n_control = [1, 1] qs_prev = Categorical(values=construct_init_qs(n_states)) qs = Categorical(values=construct_init_qs(n_states)) l_rate = 1.0 factors_to_update = [np.random.randint(len(n_states))] B = Categorical(values=np.array([ np.random.rand(ns, ns, n_control[factor]) for factor, ns in enumerate(n_states) ], dtype=object)) B.normalize() pB = Dirichlet(values=np.array([ np.ones_like(B[factor].values) for factor in range(len(n_states)) ], dtype=object)) 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=factors_to_update, return_numpy=True) validation_pB = pB.copy() for factor, _ in enumerate(n_control): validation_pB = pB[factor].copy() if factor in factors_to_update: validation_pB[:, :, action[factor]] += ( l_rate * maths.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_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 test_update_pA_single_factor_all(self): """ Test for updating prior Dirichlet parameters over sensory likelihood (pA) in the case that all observation modalities are 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 # single observation modality num_obs = [4] 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="all", return_numpy=True) validation_pA = pA + l_rate * core.spm_cross( np.eye(*num_obs)[observation], qs.values) self.assertTrue(np.all(pA_updated == validation_pA.values)) # multiple observation modalities num_obs = [3, 4] 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="all", return_numpy=True) for modality, no in enumerate(num_obs): update = core.spm_cross( np.eye(no)[observation[modality]], qs.values) validation_pA = pA[modality] + l_rate * update self.assertTrue( np.all(pA_updated[modality] == validation_pA.values))
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))