841f5efcf5
Inspired by Crucio's distributional matrix (ICSE 2026). Groups symbols by context similarity instead of first-symbol. New module: bex/distributional.py - extract_contexts(): get (left, right) context pairs - build_distribution_matrix(): symbol × context counts - cluster_symbols(): agglomerative clustering by similarity - split_by_cluster(): split sequences by cluster membership - distributional_split(): drop-in replacement for _split_by_first_symbol CLI: --cluster-method distributional (opt-in, first-symbol is default) Example: 'if' and 'while' both appear before 'return' → first-symbol: 3 groups (if, while, return) → distributional: 2 groups (if/while merged, return) 23 tests pass. Full suite: 257 tests pass.
240 lines
7.1 KiB
Python
240 lines
7.1 KiB
Python
"""Tests for distributional clustering."""
|
|
|
|
import pytest
|
|
from bex.distributional import (
|
|
extract_contexts,
|
|
build_distribution_matrix,
|
|
cosine_similarity,
|
|
jaccard_similarity,
|
|
cluster_symbols,
|
|
split_by_cluster,
|
|
split_by_distributional,
|
|
distributional_split,
|
|
)
|
|
|
|
|
|
class TestExtractContexts:
|
|
"""Test context extraction."""
|
|
|
|
def test_simple_sequence(self):
|
|
seqs = [["if", "return", "if"]]
|
|
ctxs = extract_contexts(seqs)
|
|
|
|
# "if" at position 0: (None, "return")
|
|
# "return" at position 1: ("if", "if")
|
|
# "if" at position 2: ("return", None)
|
|
assert ctxs["if"] == [(None, "return"), ("return", None)]
|
|
assert ctxs["return"] == [("if", "if")]
|
|
|
|
def test_empty_sequence(self):
|
|
seqs = [[]]
|
|
ctxs = extract_contexts(seqs)
|
|
assert ctxs == {}
|
|
|
|
def test_single_symbol(self):
|
|
seqs = [["return"]]
|
|
ctxs = extract_contexts(seqs)
|
|
assert ctxs["return"] == [(None, None)]
|
|
|
|
def test_multiple_sequences(self):
|
|
seqs = [
|
|
["if", "return"], # if at pos 0: (None, return), return at pos 1: (if, None)
|
|
["return", "if"], # return at pos 0: (None, if), if at pos 1: (return, None)
|
|
]
|
|
ctxs = extract_contexts(seqs)
|
|
|
|
# if appears at: pos 0 in seq1, pos 1 in seq2
|
|
assert (None, "return") in ctxs["if"] # from seq1
|
|
assert ("return", None) in ctxs["if"] # from seq2
|
|
|
|
# return appears at: pos 1 in seq1, pos 0 in seq2
|
|
assert ("if", None) in ctxs["return"] # from seq1
|
|
assert (None, "if") in ctxs["return"] # from seq2
|
|
|
|
|
|
class TestBuildDistributionMatrix:
|
|
"""Test distribution matrix construction."""
|
|
|
|
def test_simple(self):
|
|
seqs = [
|
|
["if", "return", "if", "return"],
|
|
["return", "if", "return", "if"],
|
|
]
|
|
syms, ctxs, mat = build_distribution_matrix(seqs, min_occurrences=1)
|
|
|
|
assert len(syms) == 2
|
|
assert "if" in syms
|
|
assert "return" in syms
|
|
assert len(ctxs) > 0
|
|
assert len(mat) == 2
|
|
|
|
def test_filter_rare(self):
|
|
seqs = [
|
|
["if", "return"],
|
|
["rare"], # rare symbol
|
|
]
|
|
syms, ctxs, mat = build_distribution_matrix(seqs, min_occurrences=2)
|
|
|
|
# "rare" should be filtered out
|
|
assert "rare" not in syms
|
|
|
|
def test_empty(self):
|
|
syms, ctxs, mat = build_distribution_matrix([], min_occurrences=1)
|
|
assert syms == []
|
|
|
|
|
|
class TestSimilarity:
|
|
"""Test similarity measures."""
|
|
|
|
def test_cosine_identical(self):
|
|
assert cosine_similarity([1, 2, 3], [1, 2, 3]) == pytest.approx(1.0)
|
|
|
|
def test_cosine_orthogonal(self):
|
|
assert cosine_similarity([1, 0], [0, 1]) == pytest.approx(0.0)
|
|
|
|
def test_cosine_similar(self):
|
|
sim = cosine_similarity([1, 2, 3], [1, 2, 4])
|
|
assert sim > 0.9
|
|
|
|
def test_jaccard_identical(self):
|
|
assert jaccard_similarity({1, 2}, {1, 2}) == 1.0
|
|
|
|
def test_jaccard_disjoint(self):
|
|
assert jaccard_similarity({1}, {2}) == 0.0
|
|
|
|
def test_jaccard_partial(self):
|
|
assert jaccard_similarity({1, 2}, {2, 3}) == pytest.approx(1/3)
|
|
|
|
|
|
class TestClusterSymbols:
|
|
"""Test symbol clustering."""
|
|
|
|
def test_similar_symbols(self):
|
|
# Two symbols with identical context distributions
|
|
symbols = ["a", "b"]
|
|
matrix = [
|
|
[1, 2, 0], # a
|
|
[1, 2, 0], # b (same as a)
|
|
]
|
|
clusters = cluster_symbols(symbols, matrix, threshold=0.5)
|
|
|
|
# Should be in same cluster
|
|
assert clusters["a"] == clusters["b"]
|
|
|
|
def test_dissimilar_symbols(self):
|
|
# Two symbols with different contexts
|
|
symbols = ["a", "b"]
|
|
matrix = [
|
|
[1, 0, 0], # a
|
|
[0, 0, 1], # b (different from a)
|
|
]
|
|
clusters = cluster_symbols(symbols, matrix, threshold=0.5)
|
|
|
|
# Should be in different clusters
|
|
assert clusters["a"] != clusters["b"]
|
|
|
|
def test_empty(self):
|
|
clusters = cluster_symbols([], [], threshold=0.5)
|
|
assert clusters == {}
|
|
|
|
|
|
class TestSplitByCluster:
|
|
"""Test sequence splitting by cluster."""
|
|
|
|
def test_basic(self):
|
|
seqs = [
|
|
["if", "return"],
|
|
["if", "return"],
|
|
["return", "if"],
|
|
["return", "if"],
|
|
]
|
|
clusters = {"if": 0, "return": 1}
|
|
groups = split_by_cluster(seqs, clusters, min_cluster_size=1)
|
|
|
|
assert 0 in groups
|
|
assert 1 in groups
|
|
assert len(groups[0]) == 2
|
|
assert len(groups[1]) == 2
|
|
|
|
def test_empty_sequences(self):
|
|
seqs = [
|
|
[],
|
|
["if", "return"],
|
|
]
|
|
clusters = {"if": 0}
|
|
groups = split_by_cluster(seqs, clusters, min_cluster_size=1)
|
|
|
|
assert "__empty__" in groups
|
|
assert len(groups["__empty__"]) == 1
|
|
|
|
def test_filter_small(self):
|
|
seqs = [
|
|
["if", "return"],
|
|
["return", "if"],
|
|
]
|
|
clusters = {"if": 0, "return": 1}
|
|
groups = split_by_cluster(seqs, clusters, min_cluster_size=2)
|
|
|
|
# Both clusters have only 1 sequence, should be filtered
|
|
assert 0 not in groups
|
|
assert 1 not in groups
|
|
|
|
|
|
class TestDistributionalSplit:
|
|
"""Test high-level distributional split."""
|
|
|
|
def test_similar_first_symbols(self):
|
|
# "if" and "while" both appear before "return"
|
|
# They should cluster together
|
|
seqs = [
|
|
["if", "return", "if", "return"],
|
|
["if", "return"],
|
|
["while", "return", "while", "return"],
|
|
["while", "return"],
|
|
["return", "if"],
|
|
["return", "if"],
|
|
]
|
|
groups = split_by_distributional(seqs, threshold=0.3, min_cluster_size=2)
|
|
|
|
# Should have fewer groups than first-symbol split
|
|
# (if and while should merge)
|
|
assert len(groups) < 6
|
|
|
|
def test_different_first_symbols(self):
|
|
# "if" and "return" have different contexts
|
|
seqs = [
|
|
["if", "return", "if"],
|
|
["if", "return"],
|
|
["return", "if", "return"],
|
|
["return", "if"],
|
|
]
|
|
groups = split_by_distributional(seqs, threshold=0.5, min_cluster_size=2)
|
|
|
|
# Should still have 2 groups
|
|
assert len(groups) >= 2
|
|
|
|
def test_empty(self):
|
|
groups = split_by_distributional([], threshold=0.5)
|
|
assert groups == {}
|
|
|
|
|
|
class TestDropInReplacement:
|
|
"""Test that distributional_split is a drop-in replacement."""
|
|
|
|
def test_same_signature(self):
|
|
# Should work like _split_by_first_symbol but smarter
|
|
seqs = [
|
|
["if", "return", "if"],
|
|
["if", "return"],
|
|
["return", "if"],
|
|
["return", "if"],
|
|
]
|
|
groups = distributional_split(seqs, threshold=0.5)
|
|
|
|
# Should return dict of groups
|
|
assert isinstance(groups, dict)
|
|
assert len(groups) > 0
|
|
|
|
# All sequences should be in some group
|
|
total = sum(len(v) for v in groups.values())
|
|
assert total == len(seqs)
|