fix test and output

Author: papajohn

datascience/util.py

@@ -136,5 +136,9 @@ def wrapper(args):
         return f(*args)
 
     result = optimize.minimize(wrapper, start, **vargs)
-    return np.round(result.x, 7)
+    rounded = np.round(result.x, 7)
+    if len(start) == 1:
+        return rounded.item(0)
+    else:
+        return rounded
 

tests/test_util.py

@@ -42,5 +42,5 @@ def test_table_apply():
     assert all(newtab['b'] == tab['b'] + 1)
 
 def test_minimize():
-    assert 2 == ds.minimize(lambda x: (x-2)**2)
+    assert (2 == ds.minimize(lambda x: (x-2)**2)) == True
     assert [2, 1] == list(ds.minimize(lambda x, y: (x-2)**2 + (y-1)**2))