1 files changed, 109 insertions, 0 deletions
diff --git a/tests/test_currents.py b/tests/test_currents.py
new file mode 100644
index 0000000..2b01a43
--- /dev/null
+++ b/tests/test_currents.py
@@ -0,0 +1,109 @@
+import unittest
+
+import numpy as np
+
+from tests.helpers import line_world, make_history, make_race
+from worldhistory.habitat import environment, suitability
+from worldhistory.migration import drift_and_bud, long_jumps
+from worldhistory.predicates import evaluate
+from worldhistory.state import new_state
+
+MP = make_history().migration
+
+
+def east(w, speed=1.0):
+    """Uniform current along the line (the line runs along the equator, eastward = +lon)."""
+    lon = np.radians(w.lon)
+    return speed * np.stack([-np.sin(lon), np.cos(lon), np.zeros(w.n)], 1)
+
+
+def sea_world(n, spacing_km=100.0, current=True, **over):
+    w = line_world(n, spacing_km=spacing_km, ocean=True, elevation_m=-3000.0, **over)
+    if current:
+        w.base["current"] = east(w)
+    return w
+
+
+def swimmer(bias):
+    return make_race(demography={"founder": 2.0, "allee": 1.0, "jump_path": "sea", "current_bias": bias,
+                                 "mobility": 0.2},
+                     habitat={"realm": "sea", "terms": [{"p": "sea", "w": 1.0}]})
+
+
+class SailCostTest(unittest.TestCase):
+    def test_downstream_cheaper_than_still_cheaper_than_upstream(self):
+        w = sea_world(3)
+        c = w.sail_cost(5.0)
+        k_east = list(w.nb[1]).index(2)
+        k_west = list(w.nb[1]).index(0)
+        edge_h = 100.0 * 1000.0 / 3600.0                                # 100 km in hours per (m/s)
+        self.assertAlmostEqual(c[1, k_east], edge_h / 6.0, delta=0.01)  # 5 + 1 m/s
+        self.assertAlmostEqual(c[1, k_west], edge_h / 4.0, delta=0.01)  # 5 − 1 m/s
+        still = sea_world(3, current=False).sail_cost(5.0)
+        self.assertAlmostEqual(still[1, k_east], edge_h / 5.0, delta=0.01)
+        self.assertTrue(np.isinf(c[0, list(w.nb[0]).index(-1)]))        # no neighbour: no edge
+
+    def test_land_edges_are_impassable(self):
+        w = line_world(3, ocean=[True, True, False])
+        c = w.sail_cost(5.0)
+        self.assertTrue(np.isinf(c[1, list(w.nb[1]).index(2)]))
+
+
+class FieldsTest(unittest.TestCase):
+    def test_missing_current_is_still_water(self):
+        w = sea_world(3, current=False)
+        self.assertEqual(w.current.shape, (3, 3))
+        self.assertTrue(np.all(w.current == 0))
+
+    def test_productivity_predicate(self):
+        w = sea_world(3, productivity=np.array([0.0, 0.5, 1.0]))
+        np.testing.assert_allclose(evaluate("productivity", w), [0.0, 0.5, 1.0])
+        np.testing.assert_allclose(evaluate("productivity", sea_world(3)), [0.0, 0.0, 0.0])   # absent → 0
+
+
+class JumpBiasTest(unittest.TestCase):
+    def landings(self, bias, current=True):
+        w = sea_world(400, spacing_km=10.0, current=current)
+        r = swimmer(bias)
+        env, suit = environment(w), suitability(w, r)[None]
+        east_n = west_n = 0
+        for seed in range(150):
+            st = new_state(1, w.n, np.zeros((1, 6)), seed=0)
+            st.P[0, 200] = 1e6
+            long_jumps(st, w, [r], np.full((1, w.n), 1e6), np.ones((1, w.n), bool), np.random.default_rng(seed),
+                       {**MP, "jump_rate": 1.0}, np.ones((1, w.n)), env=env, suit=suit, f_min=0.0)
+            hit = np.flatnonzero(st.P[0] > 0)
+            east_n += int(np.sum(hit > 200))
+            west_n += int(np.sum(hit < 200))
+        return east_n, west_n
+
+    def test_downstream_jumps_dominate_with_bias(self):
+        e, w = self.landings(3.0)
+        self.assertGreater(e, 3 * max(w, 1))
+
+    def test_no_bias_is_symmetric(self):
+        e, w = self.landings(0.0)
+        self.assertLess(abs(e - w), 0.35 * (e + w) + 5)
+
+    def test_no_current_field_ignores_bias(self):
+        e, w = self.landings(3.0, current=False)
+        self.assertLess(abs(e - w), 0.35 * (e + w) + 5)
+
+
+class DriftBiasTest(unittest.TestCase):
+    def drift(self, bias):
+        w = sea_world(3)
+        r = swimmer(bias)
+        st = new_state(1, 3, np.zeros((1, 6)), seed=0)
+        st.P[0] = [100.0, 1000.0, 100.0]                          # neighbours settled: drift, not founding
+        drift_and_bud(st, w, [r], np.full((1, 3), 1e5), np.ones((1, 3), bool), np.zeros((1, 3)),
+                      np.zeros((1, 3)), np.ones((1, 3)), np.random.default_rng(0), MP)
+        return st.P[0]
+
+    def test_drift_follows_the_current(self):
+        P = self.drift(2.0)
+        self.assertGreater(P[2] - 100.0, 3.0 * max(P[0] - 100.0, 1e-9))   # the downstream neighbour gains most
+
+    def test_no_bias_drifts_evenly(self):
+        P = self.drift(0.0)
+        self.assertAlmostEqual(P[0], P[2], delta=1e-6 * P.sum())