hippocampaif/neocortex/predictive_coding.py

"""
Predictive Coding — Hierarchical Prediction-Error Minimization

Implements Friston's Box 3: the canonical microcircuit for hierarchical
prediction error processing:
- Superficial Pyramidal (SG): prediction errors ε (bottom-up)
- Layer 4 (L4): state estimation
- Deep Pyramidal (IG): predictions μ (top-down)

Recognition dynamics: perception as gradient descent on free energy.
μ̇ = Dμ − ∂F/∂μ  (internal state update)
ε = observed − predicted  (prediction error)

Reference: Friston (2009) Box 3, Figure I
Author: Algorembrant, Rembrant Oyangoren Albeos (2026)
"""

import numpy as np
from typing import Optional


class CorticalColumn:
    """
    A single cortical column implementing the canonical microcircuit.
    
    Contains:
    - Superficial granular (SG): prediction error neurons
    - Layer 4 (L4): state representation neurons
    - Infragranular (IG): prediction neurons (deep pyramidal)
    """
    
    def __init__(self, size: int, level: int = 0):
        """
        Args:
            size: Number of units in this column.
            level: Hierarchical level (0 = lowest/sensory).
        """
        self.size = size
        self.level = level
        
        # State variables (generalized coordinates of motion)
        self.mu = np.zeros(size)          # Internal state estimates (expectations)
        self.mu_dot = np.zeros(size)      # Velocity of expectations
        self.epsilon = np.zeros(size)     # Prediction errors
        self.prediction = np.zeros(size)  # Top-down predictions to level below
        
        # Precision (inverse variance) — controls gain on prediction errors
        self.precision = np.ones(size)    # π = 1/σ² (higher = more confident)
        
        # Connection weights
        self.forward_weights = None   # Bottom-up: from level below
        self.backward_weights = None  # Top-down: to level below
        self.lateral_weights = None   # Within level
        
    def initialize_connections(self, input_size: int, output_size: Optional[int] = None):
        """Initialize synaptic weights for this column's connections."""
        self.forward_weights = np.random.randn(self.size, input_size) * 0.1
        if output_size is not None:
            self.backward_weights = np.random.randn(output_size, self.size) * 0.1
        self.lateral_weights = np.random.randn(self.size, self.size) * 0.01
        np.fill_diagonal(self.lateral_weights, 0)


class PredictiveCodingHierarchy:
    """
    Full hierarchical predictive coding network.
    
    Implements the hierarchical generative model from Friston Box 3:
    - Each level generates predictions for the level below
    - Prediction errors propagate upward (superficial pyramidal)
    - Predictions propagate downward (deep pyramidal)
    - Recognition dynamics minimize free energy via gradient descent
    
    This is NOT variational inference in the ML sense — this is 
    biological free-energy minimization via neural dynamics.
    """
    
    def __init__(self, layer_sizes: list[int], learning_rate: float = 0.01,
                 dt: float = 0.1, n_iterations: int = 10):
        """
        Args:
            layer_sizes: Sizes of each hierarchical level [sensory, ..., abstract].
            learning_rate: Step size for recognition dynamics.
            dt: Integration time step.
            n_iterations: Number of iterations per perception step.
        """
        self.n_levels = len(layer_sizes)
        self.learning_rate = learning_rate
        self.dt = dt
        self.n_iterations = n_iterations
        
        # Build cortical columns at each level
        self.columns: list[CorticalColumn] = []
        for i, size in enumerate(layer_sizes):
            col = CorticalColumn(size, level=i)
            self.columns.append(col)
        
        # Initialize inter-level connections
        for i in range(1, self.n_levels):
            self.columns[i].initialize_connections(
                input_size=layer_sizes[i-1],
                output_size=layer_sizes[i-1]
            )
        # Level 0 has lateral connections only
        self.columns[0].lateral_weights = np.random.randn(
            layer_sizes[0], layer_sizes[0]
        ) * 0.01
        np.fill_diagonal(self.columns[0].lateral_weights, 0)
    
    def _generate_prediction(self, level: int) -> np.ndarray:
        """
        Generate top-down prediction from level i to level i-1.
        
        g(μ⁽ⁱ⁾) — the generative model mapping from higher to lower.
        """
        col = self.columns[level]
        if col.backward_weights is not None:
            # Nonlinear generative mapping (sigmoid for bounded predictions)
            hidden = np.tanh(col.mu)
            return np.dot(col.backward_weights, hidden)
        return np.zeros(self.columns[level - 1].size if level > 0 else col.size)
    
    def _compute_prediction_errors(self, sensory_input: np.ndarray):
        """
        Compute prediction errors at each level of the hierarchy.
        
        ε⁽ⁱ⁾ = μ⁽ⁱ⁻¹⁾ − g(μ⁽ⁱ⁾)
        
        Prediction error = what I observe − what I predicted.
        """
        # Level 0: error between sensory input and level 1's prediction
        if self.n_levels > 1:
            prediction_from_above = self._generate_prediction(1)
            self.columns[0].epsilon = sensory_input - prediction_from_above
        else:
            self.columns[0].epsilon = sensory_input - self.columns[0].mu
        
        # Higher levels: error between current state and prediction from above
        for i in range(1, self.n_levels - 1):
            prediction_from_above = self._generate_prediction(i + 1)
            self.columns[i].epsilon = self.columns[i].mu - prediction_from_above
    
    def _recognition_dynamics(self):
        """
        Recognition dynamics — gradient descent on free energy.
        
        μ̇⁽ⁱ⁾ = Dμ⁽ⁱ⁾ − ∂F/∂μ⁽ⁱ⁾
        
        The internal states update to minimize prediction error,
        weighted by precision. This IS perception in the Fristonian framework.
        """
        for i in range(self.n_levels):
            col = self.columns[i]
            
            # Gradient of free energy w.r.t. internal states
            # ∂F/∂μ = precision-weighted prediction error + prior gradient
            
            dF_dmu = np.zeros(col.size)
            
            # Bottom-up: precision-weighted error from level below
            if i > 0:
                lower_col = self.columns[i - 1]
                if col.forward_weights is not None:
                    # Error signal from lower level, weighted by precision
                    weighted_error = lower_col.precision * lower_col.epsilon
                    dF_dmu -= np.dot(col.forward_weights, weighted_error)
            
            # Top-down: prediction error at this level
            if i < self.n_levels - 1:
                dF_dmu += col.precision * col.epsilon
            
            # Lateral dynamics (recurrent processing within level)
            if col.lateral_weights is not None:
                lateral = np.dot(col.lateral_weights, col.mu)
                dF_dmu -= 0.1 * lateral
            
            # Update internal states (gradient descent on F)
            col.mu_dot = -self.learning_rate * dF_dmu
            col.mu += col.mu_dot * self.dt
            
            # Bounded activation
            col.mu = np.clip(col.mu, -5.0, 5.0)
    
    def process(self, sensory_input: np.ndarray) -> dict:
        """
        Run perception: minimize free energy given sensory input.
        
        This is "seeing" — the brain settling into an interpretation
        of sensory data that minimizes surprise (prediction error).
        
        Args:
            sensory_input: Raw sensory data (preprocessed by retina/V1).
            
        Returns:
            Dict with internal states, prediction errors, and free energy.
        """
        free_energy_history = []
        
        for iteration in range(self.n_iterations):
            # 1. Compute prediction errors at all levels
            self._compute_prediction_errors(sensory_input)
            
            # 2. Run recognition dynamics (update internal states)
            self._recognition_dynamics()
            
            # 3. Update predictions (generative model output)
            for i in range(1, self.n_levels):
                self.columns[i].prediction = self._generate_prediction(i)
            
            # 4. Compute total free energy (should decrease over iterations)
            F = self._compute_free_energy()
            free_energy_history.append(F)
        
        return {
            'states': [col.mu.copy() for col in self.columns],
            'errors': [col.epsilon.copy() for col in self.columns],
            'predictions': [col.prediction.copy() for col in self.columns],
            'free_energy': free_energy_history,
            'final_F': free_energy_history[-1] if free_energy_history else 0.0
        }
    
    def _compute_free_energy(self) -> float:
        """
        Compute variational free energy across the hierarchy.
        
        F ≈ Σᵢ (εᵢ)ᵀ Πᵢ εᵢ  (precision-weighted sum of squared errors)
        
        This is the Laplace approximation to true free energy.
        """
        F = 0.0
        for col in self.columns:
            # Precision-weighted prediction error (energy term)
            F += 0.5 * np.sum(col.precision * col.epsilon**2)
        return float(F)
    
    def learn(self, learning_rate: float = 0.001):
        """
        Update generative model parameters (synaptic plasticity).
        
        This slowly adjusts the top-down generative model to better
        predict sensory input. Corresponds to synaptic plasticity
        in the brain (much slower timescale than recognition dynamics).
        """
        for i in range(1, self.n_levels):
            col = self.columns[i]
            lower = self.columns[i - 1]
            
            if col.backward_weights is not None:
                # Gradient of F w.r.t. backward weights
                # ΔW = -lr * ε * ∂g/∂W
                hidden = np.tanh(col.mu)
                dg_dW = np.outer(lower.epsilon, hidden)
                col.backward_weights += learning_rate * dg_dW
            
            if col.forward_weights is not None:
                # Forward weights learn from prediction errors
                dW = np.outer(col.mu, lower.epsilon * lower.precision)
                col.forward_weights += learning_rate * dW
    
    def get_representation(self, level: int) -> np.ndarray:
        """Get the internal state representation at a given hierarchical level."""
        return self.columns[level].mu.copy()
    
    def get_total_surprise(self) -> float:
        """Total surprise = total precision-weighted prediction error."""
        return self._compute_free_energy()