<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,
                                                     int top){
        return topK(x, classProbEstimator, labelCalibrator, vectorCalibrator,predictionFeatureExtractor,0,Integer.MAX_VALUE,top);
    }


    public static List<Pair<MultiLabel,Double>> topK(Vector x, MultiLabelClassifier.ClassProbEstimator classProbEstimator, LabelCalibrator labelCalibrator,
                                                     VectorCalibrator vectorCalibrator, PredictionFeatureExtractor predictionFeatureExtractor,
                                                     int minSetSize, int maxSetSize,
                                                     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);

                                                     VectorCalibrator vectorCalibrator, PredictionVectorizer predictionVectorizer,
                                                     List<MultiLabel> support,
                                                     int top){
        Set<MultiLabel> supportSets = new HashSet<>(support);
        List<Pair<MultiLabel,Double>> list = new ArrayList<>();
        double[] marginals = labelCalibrator.calibratedClassProbs(cbm.predictClassProbs(x));
        DynamicProgramming dynamicProgramming = new DynamicProgramming(marginals);
        BMDistribution bmDistribution = cbm.computeBM(x,0.001);
        int count = 0;
        for (int i=0;i<Integer.MAX_VALUE;i++) {
            MultiLabel candidate = dynamicProgramming.nextHighestVector();
            if (supportSets.contains(candidate)) {
                Map<MultiLabel,Integer> positionMap = new HashMap<>();
                positionMap.put(candidate,i);
                Vector yFeature = predictionVectorizer.feature(bmDistribution,candidate,marginals, Optional.of(positionMap));
                double pro = vectorCalibrator.calibrate(yFeature);
                list.add(new Pair<>(candidate,pro));
                count += 1;
            }
            if (count==top){
                break;
            }
        }

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


            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> bug fix

<DFF> @@ -40,29 +40,15 @@ public class TopKFinder {
                                                      VectorCalibrator vectorCalibrator, PredictionVectorizer predictionVectorizer,
                                                      List<MultiLabel> support,
                                                      int top){
-        Set<MultiLabel> supportSets = new HashSet<>(support);
-        List<Pair<MultiLabel,Double>> list = new ArrayList<>();
         double[] marginals = labelCalibrator.calibratedClassProbs(cbm.predictClassProbs(x));
-        DynamicProgramming dynamicProgramming = new DynamicProgramming(marginals);
+        Map<MultiLabel,Integer> positionMap = predictionVectorizer.positionMap(marginals);
         BMDistribution bmDistribution = cbm.computeBM(x,0.001);
-        int count = 0;
-        for (int i=0;i<Integer.MAX_VALUE;i++) {
-            MultiLabel candidate = dynamicProgramming.nextHighestVector();
-            if (supportSets.contains(candidate)) {
-                Map<MultiLabel,Integer> positionMap = new HashMap<>();
-                positionMap.put(candidate,i);
-                Vector yFeature = predictionVectorizer.feature(bmDistribution,candidate,marginals, Optional.of(positionMap));
-                double pro = vectorCalibrator.calibrate(yFeature);
-                list.add(new Pair<>(candidate,pro));
-                count += 1;
-            }
-            if (count==top){
-                break;
-            }
-        }
-
         Comparator<Pair<MultiLabel,Double>> comparator = Comparator.comparing(pair->pair.getSecond());
-        return list.stream().sorted(comparator.reversed()).collect(Collectors.toList());
+        return support.stream().map(multiLabel -> {
+            Vector yFeature = predictionVectorizer.feature(bmDistribution,multiLabel,marginals, Optional.of(positionMap));
+            double pro = vectorCalibrator.calibrate(yFeature);
+            return new Pair<>(multiLabel,pro);
+        }).sorted(comparator.reversed()).limit(top).collect(Collectors.toList());
     }
 
 

bug fix