CelerisLab/tests/unit/test_union_refinement.py

"""CPU tests: union tables vs current FixedPatchSet2D; numpy + CelerisLab on PYTHONPATH.
Optional numerics: ADAPTER_DIR=<adapter bundle dir> [ADAPTER_FILE=adapter.py or
multilevel_numeric_adapter.py (default: whichever exists)]."""
from dataclasses import replace
import importlib.util
import os
import unittest

import numpy as np
from CelerisLab.config import LBMConfig
from CelerisLab.lbm.refinement import (ACTIVE, OUTER_GHOST, FixedPatchSet2D, _middle_config,
                                       _nested_kind, _transfer_footprint)
from CelerisLab.lbm.refinement_tables import blocked_kind

from CelerisLab.lbm.refinement import (FixedPatchSet2D as UnionPatchSet2D,
    _check_parent_pull_closure as check_parent_pull_closure, _dilate as dilate)

def table_sets(pair):
    return (pair.fine_kind.tobytes(), pair.covered.tobytes(),
            frozenset(zip(map(int,pair.ghost_fine),map(int,pair.ghost_coarse))),
            frozenset((int(t),tuple(map(int,c))) for t,c in zip(pair.coalesce_coarse,pair.coalesce_fine)))

CFG = LBMConfig(nx=30, ny=30, viscosity=.1)
A, B, U = (8, 8, 16, 20), (16, 8, 24, 20), (8, 8, 24, 20)


def union(*rects, cfg=CFG):
    return UnionPatchSet2D.from_config(cfg, rects)


class Equivalence(unittest.TestCase):
    def test_abutting_boxes_equal_their_union_rectangle(self):
        reference = table_sets(FixedPatchSet2D.from_config(CFG, [U]))
        for rects in ((A, B), (B, A), (U,)):
            with self.subTest(rects=rects):
                self.assertEqual(table_sets(union(*rects)), reference)

    def test_vertical_abutting(self):
        self.assertEqual(table_sets(union((8, 8, 20, 14), (8, 14, 20, 22))),
                         table_sets(FixedPatchSet2D.from_config(CFG, [(8, 8, 20, 22)])))

    def test_separated_rectangles_reproduce_current_tables(self):
        rects = [(4, 4, 10, 10), (16, 16, 26, 26)]  # separated rectangles; union dependencies determine legality
        self.assertEqual(table_sets(union(*rects)),
                         table_sets(FixedPatchSet2D.from_config(CFG, rects)))

    def test_layout_tiles_are_free_of_rectangle_alignment(self):
        pair = union(A, B)
        support = dilate(pair.fine_kind == ACTIVE, 4)
        reference = FixedPatchSet2D.from_config(CFG, [U]).fine_layout((4, 4))
        np.testing.assert_array_equal(pair.fine_layout((4, 4)).owned, reference.owned)
        layout = pair.fine_layout((6, 10))  # not aligned to the rectangles
        self.assertTrue(layout.owned[support].all())
        kind = blocked_kind(layout, pair.fine_kind)
        owned = layout.owned
        np.testing.assert_array_equal(kind[layout.owner[owned]], pair.fine_kind[owned])


class NewlyLegal(unittest.TestCase):
    def check(self, pair):
        _transfer_footprint(pair)
        check_parent_pull_closure(pair)
        ghosts = pair.ghost_fine.astype(np.intp)
        self.assertFalse((pair.fine_kind.ravel()[ghosts] == ACTIVE).any())
        self.assertFalse(pair.covered.ravel()[pair.ghost_coarse.astype(np.intp)].any())

    def test_small_gaps_and_l_shape(self):
        for rects in (((8, 8, 16, 20), (17, 8, 24, 20)),   # gap 1
                      ((8, 8, 16, 20), (18, 8, 24, 20)),   # gap 2
                      ((8, 8, 20, 14), (8, 14, 14, 22))):  # L-shape
            with self.subTest(rects=rects):
                self.check(union(*rects))

    def test_corner_contact_only_by_dependency_checks(self):
        # Policy: no geometric ban; accepted iff the direction-specific checks pass.
        try:
            pair = union((8, 8, 16, 16), (16, 16, 24, 24))
        except ValueError as error:
            self.assertRegex(str(error), "outer ghost|allocation|INNER|surviving|coalesce")
        else:
            self.check(pair)


class Negatives(unittest.TestCase):
    def test_inputs_are_validated_before_conversion(self):
        for bad in ((8.5, 8, 16, 20), ("8", 8, 16, 20), (True, 8, 16, 20), "abcd", (8, 8, 16)):
            with self.subTest(bad=bad), self.assertRaises(ValueError):
                union(bad, B)
        self.assertEqual(union(tuple(np.int64(v) for v in A), B).bounds, (A, B))
        pair = union(A, B)
        for block in ((4.0, 4), (True, 4), (0, 4), (4,), ("4", 4)):
            with self.subTest(block=block), self.assertRaises(ValueError):
                pair.fine_layout(block)

    def test_missing_donor_at_domain_margin(self):
        with self.assertRaisesRegex(ValueError, "margin"):
            union((1, 8, 10, 20), (10, 8, 20, 20))

    def test_missing_or_duplicate_coalesce_row(self):
        pair = union(A, B)
        with self.assertRaisesRegex(ValueError, "lacks its coalesce row"):
            check_parent_pull_closure(replace(pair, coalesce_coarse=pair.coalesce_coarse[1:],
                                              coalesce_fine=pair.coalesce_fine[1:]))
        doubled = replace(pair, coalesce_coarse=np.r_[pair.coalesce_coarse, pair.coalesce_coarse[:1]],
                          coalesce_fine=np.r_[pair.coalesce_fine, pair.coalesce_fine[:1]])
        with self.assertRaisesRegex(ValueError, "duplicate"):
            check_parent_pull_closure(doubled)

    def test_direction_specific_dependency_on_outer_cell(self):
        pair = union(A, B)
        active = pair.fine_kind == ACTIVE
        kind = pair.fine_kind.copy()
        kind[dilate(active, 2) & ~dilate(active, 1)] = OUTER_GHOST  # thin the INNER band
        with self.assertRaisesRegex(ValueError, "outer ghost"):
            _transfer_footprint(replace(pair, fine_kind=kind))

    def test_nested_union_child_donors_must_be_parent_active(self):
        parent = FixedPatchSet2D.from_config(CFG, [(6, 6, 24, 24)])  # L1 ACTIVE 12..48
        middle = _middle_config(CFG, parent)
        with self.assertRaisesRegex(ValueError, "donors/receivers"):
            _nested_kind(parent, UnionPatchSet2D.from_config(middle, [(12, 20, 30, 40), (30, 20, 40, 40)]))
        child = UnionPatchSet2D.from_config(middle, [(16, 20, 30, 40), (30, 20, 40, 40)])
        reference = FixedPatchSet2D.from_config(middle, [(16, 20, 40, 40)])
        self.assertEqual(table_sets(child), table_sets(reference))
        np.testing.assert_array_equal(_nested_kind(parent, child), _nested_kind(parent, reference))


@unittest.skipUnless(os.environ.get("ADAPTER_DIR"), "set ADAPTER_DIR for the CPU D2Q9 adapter run")
class AdapterNumerics(unittest.TestCase):
    def test_eight_steps_bitwise(self):
        folder = os.environ["ADAPTER_DIR"]
        name = os.environ.get("ADAPTER_FILE") or next(
            n for n in ("multilevel_numeric_adapter.py", "adapter.py")
            if os.path.exists(os.path.join(folder, n)))
        path = os.path.join(folder, name)
        spec = importlib.util.spec_from_file_location("multilevel_numeric_adapter", path)
        adapter = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(adapter)
        results = []
        for pair in (union(A, B), FixedPatchSet2D.from_config(CFG, [U])):
            kinds = [np.where(pair.covered, OUTER_GHOST, ACTIVE).astype(np.uint8), pair.fine_kind]
            taus = [pair.tau_coarse, pair.tau_fine]
            initial = []
            for level, kind in enumerate(kinds):
                ny, nx = kind.shape
                y = (np.arange(ny) + .5) / 2 ** level
                ux = 1e-3 * np.sin(2 * np.pi * y / 30)[:, None] * np.ones((1, nx))
                cx = np.array([0, 1, -1, 0, 0, 1, -1, 1, -1], float)[:, None, None]
                w = np.array([4/9] + [1/9]*4 + [1/36]*4)[:, None, None]
                g = w * (1 + 3*cx*ux + 4.5*(cx*ux)**2 - 1.5*ux*ux)
                g[:, kind != ACTIVE] = np.nan
                initial.append(g)
            run = adapter.make_scheduler(kinds, taus, [pair], initial)
            run.step(8)
            results.append([level.current.copy() for level in run.levels])
        for a, b in zip(*results):
            np.testing.assert_array_equal(a, b)


if __name__ == "__main__":
    unittest.main()