Compute the kinetic and potential energy of the planetary system

def energy(q, v):
    """Compute the kinetic and potential energy of the planetary system"""
    
    # Number of points
    N: int = len(q)
    
    # Initialize arrays to zero of the correct size
    T: np.ndarray = np.zeros(N)
    U: np.ndarray = np.zeros(N)    
    
    # Add up kinetic energy of each body
    for i in range(B):
        # Kinetic energy is 1/2 mv^2
        m = mass[i]
        vi = v[:, slices[i]]
        T += 0.5 * m * np.sum(vi * vi, axis=1)
    
    # Add up potential energy of each pair of bodies
    for i in range(B):
        for j in range(i+1, B):
            # Masses of these two bodies
            mi = mass[i]
            mj = mass[j]
        
            # Positions of body i and j
            qi: np.ndarray = q[:, slices[i]]
            qj: np.ndarray = q[:, slices[j]]
        
            # Potential energy is -G m1 m2  / r
            dv_ij = qj - qi
            r_ij = np.linalg.norm(dv_ij, axis=1)
            U -= G * mi * mj * 1.0 / r_ij
    
    # Total energy H = T + U
    H = T + U
    
    return H, T, U