f_beta.py

import numpy as np
from numba import njit, prange


@njit(parallel=True)
def _fast_f_beta_score_by_row(y_pred: np.ndarray, y_true: np.ndarray, beta: float) -> float:
    num_rows: int = y_true.shape[0]
    num_cols: int = y_true.shape[1]
    score: float = 0
    b2 = beta * beta

    for r in prange(num_rows):
        tp: int = 0
        tn: int = 0
        fp: int = 0
        fn: int = 0
        _precision: float = 1.0
        _recall: float = 0.0
        _score: float = 0.0
        for i in prange(num_cols):
            if y_pred[r][i] == 1 and y_true[r][i] == 1:
                tp += 1
            if y_pred[r][i] == 1 and y_true[r][i] == 0:
                fp += 1
            if y_pred[r][i] == 0 and y_true[r][i] == 1:
                fn += 1
            if y_pred[r][i] == 0 and y_true[r][i] == 0:
                tn += 1
        if tp + fp != 0.0:
            _precision = tp / (tp + fp)
        if tp + fn != 0.0:
            _recall = tp / (tp + fn)
        _delim = (b2 * _precision) + _recall
        if _delim != 0.0:
            _score = (1 + b2) * _precision * _recall / _delim
        score = score + _score / num_rows
    return score