FlowTracker/src/main/java/edu/nju/ics/frontier/learning/TrainAndTestSession.java

package edu.nju.ics.frontier.learning;

import edu.nju.ics.frontier.configuration.Configuration;
import edu.nju.ics.frontier.util.Assertion;

import java.io.File;
import java.util.*;

public abstract class TrainAndTestSession {
    protected String name;
    protected int fold;
    protected double validRate;
    protected boolean isPrintInfo;

    public TrainAndTestSession(String name, int fold, double validRate, boolean isPrintInfo) {
        this.name = name;
        this.fold = fold;
        this.validRate = validRate;
        this.isPrintInfo = isPrintInfo;
    }

    public String getName() {
        return name;
    }

    public int getFold() {
        return fold;
    }

    public void setFold(int fold) {
        this.fold = fold;
    }

    public double getValidRate() {
        return validRate;
    }

    public void setValidRate(double validRate) {
        this.validRate = validRate;
    }

    public boolean isPrintInfo() {
        return isPrintInfo;
    }

    public void setPrintInfo(boolean printInfo) {
        isPrintInfo = printInfo;
    }

    @Override
    public String toString() {
        return this.name;
    }

    public Map<String, Evaluation> crossValidation(String root) throws Exception {
        Configuration.out.println(toString());

        // attempt to load training set, validation set, and test set from local file system
        boolean isLoadDataFromLocalFileSystem;
        if (root != null) {
            isLoadDataFromLocalFileSystem = true;
            for (int foldIndex = 0; foldIndex < this.fold; foldIndex++) {
                File trainFile = new File(root + File.separator + "fold_" + foldIndex + File.separator + "train.json");
                if (!trainFile.exists()) {
                    isLoadDataFromLocalFileSystem = false;
                    break;
                }
                File testFile = new File(root + File.separator + "fold_" + foldIndex + File.separator + "test.json");
                if (!testFile.exists()) {
                    isLoadDataFromLocalFileSystem = false;
                    break;
                }
            }
        } else {
            isLoadDataFromLocalFileSystem = false;
        }

        // container of confusion matrices generated in each fold
        Map<String, List<int[]>> tmp = new HashMap<String, List<int[]>>();
        int[] classes = null;
        if (isLoadDataFromLocalFileSystem) {
            System.out.println("load data from local file system: " + root);
            // cross validation
            for (int foldIndex = 0; foldIndex < this.fold; foldIndex++) {
                Configuration.out.printf(" > %s:fold %d\n", this.name, foldIndex);
                // get training and test sets
                File trainFile = new File(root + File.separator + "fold_" + foldIndex + File.separator + "train.json");
                TimeDatabase trainDb = TimeDatabase.deserialize(trainFile.getAbsolutePath());
                File testFile = new File(root + File.separator + "fold_" + foldIndex + File.separator + "test.json");
                TimeDatabase testDb = TimeDatabase.deserialize(testFile.getAbsolutePath());
                File validFile = new File(root + File.separator + "fold_" + foldIndex + File.separator + "valid.json");
                TimeDatabase validDb = validFile.exists() ? TimeDatabase.deserialize(validFile.getAbsolutePath()) : null;
                // initial classes
                if (classes == null) {
                    classes = trainDb.getClasses();
                }
                // log
                if (this.isPrintInfo) {
                    Configuration.out.println(trainDb.simpleDescribe(String.format("%s:training set", this.name)));
                    Configuration.out.println(validDb == null ? "null" : validDb.simpleDescribe(String.format("%s:validation set", this.name)));
                    Configuration.out.println(testDb.simpleDescribe(String.format("%s:test set", this.name)));
                }
                // do validation
                Map<String, List<int[]>> truePredLabelPairsMap = validation(foldIndex, trainDb, validDb, testDb);
                for (Map.Entry<String, List<int[]>> entry : truePredLabelPairsMap.entrySet()) {
                    String key = entry.getKey();
                    List<int[]> truePredLabelPairs = entry.getValue();
                    List<int[]> result = tmp.get(key);
                    if (result == null) {
                        result = new ArrayList<int[]>();
                    }
                    result.addAll(truePredLabelPairs);
                    tmp.put(key, result);
                }
            }
        } else {
            System.out.println("init data from scratch");
            // load dataset
            TimeDatabase db = DataSource.initDatabase();
            // log
            if (this.isPrintInfo) {
                Configuration.out.println(db.describe(this.name + ":dataset"));
            }
            // split training and test sets
            int tpNum = db.getNumberOfTimePoints();
            Assertion.assertPositive(tpNum);
            int[][][] trainAndTestFolds = DataSource.splitTrainAndTestByKFold(tpNum, this.fold);
            // cross validation
            for (int foldIndex = 0; foldIndex < trainAndTestFolds.length; foldIndex++) {
                Configuration.out.printf(" > %s:fold %d\n", this.name, foldIndex);
                // get training and test sets
                int[] trainIndices = trainAndTestFolds[foldIndex][0];
                int[] testIndices = trainAndTestFolds[foldIndex][1];
                TimeDatabase[] trainValidAndTestDb = prepareTrainValidAndTestSets(db.slice(trainIndices), db.slice(testIndices));
                TimeDatabase trainDb = trainValidAndTestDb[0];
                TimeDatabase validDb = trainValidAndTestDb[1];
                TimeDatabase testDb = trainValidAndTestDb[2];
                // initial classes
                if (classes == null) {
                    classes = trainDb.getClasses();
                }
                // serialize
                TimeDatabase.serialize(trainDb, root + File.separator + "fold_" + foldIndex + File.separator + "train.json");
                TimeDatabase.serialize(testDb, root + File.separator + "fold_" + foldIndex + File.separator + "test.json");
                if (validDb != null) {
                    TimeDatabase.serialize(validDb, root + File.separator + "fold_" + foldIndex + File.separator + "valid.json");
                }
                // log
                if (this.isPrintInfo) {
                    Configuration.out.println(trainDb.simpleDescribe(String.format("%s:training set", this.name)));
                    Configuration.out.println(validDb == null ? "null" : validDb.simpleDescribe(String.format("%s:validation set", this.name)));
                    Configuration.out.println(testDb.simpleDescribe(String.format("%s:test set", this.name)));
                }
                // do validation
                Map<String, List<int[]>> truePredLabelPairsMap = validation(foldIndex, trainDb, validDb, testDb);
                for (Map.Entry<String, List<int[]>> entry : truePredLabelPairsMap.entrySet()) {
                    String key = entry.getKey();
                    List<int[]> truePredLabelPairs = entry.getValue();
                    List<int[]> result = tmp.get(key);
                    if (result == null) {
                        result = new ArrayList<int[]>();
                    }
                    result.addAll(truePredLabelPairs);
                    tmp.put(key, result);
                }
            }
        }

        List<String> keys = new ArrayList<String>(tmp.keySet());
        Collections.sort(keys);
        Map<String, Evaluation> result = new HashMap<String, Evaluation>();
        for (String key : keys) {
            List<int[]> value = tmp.get(key);
            Evaluation evaluation = Evaluation.evaluate(key, classes, value, true);
            result.put(key, evaluation);
        }
        return result;
    }

    protected abstract TimeDatabase[] prepareTrainValidAndTestSets(TimeDatabase trainDb, TimeDatabase testDb);

    protected abstract Map<String, List<int[]>> validation(int foldIndex, TimeDatabase trainDb, TimeDatabase validDb, TimeDatabase testDb) throws Exception;
}