docs: def solve_mbar_for_all_states(u_kn, N_k, f_k, solver_protocol):
      """Solve for free energies of states with samples, then calculate for
      empty states.
  
      Parameters
      ----------
      u_kn : np.ndarray, shape=(n_states, n_samples), dtype='float'
          The reduced potential energies, i.e. -log unnormalized probabilities
      N_k : np.ndarray, shape=(n_states), dtype='int'
          The number of samples in each state
      f_k : np.ndarray, shape=(n_states), dtype='float'
          The reduced free energies of each state
      solver_protocol: tuple(dict()), optional, default=None
          Sequence of dictionaries of steps in solver protocol for final
          stage of refinement.
  
      Returns
      -------
      f_k : np.ndarray, shape=(n_states), dtype='float'
          The free energies of states
      """
      states_with_samples = np.where(N_k > 0)[0]
  
      if len(states_with_samples) == 1:
          f_k_nonzero = np.array([0.0])
      else:
          f_k_nonzero, all_results = solve_mbar(u_kn[states_with_samples], N_k[states_with_samples],
                                                f_k[states_with_samples], solver_protocol=solver_protocol)
  
      f_k[states_with_samples] = f_k_nonzero
  
      # Update all free energies because those from states with zero samples are not correctly computed by solvers.
      f_k = self_consistent_update(u_kn, N_k, f_k)
      # This is necessary because state 0 might have had zero samples,
      # but we still want that state to be the reference with free energy 0.
      f_k -= f_k[0]
  
      return f_k