def R_score_v2()

in neural/linear/stats.py [0:0]

• 21 lines of code
• 9 McCabe index (conditional complexity)

def R_score_v2(Y_true, Y_pred, score="r", avg_out="times", start=0):
    """
    Y_true: numpy or torch (B, T, C)
    Y_pred: numpy or torch (B, T, C)
    """

    if type(Y_true) is not np.ndarray:
        Y_true = torch.from_numpy(Y_true)
    if type(Y_pred) is not np.ndarray:
        Y_pred = torch.from_numpy(Y_pred)

    if avg_out == "epochs":
        dim = 0

    elif avg_out == "times":
        dim = 1
        Y_pred, Y_true = Y_pred[:, start:, :], Y_true[:, start:, :]

    if score == "r":
        cov = (Y_true * Y_pred).mean(dim)
        na, nb = [(i**2).mean(dim)**0.5 for i in [Y_true, Y_pred]]
        norms = na * nb
        R_matrix = cov / norms

    if score == "relativemse":
        Y_err = Y_pred - Y_true
        R_matrix = (Y_err**2).mean(dim) / (Y_true**2).mean(dim)  # rename this score matrix!!!

    if type(R_matrix) is not np.ndarray:
        R_matrix = R_matrix.cpu().numpy()

    return R_matrix