""" Stage 2: Concept Discovery Discover semantic concepts that characterize regions of the program space. Three complementary discovery mechanisms: 1. BehavioralConceptDiscovery (TRACED, 2306.07487): Concepts from execution traces — branch coverage, variable quantiles, output signatures 2. SAEConceptDiscovery (CB-SAE, 2512.10805 + DN-CBM, 2407.14499): Train Sparse Autoencoder on LLM hidden states; auto-name neurons as concepts 3. AbstractionConceptDiscovery (LILO, 2310.19791 / ReGAL, 2401.16467): Discover shared program fragments and name them via LLM """ from __future__ import annotations import ast import hashlib import logging from abc import ABC, abstractmethod from dataclasses import dataclass, field from typing import Any, Optional, Callable from collections import defaultdict, Counter import numpy as np from reason_first_program.program_space import Program, ProgramSpace logger = logging.getLogger(__name__) @dataclass class Concept: """ A semantic concept that characterizes a region of the program space. A concept is defined by: - name: human-readable label - description: natural language description - type: 'behavioral', 'representational', or 'structural' - detector: a function Program -> float (concept score in [0, 1]) - vector: optional embedding vector for this concept (for steering) - programs: set of program_ids that exhibit this concept """ name: str description: str type: str # 'behavioral', 'representational', 'structural' detector: Optional[Callable[[Program], float]] = None vector: Optional[np.ndarray] = None programs: set[str] = field(default_factory=set) metadata: dict[str, Any] = field(default_factory=dict) @property def concept_id(self) -> str: return hashlib.sha256( f"{self.name}:{self.type}".encode() ).hexdigest()[:12] def score(self, program: Program) -> float: """Score how strongly a program exhibits this concept.""" if self.detector is not None: return self.detector(program) return 1.0 if program.program_id in self.programs else 0.0 def to_dict(self) -> dict[str, Any]: return { "concept_id": self.concept_id, "name": self.name, "description": self.description, "type": self.type, "n_programs": len(self.programs), "has_vector": self.vector is not None, "metadata": self.metadata, } class ConceptSet: """A collection of concepts discovered for a program space.""" def __init__(self): self._concepts: dict[str, Concept] = {} def add(self, concept: Concept) -> None: self._concepts[concept.concept_id] = concept def get(self, concept_id: str) -> Optional[Concept]: return self._concepts.get(concept_id) def get_by_name(self, name: str) -> Optional[Concept]: for c in self._concepts.values(): if c.name == name: return c return None @property def concepts(self) -> list[Concept]: return list(self._concepts.values()) @property def names(self) -> list[str]: return [c.name for c in self._concepts.values()] def score_program(self, program: Program) -> dict[str, float]: """Score a program against all concepts. Returns {concept_name: score}.""" return {c.name: c.score(program) for c in self._concepts.values()} def score_matrix(self, programs: list[Program]) -> np.ndarray: """ Compute concept score matrix for a list of programs. Returns shape (n_programs, n_concepts). """ matrix = np.zeros((len(programs), len(self._concepts))) for i, program in enumerate(programs): for j, concept in enumerate(self._concepts.values()): matrix[i, j] = concept.score(program) return matrix def concept_lattice(self) -> list[tuple[frozenset[str], frozenset[str]]]: """ Build the Formal Concept Analysis (FCA) lattice. Each node = (extent: set of programs, intent: set of concepts). The Galois connection: extent(B) = {p ∈ P : ∀c ∈ B, c.score(p) > 0.5} intent(A) = {c ∈ C : ∀p ∈ A, c.score(p) > 0.5} """ all_programs = set() for c in self._concepts.values(): all_programs |= c.programs lattice = [] # For each subset of concepts, compute its extent concept_list = list(self._concepts.values()) n = len(concept_list) # For tractability, we compute single-concept and pairwise nodes # Full lattice enumeration is exponential for i in range(n): ci = concept_list[i] intent = frozenset([ci.name]) extent = frozenset(ci.programs) if extent: lattice.append((extent, intent)) for j in range(i + 1, n): cj = concept_list[j] joint_extent = frozenset(ci.programs & cj.programs) joint_intent = frozenset([ci.name, cj.name]) if joint_extent: lattice.append((joint_extent, joint_intent)) return lattice def to_dict(self) -> list[dict[str, Any]]: return [c.to_dict() for c in self._concepts.values()] class ConceptDiscovery(ABC): """Abstract base for concept discovery methods.""" @abstractmethod def discover(self, space: ProgramSpace) -> ConceptSet: """Discover concepts from a program space.""" ... class BehavioralConceptDiscovery(ConceptDiscovery): """ Discover concepts from execution behavior. Based on TRACED (2306.07487): execution traces reveal semantic properties that are independent of syntactic implementation. Concepts discovered: - Algorithmic pattern (iterative, recursive, divide-and-conquer) - Space complexity class (constant, linear, quadratic) - Time complexity class (from execution time quantiles) - Error handling strategy (try/except, assertions, guard clauses) - Data structure usage (list, dict, set, heap, etc.) - Mutation pattern (in-place, copy-on-write, immutable) - Control flow pattern (early return, single return, generator) """ # AST-based detectors for syntactic behavioral proxies SYNTACTIC_DETECTORS = { "uses_recursion": "Implements the solution using recursive function calls", "uses_iteration": "Uses explicit loops (for/while) for the main computation", "uses_list_comprehension": "Uses list/dict/set comprehensions", "uses_generator": "Uses yield or generator expressions", "uses_try_except": "Handles errors with try/except blocks", "uses_assertion": "Uses assert statements for validation", "uses_early_return": "Returns early on edge cases (guard clauses)", "uses_sorting": "Uses sort() or sorted() built-in", "uses_dict": "Uses dictionaries as primary data structure", "uses_set": "Uses sets for membership testing or deduplication", "uses_heap": "Uses heapq module for priority queue operations", "uses_functional": "Uses map/filter/reduce functional patterns", "uses_lambda": "Uses lambda expressions", "uses_mutation": "Modifies data structures in-place", "single_return": "Has exactly one return statement at the end", "uses_defaultdict": "Uses collections.defaultdict", "uses_counter": "Uses collections.Counter", "uses_enumerate": "Uses enumerate for index-value iteration", "uses_zip": "Uses zip to iterate multiple sequences", "uses_slicing": "Uses list/string slicing operations", } def _detect_uses_recursion(self, source: str, func_name: str) -> bool: try: tree = ast.parse(source) for node in ast.walk(tree): if isinstance(node, ast.Call): if isinstance(node.func, ast.Name) and node.func.id == func_name: return True return False except SyntaxError: return False def _detect_uses_iteration(self, source: str) -> bool: try: tree = ast.parse(source) for node in ast.walk(tree): if isinstance(node, (ast.For, ast.While)): return True return False except SyntaxError: return False def _detect_uses_list_comprehension(self, source: str) -> bool: try: tree = ast.parse(source) for node in ast.walk(tree): if isinstance(node, (ast.ListComp, ast.DictComp, ast.SetComp)): return True return False except SyntaxError: return False def _detect_uses_generator(self, source: str) -> bool: try: tree = ast.parse(source) for node in ast.walk(tree): if isinstance(node, (ast.Yield, ast.YieldFrom, ast.GeneratorExp)): return True return False except SyntaxError: return False def _detect_uses_try_except(self, source: str) -> bool: try: tree = ast.parse(source) for node in ast.walk(tree): if isinstance(node, ast.Try): return True return False except SyntaxError: return False def _detect_uses_assertion(self, source: str) -> bool: try: tree = ast.parse(source) for node in ast.walk(tree): if isinstance(node, ast.Assert): return True return False except SyntaxError: return False def _detect_uses_early_return(self, source: str) -> bool: try: tree = ast.parse(source) for node in ast.walk(tree): if isinstance(node, ast.FunctionDef): returns = [n for n in ast.walk(node) if isinstance(n, ast.Return)] if len(returns) > 1: return True return False except SyntaxError: return False def _detect_uses_sorting(self, source: str) -> bool: return "sorted(" in source or ".sort(" in source def _detect_uses_dict(self, source: str) -> bool: try: tree = ast.parse(source) for node in ast.walk(tree): if isinstance(node, ast.Dict): return True if isinstance(node, ast.Call): if isinstance(node.func, ast.Name) and node.func.id == "dict": return True return False except SyntaxError: return False def _detect_uses_set(self, source: str) -> bool: try: tree = ast.parse(source) for node in ast.walk(tree): if isinstance(node, ast.Set): return True if isinstance(node, ast.Call): if isinstance(node.func, ast.Name) and node.func.id == "set": return True return False except SyntaxError: return False def _detect_uses_heap(self, source: str) -> bool: return "heapq" in source or "heappush" in source or "heappop" in source def _detect_uses_functional(self, source: str) -> bool: return any(f in source for f in ["map(", "filter(", "reduce(", "functools"]) def _detect_uses_lambda(self, source: str) -> bool: try: tree = ast.parse(source) for node in ast.walk(tree): if isinstance(node, ast.Lambda): return True return False except SyntaxError: return False def _detect_uses_mutation(self, source: str) -> bool: """Detect in-place mutation: .append, .extend, .pop, .remove, augmented assign.""" mutation_markers = [".append(", ".extend(", ".pop(", ".remove(", ".insert(", ".update(", ".add(", ".discard("] if any(m in source for m in mutation_markers): return True try: tree = ast.parse(source) for node in ast.walk(tree): if isinstance(node, ast.AugAssign): return True return False except SyntaxError: return False def _detect_single_return(self, source: str) -> bool: try: tree = ast.parse(source) for node in ast.walk(tree): if isinstance(node, ast.FunctionDef): returns = [n for n in ast.walk(node) if isinstance(n, ast.Return)] return len(returns) == 1 return False except SyntaxError: return False def _detect_uses_defaultdict(self, source: str) -> bool: return "defaultdict" in source def _detect_uses_counter(self, source: str) -> bool: return "Counter(" in source def _detect_uses_enumerate(self, source: str) -> bool: return "enumerate(" in source def _detect_uses_zip(self, source: str) -> bool: return "zip(" in source def _detect_uses_slicing(self, source: str) -> bool: try: tree = ast.parse(source) for node in ast.walk(tree): if isinstance(node, ast.Subscript): if isinstance(node.slice, ast.Slice): return True return False except SyntaxError: return False def discover(self, space: ProgramSpace) -> ConceptSet: """ Discover behavioral concepts from the program space. Uses AST analysis + execution trace analysis. """ concept_set = ConceptSet() programs = space.valid_programs if space.valid_programs else space.programs func_name = space.stub.name # --- AST-based behavioral concepts --- detector_methods = { "uses_recursion": lambda s: self._detect_uses_recursion(s, func_name), "uses_iteration": self._detect_uses_iteration, "uses_list_comprehension": self._detect_uses_list_comprehension, "uses_generator": self._detect_uses_generator, "uses_try_except": self._detect_uses_try_except, "uses_assertion": self._detect_uses_assertion, "uses_early_return": self._detect_uses_early_return, "uses_sorting": self._detect_uses_sorting, "uses_dict": self._detect_uses_dict, "uses_set": self._detect_uses_set, "uses_heap": self._detect_uses_heap, "uses_functional": self._detect_uses_functional, "uses_lambda": self._detect_uses_lambda, "uses_mutation": self._detect_uses_mutation, "single_return": self._detect_single_return, "uses_defaultdict": self._detect_uses_defaultdict, "uses_counter": self._detect_uses_counter, "uses_enumerate": self._detect_uses_enumerate, "uses_zip": self._detect_uses_zip, "uses_slicing": self._detect_uses_slicing, } for concept_name, detect_fn in detector_methods.items(): description = self.SYNTACTIC_DETECTORS[concept_name] positive_programs: set[str] = set() for program in programs: source = program.full_source try: if detect_fn(source): positive_programs.add(program.program_id) except Exception: continue # Only create concept if it's discriminative (not all or nothing) if 0 < len(positive_programs) < len(programs): def make_detector(pids): def detector(p: Program) -> float: return 1.0 if p.program_id in pids else 0.0 return detector concept = Concept( name=concept_name, description=description, type="behavioral", detector=make_detector(positive_programs), programs=positive_programs, metadata={ "coverage": len(positive_programs) / len(programs), "n_positive": len(positive_programs), "n_total": len(programs), }, ) concept_set.add(concept) # --- Execution-based behavioral concepts --- if programs and programs[0].traces: # Time complexity proxy: fast vs slow programs exec_times = [] for p in programs: avg_time = np.mean([t.execution_time_ms for t in p.traces if t.succeeded]) exec_times.append((p.program_id, avg_time)) if exec_times: times = [t for _, t in exec_times] median_time = np.median(times) fast_programs = {pid for pid, t in exec_times if t < median_time} slow_programs = {pid for pid, t in exec_times if t >= median_time} if fast_programs and slow_programs: concept_set.add(Concept( name="fast_execution", description="Executes faster than median (lower time complexity or better constants)", type="behavioral", programs=fast_programs, metadata={"median_time_ms": float(median_time)}, )) logger.info(f"Discovered {len(concept_set.concepts)} behavioral concepts") return concept_set class SAEConceptDiscovery(ConceptDiscovery): """ Discover concepts via Sparse Autoencoder on LLM hidden states. Based on CB-SAE (2512.10805) and DN-CBM (2407.14499). Pipeline: 1. Pass programs through a code LLM, extract hidden states 2. Train SAE: z = TopK(E_sae(v - b)), v_hat = D_sae @ z + b 3. Auto-name neurons by cosine similarity with concept vocabulary 4. Prune neurons with low interpretability/steerability """ def __init__( self, model_name: str = "deepseek-ai/deepseek-coder-1.3b-instruct", layer: int = 12, expansion_factor: int = 32, top_k: int = 32, n_training_steps: int = 5000, learning_rate: float = 3e-4, device: str = "auto", ): self.model_name = model_name self.layer = layer self.expansion_factor = expansion_factor self.top_k = top_k self.n_training_steps = n_training_steps self.learning_rate = learning_rate self.device = device # SAE parameters (initialized during training) self.encoder = None # E_sae ∈ R^{ω×d} self.decoder = None # D_sae ∈ R^{d×ω} self.bias = None # b ∈ R^d self.concept_names: dict[int, str] = {} self.concept_scores: dict[int, float] = {} # interpretability score # Default concept vocabulary for code CODE_CONCEPT_VOCABULARY = [ "recursion", "iteration", "memoization", "dynamic_programming", "sorting", "searching", "hashing", "tree_traversal", "graph_search", "backtracking", "greedy", "divide_and_conquer", "two_pointers", "sliding_window", "binary_search", "stack_based", "queue_based", "heap_based", "linked_list", "array_manipulation", "string_processing", "bit_manipulation", "mathematical_formula", "brute_force", "optimization", "caching", "lazy_evaluation", "eager_evaluation", "immutable_operations", "mutable_operations", "error_handling", "input_validation", "edge_case_handling", "functional_style", "imperative_style", "object_oriented", "generator_pattern", "accumulator_pattern", "map_reduce_pattern", "fold_pattern", "filter_pattern", "pipeline_pattern", "early_termination", "exhaustive_search", "probabilistic", "space_efficient", "time_efficient", "readable_code", "compact_code", "verbose_code", "pythonic_idiom", ] def extract_hidden_states( self, programs: list[Program] ) -> Optional[np.ndarray]: """ Extract hidden states from a code LLM for each program. Returns shape (n_programs, hidden_dim). """ try: import torch from transformers import AutoTokenizer, AutoModel except ImportError: logger.error("transformers and torch required for SAE concept discovery") return None logger.info(f"Loading model {self.model_name} for hidden state extraction...") tokenizer = AutoTokenizer.from_pretrained(self.model_name, trust_remote_code=True) model = AutoModel.from_pretrained( self.model_name, trust_remote_code=True, output_hidden_states=True, ) if self.device == "auto": model = model.to("cuda" if torch.cuda.is_available() else "cpu") else: model = model.to(self.device) model.eval() hidden_states = [] for program in programs: inputs = tokenizer( program.full_source, return_tensors="pt", truncation=True, max_length=512, ).to(model.device) with torch.no_grad(): outputs = model(**inputs) # Extract hidden state from target layer, mean-pool over tokens layer_output = outputs.hidden_states[self.layer] pooled = layer_output.mean(dim=1).cpu().numpy() hidden_states.append(pooled[0]) return np.array(hidden_states) def train_sae(self, hidden_states: np.ndarray) -> None: """ Train Sparse Autoencoder on hidden states. SAE objective: ||v - v_hat||_2^2 + λ||z||_1 Using TopK activation (2512.10805 recommends TopK over L1). """ try: import torch import torch.nn as nn import torch.optim as optim except ImportError: logger.error("torch required for SAE training") return d = hidden_states.shape[1] omega = d * self.expansion_factor # Initialize SAE parameters device = "cuda" if torch.cuda.is_available() else "cpu" encoder_w = torch.randn(omega, d, device=device) * 0.01 decoder_w = torch.randn(d, omega, device=device) * 0.01 bias = torch.zeros(d, device=device) encoder_w.requires_grad_(True) decoder_w.requires_grad_(True) bias.requires_grad_(True) optimizer = optim.Adam([encoder_w, decoder_w, bias], lr=self.learning_rate) data = torch.tensor(hidden_states, dtype=torch.float32, device=device) n = data.shape[0] logger.info(f"Training SAE: d={d}, ω={omega}, TopK={self.top_k}") for step in range(self.n_training_steps): # Sample batch idx = torch.randint(0, n, (min(64, n),), device=device) v = data[idx] # Encode: z = TopK(E_sae(v - b)) pre_act = (v - bias) @ encoder_w.T # (batch, omega) # TopK activation topk_vals, topk_idx = torch.topk(pre_act, self.top_k, dim=-1) z = torch.zeros_like(pre_act) z.scatter_(1, topk_idx, topk_vals.clamp(min=0)) # Decode: v_hat = D_sae @ z + b v_hat = z @ decoder_w.T + bias # Loss: reconstruction loss = ((v - v_hat) ** 2).mean() optimizer.zero_grad() loss.backward() optimizer.step() if step % 1000 == 0: logger.info(f"SAE step {step}: loss={loss.item():.6f}") # Store trained parameters self.encoder = encoder_w.detach().cpu().numpy() self.decoder = decoder_w.detach().cpu().numpy() self.bias = bias.detach().cpu().numpy() logger.info("SAE training complete") def auto_name_neurons( self, concept_vocabulary: Optional[list[str]] = None, ) -> dict[int, str]: """ Auto-name SAE neurons using cosine similarity with concept vocabulary. Based on CLIP-Dissect approach from CB-SAE (2512.10805). """ if self.decoder is None: logger.error("SAE not trained yet") return {} vocab = concept_vocabulary or self.CODE_CONCEPT_VOCABULARY try: from sentence_transformers import SentenceTransformer except ImportError: logger.warning( "sentence-transformers not available; using hash-based naming" ) # Fallback: name by neuron index omega = self.decoder.shape[1] return {i: f"concept_{i}" for i in range(omega)} # Embed concept vocabulary embedder = SentenceTransformer("all-MiniLM-L6-v2") vocab_embeddings = embedder.encode(vocab) vocab_embeddings = vocab_embeddings / np.linalg.norm( vocab_embeddings, axis=1, keepdims=True ) # For each SAE neuron, compute cosine similarity with vocab # Decoder columns = concept directions in activation space decoder_cols = self.decoder.T # (omega, d) # Project decoder columns to same dimension as vocab embeddings # Use PCA or random projection if dimensions don't match if decoder_cols.shape[1] != vocab_embeddings.shape[1]: # Simple approach: encode the neuron's "meaning" via its top-activating pattern logger.info("Dimension mismatch; using neuron activation patterns for naming") omega = decoder_cols.shape[0] self.concept_names = {i: f"sae_concept_{i}" for i in range(omega)} return self.concept_names decoder_normed = decoder_cols / ( np.linalg.norm(decoder_cols, axis=1, keepdims=True) + 1e-8 ) similarities = decoder_normed @ vocab_embeddings.T # (omega, |vocab|) self.concept_names = {} self.concept_scores = {} for i in range(similarities.shape[0]): best_idx = similarities[i].argmax() best_score = similarities[i, best_idx] self.concept_names[i] = vocab[best_idx] self.concept_scores[i] = float(best_score) return self.concept_names def discover(self, space: ProgramSpace) -> ConceptSet: """Full SAE-based concept discovery pipeline.""" programs = space.valid_programs if space.valid_programs else space.programs if not programs: return ConceptSet() # Step 1: Extract hidden states hidden_states = self.extract_hidden_states(programs) if hidden_states is None: return ConceptSet() # Step 2: Train SAE self.train_sae(hidden_states) # Step 3: Auto-name neurons self.auto_name_neurons() # Step 4: Get SAE activations for all programs if self.encoder is not None and self.bias is not None: v = hidden_states - self.bias pre_act = v @ self.encoder.T # TopK activations = np.zeros_like(pre_act) for i in range(pre_act.shape[0]): topk_idx = np.argsort(pre_act[i])[-self.top_k:] activations[i, topk_idx] = np.maximum(pre_act[i, topk_idx], 0) # Step 5: Build concepts from top neurons concept_set = ConceptSet() omega = activations.shape[1] # Find discriminative neurons (high variance in activation across programs) neuron_variance = activations.var(axis=0) top_neurons = np.argsort(neuron_variance)[-50:] # Top 50 most discriminative for neuron_idx in top_neurons: name = self.concept_names.get(neuron_idx, f"sae_concept_{neuron_idx}") neuron_acts = activations[:, neuron_idx] median_act = np.median(neuron_acts) # Programs where this neuron is active above median active_programs = set() for i, program in enumerate(programs): if neuron_acts[i] > median_act: active_programs.add(program.program_id) if 0 < len(active_programs) < len(programs): concept = Concept( name=f"{name}", description=f"SAE neuron {neuron_idx}: represents '{name}' pattern", type="representational", programs=active_programs, vector=self.decoder[:, neuron_idx] if self.decoder is not None else None, metadata={ "neuron_idx": int(neuron_idx), "variance": float(neuron_variance[neuron_idx]), "interpretability_score": self.concept_scores.get(neuron_idx, 0.0), }, ) concept_set.add(concept) logger.info( f"Discovered {len(concept_set.concepts)} SAE-based concepts " f"from {omega} neurons" ) return concept_set class AbstractionConceptDiscovery(ConceptDiscovery): """ Discover concepts as shared program fragments (abstractions). Based on LILO (2310.19791) / ReGAL (2401.16467). Pipeline: 1. Parse all programs into ASTs 2. Find common subtree patterns (frequent subgraph mining) 3. Name patterns via LLM (AutoDoc) 4. Each named pattern = a structural concept """ def __init__(self, min_frequency: float = 0.1, max_frequency: float = 0.9): """ Args: min_frequency: Minimum fraction of programs that must contain the pattern max_frequency: Maximum fraction (patterns in all programs aren't discriminative) """ self.min_frequency = min_frequency self.max_frequency = max_frequency def _extract_ast_patterns(self, source: str) -> list[str]: """Extract structural patterns from AST.""" patterns = [] try: tree = ast.parse(source) except SyntaxError: return patterns for node in ast.walk(tree): # Pattern: control flow structure if isinstance(node, ast.For): # Check if it's a for-in-range, for-in-enumerate, etc. if isinstance(node.iter, ast.Call): if isinstance(node.iter.func, ast.Name): patterns.append(f"for_{node.iter.func.id}") elif isinstance(node.iter.func, ast.Attribute): patterns.append(f"for_{node.iter.func.attr}") else: patterns.append("for_iter") elif isinstance(node, ast.While): patterns.append("while_loop") elif isinstance(node, ast.If): # Nested if depth depth = 0 n = node while isinstance(n, ast.If) and n.orelse: if len(n.orelse) == 1 and isinstance(n.orelse[0], ast.If): depth += 1 n = n.orelse[0] else: break patterns.append(f"if_chain_{depth}" if depth > 0 else "if_simple") elif isinstance(node, ast.FunctionDef): # Nested function definition if any(isinstance(child, ast.FunctionDef) for child in ast.walk(node)): inner_fns = [ child for child in ast.walk(node) if isinstance(child, ast.FunctionDef) and child is not node ] if inner_fns: patterns.append("nested_function") elif isinstance(node, ast.ListComp): patterns.append("list_comprehension") # Nested comprehension if len(node.generators) > 1: patterns.append("nested_comprehension") elif isinstance(node, ast.DictComp): patterns.append("dict_comprehension") elif isinstance(node, ast.Call): if isinstance(node.func, ast.Name): patterns.append(f"call_{node.func.id}") elif isinstance(node.func, ast.Attribute): patterns.append(f"method_{node.func.attr}") return patterns def discover(self, space: ProgramSpace) -> ConceptSet: """Discover structural abstraction concepts.""" programs = space.valid_programs if space.valid_programs else space.programs if not programs: return ConceptSet() n = len(programs) min_count = max(1, int(n * self.min_frequency)) max_count = int(n * self.max_frequency) # Extract patterns from all programs program_patterns: dict[str, list[str]] = {} pattern_counter: Counter = Counter() for program in programs: patterns = self._extract_ast_patterns(program.full_source) program_patterns[program.program_id] = patterns pattern_counter.update(set(patterns)) # Count each pattern once per program # Find discriminative patterns concept_set = ConceptSet() # Human-readable descriptions for common patterns pattern_descriptions = { "for_range": "Uses for-range loop (index-based iteration)", "for_enumerate": "Uses enumerate for index-value pairs", "for_zip": "Uses zip to iterate multiple sequences in parallel", "for_iter": "Uses direct iterable iteration", "while_loop": "Uses while loop (condition-based iteration)", "if_simple": "Uses simple conditional branching", "if_chain_1": "Uses if-elif chain (multiple conditions)", "if_chain_2": "Uses deeply nested if-elif chains", "nested_function": "Defines helper function(s) inside main function", "list_comprehension": "Uses list comprehension for concise collection building", "dict_comprehension": "Uses dict comprehension", "nested_comprehension": "Uses nested comprehension (multiple generators)", "call_sorted": "Calls sorted() built-in", "call_len": "Calls len() built-in", "call_sum": "Calls sum() built-in", "call_max": "Calls max() built-in", "call_min": "Calls min() built-in", "method_append": "Uses list.append() for accumulation", "method_items": "Uses dict.items() for key-value iteration", "method_keys": "Uses dict.keys()", "method_values": "Uses dict.values()", "method_join": "Uses str.join() for string building", "method_split": "Uses str.split() for tokenization", } for pattern, count in pattern_counter.items(): if min_count <= count <= max_count: # Find programs that have this pattern positive_programs = { pid for pid, patterns in program_patterns.items() if pattern in patterns } description = pattern_descriptions.get( pattern, f"Uses {pattern.replace('_', ' ')} pattern", ) concept = Concept( name=f"struct_{pattern}", description=description, type="structural", programs=positive_programs, metadata={ "frequency": count / n, "n_programs": count, "pattern": pattern, }, ) concept_set.add(concept) logger.info(f"Discovered {len(concept_set.concepts)} structural concepts") return concept_set class UnifiedConceptDiscovery: """ Combines all three concept discovery methods into a single pipeline. Merges behavioral, representational, and structural concepts. """ def __init__( self, behavioral: Optional[BehavioralConceptDiscovery] = None, sae: Optional[SAEConceptDiscovery] = None, abstraction: Optional[AbstractionConceptDiscovery] = None, ): self.behavioral = behavioral or BehavioralConceptDiscovery() self.sae = sae # Optional: requires GPU self.abstraction = abstraction or AbstractionConceptDiscovery() def discover(self, space: ProgramSpace) -> ConceptSet: """Run all discovery methods and merge results.""" merged = ConceptSet() # Always run behavioral discovery logger.info("Running behavioral concept discovery...") behavioral_concepts = self.behavioral.discover(space) for c in behavioral_concepts.concepts: merged.add(c) # Run SAE discovery if available if self.sae is not None: logger.info("Running SAE concept discovery...") try: sae_concepts = self.sae.discover(space) for c in sae_concepts.concepts: merged.add(c) except Exception as e: logger.warning(f"SAE discovery failed: {e}") # Always run abstraction discovery logger.info("Running abstraction concept discovery...") abstraction_concepts = self.abstraction.discover(space) for c in abstraction_concepts.concepts: merged.add(c) logger.info( f"Unified discovery: {len(merged.concepts)} total concepts " f"({len(behavioral_concepts.concepts)} behavioral, " f"{len(abstraction_concepts.concepts)} structural)" ) return merged