hippocampaif/core_knowledge/geometry_system.py

"""
Geometry System — Core Knowledge of Geometry

Implements innate geometric/spatial reasoning:
1. Spatial relations: left, right, above, below, inside, outside
2. Distance metrics with smooth deformations (Distortable Canvas integration)
3. Surface layout representations (navigable surfaces)
4. Shape primitives for recognition

Boosted by the Distortable Canvas paper: images as smooth functions
on elastic 2D canvas with deformation fields.

Reference: Spelke & Lee (2012), oneandtrulyone Distortable Canvas paper
Author: Algorembrant, Rembrant Oyangoren Albeos (2026)
"""

import numpy as np
from scipy.ndimage import gaussian_filter


class GeometrySystem:
    """
    Innate spatial and geometric reasoning.
    
    Provides core geometric computations including spatial relations,
    smooth deformation fields (from Distortable Canvas), and surface
    layout representations for navigation and object reasoning.
    """
    
    def __init__(self, canvas_resolution: int = 28):
        """
        Args:
            canvas_resolution: Default canvas size for deformation operations.
        """
        self.canvas_resolution = canvas_resolution
    
    # ----- Spatial Relations (innate categorical distinctions) -----
    
    def spatial_relation(self, pos_a: np.ndarray, pos_b: np.ndarray) -> dict:
        """
        Compute innate categorical spatial relations between two positions.
        
        Infants distinguish these before learning language labels for them.
        
        Args:
            pos_a: Reference position (ndarray, at least 2D).
            pos_b: Target position.
            
        Returns:
            Dict with boolean spatial relations and continuous distances.
        """
        a = np.asarray(pos_a, dtype=np.float64)
        b = np.asarray(pos_b, dtype=np.float64)
        diff = b - a
        
        relations = {
            'distance': float(np.linalg.norm(diff)),
            'direction': diff / (np.linalg.norm(diff) + 1e-8),
        }
        
        if len(diff) >= 2:
            relations['right_of'] = bool(diff[0] > 0)
            relations['left_of'] = bool(diff[0] < 0)
            relations['above'] = bool(diff[1] < 0)  # Assuming y-axis points down (image coords)
            relations['below'] = bool(diff[1] > 0)
            
        return relations
    
    def is_inside(self, point: np.ndarray, bbox_min: np.ndarray, 
                  bbox_max: np.ndarray) -> bool:
        """
        Check if a point is inside a bounding region.
        
        Containment is a core geometric concept — infants reason about
        "inside" and "outside" from very early on.
        """
        p = np.asarray(point, dtype=np.float64)
        return bool(np.all(p >= bbox_min) and np.all(p <= bbox_max))
    
    # ----- Smooth Deformation Fields (Distortable Canvas) -----
    
    def create_deformation_field(self, shape: tuple[int, int],
                                  smoothness: float = 3.0,
                                  magnitude: float = 2.0) -> tuple[np.ndarray, np.ndarray]:
        """
        Create a smooth random deformation field.
        
        From the Distortable Canvas paper: images live on an elastic 2D canvas
        that can be smoothly warped. The deformation field u(x,y), v(x,y) defines
        how each pixel coordinate shifts.
        
        Args:
            shape: (H, W) shape of the canvas.
            smoothness: Gaussian sigma for smoothness regularization.
                       Higher = smoother/more rigid deformation.
            magnitude: Maximum displacement magnitude.
            
        Returns:
            Tuple of (u_field, v_field), each of shape (H, W).
        """
        H, W = shape
        # Random initial displacements
        u = np.random.randn(H, W) * magnitude
        v = np.random.randn(H, W) * magnitude
        
        # Smooth with Gaussian filter (biological smoothness constraint)
        u = gaussian_filter(u, sigma=smoothness)
        v = gaussian_filter(v, sigma=smoothness)
        
        return u, v
    
    def apply_deformation(self, image: np.ndarray,
                          u_field: np.ndarray,
                          v_field: np.ndarray) -> np.ndarray:
        """
        Apply a smooth deformation field to an image.
        
        This is the core operation from the Distortable Canvas paper:
        warp the canvas to align one image to another.
        
        Args:
            image: 2D image array (H, W).
            u_field: Horizontal displacement field (H, W).
            v_field: Vertical displacement field (H, W).
            
        Returns:
            Warped image.
        """
        H, W = image.shape[:2]
        # Create coordinate grids
        y_coords, x_coords = np.mgrid[0:H, 0:W].astype(np.float64)
        
        # Apply deformation
        new_x = x_coords + u_field
        new_y = y_coords + v_field
        
        # Clamp to valid range
        new_x = np.clip(new_x, 0, W - 1)
        new_y = np.clip(new_y, 0, H - 1)
        
        # Bilinear interpolation
        x0 = np.floor(new_x).astype(int)
        x1 = np.minimum(x0 + 1, W - 1)
        y0 = np.floor(new_y).astype(int)
        y1 = np.minimum(y0 + 1, H - 1)
        
        wx = new_x - x0
        wy = new_y - y0
        
        result = (image[y0, x0] * (1 - wx) * (1 - wy) +
                  image[y1, x0] * (1 - wx) * wy +
                  image[y0, x1] * wx * (1 - wy) +
                  image[y1, x1] * wx * wy)
        
        return result
    
    def canvas_distance(self, u_field: np.ndarray, v_field: np.ndarray) -> float:
        """
        Compute the canvas distortion energy (from Distortable Canvas paper).
        
        This measures how much geometric warping is needed. The Jacobian
        penalty ensures the deformation is smooth and doesn't tear/fold.
        
        Energy = sum of squared gradients of the deformation field.
        """
        # Gradient of deformation field (Jacobian components)
        du_dx = np.gradient(u_field, axis=1)
        du_dy = np.gradient(u_field, axis=0)
        dv_dx = np.gradient(v_field, axis=1)
        dv_dy = np.gradient(v_field, axis=0)
        
        # Frobenius norm of the Jacobian of the displacement
        jacobian_energy = np.sum(du_dx**2 + du_dy**2 + dv_dx**2 + dv_dy**2)
        
        return float(jacobian_energy)
    
    def color_distance(self, image1: np.ndarray, image2: np.ndarray) -> float:
        """
        Pixel-wise intensity distance between two images.
        
        From Distortable Canvas: the "color distortion" component
        of the dual distance metric.
        """
        return float(np.sum((image1.astype(np.float64) - image2.astype(np.float64))**2))
    
    def dual_distance(self, image1: np.ndarray, image2: np.ndarray,
                      u_field: np.ndarray, v_field: np.ndarray,
                      lambda_weight: float = 0.1) -> float:
        """
        Compute the dual distance from the Distortable Canvas paper.
        
        dual_distance = color_distance + lambda * canvas_distance
        
        This balances pixel-level similarity against geometric warping cost.
        """
        # Warp image1 toward image2
        warped = self.apply_deformation(image1, u_field, v_field)
        
        color_dist = self.color_distance(warped, image2)
        canvas_dist = self.canvas_distance(u_field, v_field)
        
        return color_dist + lambda_weight * canvas_dist
    
    # ----- Shape Primitives -----
    
    def compute_centroid(self, points: np.ndarray) -> np.ndarray:
        """Compute the centroid (center of mass) of a set of points."""
        return np.mean(points, axis=0)
    
    def compute_extent(self, points: np.ndarray) -> dict:
        """Compute spatial extent (bounding box, spread) of a point set."""
        mins = np.min(points, axis=0)
        maxs = np.max(points, axis=0)
        return {
            'min': mins,
            'max': maxs,
            'extent': maxs - mins,
            'center': (mins + maxs) / 2.0
        }
    
    def angular_relation(self, center: np.ndarray, point: np.ndarray) -> float:
        """
        Compute the angle from center to point (in radians).
        Used for encoding relative angular position.
        """
        diff = np.asarray(point, dtype=np.float64) - np.asarray(center, dtype=np.float64)
        return float(np.arctan2(diff[1], diff[0]))