Merge pull request #3 from tylerbessire/codex/continue-implementation-for-sota-model

Implement numerical features and episodic memory

Author: tylerbessire

arc_solver/features.py

@@ -18,51 +18,57 @@
 
 def extract_task_features(train_pairs: List[Tuple[Array, Array]]) -> Dict[str, Any]:
     """Extract a comprehensive feature vector from training pairs.
-    
+
     These features capture task-level properties that can help predict which
     DSL operations are likely to be relevant for solving the task.
     """
+
+    # Ensure arrays are integer typed for canonicalisation
     try:
-        train_pairs = [
-            (canonicalize_D4(inp), canonicalize_D4(out))
+        original_pairs = [
+            (np.asarray(inp, dtype=int), np.asarray(out, dtype=int))
             for inp, out in train_pairs
         ]
-    except TypeError as exc:
+        canonical_pairs = [
+            (canonicalize_D4(inp), canonicalize_D4(out))
+            for inp, out in original_pairs
+        ]
+    except Exception as exc:
         raise ValueError(f"invalid grid in train_pairs: {exc}") from exc
 
     features: Dict[str, Any] = {}
-    
-    # Basic grid statistics
-    input_shapes = [inp.shape for inp, _ in train_pairs]
-    output_shapes = [out.shape for _, out in train_pairs]
-    
+
+    # Basic grid statistics using original shapes
+    input_shapes = [inp.shape for inp, _ in original_pairs]
+    output_shapes = [out.shape for _, out in original_pairs]
+
     features.update({
-        'num_train_pairs': len(train_pairs),
+        'num_train_pairs': len(original_pairs),
         'input_height_mean': np.mean([s[0] for s in input_shapes]),
         'input_width_mean': np.mean([s[1] for s in input_shapes]),
         'output_height_mean': np.mean([s[0] for s in output_shapes]),
         'output_width_mean': np.mean([s[1] for s in output_shapes]),
-        'shape_preserved': all(inp.shape == out.shape for inp, out in train_pairs),
+        'shape_preserved': all(inp.shape == out.shape for inp, out in original_pairs),
         'size_ratio_mean': np.mean([
             (out.shape[0] * out.shape[1]) / (inp.shape[0] * inp.shape[1])
-            for inp, out in train_pairs
+            for inp, out in original_pairs
         ]),
     })
-    
-    # Color analysis
-    input_colors = []
-    output_colors = []
-    color_mappings = []
-    
-    for inp, out in train_pairs:
+
+    # Color analysis on canonical pairs
+    input_colors: List[int] = []
+    output_colors: List[int] = []
+    color_mappings: List[int] = []
+
+    for inp, out in canonical_pairs:
         inp_hist = histogram(inp)
         out_hist = histogram(out)
         input_colors.append(len(inp_hist))
         output_colors.append(len(out_hist))
-        
+
         # Try to detect color mappings
         if inp.shape == out.shape:
-            mapping = {}
+            mapping: Dict[int, int] = {}
             valid_mapping = True
             for i_val, o_val in zip(inp.flatten(), out.flatten()):
                 if i_val in mapping and mapping[i_val] != o_val:
@@ -71,49 +77,51 @@ def extract_task_features(train_pairs: List[Tuple[Array, Array]]) -> Dict[str, A
                 mapping[i_val] = o_val
             if valid_mapping:
                 color_mappings.append(len(mapping))
-    
+
     features.update({
         'input_colors_mean': np.mean(input_colors),
         'output_colors_mean': np.mean(output_colors),
-        'background_color_consistent': len(set(bg_color(inp) for inp, _ in train_pairs)) == 1,
+        'background_color_consistent': len(set(bg_color(inp) for inp, _ in canonical_pairs)) == 1,
         'has_color_mapping': len(color_mappings) > 0,
         'color_mapping_size': np.mean(color_mappings) if color_mappings else 0,
     })
-    
-    # Object analysis
-    input_obj_counts = []
-    output_obj_counts = []
-    
-    for inp, out in train_pairs:
+
+    # Object analysis on canonical pairs
+    input_obj_counts: List[int] = []
+    output_obj_counts: List[int] = []
+
+    for inp, out in canonical_pairs:
         inp_objects = connected_components(inp)
         out_objects = connected_components(out)
         input_obj_counts.append(len(inp_objects))
         output_obj_counts.append(len(out_objects))
-    
+
     features.update({
         'input_objects_mean': np.mean(input_obj_counts),
         'output_objects_mean': np.mean(output_obj_counts),
-        'object_count_preserved': np.mean([
+        'object_count_preserved': all(
             len(connected_components(inp)) == len(connected_components(out))
-            for inp, out in train_pairs
-        ]),
+            for inp, out in canonical_pairs
+        ),
     })
-    
-    # Transformation hints
+
+    # Transformation hints from original pairs
     features.update({
-        'likely_rotation': _detect_rotation_patterns(train_pairs),
-        'likely_reflection': _detect_reflection_patterns(train_pairs),
-        'likely_translation': _detect_translation_patterns(train_pairs),
-        'likely_recolor': _detect_recolor_patterns(train_pairs),
-        'likely_crop': _detect_crop_patterns(train_pairs),
-        'likely_pad': _detect_pad_patterns(train_pairs),
+        'likely_rotation': _detect_rotation_patterns(original_pairs),
+        'likely_reflection': _detect_reflection_patterns(original_pairs),
+        'likely_translation': _detect_translation_patterns(original_pairs),
+        'likely_recolor': _detect_recolor_patterns(original_pairs),
+        'likely_crop': _detect_crop_patterns(original_pairs),
+        'likely_pad': _detect_pad_patterns(original_pairs),
     })
-    
+
     return features
 
 
 def _detect_rotation_patterns(train_pairs: List[Tuple[Array, Array]]) -> float:
     """Detect if rotation transformations are likely."""
+    if not train_pairs:
+        return 0.0
     rotation_score = 0.0
     for inp, out in train_pairs:
         if inp.shape[0] == inp.shape[1] and out.shape[0] == out.shape[1]:
@@ -127,6 +135,8 @@ def _detect_rotation_patterns(train_pairs: List[Tuple[Array, Array]]) -> float:
 
 def _detect_reflection_patterns(train_pairs: List[Tuple[Array, Array]]) -> float:
     """Detect if reflection transformations are likely."""
+    if not train_pairs:
+        return 0.0
     reflection_score = 0.0
     for inp, out in train_pairs:
         if inp.shape == out.shape:
@@ -139,7 +149,7 @@ def _detect_reflection_patterns(train_pairs: List[Tuple[Array, Array]]) -> float
 
 def _detect_translation_patterns(train_pairs: List[Tuple[Array, Array]]) -> float:
     """Detect if translation transformations are likely."""
-    if not all(inp.shape == out.shape for inp, out in train_pairs):
+    if not train_pairs or not all(inp.shape == out.shape for inp, out in train_pairs):
         return 0.0
 
     translation_score = 0.0
@@ -156,6 +166,8 @@ def _detect_translation_patterns(train_pairs: List[Tuple[Array, Array]]) -> floa
 
 def _detect_recolor_patterns(train_pairs: List[Tuple[Array, Array]]) -> float:
     """Detect if recoloring transformations are likely."""
+    if not train_pairs:
+        return 0.0
     recolor_score = 0.0
     for inp, out in train_pairs:
         if inp.shape == out.shape:
@@ -174,6 +186,8 @@ def _detect_recolor_patterns(train_pairs: List[Tuple[Array, Array]]) -> float:
 
 def _detect_crop_patterns(train_pairs: List[Tuple[Array, Array]]) -> float:
     """Detect if cropping transformations are likely."""
+    if not train_pairs:
+        return 0.0
     crop_score = 0.0
     for inp, out in train_pairs:
         if (inp.shape[0] > out.shape[0] or inp.shape[1] > out.shape[1]):
@@ -183,6 +197,8 @@ def _detect_crop_patterns(train_pairs: List[Tuple[Array, Array]]) -> float:
 
 def _detect_pad_patterns(train_pairs: List[Tuple[Array, Array]]) -> float:
     """Detect if padding transformations are likely."""
+    if not train_pairs:
+        return 0.0
     pad_score = 0.0
     for inp, out in train_pairs:
         if (inp.shape[0] < out.shape[0] or inp.shape[1] < out.shape[1]):
@@ -225,3 +241,58 @@ def _operation_hints(features: Dict[str, Any]) -> str:
         hints.append('P')
 
     return "".join(hints) if hints else "U"  # U for unknown
+
+
+def compute_numerical_features(train_pairs: List[Tuple[Array, Array]]) -> np.ndarray:
+    """Convert task features to a numerical vector.
+
+    This utility is primarily used by learning components that expect a fixed
+    numeric representation.  The order of features is deterministic to ensure
+    reproducibility across runs.
+
+    Args:
+        train_pairs: List of training input/output grid pairs.
+
+    Returns:
+        A 1-D numpy array of feature values. Boolean features are encoded as
+        ``0.0`` or ``1.0``.
+    """
+
+    features = extract_task_features(train_pairs)
+
+    numerical_keys = [
+        'num_train_pairs',
+        'input_height_mean',
+        'input_width_mean',
+        'output_height_mean',
+        'output_width_mean',
+        'shape_preserved',
+        'size_ratio_mean',
+        'input_colors_mean',
+        'output_colors_mean',
+        'background_color_consistent',
+        'has_color_mapping',
+        'color_mapping_size',
+        'input_objects_mean',
+        'output_objects_mean',
+        'object_count_preserved',
+        'likely_rotation',
+        'likely_reflection',
+        'likely_translation',
+        'likely_recolor',
+        'likely_crop',
+        'likely_pad',
+    ]
+
+    values: List[float] = []
+    for key in numerical_keys:
+        val = features.get(key, 0)
+        if isinstance(val, bool):
+            values.append(1.0 if val else 0.0)
+        else:
+            try:
+                values.append(float(val))
+            except (TypeError, ValueError):  # pragma: no cover - defensive path
+                values.append(0.0)
+
+    return np.array(values, dtype=float)

arc_solver/guidance.py

@@ -38,20 +38,22 @@ def __init__(self, input_dim: int, hidden_dim: int = 32):
 
     def forward(self, x: np.ndarray) -> np.ndarray:
         """Forward pass through the network."""
+        if x.ndim == 1:
+            x = x.reshape(1, -1)
         # First layer
         h = np.maximum(0, np.dot(x, self.weights1) + self.bias1)  # ReLU
         # Output layer with sigmoid
         out = 1.0 / (1.0 + np.exp(-(np.dot(h, self.weights2) + self.bias2)))
-        return out
+        return out.squeeze()
 
     def predict_operations(self, features: Dict[str, Any], threshold: float = 0.5) -> List[str]:
         """Predict which operations are likely relevant."""
         feature_vector = self._features_to_vector(features)
-        probabilities = self.forward(feature_vector)
-        
+        probabilities = self.forward(feature_vector).ravel()
+
         relevant_ops = []
         for i, prob in enumerate(probabilities):
-            if prob > threshold:
+            if float(prob) > threshold:
                 relevant_ops.append(self.operations[i])
 
         return relevant_ops if relevant_ops else ['identity']