Improve functions and add tests for generic cases

logistic.py

@@ -1 +1,28 @@
-# Your code goes here
+from matplotlib import pyplot as plt
+
+
+def logistic_step(x, r):
+    """ Compute the logistic map for a given value of x and r"""
+    constant_val = 1
+    return r * x * (constant_val - x)
+
+def run_iterations(start, r, n_iter):
+    """ Run iterations of logistic function """
+    val = [start]
+    for i_iter in range(n_iter):
+        val.append(logistic_step(val[i_iter], r))
+    return val
+
+if __name__=='__main__':
+    start_vals = [i/10 for i in range(1, 6)]
+    r = 1.2
+    n_iter = 20
+    fig_name = "test_trajectory"
+
+    fig, ax = plt.subplots(figsize=(10, 5))
+    for i_start_val in start_vals:
+        trajectory = run_iterations(i_start_val, r, n_iter)
+        ax.plot(list(range(n_iter+1)), trajectory, label=str(i_start_val))
+    fig.suptitle(f'Logistic function: r={r}, n_iter={n_iter}')
+
+    fig.savefig(fig_name)

logistic_fit.py

@@ -1,6 +1,6 @@
 import numpy as np
 
-from logistic import iterate_f
+from logistic import run_iterations
 
 
 def fit_r(xs):
@@ -23,10 +23,10 @@ def fit_r(xs):
     x0 = xs[0]
     it = len(xs) - 1
 
-    def error(r):
-        return np.linalg.norm(xs - iterate_f(it, x0, r))
+    def compute_error(r):
+        return np.linalg.norm(xs - run_iterations(x0, r, it))
 
     errors = []
     for r in np.linspace(0, 4, 4001):
-        errors.append((r, error(r)))
+        errors.append((r, compute_error(r)))
     return min(errors, key=lambda x: x[1])[0]

test_logistic.py

@@ -1,16 +1,65 @@
 from numpy.testing import assert_allclose
+import pytest
+from logistic import logistic_step, run_iterations
+from logistic_fit import fit_r
+from itertools import product
+import numpy as np
 
-from logistic import f
 
-# Add here your test for the logistic map
+SEED = np.random.randint(0, 2**31)
 
 
-def test_f_corner_cases():
-    # Test cases are (x, r, expected)
-    cases = [
+@pytest.fixture
+def random_state():
+    print(f'Using seed {SEED}')
+    random_state = np.random.RandomState(SEED)
+    return random_state
+
+
+@pytest.mark.parametrize('x, r, expected', [
         (0, 1.1, 0),
         (1, 3.7, 0),
     ]
-    for x, r, expected in cases:
-        result = f(x, r)
-        assert_allclose(result, expected)
+    )
+def test_logistic_step_corner_cases(x, r, expected):
+    result = logistic_step(x, r)
+    assert_allclose(result, expected)
+
+@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_logistic_step_generic_cases(x, r, expected):
+    result = logistic_step(x, r)
+    assert_allclose(result, expected)
+
+@pytest.mark.parametrize('x, r, n_iter, expected', [
+        (0.1, 2.2, 1, [0.1, 0.198]),
+        (0.2, 3.4, 4, [0.2, 0.544, 0.843418, 0.449019, 0.841163]),
+        (0.5, 2, 3, [0.5, 0.5, 0.5, 0.5]),
+    ]
+    )
+def test_logistic_iterations(x, r, n_iter, expected):
+    result = run_iterations(x, r, n_iter)
+    assert_allclose(result, expected, rtol=1e-4)
+
+
+
+
+@pytest.mark.parametrize('x, r, n_iter', [(i/10,f/10,20) for i,f in product(range(1,10,3),range(10,30,4))]
+    )
+def test_logistic_fit_r(x, r, n_iter):
+    result = fit_r(run_iterations(x, r, n_iter))
+    assert_allclose(result, r, rtol=1e-3)
+
+
+def test_logistic_convergence(random_state):
+    r = 1.5
+    n_iter = 100
+    expected = 1/3
+    for _ in range(100):
+        x = random_state.uniform(0.0001, 0.9999)
+        result = run_iterations(x, r, n_iter)
+        assert_allclose(result[-1], expected, rtol=1e-4)