diff --git a/.gitignore b/.gitignore index 778d69dd..5f24acf9 100644 --- a/.gitignore +++ b/.gitignore @@ -7,4 +7,5 @@ __pycache__ .ipynb_checkpoints/ .pytest_cache env/ -pymdp.egg-info \ No newline at end of file +pymdp.egg-info +inferactively_pymdp.egg-info diff --git a/pymdp/utils.py b/pymdp/utils.py index 005bd4ef..f8f69956 100644 --- a/pymdp/utils.py +++ b/pymdp/utils.py @@ -16,6 +16,27 @@ EPS_VAL = 1e-16 # global constant for use in norm_dist() +class Dimensions(object): + """ + The Dimensions class stores all data related to the size and shape of a model. + """ + def __init__( + self, + num_observations=None, + num_observation_modalities=0, + num_states=None, + num_state_factors=0, + num_controls=None, + num_control_factors=0, + ): + self.num_observations=num_observations + self.num_observation_modalities=num_observation_modalities + self.num_states=num_states + self.num_state_factors=num_state_factors + self.num_controls=num_controls + self.num_control_factors=num_control_factors + + def sample(probabilities): probabilities = probabilities.squeeze() if len(probabilities) > 1 else probabilities sample_onehot = np.random.multinomial(1, probabilities) @@ -196,22 +217,22 @@ def get_model_dimensions(A=None, B=None): def get_model_dimensions_from_labels(model_labels): modalities = model_labels['observations'] - num_modalities = len(modalities.keys()) - num_obs = [len(modalities[modality]) for modality in modalities.keys()] - factors = model_labels['states'] - num_factors = len(factors.keys()) - num_states = [len(factors[factor]) for factor in factors.keys()] - if 'actions' in model_labels.keys(): + res = Dimensions( + num_observations=[len(modalities[modality]) for modality in modalities.keys()], + num_observation_modalities=len(modalities.keys()), + num_states=[len(factors[factor]) for factor in factors.keys()], + num_state_factors=len(factors.keys()), + ) + if 'actions' in model_labels.keys(): controls = model_labels['actions'] - num_control_fac = len(controls.keys()) - num_controls = [len(controls[cfac]) for cfac in controls.keys()] + res.num_controls=[len(controls[cfac]) for cfac in controls.keys()] + res.num_control_factors=len(controls.keys()) + + return res - return num_obs, num_modalities, num_states, num_factors, num_controls, num_control_fac - else: - return num_obs, num_modalities, num_states, num_factors def norm_dist(dist): @@ -449,21 +470,18 @@ def construct_full_a(A_reduced, original_factor_idx, num_states): def create_A_matrix_stub(model_labels): - num_obs, _, num_states, _= get_model_dimensions_from_labels(model_labels) + dimensions = get_model_dimensions_from_labels(model_labels) obs_labels, state_labels = model_labels['observations'], model_labels['states'] state_combinations = pd.MultiIndex.from_product(list(state_labels.values()), names=list(state_labels.keys())) - num_state_combos = np.prod(num_states) - # num_rows = (np.array(num_obs) * num_state_combos).sum() - num_rows = sum(num_obs) + num_rows = sum(dimensions.num_observations) cell_values = np.zeros((num_rows, len(state_combinations))) obs_combinations = [] for modality in obs_labels.keys(): levels_to_combine = [[modality]] + [obs_labels[modality]] - # obs_combinations += num_state_combos * list(itertools.product(*levels_to_combine)) obs_combinations += list(itertools.product(*levels_to_combine)) @@ -475,7 +493,7 @@ def create_A_matrix_stub(model_labels): def create_B_matrix_stubs(model_labels): - _, _, num_states, _, num_controls, _ = get_model_dimensions_from_labels(model_labels) + dimensions = get_model_dimensions_from_labels(model_labels) state_labels = model_labels['states'] action_labels = model_labels['actions'] @@ -489,9 +507,9 @@ def create_B_matrix_stubs(model_labels): prev_state_action_combos = pd.MultiIndex.from_product(factor_list, names=[factor, list(action_labels.keys())[f_idx]]) - num_state_action_combos = num_states[f_idx] * num_controls[f_idx] + num_state_action_combos = dimensions.num_states[f_idx] * dimensions.num_controls[f_idx] - num_rows = num_states[f_idx] + num_rows = dimensions.num_states[f_idx] cell_values = np.zeros((num_rows, num_state_action_combos)) @@ -544,13 +562,12 @@ def convert_A_stub_to_ndarray(A_stub, model_labels): This function converts a multi-index pandas dataframe `A_stub` into an object array of different A matrices, one per observation modality. """ + dimensions = get_model_dimensions_from_labels(model_labels) - num_obs, num_modalities, num_states, num_factors = get_model_dimensions_from_labels(model_labels) - - A = obj_array(num_modalities) + A = obj_array(dimensions.num_observation_modalities) for g, modality_name in enumerate(model_labels['observations'].keys()): - A[g] = A_stub.loc[modality_name].to_numpy().reshape(num_obs[g], *num_states) + A[g] = A_stub.loc[modality_name].to_numpy().reshape(dimensions.num_observations[g], *dimensions.num_states) assert (A[g].sum(axis=0) == 1.0).all(), 'A matrix not normalized! Check your initialization....\n' return A @@ -561,13 +578,13 @@ def convert_B_stubs_to_ndarray(B_stubs, model_labels): of different B matrices, one per hidden state factor """ - _, _, num_states, num_factors, num_controls, num_control_fac = get_model_dimensions_from_labels(model_labels) + dimensions = get_model_dimensions_from_labels(model_labels) - B = obj_array(num_factors) + B = obj_array(dimensions.num_control_factors) for f, factor_name in enumerate(B_stubs.keys()): - B[f] = B_stubs[factor_name].to_numpy().reshape(num_states[f], num_states[f], num_controls[f]) + B[f] = B_stubs[factor_name].to_numpy().reshape(dimensions.num_states[f], dimensions.num_states[f], dimensions.num_controls[f]) assert (B[f].sum(axis=0) == 1.0).all(), 'B matrix not normalized! Check your initialization....\n' return B diff --git a/test/test_wrappers.py b/test/test_wrappers.py index db25c984..254d902b 100644 --- a/test/test_wrappers.py +++ b/test/test_wrappers.py @@ -1,15 +1,11 @@ import os import unittest from pathlib import Path -import shutil -import tempfile -import numpy as np -import itertools import pandas as pd from pandas.testing import assert_frame_equal -from pymdp.utils import create_A_matrix_stub, read_A_matrix, create_B_matrix_stubs, read_B_matrices +from pymdp.utils import Dimensions, get_model_dimensions_from_labels, create_A_matrix_stub, read_A_matrix, create_B_matrix_stubs, read_B_matrices tmp_path = Path('tmp_dir') @@ -18,11 +14,62 @@ class TestWrappers(unittest.TestCase): + def test_get_model_dimensions_from_labels(self): + """ + Tests model dimension extraction from labels including observations, states and actions. + """ + model_labels = { + "observations": { + "species_observation": [ + "absent", + "present", + ], + "budget_observation": [ + "high", + "medium", + "low", + ], + }, + "states": { + "species_state": [ + "extant", + "extinct", + ], + }, + "actions": { + "conservation_action": [ + "manage", + "survey", + "stop", + ], + }, + } + + want = Dimensions( + num_observations=[2, 3], + num_observation_modalities=2, + num_states=[2], + num_state_factors=1, + num_controls=[3], + num_control_factors=1, + ) + + got = get_model_dimensions_from_labels(model_labels) + + self.assertEqual(want.num_observations, got.num_observations) + self.assertEqual(want.num_observation_modalities, got.num_observation_modalities) + self.assertEqual(want.num_states, got.num_states) + self.assertEqual(want.num_state_factors, got.num_state_factors) + self.assertEqual(want.num_controls, got.num_controls) + self.assertEqual(want.num_control_factors, got.num_control_factors) + def test_A_matrix_stub(self): """ This tests the construction of a 2-modality, 2-hidden state factor pandas MultiIndex dataframe using the `model_labels` dictionary, which contains the modality- and factor-specific levels, labeled with string - identifiers + identifiers. + + Note: actions are ignored when creating an A matrix stub """ model_labels = { @@ -41,6 +88,9 @@ def test_A_matrix_stub(self): "weather_state": ["raining", "clear"], "sprinkler_state": ["on", "off"], }, + "actions": { + "actions": ["something", "nothing"], + } } num_hidden_state_factors = len(model_labels["states"])