Arcgentica

PROMPT STRUCTURE
[SYSTEM MESSAGE]
  - Role: Expert Python programmer specializing in grid transformations
  - Constraints: Pure function, deterministic, numpy-compatible
  - Output format: Complete Python function

[TASK CONTEXT]
  - Task description with demonstration input/output grid pairs
  - Grid dimensions and color palette information

[PARENT PROGRAM]
  - Current candidate source code being evolved
  - Its fitness score on training examples

[RUNTIME TRACES]  <-- THE KEY INNOVATION
  - For each training example:
    - Input grid (formatted)
    - Expected output grid
    - Actual output produced by parent program
    - Intermediate variable states during execution
    - Any exceptions or errors encountered
    - Diff between expected and actual output (highlighted cells)

[MUTATION INSTRUCTION]
  - Specific guidance on what kind of change to attempt
  - Hints from trace analysis (e.g., "the rotation logic fails on non-square grids")

bash
# Clone the repository
git clone https://github.com/symbolica-ai/arcgentica.git
cd arcgentica

# Install dependencies
pip install -r requirements.txt

# Configure API keys
export ANTHROPIC_API_KEY="your-key-here"
export OPENAI_API_KEY="your-key-here"

# Run on ARC-AGI-2 evaluation set
python run.py --config configs/default.yaml --tasks arc-agi-2

# Run on a single task for testing
python run.py --task-id <task-id> --generations 50 --population-size 20

+=====================================================================+
|                    ARCGENTICA SYSTEM ARCHITECTURE                    |
+=====================================================================+
[ ARC-AGI-2 Task ]
                    (input/output grid pairs)
                              |
                              v
               +---------------------------+
|   TASK PARSER & ENCODER   |
|  - Parse grid pairs       |
|  - Encode as numpy arrays  |
|  - Extract task features   |
+---------------------------+
                              |
                              v
               +---------------------------+
|   INITIAL POPULATION      |
|   SEEDER (Generation 0)   |
|  - LLM zero-shot gen      |
|  - Template programs       |
|  - Random initialization   |
+---------------------------+
                              |
           +==================|==================+
           |        EVOLUTIONARY LOOP              |
           |                  v                  |
           |  +-----------------------------+    |
           |  |    EXECUTION ENGINE         |    |
           |  |  - Sandboxed Python exec    |    |
           |  |  - Timeout enforcement      |    |
           |  |  - Runtime trace capture    |    |
           |  |    * Variable snapshots     |    |
           |  |    * Intermediate grids     |    |
           |  |    * Exception traces       |    |
           |  |    * Control flow path      |    |
           |  +-----------------------------+    |
           |                  |                  |
           |                  v                  |
           |  +-----------------------------+    |
           |  |    FITNESS EVALUATOR        |    |
           |  |  - Grid comparison          |    |
           |  |  - Cell-level accuracy      |    |
           |  |  - Partial credit scoring   |    |
           |  |  - Shape/pattern matching   |    |
           |  +-----------------------------+    |
           |                  |                  |
           |                  v                  |
           |  +-----------------------------+    |
           |  |    PARENT SELECTION         |    |
           |  |  - Tournament selection     |    |
           |  |  - Fitness-proportionate    |    |
           |  |  - Diversity bonus          |    |
           |  |  - Elite preservation       |    |
           |  +-----------------------------+    |
           |                  |                  |
           |                  v                  |
           |  +-----------------------------+    |
           |  |  CONTEXT BUILDER            |    |
           |  |  - Format parent code       |    |
           |  |  - Attach runtime traces    |    |
           |  |  - Compute output diffs     |    |
           |  |  - Add mutation hints       |    |
           |  |  - Compress to fit context  |    |
           |  +-----------------------------+    |
           |                  |                  |
           |                  v                  |
           |  +-----------------------------+    |
           |  |  LLM MUTATION ENGINE        |    |
           |  |  - Send prompt to LLM      |    |
           |  |  - Parse generated code     |    |
           |  |  - Validate syntax          |    |
           |  |  - Inject into population   |    |
           |  +-----------------------------+    |
           |                  |                  |
           |                  v                  |
           |  +-----------------------------+    |
           |  |  POPULATION MANAGER         |    |
           |  |  - Add offspring            |    |
           |  |  - Remove weakest           |    |
           |  |  - Maintain diversity       |    |
           |  |  - Check termination        |    |
           |  +-----------------------------+    |
           |                  |                  |
           |         [Loop if not converged]     |
           +=====================================+
                              |
                              v
               +---------------------------+
|   SOLUTION EXTRACTOR      |
|  - Select best program    |
|  - Apply to test input    |
|  - Format submission      |
+---------------------------+
                              |
                              v
                    [ Predicted Output Grid ]

Task I/O Pairs -----> Execution Engine -----> Runtime Traces
       |                       ^                        |
       |                       |                        v
       |              Candidate Context Builder
       |              Programs                       |
       |                  ^                             v
       |                  |                     LLM Mutation
       |                  |                     Engine
       |                  |                        |    |
       |            Population  <--- New <--------+    |
       |            Manager Offspring          |
       |                  |                             |
       +---> Fitness <----+                 API Calls to
             Evaluator Claude / GPT-4 /
Gemini / etc.

python
class TaskEncoder:
    """Encodes ARC-AGI-2 tasks into internal representations."""

    def encode(self, task_json: dict) -> Task:
        train_pairs = [
            IOPair(
                input=np.array(ex["input"]),
                output=np.array(ex["output"])
            )
            for ex in task_json["train"]
        ]
        test_inputs = [
            np.array(ex["input"])
            for ex in task_json["test"]
        ]
        # Extract meta-features for search strategy hints
        features = self._extract_features(train_pairs)
        return Task(
            train=train_pairs,
            test=test_inputs,
            features=features
        )

    def _extract_features(self, pairs: list[IOPair]) -> TaskFeatures:
        """Extract structural features: grid sizes, color palettes,
        symmetry, size changes, etc."""
        return TaskFeatures(
            input_sizes=[(p.input.shape) for p in pairs],
            output_sizes=[(p.output.shape) for p in pairs],
            colors_used=set().union(*[set(p.input.flatten()) | set(p.output.flatten()) for p in pairs]),
            size_changes=any(p.input.shape != p.output.shape for p in pairs),
            has_symmetry=self._detect_symmetry(pairs),
        )

python
@dataclass
class Individual:
    source_code: str
    fitness: float           # Aggregate fitness across training examples
    per_example_scores: list[float]  # Per-example fitness breakdown
    generation: int
    parent_id: Optional[str]
    mutation_type: str       # e.g., "llm_mutation", "crossover", "seed"
    runtime_traces: list[RuntimeTrace]
    id: str = field(default_factory=lambda: uuid4().hex[:8])

class PopulationManager:
    def __init__(self, max_size: int = 50, elite_fraction: float = 0.1):
        self.population: list[Individual] = []
        self.max_size = max_size
        self.elite_count = max(1, int(max_size * elite_fraction))
        self.archive: list[Individual] = []  # Hall of fame

    def add(self, individual: Individual):
        self.population.append(individual)
        if len(self.population) > self.max_size:
            self._cull()

    def _cull(self):
        """Remove worst individuals while preserving elites and diversity."""
        # Sort by fitness descending
        self.population.sort(key=lambda x: x.fitness, reverse=True)
        # Always keep elites
        elites = self.population[:self.elite_count]
        rest = self.population[self.elite_count:]
        # Apply diversity-aware culling to the rest
        diverse_rest = self._diversity_filter(rest, self.max_size - self.elite_count)
        self.population = elites + diverse_rest

python
class ExecutionEngine:
    """Sandboxed execution with runtime trace capture."""

    def __init__(self, timeout_seconds: float = 5.0):
        self.timeout = timeout_seconds

    def execute_with_trace(
        self, source_code: str, input_grid: np.ndarray
    ) -> ExecutionResult:
        """Execute candidate program and capture runtime trace."""
        trace = RuntimeTrace()
        try:
            # Compile the source code
            compiled = compile(source_code, "", "exec")

            # Set up instrumented namespace
            namespace = {
                "np": np,
                "numpy": np,
                "__trace__": trace,
            }

            # Inject tracing hooks
            instrumented_code = self._instrument(source_code)

            # Execute with timeout
            with timeout(self.timeout):
                exec(compiled, namespace)
                transform_fn = namespace.get("transform")
                if transform_fn is None:
                    return ExecutionResult(
                        success=False,
                        error="No 'transform' function defined",
                        trace=trace
                    )
                output = transform_fn(input_grid.copy())

            trace.final_output = np.array(output)
            return ExecutionResult(
                success=True,
                output=np.array(output),
                trace=trace
            )
        except TimeoutError:
            trace.error = "Execution timed out"
            return ExecutionResult(success=False, error="Timeout", trace=trace)
        except Exception as e:
            trace.error = str(e)
            trace.traceback = traceback.format_exc()
            return ExecutionResult(success=False, error=str(e), trace=trace)

python
@dataclass
class RuntimeTrace:
    """Captures execution dynamics of a candidate program."""
    variable_snapshots: list[dict] = field(default_factory=list)
    intermediate_grids: list[np.ndarray] = field(default_factory=list)
    loop_iterations: dict[str, int] = field(default_factory=dict)
    function_calls: list[str] = field(default_factory=list)
    final_output: Optional[np.ndarray] = None
    error: Optional[str] = None
    traceback: Optional[str] = None
    execution_time_ms: float = 0.0

    def snapshot(self, label: str, local_vars: dict):
        """Capture current variable state at a program point."""
        snapshot = {"label": label}
        for name, value in local_vars.items():
            if isinstance(value, np.ndarray):
                snapshot[name] = value.tolist()
            elif isinstance(value, (int, float, str, bool, list, tuple)):
                snapshot[name] = value
        self.variable_snapshots.append(snapshot)

    def format_for_llm(self, max_tokens: int = 2000) -> str:
        """Format trace for inclusion in LLM prompt."""
        parts = []
        if self.error:
            parts.append(f"ERROR: {self.error}")
            if self.traceback:
                # Include last 5 lines of traceback
                tb_lines = self.traceback.strip().split("\n")
                parts.append("Traceback (last 5 lines):")
                parts.extend(tb_lines[-5:])

        if self.final_output is not None:
            parts.append(f"Output shape: {np.array(self.final_output).shape}")
            parts.append(f"Output:\n{self._format_grid(self.final_output)}")

        if self.variable_snapshots:
            parts.append("Key variable states:")
            # Include most informative snapshots
            for snap in self._select_informative(self.variable_snapshots):
                parts.append(f"  At '{snap.get('label', '?')}': {self._summarize(snap)}")

        result = "\n".join(parts)
        return self._truncate(result, max_tokens)

python
class FitnessEvaluator:
    """Multi-criteria fitness evaluation for ARC programs."""

    def evaluate(self, individual: Individual, task: Task) -> float:
        """Evaluate fitness across all training examples."""
        scores = []
        for pair in task.train:
            result = self.engine.execute_with_trace(
                individual.source_code, pair.input
            )
            if not result.success:
                scores.append(0.0)  # Crashed programs get zero fitness
            else:
                score = self._compare_grids(result.output, pair.output)
                scores.append(score)
            individual.runtime_traces.append(result.trace)

        individual.per_example_scores = scores
        individual.fitness = self._aggregate(scores)
        return individual.fitness

    def _compare_grids(self, actual: np.ndarray, expected: np.ndarray) -> float:
        """Compute similarity between actual and expected grids."""
        if actual.shape != expected.shape:
            # Partial credit for correct dimensions on one axis
            shape_score = sum(
                a == e for a, e in
                zip(actual.shape, expected.shape)
            ) / len(expected.shape) * 0.1
            return shape_score

        # Cell-level accuracy
        correct_cells = np.sum(actual == expected)
        total_cells = expected.size
        cell_accuracy = correct_cells / total_cells

        # Exact match bonus
        if cell_accuracy == 1.0:
            return 1.0

        return cell_accuracy * 0.9  # Reserve 0.1 for exact match bonus

    def _aggregate(self, scores: list[float]) -> float:
        """Aggregate per-example scores. Use minimum to encourage
        solutions that generalize across all examples."""
        if not scores:
            return 0.0
        # Weighted: emphasize worst-case performance
        return 0.4 * min(scores) + 0.6 * (sum(scores) / len(scores))

python
def capture_runtime_context(program: str, input_grid: np.ndarray,
                            expected_output: np.ndarray) -> dict:
    """
    Full runtime context capture pipeline.
    Returns structured context for LLM consumption.
    """
    result = execution_engine.execute_with_trace(program, input_grid)

    context = {
        "program_source": program,
        "input_grid": format_grid(input_grid),
        "expected_output": format_grid(expected_output),
        "execution_successful": result.success,
    }

    if result.success:
        context["actual_output"] = format_grid(result.output)
        context["output_diff"] = compute_grid_diff(result.output, expected_output)
        context["cell_accuracy"] = compute_cell_accuracy(result.output, expected_output)
        context["shape_match"] = result.output.shape == expected_output.shape

        # Key: include intermediate variable states
        context["intermediate_states"] = result.trace.format_for_llm()
    else:
        context["error_message"] = result.trace.error
        context["error_traceback"] = result.trace.traceback
        # Even on error, partial traces may exist
        context["partial_trace"] = result.trace.format_for_llm()

    return context

python
def compute_grid_diff(actual: np.ndarray, expected: np.ndarray) -> str:
    """
    Produce a human-readable diff between actual and expected grids.
    Cells that match are shown normally; mismatches are highlighted.
    """
    if actual.shape != expected.shape:
        return (f"Shape mismatch: got {actual.shape}, "
                f"expected {expected.shape}")

    lines = []
    for i in range(expected.shape[0]):
        row = []
        for j in range(expected.shape[1]):
            if actual[i, j] == expected[i, j]:
                row.append(f" {actual[i,j]} ")
            else:
                row.append(f"[{actual[i,j]}->{expected[i,j]}]")
        lines.append(" ".join(row))
    return "\n".join(lines)

# Example output:
#  0   0   0   0   0
#  0  [3->7]  0   0   0
#  0   0   2   0   0
#  0   0   0  [3->7]  0
#  0   0   0   0   0

python
def construct_mutation_prompt(
    parent: Individual,
    task: Task,
    mutation_type: str,
    traces: list[dict]
) -> str:
    """Build the LLM prompt for a specific mutation type."""

    prompt_parts = [
        SYSTEM_PROMPT,
        format_task_description(task),
        f"\n## Current Program (fitness: {parent.fitness:.3f})\n",
        f"```python\n{parent.source_code}\n```\n",
    ]

    # Add runtime context for each training example
    prompt_parts.append("\n## Runtime Analysis\n")
    for i, (trace, pair) in enumerate(zip(traces, task.train)):
        prompt_parts.append(f"\n### Training Example {i+1}:")
        prompt_parts.append(f"Input ({pair.input.shape}):\n{format_grid(pair.input)}")
        prompt_parts.append(f"Expected output ({pair.output.shape}):\n{format_grid(pair.output)}")

        if trace["execution_successful"]:
            prompt_parts.append(f"Actual output:\n{trace['actual_output']}")
            prompt_parts.append(f"Cell accuracy: {trace['cell_accuracy']:.1%}")
            prompt_parts.append(f"Diff (mismatches in [actual->expected]):\n{trace['output_diff']}")
            if trace.get("intermediate_states"):
                prompt_parts.append(f"Intermediate states:\n{trace['intermediate_states']}")
        else:
            prompt_parts.append(f"EXECUTION FAILED: {trace['error_message']}")
            if trace.get("error_traceback"):
                prompt_parts.append(f"Traceback:\n{trace['error_traceback']}")

    # Mutation-type-specific instructions
    mutation_instructions = {
        "targeted_fix": "Fix the specific errors shown in the diff above. Focus on the cells that are wrong.",
        "structural_rewrite": "The current approach has fundamental issues. Rewrite the core algorithm while keeping any parts that work correctly.",
        "strategy_shift": "Try a completely different approach to solve this task. Analyze the input/output patterns from scratch.",
        "refinement": "The program is nearly correct. Make minimal, precise changes to handle the remaining edge cases.",
    }

    prompt_parts.append(f"\n## Instruction\n{mutation_instructions.get(mutation_type, 'Improve the program.')}")
    prompt_parts.append("\nProvide the complete improved Python function. Start with `def transform(grid):`")

    return "\n".join(prompt_parts)

python
class ParentSelector:
    """Multi-strategy parent selection for evolutionary search."""

    def select(self, population: list[Individual],
               strategy: str = "mixed") -> Individual:
        if strategy == "tournament":
            return self._tournament_select(population, k=3)
        elif strategy == "fitness_proportionate":
            return self._fitness_proportionate(population)
        elif strategy == "diversity_weighted":
            return self._diversity_weighted(population)
        elif strategy == "mixed":
            # Probabilistic mix of strategies
            r = random.random()
            if r < 0.4:
                return self._tournament_select(population, k=3)
            elif r < 0.7:
                return self._fitness_proportionate(population)
            elif r < 0.9:
                return self._diversity_weighted(population)
            else:
                return random.choice(population)  # Pure random for exploration

    def _tournament_select(self, pop: list[Individual], k: int) -> Individual:
        """Select best from k random individuals."""
        contestants = random.sample(pop, min(k, len(pop)))
        return max(contestants, key=lambda x: x.fitness)

    def _fitness_proportionate(self, pop: list[Individual]) -> Individual:
        """Roulette wheel selection proportional to fitness."""
        total = sum(ind.fitness + 0.01 for ind in pop)  # +epsilon to avoid zero
        r = random.uniform(0, total)
        cumulative = 0
        for ind in pop:
            cumulative += ind.fitness + 0.01
            if cumulative >= r:
                return ind
        return pop[-1]

    def _diversity_weighted(self, pop: list[Individual]) -> Individual:
        """Select considering both fitness and uniqueness."""
        scores = []
        for ind in pop:
            uniqueness = self._compute_uniqueness(ind, pop)
            combined = 0.5 * ind.fitness + 0.5 * uniqueness
            scores.append(combined)
        idx = random.choices(range(len(pop)), weights=scores, k=1)[0]
        return pop[idx]

    def _compute_uniqueness(self, ind: Individual,
                            pop: list[Individual]) -> float:
        """Measure how different this individual is from the population."""
        if len(pop) <= 1:
            return 1.0
        similarities = []
        for other in pop:
            if other.id == ind.id:
                continue
            sim = self._code_similarity(ind.source_code, other.source_code)
            similarities.append(sim)
        avg_similarity = sum(similarities) / len(similarities)
        return 1.0 - avg_similarity

Gen 0: [Seed1] [Seed2] [Seed3] ... [SeedN]   (LLM zero-shot generation)
         |       |       |           |
         v       v       v           v
Gen 1: [Eval]  [Eval]  [Eval]  ... [Eval]    (Fitness evaluation)
         |       |       |           |
         v       v       v           v
       [Select parents] --> [Mutate via LLM] --> [New offspring]
         |                                         |
         v                                         v
Gen 2: [Elites preserved] + [New offspring] = [New population]
         |
         v
       [Repeat until convergence or budget exhaustion]

python
def code_similarity(code_a: str, code_b: str) -> float:
    """Compute similarity between two programs using multiple metrics."""

    # 1. Token-level Jaccard similarity
    tokens_a = set(tokenize(code_a))
    tokens_b = set(tokenize(code_b))
    jaccard = len(tokens_a & tokens_b) / max(len(tokens_a | tokens_b), 1)

    # 2. AST structural similarity
    try:
        ast_a = ast.parse(code_a)
        ast_b = ast.parse(code_b)
        ast_sim = ast_structural_similarity(ast_a, ast_b)
    except SyntaxError:
        ast_sim = jaccard  # Fallback

    # 3. Behavioral similarity (same outputs on same inputs?)
    # This is captured separately in the fitness evaluator

    return 0.4 * jaccard + 0.6 * ast_sim

python
def behavioral_fingerprint(individual: Individual, task: Task) -> tuple:
    """Create a hashable fingerprint of program behavior."""
    outputs = []
    for pair in task.train:
        result = execute(individual.source_code, pair.input)
        if result.success:
            outputs.append(tuple(result.output.flatten()))
        else:
            outputs.append(("ERROR", result.trace.error))
    return tuple(outputs)

python
class LLMOrchestrator:
    """Manages concurrent LLM calls for mutation."""

    def __init__(self, config: OrchestratorConfig):
        self.clients = {
            "anthropic": AsyncAnthropic(api_key=config.anthropic_key),
            "openai": AsyncOpenAI(api_key=config.openai_key),
        }
        self.semaphore = asyncio.Semaphore(config.max_concurrent_calls)
        self.rate_limiter = TokenBucketRateLimiter(
            tokens_per_second=config.max_tokens_per_second
        )

    async def generate_mutations(
        self,
        parents: list[Individual],
        task: Task,
        num_mutations: int
    ) -> list[Individual]:
        """Generate multiple mutations concurrently."""
        tasks = []
        for _ in range(num_mutations):
            parent = self.selector.select(parents)
            mutation_type = self._choose_mutation_type(parent)
            model = self._choose_model(mutation_type)
            tasks.append(
                self._generate_single_mutation(parent, task, mutation_type, model)
            )
        results = await asyncio.gather(*tasks, return_exceptions=True)
        return [r for r in results if isinstance(r, Individual)]

    async def _generate_single_mutation(
        self, parent: Individual, task: Task,
        mutation_type: str, model: str
    ) -> Individual:
        async with self.semaphore:
            await self.rate_limiter.acquire()
            prompt = construct_mutation_prompt(parent, task, mutation_type,
                                               parent.runtime_traces)
            response = await self._call_llm(model, prompt)
            source_code = extract_python_code(response)
            return Individual(
                source_code=source_code,
                fitness=0.0,  # Will be evaluated later
                generation=parent.generation + 1,
                parent_id=parent.id,
                mutation_type=mutation_type,
                runtime_traces=[],
            )

    def _choose_model(self, mutation_type: str) -> str:
        """Select model based on mutation type and budget."""
        if mutation_type == "strategy_shift":
            return "claude-3-5-sonnet"  # Strong reasoning for novel approaches
        elif mutation_type == "targeted_fix":
            return "gpt-4o"  # Fast and precise for small fixes
        else:
            return random.choice(["claude-3-5-sonnet", "gpt-4o"])

python
class AdaptiveSearchStrategy:
    """Dynamically select search strategy based on population state."""

    def __init__(self):
        self.stagnation_counter = 0
        self.best_fitness_history = []

    def select_strategy(self, population: list[Individual]) -> str:
        best_fitness = max(ind.fitness for ind in population)
        self.best_fitness_history.append(best_fitness)

        # Detect stagnation
        if len(self.best_fitness_history) >= 5:
            recent = self.best_fitness_history[-5:]
            if max(recent) - min(recent) < 0.01:
                self.stagnation_counter += 1
            else:
                self.stagnation_counter = 0

        # Strategy selection logic
        if best_fitness > 0.95:
            return "intensification"  # Almost there, focus!
        elif best_fitness > 0.7:
            return "exploitation"     # Good programs exist, refine them
        elif self.stagnation_counter > 3:
            if best_fitness < 0.1:
                return "restart"      # Nothing works, start over
            else:
                return "exploration"  # Stuck, try new approaches
        else:
            return "balanced"         # Default mixed strategy

python
def composite_fitness(
    program: str,
    task: Task,
    engine: ExecutionEngine,
    weights: tuple = (0.4, 0.3, 0.2, 0.1)
) -> float:
    """
    Composite fitness function with gradient signal for partial solutions.

    Components:
      1. Worst-case cell accuracy (min across examples)
      2. Average cell accuracy
      3. Shape correctness
      4. Execution success
    """
    w_min, w_avg, w_shape, w_exec = weights

    cell_accuracies = []
    shape_matches = []
    exec_successes = []

    for pair in task.train:
        result = engine.execute_with_trace(program, pair.input)

        if not result.success:
            cell_accuracies.append(0.0)
            shape_matches.append(0.0)
            exec_successes.append(0.0)
            continue

        exec_successes.append(1.0)

        if result.output.shape != pair.output.shape:
            shape_matches.append(0.0)
            # Can still compute partial cell accuracy via padding/cropping
            cell_acc = _padded_cell_accuracy(result.output, pair.output)
            cell_accuracies.append(cell_acc * 0.5)  # Penalty for wrong shape
        else:
            shape_matches.append(1.0)
            cell_acc = np.mean(result.output == pair.output)
            cell_accuracies.append(float(cell_acc))

    if not cell_accuracies:
        return 0.0

    fitness = (
        w_min * min(cell_accuracies) +
        w_avg * (sum(cell_accuracies) / len(cell_accuracies)) +
        w_shape * (sum(shape_matches) / len(shape_matches)) +
        w_exec * (sum(exec_successes) / len(exec_successes))
    )

    # Bonus for perfect solution
    if all(ca == 1.0 for ca in cell_accuracies):
        fitness = 1.0

    return fitness

text
SYSTEM PROMPT:
You are an expert Python programmer specializing in grid transformation problems.
You write clean, efficient numpy code.

Rules:
1. Write a single function: def transform(grid: np.ndarray) -> np.ndarray
2. Use only numpy operations and Python standard library
3. The function must be deterministic (no randomness)
4. Handle grids of any reasonable size
5. Do not use external libraries beyond numpy
6. Return a numpy array with integer values (0-9 representing colors)
7. Analyze the runtime traces carefully to understand what the current
   program does wrong, then fix it precisely

python
def format_grid(grid: np.ndarray) -> str:
    """Format a grid for LLM consumption with color names."""
    COLOR_MAP = {
        0: ".", 1: "B", 2: "R", 3: "G", 4: "Y",
        5: "A", 6: "M", 7: "O", 8: "C", 9: "W"
    }
    lines = []
    for row in grid:
        line = " ".join(COLOR_MAP.get(int(c), str(c)) for c in row)
        lines.append(line)
    return "\n".join(lines)

python
MUTATION_TEMPLATES = {
    "targeted_fix": """
The program above is close to correct but has specific errors.
Look at the diff between actual and expected outputs:
- Cells marked [actual->expected] show mismatches
- Focus on WHY these specific cells are wrong
- Make the minimal change needed to fix these errors
- Do not change parts of the code that produce correct results
""",
    "structural_rewrite": """
The current approach has fundamental limitations.
Analyze the input/output pairs to understand the transformation rule:
- What patterns exist in the inputs?
- How do the outputs relate to the inputs?
- What spatial/color operations are involved?
Rewrite the program using a more appropriate algorithm.
""",
    "strategy_shift": """
Forget the current program entirely.
Study the input/output examples carefully and describe the transformation:
1. What objects/patterns appear in the input?
2. How are they transformed in the output?
3. What rules govern the transformation?
Then write a completely new program implementing your analysis.
""",
    "refinement": """
The program is nearly correct (fitness: {fitness:.1%}).
It fails on specific edge cases. Analyze the runtime traces to find:
- Which examples fail and why
- Whether the issue is boundary conditions, special cases, or off-by-one errors
Make precise, minimal fixes. Do not restructure working code.
"""
}

python
class CostController:
    """Manage LLM API spending across the search."""

    def __init__(self, config: CostConfig):
        self.per_task_budget = config.per_task_budget  # e.g., $5.00
        self.global_budget = config.global_budget      # e.g., $500.00
        self.spent_per_task: dict[str, float] = {}
        self.total_spent = 0.0
        # Model costs per 1M tokens (approximate)
        self.model_costs = {
            "claude-3-5-sonnet": {"input": 3.0, "output": 15.0},
            "claude-3-opus":     {"input": 15.0, "output": 75.0},
            "gpt-4o":            {"input": 2.5, "output": 10.0},
            "gpt-4o-mini":       {"input": 0.15, "output": 0.60},
        }

    def can_spend(self, task_id: str, estimated_tokens: int,
                  model: str) -> bool:
        """Check if we can afford this API call."""
        cost = self._estimate_cost(estimated_tokens, model)
        task_spent = self.spent_per_task.get(task_id, 0.0)
        return (
            task_spent + cost <= self.per_task_budget and
            self.total_spent + cost <= self.global_budget
        )

    def record(self, task_id: str, input_tokens: int,
               output_tokens: int, model: str):
        """Record actual cost after API call."""
        rates = self.model_costs.get(model, {"input": 5.0, "output": 15.0})
        cost = (
            input_tokens * rates["input"] / 1_000_000 +
            output_tokens * rates["output"] / 1_000_000
        )
        self.spent_per_task[task_id] = self.spent_per_task.get(task_id, 0.0) + cost
        self.total_spent += cost

    def suggest_model(self, task_id: str, remaining_generations: int) -> str:
        """Suggest a model based on remaining budget."""
        remaining = self.per_task_budget - self.spent_per_task.get(task_id, 0.0)
        avg_cost_per_gen = self.spent_per_task.get(task_id, 1.0) / max(1, ...)
        if remaining < avg_cost_per_gen * remaining_generations:
            return "gpt-4o-mini"  # Switch to cheaper model
        return "claude-3-5-sonnet"  # Use best model

python
import numpy as np
from typing import Callable

# Eq. 2: Composite fitness function
def fitness(
    program_fn: Callable,
    examples: list[tuple[np.ndarray, np.ndarray]],
    w1: float = 0.4,
    w2: float = 0.6
) -> float:
    """
    F(p, T) = w1 * min_i f(p, x_i, y_i) + w2 * mean_i f(p, x_i, y_i)

    where f(p, x, y) = |{(i,j) : p(x)[i,j] == y[i,j]}| / |y|
    """
    scores = []
    for x, y in examples:
        try:
            pred = program_fn(x.copy())
            if pred.shape != y.shape:
                scores.append(0.0)
            else:
                scores.append(float(np.mean(pred == y)))
        except Exception:
            scores.append(0.0)

    if not scores:
        return 0.0
    return w1 * min(scores) + w2 * np.mean(scores)


# Eq. 6: Expected fitness improvement per generation
def expected_improvement(
    population_fitness: list[float],
    selection_intensity: float = 1.5,  # ~ top 20% selection
    runtime_context_quality: float = 0.8,  # 0-1, how good are traces
) -> float:
    """
    E[Delta F] = alpha * sigma_F * I / sqrt(N)

    Parameters:
      population_fitness: list of fitness values for current population
      selection_intensity: alpha, depends on selection pressure
      runtime_context_quality: I, depends on trace informativeness
    """
    sigma_f = np.std(population_fitness)
    n = len(population_fitness)
    delta_f = selection_intensity * sigma_f * runtime_context_quality / np.sqrt(n)
    return delta_f


# Behavioral diversity metric for population
def population_diversity(
    behavioral_fingerprints: list[tuple]
) -> float:
    """
    Diversity = |unique fingerprints| / |population|

    Measures fraction of behaviorally distinct programs.
    """
    unique = len(set(behavioral_fingerprints))
    total = len(behavioral_fingerprints)
    return unique / total if total > 0 else 0.0


# Information gain from runtime context
def runtime_context_information_gain(
    accuracy_with_context: list[float],
    accuracy_without_context: list[float]
) -> float:
    """
    Approximate information gain from runtime traces.

    IG = E[F_with_context] - E[F_without_context]

    Based on empirical measurements of mutation quality with
    and without runtime traces.
    """
    mean_with = np.mean(accuracy_with_context)
    mean_without = np.mean(accuracy_without_context)
    return mean_with - mean_without


# Grid cell-level accuracy (core fitness primitive)
def cell_accuracy(actual: np.ndarray, expected: np.ndarray) -> float:
    """
    f(p, x, y) = (1/|y|) * sum_{i,j} 1[p(x)[i,j] == y[i,j]]

    Returns fraction of correctly predicted cells.
    Handles shape mismatches with zero score.
    """
    if actual.shape != expected.shape:
        return 0.0
    return float(np.sum(actual == expected)) / expected.size


# Tournament selection probability
def tournament_selection_prob(
    fitness_rank: int,
    population_size: int,
    tournament_size: int
) -> float:
    """
    P(rank k selected) = (1 - ((N-k)/N)^T) - (1 - ((N-k+1)/N)^T)

    where N = population_size, k = rank (1=best), T = tournament_size.
    Approximation for large N.
    """
    N = population_size
    T = tournament_size
    k = fitness_rank
    # Probability of being the best in a tournament of size T
    # when your rank is k out of N
    p_not_selected = ((k - 1) / N) ** T
    p_not_selected_or_better = (k / N) ** T
    return p_not_selected_or_better - p_not_selected

python
class PopulationMemory:
    """Long-term memory across evolutionary generations."""

    def __init__(self):
        self.archive: list[Individual] = []  # Best individuals ever seen
        self.strategy_log: list[dict] = []   # What strategies worked
        self.stagnation_history: list[int] = []  # Track stagnation events

    def update(self, population: list[Individual], generation: int):
        best = max(population, key=lambda x: x.fitness)
        if not self.archive or best.fitness > self.archive[-1].fitness:
            self.archive.append(best)
            self.strategy_log.append({
                "generation": generation,
                "fitness": best.fitness,
                "mutation_type": best.mutation_type,
                "parent_fitness": self._get_parent_fitness(best),
            })

    def get_effective_strategies(self) -> list[str]:
        """Return mutation types that historically produced improvements."""
        improvements = [
            s for s in self.strategy_log
            if s["fitness"] > s.get("parent_fitness", 0)
        ]
        return [s["mutation_type"] for s in improvements]