1 files changed, 80 insertions, 0 deletions
diff --git a/tests/test_predicates.py b/tests/test_predicates.py
new file mode 100644
index 0000000..542b67f
--- /dev/null
+++ b/tests/test_predicates.py
@@ -0,0 +1,80 @@
+import unittest
+
+import numpy as np
+
+from tests.helpers import fields, line_world
+from worldhistory.predicates import check, depth, evaluate, parse
+from worldhistory.regions import Regions, region_mask
+
+
+class PredicateTest(unittest.TestCase):
+    def setUp(self):
+        # cells: 0 forest land, 1 desert land by a river, 2 mountain land, 3 shelf sea, 4 deep sea with a vent
+        self.w = line_world(5, ocean=[False, False, False, True, True], holdridge=[20, 23, 12, 0, 0],
+                            landform=[1, 1, 3, 0, 0], river=[False, True, False, False, False],
+                            strahler=[0, 4, 0, 0, 0], elevation_m=[100, 50, 2000, -150, -3000],
+                            vent_potential=[0, 0, 0, 0, 0.8], T_mean=[10, 30, 0, 12, 5], lithology=[1, 3, 4, 0, 0])
+
+    def test_parse(self):
+        self.assertEqual(parse("land"), ("land", []))
+        self.assertEqual(parse("warm:5:20"), ("warm", [5.0, 20.0]))
+        self.assertEqual(parse("lithology:3,4"), ("lithology", [[3, 4]]))
+        self.assertEqual(parse("above:o2_fraction:0.3"), ("above", ["o2_fraction", 0.3]))
+
+    def test_basic(self):
+        e = lambda p: evaluate(p, self.w).tolist()
+        self.assertEqual(e("land"), [1, 1, 1, 0, 0])
+        self.assertEqual(e("forest"), [1, 0, 0, 0, 0])
+        self.assertEqual(e("desert"), [0, 1, 0, 0, 0])
+        self.assertEqual(e("mountain"), [0, 0, 1, 0, 0])
+        self.assertEqual(e("coast"), [0, 0, 1, 0, 0])
+        self.assertEqual(e("shelf"), [0, 0, 0, 1, 0])
+        self.assertEqual(e("deep_sea"), [0, 0, 0, 0, 1])
+        self.assertAlmostEqual(e("water")[1], 1.0)              # 0.4 + 0.15*4 = 1.0
+        self.assertAlmostEqual(e("vent")[4], 0.8)
+        self.assertEqual(e("lithology:3,4"), [0, 1, 1, 0, 0])
+
+    def test_ranges_and_products(self):
+        self.assertEqual(evaluate("warm:5:25", self.w).tolist(), [0.25, 1.0, 0.0, 0.35, 0.0])
+        self.assertEqual(evaluate(["land", "above:T_mean:20"], self.w).tolist(), [0, 1, 0, 0, 0])
+        np.testing.assert_array_equal(depth(self.w.fields), [0, 0, 0, 150, 3000])
+
+    def test_check(self):
+        check(["land", "warm:1:2", "lithology:1"])
+        for bad in ("nosuch", "warm:1", "above:nofield:1"):
+            with self.assertRaises(ValueError):
+                check(bad, field_names=set(fields(1)))
+
+
+class RegionTest(unittest.TestCase):
+    def test_kinds(self):
+        w = line_world(4, ocean=[False, False, False, True], plate=[1, 1, 2, 2], gravity_g=1.0)
+        w.eras["late"] = fields(4, ocean=[False, True, False, True], gravity_g=[1.0, 0.35, 0.35, 0.35])
+        self.assertEqual(region_mask(w, {"kind": "all"}).sum(), 4)
+        self.assertEqual(region_mask(w, {"kind": "plate", "ids": [2]}).tolist(), [False, False, True, True])
+        self.assertEqual(region_mask(w, {"kind": "changed", "era": "late", "fields": ["gravity_g"]}).tolist(),
+                         [False, True, True, True])
+        self.assertEqual(region_mask(w, {"kind": "changed", "era": "late", "fields": ["gravity_g"], "land": "base"})
+                         .tolist(), [False, True, True, False])
+        box = region_mask(w, {"kind": "box", "lat": [-1, 1], "lon": [0.5, 10]})
+        self.assertEqual(box.tolist(), [False, True, True, True])
+        self.assertEqual(region_mask(w, {"kind": "predicate", "p": "land"}).tolist(), [True, True, True, False])
+
+    def test_box_with_predicate_filter(self):
+        w = line_world(4, ocean=[False, False, True, True])
+        m = region_mask(w, {"kind": "box", "lat": [-1, 1], "lon": [0.5, 10], "p": "land"})
+        self.assertEqual(m.tolist(), [False, True, False, False])
+
+    def test_exclude_other_region(self):
+        w = line_world(4)
+        r = Regions(w, {"east": {"kind": "box", "lat": [-1, 1], "lon": [0.5, 10]},
+                        "land_not_east": {"kind": "predicate", "p": "land", "exclude": "east"}})
+        self.assertEqual(r("land_not_east").tolist(), [True, False, False, False])
+
+    def test_registry(self):
+        w = line_world(3)
+        r = Regions(w, {"west": {"kind": "box", "lat": [-1, 1], "lon": [-1, 0.5]}})
+        self.assertEqual(r("west").tolist(), [True, False, False])
+        self.assertTrue(r("all").all())
+        with self.assertRaises(KeyError):
+            r("nowhere")