docs: def computeExpectationsInner(self, A_n, u_ln, state_map,
                                   uncertainty_method=None,
                                   warning_cutoff=1.0e-10,
                                   return_theta=False):
          """
          Compute the expectations of multiple observables of phase space functions in multiple states.
  
          Compute the expectations of multiple observables of phase
          space functions [A_0(x),A_1(x),...,A_i(x)] along with the
          covariances of their estimates at multiple states.
  
          intended as an internal function to keep all the optimized and
          robust expectation code in one place, but will leave it
          open to allow for later modifications
  
          It calculates all input observables at all states which are
          specified by the list of states in the state list.
  
          Parameters
          ----------
          A_n : np.ndarray, float, shape=(I, N)
              A_in[i,n] = A_i(x_n), the value of phase observable i for configuration n
          u_ln : np.ndarray, float, shape=(L, N)
              u_ln[l,n] is the reduced potential of configuration n at state l
              if u_ln = None, we use self.u_kn
          state_map : np.ndarray, int, shape (2,NS) or shape(1,NS)
              If state_map has only one dimension where NS is the total number of states we want to simulate things
              a.  The list will be of the form ``[[0,1,2],[0,1,1]]``. This particular example indicates we want to output
              the properties of three observables total: the first property A[0] at the 0th state, the 2nd property
              A[1] at the 1th state, and the 2nd property A[1] at the 2nd state. This allows us to tailor our output to a
              large number of different situations.
          uncertainty_method : string, optional
              Choice of method used to compute asymptotic covariance method, or None to use default
              See help for computeAsymptoticCovarianceMatrix() for more information on various methods. (default: None)
          warning_cutoff : float, optional
              Warn if squared-uncertainty is negative and larger in magnitude than this number (default: 1.0e-10)
          return_theta : bool, optional
              Whether or not to return the theta matrix.  Can be useful for complicated differences of observables.
  
          Returns
          -------
          result_vals : dictionary
  
          Possible keys in the result_vals dictionary:
  
          'observables': np.ndarray, float, shape = (S)
              results_vals['observables'][i] is the estimate for the expectation of A_state_map[i](x) at the state specified by u_n[state_map[i],:]
  
          'Theta' : np.ndarray, float, shape = (K+len(state_list), K+len(state_list)) the covariance matrix of log weights.
  
          'Amin' : np.ndarray, float, shape = (S), needed for reconstructing the covariance one level up. 
  
          'f' : np.ndarray, float, shape = (K+len(state_list)), 'free energies' of the new states (i.e. ln (<A>-Amin+1)) as the log form is easier to work with.
  
          Notes
          -----
  
          Situations this will be used for :
  
          * Multiple observables, single state (called though computeMultipleExpectations)
          * Single observable, multiple states (called through computeExpectations)
  
              This has two cases: observables that don't change with state, and observables that
              do change with state.
              For example, the set of energies at state k consist in energy function of state
              1 evaluated at state 1, energies of state 2 evaluated at
              state 2, and so forth.
  
          * Computing only free energies at new states.
  
          * Would require additional work to work with potentials of mean force, because we need to ignore the functions that are zero when integrating.
  
          Examples
          --------
  
          >>> from pymbar import testsystems
          >>> (x_n, u_kn, N_k, s_n) = testsystems.HarmonicOscillatorsTestCase().sample(mode='u_kn')
          >>> mbar = MBAR(u_kn, N_k)
          >>> A_n = np.array([x_n,x_n**2,x_n**3])
          >>> u_n = u_kn[:2,:]
          >>> state_map = np.array([[0,0],[1,0],[2,0],[2,1]],int)
          >>> results = mbar.computeExpectationsInner(A_n, u_n, state_map)
  
          """
  
          # Retrieve N and K for convenience.
          mapshape = np.shape(state_map) # number of computed expectations we desire
                                                 # need to convert to matrix to be able to pick up D=1
          if len(mapshape) < 2:
              # if 1D, it's just a list of states
              state_list = state_map.copy()
              state_map = np.zeros([0,0],np.float64)
              S = 0
          else:  # if 2D, then it's a list of observables and corresponding states
              state_list = state_map[0,:]
              S = mapshape[1]
  
          # reshape arrays explicitly into 2d (even if only one state) to make it easy to manipulate
          shapeu = np.shape(u_ln)
          if len(shapeu) == 1:
              u_ln = np.reshape(u_ln,[1,shapeu[0]])
  
          shapeA = np.shape(A_n)
          if len(shapeA) == 1:
              A_n = np.reshape(A_n,[1,shapeA[0]])
  
          K = self.K
          N = self.N  # N is total number of samples
          result_vals = dict() # dictionary we will store uncertainties in
  
          # make observables all positive, allowing us to take the logarithm, which is
          # required to prevent overflow in some examples.
          # WARNING: one issue to watch for is if one of the energies is extremely
          # low (-10^10 or lower), but most of the energies of interest are much higher.
          # This could lead to roundoff problems (check with Levi N.)
  
          L_list = np.unique(state_list)
          NL = len(L_list) # number of states we need to examine
          if S > 0:
              A_list = np.unique(state_map[1,:])  # what are the unique observables
              A_min = np.zeros([len(A_list)], dtype=np.float64)
          else:
              A_list = np.zeros(0,dtype=int)
  
          for i in A_list:
              A_min[i] = np.min(A_n[i, :]) #find the minimum
              A_n[i, :] = A_n[i,:] - (A_min[i] - 1)  #all values now positive so that we can work in logarithmic scale
  
          # Augment W_nk, N_k, and c_k for q_A(x) for the observables, with one
          # row for the specified state and I rows for the observable at that
          # state.
          # log weight matrix
          msize = K + NL + S # augmented size; all of the states needed to calculate
                             # the observables, and the observables themselves.
          Log_W_nk = np.zeros([N, msize], np.float64) # log weight matrix
          N_k = np.zeros([msize], np.int64)  # counts
          f_k = np.zeros([msize], np.float64)  # free energies
  
          # <A> = A(x_n) exp[f_{k} - q_{k}(x_n)] / \sum_{k'=1}^K N_{k'} exp[f_{k'} - q_{k'}(x_n)]
          # Fill in first section of matrix with existing q_k(x) from states.
          Log_W_nk[:, 0:K] = self.Log_W_nk
          N_k[0:K] = self.N_k
          f_k[0:K] = self.f_k
  
          # Pre-calculate the log denominator: Eqns 13, 14 in MBAR paper
          states_with_samples = (self.N_k > 0)
          log_denominator_n = logsumexp(self.f_k[states_with_samples] - self.u_kn[states_with_samples].T, b=self.N_k[states_with_samples], axis=1)
          # Compute row of W_nk matrix for the extra states corresponding to u_ln
          # that the state list specifies
          for l in L_list:
              la = K+l  #l, augmented
              # Calculate log normalizing constants and log weights via Eqns 13, 14
              log_C_a = -logsumexp(-u_ln[l] - log_denominator_n)
              Log_W_nk[:, la] = log_C_a - u_ln[l] - log_denominator_n
              f_k[la] = log_C_a
  
          # Compute the remaining rows/columns of W_nk, and calculate
          # their normalizing constants c_k
          for s in range(S):
              sa = K+NL+s  # augmented s
              l = K + state_map[0,s]
              i = state_map[1,s]
              Log_W_nk[:, sa] = np.log(A_n[i, :]) + Log_W_nk[:, l]
              f_k[sa] = -logsumexp(Log_W_nk[:, sa])
              Log_W_nk[:, sa] += f_k[sa]    # normalize this row
  
          # Compute estimates of A_i[s]
          A_i = np.zeros([S], np.float64)
          for s in range(S):
              A_i[s] = np.exp(-f_k[K + NL + s])
  
          # Now that covariances are computed, add the constants back to A_i that
          # were required to enforce positivity
          for s in range(S):
              A_i[s] += (A_min[state_map[1,s]] - 1)
  
          # these values may be used outside the routine, so copy back.
          for i in A_list:
              A_n[i, :] = A_n[i,:] + (A_min[i] - 1)
  
          # expectations of the observables at these states
          if S > 0:
              result_vals['observables'] = A_i
  
          if return_theta:
              Theta_ij = self._computeAsymptoticCovarianceMatrix(
                  np.exp(Log_W_nk), N_k, method=uncertainty_method)
  
              # Note: these variances will be the same whether or not we
              # subtract a different constant from each A_i
              # for efficency, output theta in block form
              #          K*K   K*S  K*NL
              # Theta =  K*S   S*S  NL*S
              #          K*NL  NL*S NL*NL
  
              # first the observables (S of them), then the free energies (also S of them)
              if S>0:
                  si = K+NL+np.arange(S)
              else:
                  si = np.zeros(0,dtype=int)
              li = K+state_list
              i = np.concatenate((si,li))
              Theta = Theta_ij[np.ix_(i, i)]
              result_vals['Theta'] = Theta
              if S > 0:
                  # we need to return the minimum A as well
                  result_vals['Amin'] = (A_min[state_map[1,np.arange(S)]] - 1)
  
          # free energies at these new states
          result_vals['f'] =  f_k[K+state_list]
  
          # Return expectations and uncertainties.
          return result_vals