Implement logistic map function f, generic test, nan test

logistic.py

@@ -1 +1,17 @@
 # Your code goes here
+def f(x,r):
+    """
+    compute logistic map
+    """
+    return r*x*(1-x)
+
+
+def iterate_f(it, x, r):
+    """
+    run f for it iterations
+    """
+    res=[x]
+    for i in range(it):
+        x=f(x,r)
+        res.append(x)
+    return res

test_logistic.py

@@ -1,16 +1,39 @@
 from numpy.testing import assert_allclose
-
-from logistic import f
+import numpy as np
+import pytest
+from logistic import f, iterate_f
 
 # Add here your test for the logistic map
 
+@pytest.mark.parametrize('x, r, expected', [(0, 1.1, 0), (1, 3.7, 0)])
+def test_f_corner_cases(x, r, expected):
+    # Test cases are (x, r, expected)
+    result = f(x, r)
+    assert_allclose(result, expected)
+
 
-def test_f_corner_cases():
+@pytest.mark.parametrize('x, r, expected', [
+        (0.1, 2.2, 0.198),
+        (0.2, 3.4, 0.544),
+        (0.5, 2,   0.5)
+    ])
+def test_f_generic_cases(x,r,expected):
     # Test cases are (x, r, expected)
-    cases = [
-        (0, 1.1, 0),
-        (1, 3.7, 0),
-    ]
-    for x, r, expected in cases:
-        result = f(x, r)
-        assert_allclose(result, expected)
+    result = f(x, r)
+    assert_allclose(result, expected)
+
+
+@pytest.mark.parametrize('x, r, expected', [(0.1, np.nan, np.nan), (np.nan,2, np.nan)])
+def test_f_nan(x, r, expected):
+    # Test cases are (x, r, expected)
+    result = f(x, r)
+    assert_allclose(result, expected)
+
+
+
+@pytest.mark.parametrize('it, x, r, expected', [(1,0.1,2.2,[0.1,0.198]),
+                (4,0.2,3.4,[0.2,0.544,0.843418,0.449019, 0.841163]),
+                (3,0.5,2,[0.5,0.5,0.5,0.5])])
+def test_iterate_f(it, x, r, expected):
+    result = iterate_f(it, x, r)
+    assert_allclose(result, expected, atol=0.00001) #adjusted tolerance