docs: def thermal_state(omega_level, T, return_diagonal=False):
      r"""Return a thermal state for a given set of levels.
  
      INPUT:
  
      -  ``omega_level`` - The angular frequencies of each state.
      -  ``T`` - The temperature of the ensemble (in Kelvin).
      -  ``return_diagonal`` - Whether to return only the populations.
  
      >>> ground = State("Rb", 85, 5, 0, 1/Integer(2))
      >>> magnetic_states = make_list_of_states([ground], "magnetic")
      >>> omega_level = [ei.omega for ei in magnetic_states]
      >>> T = 273.15 + 20
  
      >>> print(thermal_state(omega_level, T, return_diagonal=True))
      [0.0834 0.0834 0.0834 0.0834 0.0834 0.0833 0.0833 0.0833 0.0833 0.0833
       0.0833 0.0833]
  
      """
      Ne = len(omega_level)
      E = np.array([hbar*omega_level[i] for i in range(Ne)])
      p = np.exp(-E/k_B/T)
      p = p/sum(p)
      if not return_diagonal:
          return np.diag(p)
      return p