import torch import torch.nn as nn import torch.nn.functional as F import torch.optim as optim from torchvision import datasets, transforms from torch.utils.data import DataLoader device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"Using device: {device}") # ===================================================================== # 1. THE OBSERVER NEURAL-ODE WITH LEARNABLE QUIVER ANCHOR CODEBOOK # ===================================================================== class QuiverAnchorObserverODE(nn.Module): def __init__(self, input_dim=784, hidden_dim=128, output_dim=10, num_anchors=64, steps=5, dt=0.2): super(QuiverAnchorObserverODE, self).__init__() self.steps = steps self.dt = dt self.hidden_dim = hidden_dim self.num_anchors = num_anchors # --- THE CODEBOOK MEMORY CRADLE --- # We instantiate a fixed, learnable grid of complex Euler numbers representing stationary ODE states. # Think of this as a matrix quiver plot template that the model learns over time. self.anchor_amplitudes = nn.Parameter(torch.rand(num_anchors, hidden_dim)) self.anchor_phases = nn.Parameter(torch.randn(num_anchors, hidden_dim) * 3.14159) # Router: Evaluates the input and assigns coefficients to choose/mix the trajectories self.anchor_router = nn.Sequential( nn.Linear(input_dim, hidden_dim), nn.Tanh(), nn.Linear(hidden_dim, num_anchors) ) # --- ODE VECTOR FIELD --- # Governs how the selected superposition of anchors evolves self.real_dynamics = nn.Sequential( nn.Linear(hidden_dim, hidden_dim), nn.Tanh(), nn.Linear(hidden_dim, hidden_dim) ) self.imag_dynamics = nn.Sequential( nn.Linear(hidden_dim, hidden_dim), nn.Tanh(), nn.Linear(hidden_dim, hidden_dim) ) # --- THE COLLAPSE OBSERVER --- self.decoder = nn.Linear(hidden_dim, output_dim) def forward(self, x): # Flatten MNIST input x = x.view(x.size(0), -1) batch_size = x.size(0) # 1. Materialize the Full Complex Anchor Quiver Grid # z_grid = r * (cos(theta) + i * sin(theta)) -> shape: (num_anchors, hidden_dim) grid_r = torch.sigmoid(self.anchor_amplitudes) grid_theta = torch.tanh(self.anchor_phases) * 3.14159 grid_real = grid_r * torch.cos(grid_theta) grid_imag = grid_r * torch.sin(grid_theta) # 2. CCT Question Path: Dynamic Routing via Trajectory Superposition # Compute routing coefficients alpha across the anchors routing_logits = self.anchor_router(x) routing_weights = F.softmax(routing_logits, dim=-1) # (batch_size, num_anchors) # Compute the Initial Condition C by taking a weighted superposition of our grid anchors # (batch_size, num_anchors) x (num_anchors, hidden_dim) -> (batch_size, hidden_dim) real_z = torch.matmul(routing_weights, grid_real) imag_z = torch.matmul(routing_weights, grid_imag) # 3. ODE Trajectory Flow (Evolution Loop) # Every step cross-couples real and imaginary trajectories like an active beam splitter for _ in range(self.steps): d_real = self.real_dynamics(real_z) - self.imag_dynamics(imag_z) d_imag = self.imag_dynamics(real_z) + self.real_dynamics(imag_z) real_z = real_z + self.dt * d_real imag_z = imag_z + self.dt * d_imag # 4. Superposition Collapse Observation # Extract phase-invariant amplitude probability metrics: |z|^2 probabilities = (real_z ** 2) + (imag_z ** 2) out = self.decoder(probabilities) return out # ===================================================================== # 2. DATA PROCESSING SETUP # ===================================================================== transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ]) train_dataset = datasets.MNIST(root='../data', train=True, download=True, transform=transform) test_dataset = datasets.MNIST(root='../data', train=False, download=True, transform=transform) train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True) test_loader = DataLoader(test_dataset, batch_size=1000, shuffle=False) # ===================================================================== # 3. TRAINING INFRASTRUCTURE WITH OPTIMIZATION # ===================================================================== # Instantiating 64 complex anchor trajectories model = QuiverAnchorObserverODE(input_dim=784, hidden_dim=128, output_dim=10, num_anchors=64, steps=6, dt=0.15).to(device) criterion = nn.CrossEntropyLoss() optimizer = optim.AdamW(model.parameters(), lr=0.002, weight_decay=1e-4) print("Beginning Training via Codebook Trajectory Routing...") epochs = 5 for epoch in range(1, epochs + 1): model.train() running_loss = 0.0 correct = 0 total = 0 for batch_idx, (data, target) in enumerate(train_loader): data, target = data.to(device), target.to(device) for _ in range(1 + batch_idx%2 * 2): optimizer.zero_grad() output = model(data) loss = criterion(output, target) loss.backward() optimizer.step() running_loss += loss.item() _, predicted = output.max(1) total += target.size(0) correct += predicted.eq(target).sum().item() if batch_idx % 200 == 0: print(f"Epoch {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)}] " f"Loss: {loss.item():.4f} | Running Accuracy: {100. * correct / total:.2f}%") epoch_loss = running_loss / len(train_loader) epoch_acc = 100. * correct / total print(f"--- Epoch {epoch} Complete | Avg Loss: {epoch_loss:.4f} | Accuracy: {epoch_acc:.2f}% ---\n") # ===================================================================== # 4. SYSTEM STABILITY VALIDATION (TESTING) # ===================================================================== model.eval() test_loss = 0 correct = 0 with torch.no_grad(): for data, target in test_loader: data, target = data.to(device), target.to(device) output = model(data) test_loss += criterion(output, target).item() _, predicted = output.max(1) correct += predicted.eq(target).sum().item() test_loss /= len(test_loader) test_accuracy = 100. * correct / len(test_loader.dataset) print("==========================================") print(f"ANALYTICAL CODEBOOK STABILITY EVALUATION") print(f"Average Test Loss: {test_loss:.4f}") print(f"Anchor-Aligned Test Accuracy: {test_accuracy:.2f}%") print("==========================================")