Fixed key error in h2o.balance

tgsmith61591 · Sep 7, 2016 · 8b67ac7 · 8b67ac7
1 parent 3bf6726
commit 8b67ac7
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/skutil/h2o/balance.py b/skutil/h2o/balance.py
@@ -40,7 +40,8 @@ def _validate_x_y_ratio(X, y, ratio):
 	y = _validate_target(y) # cast to string type
 
 	# generate cts. Have to get kludgier in h2o...
-	unq_vals = X[y].unique().as_data_frame(use_pandas=True)[y].values # numpy array of unique vals
+	unq_vals = X[y].unique()
+	unq_vals = unq_vals.as_data_frame(use_pandas=True)[unq_vals.columns[0]].values # numpy array of unique vals
 	unq_cts = dict([(val, X[y][X[y]==val].shape[0]) for val in unq_vals])
 
 	# validate is < max classes

diff --git a/skutil/h2o/tests/test_h2o.py b/skutil/h2o/tests/test_h2o.py
@@ -1192,15 +1192,15 @@ def balance():
 
 				# do a real undersample
 				x = Y[:60, :] # 50 zeros, 10 ones
-				b = UndersamplingClassBalancer(y='species', ratio=0.5).balance(x).as_data_frame(use_pandas=True)
+				b = H2OUndersamplingClassBalancer(target_feature='species', ratio=0.5).balance(x).as_data_frame(use_pandas=True)
 				assert b.shape[0] == 30
 				cts = b.species.value_counts()
 				assert cts[0] == 20
 				assert cts[1] == 10
 
 				# assert oversampling works
 				y = Y[:105, :]
-				d = H2OOversamplingClassBalancer(y='species', ratio=1.0).balance(y).as_data_frame(use_pandas=True)
+				d = H2OOversamplingClassBalancer(target_feature='species', ratio=1.0).balance(y).as_data_frame(use_pandas=True)
 				assert d.shape[0] == 150
 
 				cts= d.species.value_counts()