<filename>calibration/binning/general.py
import numpy as np
import anytree


class BinningTree:
    """General structure for binning samples by splitting their indices along
    an axis."""

    def __init__(self, alg):
        self.alg = alg
        self.fitted = False

    def fit(self, probs):
        """Fit binning tree to probabilities."""
        # expand one-dimensional probability vectors
        self.probs = probs if probs.ndim > 1 else np.stack(
            [probs, 1-probs], axis=-1)

        # define root node with all samples and split it with binning algorithm
        self.root = BinningNode(np.arange(self.probs.shape[0]))
        self.alg.split(self.root, self.probs)
        self.fitted = True

        return self

    @property
    def binnumbers(self):
        """Return array with bin numbers of each sample."""

        if not self.fitted:
            raise Exception("BinningTree fit needs to be called first")

        binnumbers = np.empty(self.probs.shape[0], dtype=np.int64)
        for idx, node in enumerate(
                anytree.PreOrderIter(self.root,
                                     filter_=BinningNode.leaf_filter)):
            binnumbers[node.indices] = idx

        return binnumbers

    @property
    def bins(self):
        """Return list with sample indices of each bin."""

        if not self.fitted:
            raise Exception("BinningTree fit needs to be called first")

        bins = []
        for node in anytree.PreOrderIter(self.root,
                                         filter_=BinningNode.leaf_filter):
            bins.append(node.indices)

        return bins

    def bin_data(self, data=None):
        """
        Return list with `data` split up in the bins.

        The data `data` should be an array of size `(N, ...)` where `N`
        is the number of data points to which the binning tree was fit.

        If `data` is `None` (the default), `data` is set to the probabilities
        to which the binning tree was fit.
        """
        if data is None:
            data = self.probs

        return [data[x] for x in self.bins]


class BinningNode(anytree.NodeMixin):
    """Node for saving indices and boundaries."""

    def __init__(self, indices):
        super(BinningNode, self).__init__()
        self.indices = indices

    def _post_attach(self, parent):
        """Clear parent node."""
        try:
            parent.indices = None
        except AttributeError:
            pass

    @staticmethod
    def leaf_filter(x):
        if x.is_leaf:
            try:
                return x.indices.size > 0
            except AttributeError:
                return False

        return False

<filename>calibration/binning/general.py
import numpy as np
import anytree


class BinningTree:
    """General structure for binning samples by splitting their indices along
    an axis."""

    def __init__(self, alg):
        self.alg = alg
        self.fitted = False

    def fit(self, probs):
        """Fit binning tree to probabilities."""
        # expand one-dimensional probability vectors
        self.probs = probs if probs.ndim > 1 else np.stack(
            [probs, 1-probs], axis=-1)

        # define root node with all samples and split it with binning algorithm
        self.root = BinningNode(np.arange(self.probs.shape[0]))
        self.alg.split(self.root, self.probs)
        self.fitted = True

        return self

    @property
    def binnumbers(self):
        """Return array with bin numbers of each sample."""

        if not self.fitted:
            raise Exception("BinningTree fit needs to be called first")

        binnumbers = np.empty(self.probs.shape[0], dtype=np.int64)
        for idx, node in enumerate(
                anytree.PreOrderIter(self.root,
                                     filter_=BinningNode.leaf_filter)):
            binnumbers[node.indices] = idx

        return binnumbers

    @property
    def bins(self):
        """Return list with sample indices of each bin."""

        if not self.fitted:
            raise Exception("BinningTree fit needs to be called first")

        bins = []
        for node in anytree.PreOrderIter(self.root,
                                         filter_=BinningNode.leaf_filter):
            bins.append(node.indices)

        return bins

    def bin_data(self, data=None):
        """
        Return list with `data` split up in the bins.

        The data `data` should be an array of size `(N, ...)` where `N`
        is the number of data points to which the binning tree was fit.

        If `data` is `None` (the default), `data` is set to the probabilities
        to which the binning tree was fit.
        """
        if data is None:
            data = self.probs

        return [data[x] for x in self.bins]


class BinningNode(anytree.NodeMixin):
    """Node for saving indices and boundaries."""

    def __init__(self, indices):
        super(BinningNode, self).__init__()
        self.indices = indices

    def _post_attach(self, parent):
        """Clear parent node."""
        try:
            parent.indices = None
        except AttributeError:
            pass

    @staticmethod
    def leaf_filter(x):
        if x.is_leaf:
            try:
                return x.indices.size > 0
            except AttributeError:
                return False

        return False

General structure for binning samples by splitting their indices along     an axis. Fit binning tree to probabilities. # expand one-dimensional probability vectors # define root node with all samples and split it with binning algorithm Return array with bin numbers of each sample. Return list with sample indices of each bin. Return list with `data` split up in the bins.          The data `data` should be an array of size `(N, ...)` where `N`         is the number of data points to which the binning tree was fit.          If `data` is `None` (the default), `data` is set to the probabilities         to which the binning tree was fit. Node for saving indices and boundaries. Clear parent node.