def UtoT(U):
    """
    Takes internal energy in (km/s)^2, and returns temperature in kelvin.
    Assumes mass of gas is entirely due to protons and assumes ideal gas.

    Parameters
    ----------
    U : float or np.array of dtype float
        Internal energy in (km/s)^2

    Returns
    -------
    T : float or np.array of dtype float
        Temperature in K
    """
    mass_p = 1.6726219e-27
    k_boltz = 1.38064852e-23
    temp = (2/3 * mass_p / k_boltz) * U * 1e6
    return temp