INNER CODE UNIT · Python
_mean_match_multiclass_fast
AnotherSamWilson/miceforest · miceforest/imputation_kernel.py:658
def _mean_match_multiclass_fast(
mean_match_candidates: int,
bachelor_preds: DataFrame,
random_state: np.random.RandomState,
hashed_seeds: Optional[np.ndarray],
):
"""
If mean_match_candidates is 0, choose class with highest probability.
Otherwise, randomly choose class weighted by class probabilities.
Returns a np.ndarray, because these get set to categorical later on.
"""
if mean_match_candidates == 0:
imp_values = np.argmax(bachelor_preds, axis=1)
else:
num_bachelors = bachelor_preds.shape[0]
bachelor_preds = bachelor_preds.cumsum(axis=1).to_numpy()