ability_mle

def ability_mle(dataset, difficulty, discrimination, no_estimate=np.nan):
    """Estimates the abilities for dichotomous models.

    Estimates the ability parameters (theta) for dichotomous models via
    maximum likelihood estimation.  Response sets with no variance are trimmed
    from evaluation

    Args:
        dataset: [n_items, n_participants] (2d Array) of measured responses
        difficulty: (1d Array) of difficulty parameters for each item
        discrimination: (1d Array) of disrimination parameters for each item
        no_estimate: value to use for response sets that cannot be estimated
                     defaults to numpy.nan, if a number is used then
                     -no_estimate -> 0 and no_estimate -> 1

    Returns:
        abilities: (1d array) estimated abilities
    """
    # Locations where endorsement isn't constant
    mask = np.nanvar(dataset, axis=0) > 0

    # Use only appropriate data
    valid_dataset = dataset[:, mask]

    # Call MAP with uniform distribution
    trimmed_theta = ability_map(valid_dataset, difficulty, discrimination,
                                {'distribution': uniform(-7, 14).pdf})

    # Replace no_estimate values
    thetas = np.full((dataset.shape[1],), np.abs(no_estimate), dtype='float')
    thetas[mask] = trimmed_theta

    # Convert all zeros to negative estimate
    mask2 = ~mask & (dataset.min(axis=0) == 0)
    thetas[mask2] *= -1

    return thetas