Compute the kinetic and potential energy of the earth sun system

def energy_ES(q, v):
    """Compute the kinetic and potential energy of the earth sun system"""
    # Body 0 is the sun, Body 1 is the earth
    m0 = mass[0]
    m1 = mass[1]

    # Positions of sun and earth
    q0: np.ndarray = q[:, slices[0]]
    q1: np.ndarray = q[:, slices[1]]

    # Velocities of sun and earth
    v0: np.ndarray = v[:, slices[0]]
    v1: np.ndarray = v[:, slices[1]]

    # Kinetic energy is 1/2 mv^2
    T0: np.ndarray = 0.5 * m0 * np.sum(v0 * v0, axis=1)
    T1: np.ndarray = 0.5 * m1 * np.sum(v1 * v1, axis=1)
    T: np.ndarray = T0 + T1
    
    # Potential energy is -G m1 m2  / r
    dv_01 = q1 - q0
    r_01 = np.linalg.norm(dv_01, axis=1)
    U_01: np.ndarray = -G * m0 * m1 / r_01
    U: np.ndarray = U_01
    
    # Total energy H = T + U
    H = T + U
    
    return H, T, U