code: def _prune_B(
      B, state_levels_to_prune, action_levels_to_prune, dirichlet=False
  ):
      """
      Function for pruning a transition likelihood model (with potentially multiple hidden state factors)
  
      Parameters
      -----------
      B: ``numpy.ndarray`` of ``ndim == 3`` or ``numpy.ndarray`` of dtype object
          Dynamics likelihood mapping or 'transition model', mapping from hidden states at `t` to hidden states at `t+1`, given some control state `u`.
          Each element B[f] of this object array stores a 3-D tensor for hidden state factor `f`, whose entries `B[f][s, v, u] store the probability
          of hidden state level `s` at the current time, given hidden state level `v` and action `u` at the previous time.
      state_levels_to_prune: ``list`` of ``int`` or ``list`` of ``list``
          A ``list`` of the state levels to remove. If the likelihood in question has multiple hidden state factors,
          then this will be a ``list`` of ``list``, where each sub-list within ``state_levels_to_prune`` will contain the state levels
          to remove for a particular hidden state factor
      action_levels_to_prune: ``list`` of ``int`` or ``list`` of ``list``
          A ``list`` of the control state or action levels to remove. If the likelihood in question has multiple control state factors,
          then this will be a ``list`` of ``list``, where each sub-list within ``action_levels_to_prune`` will contain the control state levels
          to remove for a particular control state factor
      dirichlet: ``Bool``, default ``False``
          A Boolean flag indicating whether the input array(s) is/are a Dirichlet distribution, and therefore should not be normalized at the end.
          @TODO: Instead, the dirichlet parameters from the pruned rows/columns should somehow be re-distributed among the remaining rows/columns
  
      Returns
      -----------
      reduced_B: ``numpy.ndarray`` of `ndim == 3` or ``numpy.ndarray`` of dtype object
          The transition model, after pruning, which lacks the hidden state levels/action levels given by the arguments ``state_levels_to_prune`` and ``action_levels_to_prune``
      """
  
      slices_to_keep_list = []
  
      if utils.is_obj_array(B):
          num_controls = [B_arr.shape[2] for _, B_arr in enumerate(B)]
  
          for c, nc in enumerate(num_controls):
              indices_c = np.array(
                  list(set(range(nc)) - set(action_levels_to_prune[c])),
                  dtype=np.intp,
              )
              slices_to_keep_list.append(indices_c)
      else:
          num_controls = B.shape[2]
          slices_to_keep = np.array(
              list(
                  set(range(num_controls)) - set(action_levels_to_prune)
              ),
              dtype=np.intp,
          )
  
      if utils.is_obj_array(
          B
      ):  # in case of multiple hidden state factors
          assert all(
              [
                  type(ns_f_levels) == list
                  for ns_f_levels in state_levels_to_prune
              ]
          )
  
          num_factors = len(B)
  
          reduced_B = utils.obj_array(num_factors)
  
          for f, B_f in enumerate(B):  # loop over modalities
              ns = B_f.shape[0]
              states_to_keep = np.array(
                  list(set(range(ns)) - set(state_levels_to_prune[f])),
                  dtype=np.intp,
              )
  
              reduced_B[f] = B_f[
                  np.ix_(
                      states_to_keep,
                      states_to_keep,
                      slices_to_keep_list[f],
                  )
              ]
  
          if not dirichlet:
              reduced_B = utils.norm_dist_obj_arr(reduced_B)
          else:
              raise (
                  NotImplementedError(
                      "Need to figure out how to re-distribute concentration parameters from pruned rows/columns, across remaining rows/columns"
                  )
              )
  
      else:  # in case of one hidden state factor
          assert all(
              [
                  type(state_level_i) == int
                  for state_level_i in state_levels_to_prune
              ]
          )
  
          ns = B.shape[0]
          states_to_keep = np.array(
              list(set(range(ns)) - set(state_levels_to_prune)),
              dtype=np.intp,
          )
  
          reduced_B = B[
              np.ix_(states_to_keep, states_to_keep, slices_to_keep)
          ]
  
          if not dirichlet:
              reduced_B = utils.norm_dist(reduced_B)
          else:
              raise (
                  NotImplementedError(
                      "Need to figure out how to re-distribute concentration parameters from pruned rows/columns, across remaining rows/columns"
                  )
              )
  
      return reduced_B
  
file_path: pymdp/pymdp/learning.py
imports: import numpy
  from pymdp import utils, maths
  import copy
label_desc: Use image node_id 6 for calling a global function with example usage: _prune_B(B, state_levels_to_prune, action_levels_to_prune, dirichlet) and returns: reduced_B