File size: 1,056 Bytes
7ee20db | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 | from layer_interactions.result import Order2Result
def test_interactions_and_selection():
# D0 = 1.0. First-order deletion costs: [0.1, 0.2, 0.05]
# Pair (0,2) has a positive interaction of +0.3, so after deleting 2,
# deleting 1 is cheaper than deleting 0.
r = Order2Result(
depth=3,
baseline_nll=1.0,
single_nll={0: 1.1, 1: 1.2, 2: 1.05},
pair_nll={
(0, 1): 1.3,
(0, 2): 1.45,
(1, 2): 1.25,
},
)
r.build_interactions().build_greedy_path(max_delete=2)
assert r.delete_order == [2, 1]
assert r.select(1)["retained_layers"] == [0]
def test_roundtrip_dict():
r = Order2Result(
depth=2,
baseline_nll=1.0,
single_nll={0: 1.1, 1: 1.2},
pair_nll={(0, 1): 1.35},
).build_interactions().build_greedy_path(max_delete=1)
rr = Order2Result.from_dict(r.to_dict())
assert rr.depth == 2
assert rr.single_nll == r.single_nll
assert rr.pair_nll == r.pair_nll
assert rr.delete_order == r.delete_order
|