docstr: Tests the refactored (Categorical-less) version of `get_expected_states`
docstr_tok: ["Tests","the","refactored","(","Categorical-less",")","version","of","`","get_expected_states","`"]
code: def test_get_expected_states(self):
      """
      Tests the refactored (Categorical-less) version of `get_expected_states`
      """
  
      """Test with single hidden state factor and single timestep"""
  
      num_states = [3]
      num_controls = [3]
  
      qs = utils.obj_array_uniform(num_states)
      B = utils.random_B_matrix(num_states, num_controls)
  
      policies = control.construct_policies(
          num_states, num_controls, policy_len=1
      )
  
      qs_pi = [
          control.get_expected_states(qs, B, policy)
          for policy in policies
      ]
  
      factor_idx = 0
      t_idx = 0
  
      for p_idx in range(len(policies)):
          self.assertTrue(
              (
                  qs_pi[p_idx][t_idx][factor_idx]
                  == B[factor_idx][
                      :, :, policies[p_idx][t_idx, factor_idx]
                  ].dot(qs[factor_idx])
              ).all()
          )
  
      """Test with single hidden state factor and multiple-timesteps"""
  
      num_states = [3]
      num_controls = [3]
  
      qs = utils.obj_array_uniform(num_states)
      B = utils.random_B_matrix(num_states, num_controls)
  
      policies = control.construct_policies(
          num_states, num_controls, policy_len=2
      )
  
      qs_pi = [
          control.get_expected_states(qs, B, policy)
          for policy in policies
      ]
  
      for p_idx in range(len(policies)):
          for t_idx in range(2):
              if t_idx == 0:
                  self.assertTrue(
                      (
                          qs_pi[p_idx][t_idx][factor_idx]
                          == B[factor_idx][
                              :, :, policies[p_idx][t_idx, factor_idx]
                          ].dot(qs[factor_idx])
                      ).all()
                  )
              else:
                  self.assertTrue(
                      (
                          qs_pi[p_idx][t_idx][factor_idx]
                          == B[factor_idx][
                              :, :, policies[p_idx][t_idx, factor_idx]
                          ].dot(qs_pi[p_idx][t_idx - 1][factor_idx])
                      ).all()
                  )
  
      """Test with multiple hidden state factors and single timestep"""
  
      num_states = [3, 3]
      num_controls = [3, 1]
  
      num_factors = len(num_states)
  
      qs = utils.obj_array_uniform(num_states)
      B = utils.random_B_matrix(num_states, num_controls)
  
      policies = control.construct_policies(
          num_states, num_controls, policy_len=1
      )
  
      qs_pi = [
          control.get_expected_states(qs, B, policy)
          for policy in policies
      ]
  
      t_idx = 0
  
      for p_idx in range(len(policies)):
          for factor_idx in range(num_factors):
              self.assertTrue(
                  (
                      qs_pi[p_idx][t_idx][factor_idx]
                      == B[factor_idx][
                          :, :, policies[p_idx][t_idx, factor_idx]
                      ].dot(qs[factor_idx])
                  ).all()
              )
  
      """Test with multiple hidden state factors and multiple timesteps"""
  
      num_states = [3, 3]
      num_controls = [3, 3]
  
      num_factors = len(num_states)
  
      qs = utils.obj_array_uniform(num_states)
      B = utils.random_B_matrix(num_states, num_controls)
  
      policies = control.construct_policies(
          num_states, num_controls, policy_len=3
      )
  
      qs_pi = [
          control.get_expected_states(qs, B, policy)
          for policy in policies
      ]
  
      for p_idx in range(len(policies)):
          for t_idx in range(3):
              for factor_idx in range(num_factors):
                  if t_idx == 0:
                      self.assertTrue(
                          (
                              qs_pi[p_idx][t_idx][factor_idx]
                              == B[factor_idx][
                                  :,
                                  :,
                                  policies[p_idx][t_idx, factor_idx],
                              ].dot(qs[factor_idx])
                          ).all()
                      )
                  else:
                      self.assertTrue(
                          (
                              qs_pi[p_idx][t_idx][factor_idx]
                              == B[factor_idx][
                                  :,
                                  :,
                                  policies[p_idx][t_idx, factor_idx],
                              ].dot(qs_pi[p_idx][t_idx - 1][factor_idx])
                          ).all()
                      )
  
imports: import os
  import unittest
  import numpy
  from pymdp import utils, maths
  from pymdp import control
label_desc: Use image node_id 1 for calling the TestControl obj's underlying member method code with example usage: obj.test_get_expected_states() without return types