resonance_transformer/test_geometric.py

import torch
from geometric_memory import GeometricEntryPoint, GeometricMemory

def verify_geometric_memory():
    print("=== VERIFYING GEOMETRIC MEMORY (PHASE 25) ===")
    
    hidden_dim = 64
    batch_size = 2
    seq_len = 10
    
    # 1. Test Entry Point
    entry_net = GeometricEntryPoint(hidden_dim)
    dummy_query = torch.randn(batch_size, seq_len, hidden_dim)
    
    entry_point = entry_net.compute_entry_hash(dummy_query)
    
    print("\n[ENTRY POINT]")
    print(f"  Theta: {entry_point['theta'].shape}")
    print(f"  Frequency (Baseline 528): {entry_point['frequency']}")
    
    # 2. Test Memory Store/Retrieve
    memory = GeometricMemory(hidden_dim)
    
    print("\n[MEMORY STORE]")
    # Store the query as a memory
    memory.store(dummy_query, entry_point)
    print(f"  Stored {len(memory.memory_map)} batches in memory.")
    
    print("\n[MEMORY RETRIEVE]")
    # Try to retrieve using the same query (should find itself)
    retrieved = memory.retrieve(dummy_query, entry_point, k=3)
    
    if retrieved is not None:
        print(f"  Retrieved Shape: {retrieved.shape}")
        # Check alignment
        # This is a self-lookup so correlation should be high
        print("  [PASS] Retrieval successful.")
    else:
        print("  [FAIL] Retrieval returned None.")

if __name__ == "__main__":
    verify_geometric_memory()