1 files changed, 81 insertions, 0 deletions
diff --git a/tests/test_climate_ocean.py b/tests/test_climate_ocean.py
new file mode 100644
index 0000000..8132c6e
--- /dev/null
+++ b/tests/test_climate_ocean.py
@@ -0,0 +1,81 @@
+import unittest
+from pathlib import Path
+
+import numpy as np
+
+from mapgen import climate as CL
+from tests.helpers import make_ctx
+
+GOLDEN = Path(__file__).parent / "data" / "climate_golden_r2.npz"
+NEW = ("current", "current_speed", "sst", "upwelling", "productivity")
+
+
+def golden_ctx(cfg=None):
+    ctx = make_ctx(2, cfg=cfg)
+    g = ctx.grid
+    land = (np.abs(g.lat) < 40) & (np.abs(g.lon) < 50)
+    ctx.data["elevation_eroded_m"] = np.where(land, np.where(np.abs(g.lon) < 10, 2500.0, 300.0),
+                                              -4000.0).astype(np.float32)
+    return ctx, land
+
+
+class OceanDisabledTest(unittest.TestCase):
+    def test_disabled_reproduces_pre_ocean_climate(self):
+        ctx, land = golden_ctx({"ocean": {"enabled": False}})
+        out = CL.run(ctx)
+        ref = np.load(GOLDEN)
+        for k in ref.files:
+            np.testing.assert_array_equal(out[k], ref[k], err_msg=k)
+        self.assertTrue(np.all(out["current"] == 0) and np.all(out["productivity"] == 0))
+        np.testing.assert_array_equal(out["sst"], out["T_mean"].astype(np.float32))
+
+
+class OceanCoupledTest(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.ctx, cls.land = golden_ctx()
+        cls.out = CL.run(cls.ctx)
+
+    def test_new_fields(self):
+        g, out = self.ctx.grid, self.out
+        for k in NEW:
+            self.assertIn(k, out)
+            self.assertEqual(out[k].dtype, np.float32, k)
+            self.assertTrue(np.all(np.isfinite(out[k])), k)
+        self.assertEqual(out["current"].shape, (g.n, 3))
+        self.assertTrue(np.all(out["current"][self.land] == 0))
+        np.testing.assert_allclose(out["current_speed"], np.linalg.norm(out["current"], axis=1), rtol=1e-5)
+        sea = ~self.land
+        self.assertTrue(np.all((out["sst"][sea] >= -1.8) & (out["sst"][sea] < 40)))            # sea water freezes at −1.8 °C
+
+    def test_currents_change_the_climate(self):
+        ref = np.load(GOLDEN)
+        self.assertGreater(np.abs(self.out["T_mean"] - ref["T_mean"]).max(), 0.5)   # the old gyre rule is gone
+
+    def test_winds_come_from_current_free_temperatures(self):
+        ctx, _ = golden_ctx({"ocean": {"relax_days": 30.0}})                      # other SST, same winds
+        np.testing.assert_array_equal(CL.run(ctx)["wind_jun"], self.out["wind_jun"])
+
+    def test_unknown_ocean_key_is_an_error(self):
+        from mapgen.config import ConfigError
+        ctx, _ = golden_ctx({"ocean": {"frictoin_days": 3.0}})
+        with self.assertRaises(ConfigError):
+            CL.run(ctx)
+
+    def test_all_land_world(self):
+        ctx = make_ctx(2)
+        ctx.data["elevation_eroded_m"] = np.full(ctx.grid.n, 500.0, np.float32)
+        out = CL.run(ctx)
+        self.assertTrue(np.all(out["current"] == 0))
+
+    def test_locked_world_has_still_ocean(self):
+        ctx, land = golden_ctx({"climate": {"lock": True}})
+        out = CL.run(ctx)
+        self.assertTrue(np.all(out["current"] == 0) and np.all(out["upwelling"] == 0))
+        np.testing.assert_array_equal(out["sst"], out["T_mean"].astype(np.float32))
+
+    def test_upwelling_and_productivity_filled(self):
+        sea = ~self.land
+        self.assertGreater(np.abs(self.out["upwelling"][sea]).max(), 1.0)
+        self.assertGreater(self.out["productivity"][sea].max(), 0.05)
+        self.assertTrue(np.all(self.out["productivity"][self.land] == 0))