tests with start values

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

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/datascience/util.py b/datascience/util.py
@@ -130,6 +130,8 @@ def minimize(f, start=None, **vargs):
         arg_count = f.__code__.co_argcount
         assert arg_count > 0, "Please pass starting values explicitly"
         start = [0] * arg_count
+    if not hasattr(start, '__len__'):
+        start = [start]
 
     @functools.wraps(f)
     def wrapper(args):

diff --git a/tests/test_util.py b/tests/test_util.py
@@ -44,3 +44,5 @@ def test_table_apply():
 def test_minimize():
     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))
+    assert (2 == ds.minimize(lambda x: (x-2)**2, 1)) == True
+    assert [2, 1] == list(ds.minimize(lambda x, y: (x-2)**2 + (y-1)**2, [1, 1]))