Source code for nlpatl.sampling.uncertainty.margin

from typing import Tuple
import numpy as np

from nlpatl.sampling import Sampling


class MarginSampling(Sampling):
	"""
		Sampling data points according to the margin confidence. Pick the lowest
			probabilies difference between the highest class and second higest class.

		:param name: Name of this sampling
		:type name: str
    """

	def __init__(self, name: str = 'margin_sampling'):	
		super().__init__(name=name)		

	def sample(self, data: np.ndarray, 
		num_sample: int) -> Tuple[np.ndarray, np.ndarray]:

		num_node = min(num_sample, len(data))

		# Calculate margin difference between first and second highest probabilties
		margin_diffs = np.partition(-data, 1, axis=1)
		margin_diffs = -margin_diffs[:, 0] + margin_diffs[:, 1]

		indices = np.argpartition(margin_diffs, num_node-1)[:num_node]
		return indices, margin_diffs[indices]