docs: def mbar_objective_and_gradient(u_kn, N_k, f_k):
      """Calculates both objective function and gradient for MBAR.
  
      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
  
  
      Returns
      -------
      obj : float
          Objective function
      grad : np.ndarray, dtype=float, shape=(n_states)
          Gradient of objective function
  
      Notes
      -----
      This objective function is essentially a doubly-summed partition function and is
      quite sensitive to precision loss from both overflow and underflow. For optimal
      results, u_kn can be preconditioned by subtracting out a `n` dependent
      vector.
  
      More optimal precision, the objective function uses math.fsum for the
      outermost sum and logsumexp for the inner sum.
      
      The gradient is equation C6 in the JCP MBAR paper; the objective
      function is its integral.
      """
      u_kn, N_k, f_k = validate_inputs(u_kn, N_k, f_k)
  
      log_denominator_n = logsumexp(f_k - u_kn.T, b=N_k, axis=1)
      log_numerator_k = logsumexp(-log_denominator_n - u_kn, axis=1)
      grad = -1 * N_k * (1.0 - np.exp(f_k + log_numerator_k))
  
      obj = math.fsum(log_denominator_n) - N_k.dot(f_k)
  
      return obj, grad