def get_key(iroot, num_states):
    """
    Get energy key for the state of interest.
    Args:
        iroot (int): state of interest
        num_states (int): total number of states
    Returns:
        key (str): energy key
    """

    # energy if only one state
    if iroot == 0 and num_states == 1:
        key = "energy"

    # otherwise energy with state suffix
    else:
        key = "energy_{}".format(iroot)
    return key