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