code: def _prune_A(
      A, obs_levels_to_prune, state_levels_to_prune, dirichlet=False
  ):
      """
      Function for pruning a observation likelihood model (with potentially multiple hidden state factors)
      :meta private:
      Parameters
      -----------
      A: ``numpy.ndarray`` with ``ndim >= 2``, or ``numpy.ndarray`` of dtype object
          Sensory likelihood mapping or 'observation model', mapping from hidden states to observations. Each element ``A[m]`` of
          stores an ``numpy.ndarray`` multidimensional array for observation modality ``m``, whose entries ``A[m][i, j, k, ...]`` store
          the probability of observation level ``i`` given hidden state levels ``j, k, ...``
      obs_levels_to_prune: ``list`` of int or ``list`` of ``list``:
          A ``list`` of the observation levels to remove. If the likelihood in question has multiple observation modalities,
          then this will be a ``list`` of ``list``, where each sub-list within ``obs_levels_to_prune`` will contain the observation levels
          to remove for a particular observation modality
      state_levels_to_prune: ``list`` of ``int``
          A ``list`` of the hidden state levels to remove (this will be the same across modalities)
      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 columns should somehow be re-distributed among the remaining columns
  
      Returns
      -----------
      reduced_A: ``numpy.ndarray`` with ndim >= 2, or ``numpy.ndarray ``of dtype object
          The observation model, after pruning, which lacks the observation or hidden state levels given by the arguments ``obs_levels_to_prune`` and ``state_levels_to_prune``
      """
  
      columns_to_keep_list = []
      if utils.is_obj_array(A):
          num_states = A[0].shape[1:]
          for f, ns in enumerate(num_states):
              indices_f = np.array(
                  list(set(range(ns)) - set(state_levels_to_prune[f])),
                  dtype=np.intp,
              )
              columns_to_keep_list.append(indices_f)
      else:
          num_states = A.shape[1]
          indices = np.array(
              list(set(range(num_states)) - set(state_levels_to_prune)),
              dtype=np.intp,
          )
          columns_to_keep_list.append(indices)
  
      if utils.is_obj_array(
          A
      ):  # in case of multiple observation modality
          assert all(
              [
                  type(o_m_levels) == list
                  for o_m_levels in obs_levels_to_prune
              ]
          )
  
          num_modalities = len(A)
  
          reduced_A = utils.obj_array(num_modalities)
  
          for m, A_i in enumerate(A):  # loop over modalities
              no = A_i.shape[0]
              rows_to_keep = np.array(
                  list(set(range(no)) - set(obs_levels_to_prune[m])),
                  dtype=np.intp,
              )
  
              reduced_A[m] = A_i[
                  np.ix_(rows_to_keep, *columns_to_keep_list)
              ]
          if not dirichlet:
              reduced_A = utils.norm_dist_obj_arr(reduced_A)
          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 observation modality
          assert all(
              [
                  type(o_levels_i) == int
                  for o_levels_i in obs_levels_to_prune
              ]
          )
  
          no = A.shape[0]
          rows_to_keep = np.array(
              list(set(range(no)) - set(obs_levels_to_prune)),
              dtype=np.intp,
          )
  
          reduced_A = A[np.ix_(rows_to_keep, *columns_to_keep_list)]
  
          if not dirichlet:
              reduced_A = utils.norm_dist(reduced_A)
          else:
              raise (
                  NotImplementedError(
                      "Need to figure out how to re-distribute concentration parameters from pruned rows/columns, across remaining rows/columns"
                  )
              )
  
      return reduced_A
  
file_path: pymdp/pymdp/learning.py
imports: import numpy
  from pymdp import utils, maths
  import copy
label_desc: Use image node_id 5 for calling a global function with example usage: _prune_A(A, obs_levels_to_prune, state_levels_to_prune, dirichlet) and returns: reduced_A