RebalanceDatasetTest.java

package water.fvec;

import static org.junit.Assert.assertEquals;

import java.io.File;

import org.junit.Test;

import water.*;

/**
 * Created by tomasnykodym on 4/1/14.
 */
public class RebalanceDatasetTest extends TestUtil {
  @Test public void testProstate(){
    Key hex = Key.make("p.hex");
    Key rebalancedKey = Key.make("rebalanced");
    Key raw = NFSFileVec.make(new File("smalldata/logreg/prostate.csv"));
    try{
      ParseDataset2.parse(hex, new Key[]{raw});
      Frame fr = UKV.get(hex);
      RebalanceDataSet rb = new RebalanceDataSet(fr,rebalancedKey,300);
      H2O.submitTask(rb);
      rb.join();
      Frame rebalanced = UKV.get(rebalancedKey);
      assertEquals(rebalanced.numRows(),fr.numRows());
      assertEquals(rebalanced.anyVec()._espc.length,301);
      for(long l = 0; l < fr.numRows(); ++l)
        for(int i = 0; i < fr.numCols(); ++i)
          assertEquals(fr.vecs()[i].at(l),rebalanced.vecs()[i].at(l),1e-8);
    } finally{
      Frame fr = UKV.get(hex);
      fr.delete();
      fr = UKV.get(rebalancedKey);
      fr.delete();
    }
    checkLeakedKeys();
  }
}