code: def _prune_prior(prior, levels_to_remove, dirichlet=False):
      """
      Function for pruning a prior Categorical distribution (e.g. C, D)
  
      Parameters
      -----------
      prior: 1D ``numpy.ndarray`` or ``numpy.ndarray`` of dtype object
          The vector(s) containing the priors over hidden states of a generative model, e.g. the prior over hidden states (``D`` vector).
      levels_to_remove: ``list`` of ``int``, ``list`` of ``list``
          A ``list`` of the levels (indices of the support) to remove. If the prior in question has multiple hidden state factors / multiple observation modalities,
          then this will be a ``list`` of ``list``, where each sub-list within ``levels_to_remove`` will contain the levels to prune for a particular hidden state factor or modality
      dirichlet: ``Bool``, default ``False``
          A Boolean flag indicating whether the input vector(s) is/are a Dirichlet distribution, and therefore should not be normalized at the end.
          @TODO: Instead, the dirichlet parameters from the pruned levels should somehow be re-distributed among the remaining levels
  
      Returns
      -----------
      reduced_prior: 1D ``numpy.ndarray`` or ``numpy.ndarray`` of dtype object
          The prior vector(s), after pruning, that lacks the hidden state or modality levels indexed by ``levels_to_remove``
      """
  
      if utils.is_obj_array(
          prior
      ):  # in case of multiple hidden state factors
          assert all(
              [type(levels) == list for levels in levels_to_remove]
          )
  
          num_factors = len(prior)
  
          reduced_prior = utils.obj_array(num_factors)
  
          factors_to_remove = []
          for f, s_i in enumerate(
              prior
          ):  # loop over factors (or modalities)
              ns = len(s_i)
              levels_to_keep = list(
                  set(range(ns)) - set(levels_to_remove[f])
              )
              if len(levels_to_keep) == 0:
                  print(
                      f"Warning... removing ALL levels of factor {f} - i.e. the whole hidden state factor is being removed\n"
                  )
                  factors_to_remove.append(f)
              else:
                  if not dirichlet:
                      reduced_prior[f] = utils.norm_dist(
                          s_i[levels_to_keep]
                      )
                  else:
                      raise (
                          NotImplementedError(
                              "Need to figure out how to re-distribute concentration parameters from pruned levels, across remaining levels"
                          )
                      )
  
          if len(factors_to_remove) > 0:
              factors_to_keep = list(
                  set(range(num_factors)) - set(factors_to_remove)
              )
              reduced_prior = reduced_prior[factors_to_keep]
  
      else:  # in case of one hidden state factor
          assert all(
              [type(level_i) == int for level_i in levels_to_remove]
          )
  
          ns = len(prior)
          levels_to_keep = list(set(range(ns)) - set(levels_to_remove))
  
          if not dirichlet:
              reduced_prior = utils.norm_dist(prior[levels_to_keep])
          else:
              raise (
                  NotImplementedError(
                      "Need to figure out how to re-distribute concentration parameters from pruned levels, across remaining levels"
                  )
              )
  
      return reduced_prior
  
file_path: pymdp/pymdp/learning.py
imports: import numpy
  from pymdp import utils, maths
  import copy
label_desc: Use image node_id 4 for calling a global function with example usage: _prune_prior(prior, levels_to_remove, dirichlet) and returns: reduced_prior