fix test and output

data-8 · Mar 8, 2016 · a7ad94e · a7ad94e
1 parent b143a40
commit a7ad94e
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/datascience/util.py b/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
 
diff --git a/tests/test_util.py b/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))