o .wiq @s>ddlmZddlZddlmZddlmZddlmZdedee ddfd d Z d edefd d Z d ede defddZ d edefddZdededededef ddZdedededeeeffddZdededededeeeeff ddZ  d'ded eded!dee deeeff d"d#Zd$eddfd%d&ZdS)()OptionalN)Tensor)Literalrank_zero_warn nan_strategynan_replace_valuereturncCsB|dvr td||dkrt|ttfstd|dSdS)NreplacedropzPArgument `nan_strategy` is expected to be one of `['replace', 'drop']`, but got r zlArgument `nan_replace` is expected to be of a type `int` or `float` when `nan_strategy = 'replace`, but got ) ValueError isinstancefloatint)rrrb/home/ubuntu/sommelier/.venv/lib/python3.10/site-packages/torchmetrics/functional/nominal/utils.py_nominal_input_validationsrconfmatcCs,|d|d}}td|||S)zDCompute the expected frequenceis from the provided confusion matrix.rz r, c -> rc)sumtorcheinsum)rmargin_sum_rowsmargin_sum_colsrrr_compute_expected_freqs#srbias_correctioncCst|}|t|j|jd}|dkrtjd|jdS|dkr=|r=||}|}||t dt || 7}t||d|S)zChi-square test of independenc of variables in a confusion matrix table. Adapted from: https://github.com/scipy/scipy/blob/v1.9.2/scipy/stats/contingency.py. rrdeviceg?) rnumelrshapendimrtensorrsignminimum ones_likeabs)rrexpected_freqsdfdiff directionrrr_compute_chi_squared)s "r-cCs,||ddk}|dd|ddkfS)aDrop all rows and columns containing only zeros. Example: >>> from torch import randint >>> from torchmetrics.functional.nominal.utils import _drop_empty_rows_and_cols >>> matrix = randint(10, size=(4, 3)) >>> matrix[1, :] = matrix[:, 1] = 0 >>> matrix tensor([[2, 0, 6], [0, 0, 0], [0, 0, 0], [3, 0, 4]]) >>> _drop_empty_rows_and_cols(matrix) tensor([[2, 6], [3, 4]]) rrN)r)rrrr_drop_empty_rows_and_cols=sr. phi_squarednum_rowsnum_cols confmat_sumcCs0ttjd|jd||d|d|dS)z#Compute bias-corrected Phi Squared.rrr)rmaxr$r)r/r0r1r2rrr_compute_phi_squared_correctedSsr4cCs8||dd|d}||dd|d}||fS)z2Compute bias-corrected number of rows and columns.rr r)r0r1r2rows_correctedcols_correctedrrr _compute_rows_and_cols_corrected`sr7cCs(t||||}t|||\}}|||fS)zBCompute bias-corrected Phi Squared and number of rows and columns.)r4r7)r/r0r1r2phi_squared_correctedr5r6rrr_compute_bias_corrected_valuesgs r9r rpredstargetr cCsD|dkr||||fSt||}||||fS)a0Handle ``NaN`` values in input data. If ``nan_strategy = 'replace'``, all ``NaN`` values are replaced with ``nan_replace_value``. If ``nan_strategy = 'drop'``, all rows containing ``NaN`` in any of two vectors are dropped. Args: preds: 1D tensor of categorical (nominal) data target: 1D tensor of categorical (nominal) data nan_strategy: Indication of whether to replace or drop ``NaN`` values nan_replace_value: Value to replace ``NaN`s when ``nan_strategy = 'replace``` Returns: Updated ``preds`` and ``target`` tensors which contain no ``Nan`` Raises: ValueError: If ``nan_strategy`` is not from ``['replace', 'drop']``. ValueError: If ``nan_strategy = replace`` and ``nan_replace_value`` is not of a type ``int`` or ``float``. r ) nan_to_numr logical_orisnan)r:r;rrrows_contain_nanrrr_handle_nan_in_datapsr@ metric_namecCstd|ddS)NzUnable to compute zG using bias correction. Please consider to set `bias_correction=False`.r)rArrr&_unable_to_use_bias_correction_warnings rB)r r)typingrrrtyping_extensionsrtorchmetrics.utilities.printsrstrrrrboolr-r.rr4tupler7r9r@rBrrrrs\      "