<NME> TopKFinder.java
<BEF> package edu.neu.ccs.pyramid.calibration;

import edu.neu.ccs.pyramid.dataset.MultiLabel;
import edu.neu.ccs.pyramid.multilabel_classification.DynamicProgramming;
import edu.neu.ccs.pyramid.multilabel_classification.MultiLabelClassifier;
import edu.neu.ccs.pyramid.multilabel_classification.cbm.BMDistribution;
import edu.neu.ccs.pyramid.multilabel_classification.cbm.CBM;
import edu.neu.ccs.pyramid.util.ArgMax;
import edu.neu.ccs.pyramid.util.ArgSort;
import edu.neu.ccs.pyramid.util.Pair;
import org.apache.mahout.math.Vector;

import java.util.*;
import java.util.stream.Collectors;


/**
 * find top K sets by calibrated probability
 */
public class TopKFinder {

    public static List<Pair<MultiLabel,Double>> topK(Vector x, MultiLabelClassifier.ClassProbEstimator classProbEstimator, LabelCalibrator labelCalibrator,
                                                     VectorCalibrator vectorCalibrator, PredictionFeatureExtractor predictionFeatureExtractor,
        DynamicProgramming dynamicProgramming = new DynamicProgramming(marginals);
        BMDistribution bmDistribution = cbm.computeBM(x,0.001);

        Map<MultiLabel,Integer> positionMap = predictionVectorizer.positionMap(marginals);
        Map<MultiLabel,Double> cdfMap = predictionVectorizer.cdfMap(marginals);


        for (int i=0;i<top;i++) {
                                                     int top){
        double[] uncalibratedMarginals = classProbEstimator.predictClassProbs(x);
        double[] marginals = labelCalibrator.calibratedClassProbs(uncalibratedMarginals);

        DynamicProgramming dynamicProgramming = new DynamicProgramming(marginals);
        //todo better
        List<Pair<MultiLabel,Double>> sparseJoint = dynamicProgramming.topK(50);

        List<MultiLabel> multiLabels = sparseJoint.stream().map(pair->pair.getFirst())
                .filter(candidate->candidate.getNumMatchedLabels() >= minSetSize && candidate.getNumMatchedLabels() <= maxSetSize)
                .collect(Collectors.toList());
        List<Pair<MultiLabel,Double>> candidates = new ArrayList<>();
                                                     List<MultiLabel> support,
                                                     int top){
        double[] marginals = labelCalibrator.calibratedClassProbs(cbm.predictClassProbs(x));
        Map<MultiLabel,Integer> positionMap = predictionVectorizer.positionMap(marginals);
        Map<MultiLabel,Double> cdfMap = predictionVectorizer.cdfMap(marginals);
        BMDistribution bmDistribution = cbm.computeBM(x,0.001);
        Comparator<Pair<MultiLabel,Double>> comparator = Comparator.comparing(pair->pair.getSecond());
        return support.stream().map(multiLabel -> {
        }

        for (MultiLabel candidate: multiLabels){
            PredictionCandidate predictionCandidate = new PredictionCandidate();
            predictionCandidate.x = x;
            predictionCandidate.labelProbs = marginals;
            predictionCandidate.multiLabel = candidate;
            predictionCandidate.sparseJoint = sparseJoint;
            Vector feature = predictionFeatureExtractor.extractFeatures(predictionCandidate);
            double score = vectorCalibrator.calibrate(feature);
            candidates.add(new Pair<>(candidate,score));
        }

        Comparator<Pair<MultiLabel,Double>> comparator = Comparator.comparing(pair->pair.getSecond());
        return candidates.stream().sorted(comparator.reversed()).limit(top).collect(Collectors.toList());
    }

    public static List<Pair<MultiLabel,Double>> topKinSupport(Vector x, MultiLabelClassifier.ClassProbEstimator classProbEstimator, LabelCalibrator labelCalibrator,
                                                              VectorCalibrator vectorCalibrator, PredictionFeatureExtractor predictionFeatureExtractor,
                                                              List<MultiLabel> support,
                                                              int top){
        double[] marginals = labelCalibrator.calibratedClassProbs(classProbEstimator.predictClassProbs(x));


        DynamicProgramming dynamicProgramming = new DynamicProgramming(marginals);
        //todo better
        List<Pair<MultiLabel,Double>> sparseJoint = dynamicProgramming.topK(50);

        List<Pair<MultiLabel,Double>> list = new ArrayList<>();

        for (MultiLabel candidate: support){
            PredictionCandidate predictionCandidate = new PredictionCandidate();
            predictionCandidate.x = x;
            predictionCandidate.labelProbs = marginals;
            predictionCandidate.multiLabel = candidate;
            predictionCandidate.sparseJoint = sparseJoint;

            Vector feature = predictionFeatureExtractor.extractFeatures(predictionCandidate);
            double pro = vectorCalibrator.calibrate(feature);
            list.add(new Pair<>(candidate,pro));
        }

        Comparator<Pair<MultiLabel,Double>> comparator = Comparator.comparing(pair->pair.getSecond());
        return list.stream().sorted(comparator.reversed()).limit(top).collect(Collectors.toList());
    }


}

<MSG> train/predict numCandidates

<DFF> @@ -24,8 +24,8 @@ public class TopKFinder {
         DynamicProgramming dynamicProgramming = new DynamicProgramming(marginals);
         BMDistribution bmDistribution = cbm.computeBM(x,0.001);
 
-        Map<MultiLabel,Integer> positionMap = predictionVectorizer.positionMap(marginals);
-        Map<MultiLabel,Double> cdfMap = predictionVectorizer.cdfMap(marginals);
+        Map<MultiLabel,Integer> positionMap = predictionVectorizer.positionMap(marginals, top);
+        Map<MultiLabel,Double> cdfMap = predictionVectorizer.cdfMap(marginals, top);
 
 
         for (int i=0;i<top;i++) {
@@ -44,8 +44,8 @@ public class TopKFinder {
                                                      List<MultiLabel> support,
                                                      int top){
         double[] marginals = labelCalibrator.calibratedClassProbs(cbm.predictClassProbs(x));
-        Map<MultiLabel,Integer> positionMap = predictionVectorizer.positionMap(marginals);
-        Map<MultiLabel,Double> cdfMap = predictionVectorizer.cdfMap(marginals);
+        Map<MultiLabel,Integer> positionMap = predictionVectorizer.positionMap(marginals, top);
+        Map<MultiLabel,Double> cdfMap = predictionVectorizer.cdfMap(marginals, top);
         BMDistribution bmDistribution = cbm.computeBM(x,0.001);
         Comparator<Pair<MultiLabel,Double>> comparator = Comparator.comparing(pair->pair.getSecond());
         return support.stream().map(multiLabel -> {

train/predict numCandidates