Merge code for test/train set into dataset object

Author: oshando

swan-pipeline/src/main/java/de/fraunhofer/iem/swan/features/FeatureSet.java

@@ -4,12 +4,11 @@
 import de.fraunhofer.iem.swan.data.Category;
 import de.fraunhofer.iem.swan.data.Method;
 import de.fraunhofer.iem.swan.features.code.CodeFeatureHandler;
-import de.fraunhofer.iem.swan.features.code.soot.SourceFileLoader;
 import de.fraunhofer.iem.swan.features.code.type.IFeature;
 import de.fraunhofer.iem.swan.features.doc.DocFeatureHandler;
 import de.fraunhofer.iem.swan.features.doc.manual.IDocFeature;
 import de.fraunhofer.iem.swan.features.doc.nlp.AnnotatedMethod;
-import de.fraunhofer.iem.swan.io.dataset.SrmList;
+import de.fraunhofer.iem.swan.io.dataset.Dataset;
 import de.fraunhofer.iem.swan.model.ModelEvaluator;
 import weka.core.Attribute;
 import weka.core.DenseInstance;
@@ -24,9 +23,8 @@ abstract class FeatureSet {
     protected Map<IFeature, Attribute> codeAttributes;
     protected final HashMap<String, Integer> instanceMap;
     protected final SwanOptions options;
-    protected SrmList trainData;
+    protected Dataset dataset;
     protected CodeFeatureHandler codeFeatureHandler;
-    protected SourceFileLoader testData;
     protected DocFeatureHandler docFeatureHandler;
     protected HashMap<String, Instances> instances;
     protected ModelEvaluator.Toolkit toolkit;
@@ -58,11 +56,10 @@ public static FeatureSet.Type getValue(String value) {
         }
     }
 
-    public FeatureSet(SrmList trainData, SourceFileLoader testData, SwanOptions options, ModelEvaluator.Toolkit toolkit) {
+    public FeatureSet(Dataset dataset, SwanOptions options, ModelEvaluator.Toolkit toolkit) {
         this.instanceMap = new HashMap<>();
         this.options = options;
-        this.trainData = trainData;
-        this.testData = testData;
+        this.dataset = dataset;
         this.toolkit = toolkit;
         instances = new HashMap<>();
     }
@@ -84,13 +81,13 @@ public List<FeatureSet.Type> initializeFeatures() {
                     break;
                 case DOC_MANUAL:
 
-                    docFeatureHandler = new DocFeatureHandler(trainData.getMethods());
+                    docFeatureHandler = new DocFeatureHandler(dataset.getTrainMethods());
                     docFeatureHandler.initialiseManualFeatureSet();
                     docFeatureHandler.evaluateManualFeatureData();
                     break;
                 case DOC_AUTO:
 
-                    docFeatureHandler = new DocFeatureHandler(trainData.getMethods());
+                    docFeatureHandler = new DocFeatureHandler(dataset.getTrainMethods());
                     docFeatureHandler.initialiseAutomaticFeatureSet();
                     docFeatureHandler.evaluateAutomaticFeatureData();
                     break;

swan-pipeline/src/main/java/de/fraunhofer/iem/swan/features/FeatureSetSelector.java

@@ -1,22 +1,21 @@
 package de.fraunhofer.iem.swan.features;
 
 import de.fraunhofer.iem.swan.cli.SwanOptions;
-import de.fraunhofer.iem.swan.features.code.soot.SourceFileLoader;
-import de.fraunhofer.iem.swan.io.dataset.SrmList;
+import de.fraunhofer.iem.swan.io.dataset.Dataset;
 import de.fraunhofer.iem.swan.model.ModelEvaluator;
 
 public class FeatureSetSelector {
 
-    public IFeatureSet select(SrmList trainData, SourceFileLoader testData, SwanOptions options) {
+    public IFeatureSet select(Dataset dataset, SwanOptions options) {
 
         switch (ModelEvaluator.Toolkit.valueOf(options.getToolkit().toUpperCase())) {
 
             case MEKA:
-                MekaFeatureSet mekaFeatureSet = new MekaFeatureSet(trainData, testData, options);
+                MekaFeatureSet mekaFeatureSet = new MekaFeatureSet(dataset, options);
                 mekaFeatureSet.createFeatures();
                 return mekaFeatureSet;
             case WEKA:
-                WekaFeatureSet wekaFeatureSet = new WekaFeatureSet(trainData, testData, options);
+                WekaFeatureSet wekaFeatureSet = new WekaFeatureSet(dataset, options);
                 wekaFeatureSet.createFeatures();
                 return wekaFeatureSet;
         }

swan-pipeline/src/main/java/de/fraunhofer/iem/swan/features/MekaFeatureSet.java

@@ -3,8 +3,7 @@
 import de.fraunhofer.iem.swan.cli.SwanOptions;
 import de.fraunhofer.iem.swan.data.Category;
 import de.fraunhofer.iem.swan.data.Method;
-import de.fraunhofer.iem.swan.features.code.soot.SourceFileLoader;
-import de.fraunhofer.iem.swan.io.dataset.SrmList;
+import de.fraunhofer.iem.swan.io.dataset.Dataset;
 import de.fraunhofer.iem.swan.model.ModelEvaluator;
 import de.fraunhofer.iem.swan.util.Util;
 import meka.filters.unsupervised.attribute.MekaClassAttributes;
@@ -16,8 +15,8 @@
 
 public class MekaFeatureSet extends FeatureSet implements IFeatureSet {
 
-    public MekaFeatureSet(SrmList trainData, SourceFileLoader testData, SwanOptions options) {
-        super(trainData, testData, options, ModelEvaluator.Toolkit.MEKA);
+    public MekaFeatureSet(Dataset dataset, SwanOptions options) {
+        super(dataset, options, ModelEvaluator.Toolkit.MEKA);
     }
 
     /**
@@ -28,15 +27,13 @@ public void createFeatures() {
         List<FeatureSet.Type> featureSets = initializeFeatures();
 
         //Create and set attributes for the train instances
-        ArrayList<Attribute> trainAttributes = createAttributes(getCategories(options.getAllClasses()), trainData.getMethods(), featureSets);
-        Instances trainInstances = createInstances(featureSets, trainAttributes, trainData.getMethods(), getCategories(options.getAllClasses()), "train-instances");
+        ArrayList<Attribute> trainAttributes = createAttributes(getCategories(options.getAllClasses()), dataset.getTrainMethods(), featureSets);
+        Instances trainInstances = createInstances(featureSets, trainAttributes, dataset.getTrainMethods(), getCategories(options.getAllClasses()), "train-instances");
         this.instances.put("train", convertToMekaInstances(trainInstances));
 
-
-        System.out.println(testData.getMethods().size());
         //Create and set attributes for the test instances.
-        ArrayList<Attribute> testAttributes = createAttributes(getCategories(options.getAllClasses()), testData.getMethods(), featureSets);
-        Instances testInstances = createInstances(featureSets, testAttributes, testData.getMethods(), getCategories(options.getAllClasses()), "test-instances");
+        ArrayList<Attribute> testAttributes = createAttributes(getCategories(options.getAllClasses()), dataset.getTestMethods(), featureSets);
+        Instances testInstances = createInstances(featureSets, testAttributes, dataset.getTestMethods(), getCategories(options.getAllClasses()), "test-instances");
         this.instances.put("test", convertToMekaInstances(testInstances));
     }
 

swan-pipeline/src/main/java/de/fraunhofer/iem/swan/features/WekaFeatureSet.java

@@ -3,8 +3,7 @@
 import de.fraunhofer.iem.swan.cli.SwanOptions;
 import de.fraunhofer.iem.swan.data.Category;
 import de.fraunhofer.iem.swan.data.Method;
-import de.fraunhofer.iem.swan.features.code.soot.SourceFileLoader;
-import de.fraunhofer.iem.swan.io.dataset.SrmList;
+import de.fraunhofer.iem.swan.io.dataset.Dataset;
 import de.fraunhofer.iem.swan.model.ModelEvaluator;
 import de.fraunhofer.iem.swan.util.Util;
 import weka.core.Attribute;
@@ -18,8 +17,8 @@
  */
 public class WekaFeatureSet extends FeatureSet implements IFeatureSet {
 
-    public WekaFeatureSet(SrmList trainData, SourceFileLoader testData, SwanOptions options) {
-        super(trainData, testData, options, ModelEvaluator.Toolkit.WEKA);
+    public WekaFeatureSet(Dataset dataset, SwanOptions options) {
+        super(dataset, options, ModelEvaluator.Toolkit.WEKA);
     }
 
     /**
@@ -32,10 +31,10 @@ public void createFeatures() {
         for (Category category : options.getAllClasses().stream().map(Category::fromText).collect(Collectors.toList())) {
 
             //Create and set attributes for the train instances
-            ArrayList<Attribute> trainAttributes = createAttributes(category, trainData.getMethods(), featureSets);
+            ArrayList<Attribute> trainAttributes = createAttributes(category, dataset.getTrainMethods(), featureSets);
 
             String instanceName = category.getId().toLowerCase() + "-train-instances";
-            Instances trainInstances = createInstances(featureSets, trainAttributes, trainData.getMethods(), Collections.singleton(category), instanceName);
+            Instances trainInstances = createInstances(featureSets, trainAttributes, dataset.getTrainMethods(), Collections.singleton(category), instanceName);
             this.instances.put(category.getId().toLowerCase(), trainInstances);
             Util.exportInstancesToArff(trainInstances);
 

swan-pipeline/src/main/java/de/fraunhofer/iem/swan/features/code/soot/SourceFileLoader.java

@@ -0,0 +1,35 @@
+package de.fraunhofer.iem.swan.io.dataset;
+
+import de.fraunhofer.iem.swan.data.Method;
+
+import java.util.Set;
+
+public class Dataset {
+
+    private SrmList train;
+    private SrmList test;
+
+    public SrmList getTrain() {
+        return train;
+    }
+
+    public void setTrain(SrmList train) {
+        this.train = train;
+    }
+
+    public SrmList getTest() {
+        return test;
+    }
+
+    public void setTest(SrmList test) {
+        this.test = test;
+    }
+
+    public Set<Method> getTrainMethods() {
+        return train.getMethods();
+    }
+
+    public Set<Method> getTestMethods() {
+        return test.getMethods();
+    }
+}

swan-pipeline/src/main/java/de/fraunhofer/iem/swan/io/dataset/Dataset.java

@@ -2,6 +2,10 @@
 
 import com.fasterxml.jackson.annotation.JsonIgnore;
 import de.fraunhofer.iem.swan.data.Method;
+import de.fraunhofer.iem.swan.util.Util;
+
+import java.io.IOException;
+import java.util.HashSet;
 import java.util.Set;
 import java.util.stream.Collectors;
 
@@ -11,8 +15,27 @@
 public class SrmList {
 
     private Set<Method> methods;
+    @JsonIgnore
+    private Set<String> testClasses;
 
     public SrmList() {
+        testClasses = new HashSet<>();
+    }
+
+    public SrmList(String sourceFileDir) {
+        try {
+            testClasses = Util.getAllClassesFromDirectory(sourceFileDir);
+        } catch (IOException e) {
+            e.printStackTrace();
+        }
+        methods = new HashSet<>();
+    }
+
+
+    public void load(final Set<Method> trainingSet) {
+
+        Util.createSubclassAnnotations(methods, "classpath");
+        methods = Util.sanityCheck(methods, trainingSet);
     }
 
     public SrmList(Set<Method> methodList) {
@@ -27,6 +50,14 @@ public void setMethods(Set<Method> methods) {
         this.methods = methods;
     }
 
+    public Set<String> getTestClasses() {
+        return testClasses;
+    }
+
+    public void setTestClasses(Set<String> testClasses) {
+        this.testClasses = testClasses;
+    }
+
     public void removeUnclassifiedMethods() {
 
         methods = methods.stream().filter(m -> m.getAllCategories().size() > 0)