import torch import torch.nn as nn import torch.optim as optim import torch.nn.functional as F from torch.utils.data import DataLoader from torchvision import datasets, transforms import numpy as np from tqdm import tqdm from torchdiffeq import odeint_adjoint as odeint # ============================================================ # XYFLOW Classifier: Vector Field + Attractor Basins # ============================================================ class VectorField(nn.Module): def __init__(self, latent_dim): super().__init__() self.net = nn.Sequential( nn.Linear(latent_dim, 128), nn.Tanh(), nn.Linear(128, 128), nn.Tanh(), nn.Linear(128, latent_dim) ) def forward(self, t, z=None): if z is None: # Called as model.field(centers) in training/loss computation return self.net(t) # Called as model.field(t, z) by odeint return self.net(z) class XYFlowClassifier(nn.Module): def __init__(self, latent_dim=64, num_classes=10, int_time=2.0, lr=1e-3): """ latent_dim : dimension of the phase space num_classes: number of attractors (one per class) int_time : integration time for the ODE flow """ super().__init__() self.latent_dim = latent_dim self.num_classes = num_classes self.int_time = int_time # ------ 1. Embedding network: image -> initial point z0 in phase space ------ self.embed = nn.Sequential( nn.Conv2d(1, 16, 3, stride=2, padding=1), nn.ReLU(), nn.Conv2d(16, 32, 3, stride=2, padding=1), nn.ReLU(), nn.Conv2d(32, 64, 3, stride=2, padding=1), nn.ReLU(), nn.Flatten(), nn.Linear(64 * 4 * 4, latent_dim) # MNIST 28x28 -> 4x4 after 3 strides ) # ------ 2. Vector field: the "source code" of the classifier ------ # It defines how every point moves (dz/dt = f(z)). self.field = VectorField(latent_dim) # ------ 3. Class attractors: fixed points we want trajectories to converge to ------ # Each attractor corresponds to a digit (0-9). self.centers = nn.Parameter(torch.randn(num_classes, latent_dim)) def forward(self, x, return_trajectory=False): """ x: images [B,1,28,28] returns: logits (negative squared distance to each attractor center) """ B = x.shape[0] z0 = self.embed(x) # map image to phase space # Integrate dz/dt = field(z) from t=0 to t=self.int_time # t = [0, int_time] (at least two points for odeint) t = torch.linspace(0, self.int_time, 2).to(x.device) z_t = odeint(self.field, z0, t, method='dopri5') # shape [2, B, D] z_final = z_t[-1] # final state after flow # Classification: which attractor are we closest to? # logits = -||z_final - center_c||^2 (higher = more attracted) diff = z_final.unsqueeze(1) - self.centers.unsqueeze(0) # [B, 10, D] dist2 = (diff ** 2).sum(dim=2) # squared Euclidean distance logits = -dist2 # negative distance -> logits if return_trajectory: return logits, z_t, z0 return logits def classify(self, x): """Return predicted class indices.""" logits = self.forward(x) return logits.argmax(dim=1) # ============================================================ # Training Setup with XYFLOW-Specific Losses # ============================================================ def train_epoch(model, loader, optimizer, device): model.train() total_loss = 0 correct = 0 total = 0 for x, y in tqdm(loader, desc='Training'): x, y = x.to(device), y.to(device) optimizer.zero_grad() # --- Main classification loss --- logits = model(x) ce_loss = F.cross_entropy(logits, y) # --- Fixed-point loss: ensure centers are real fixed points (f(center) ≈ 0) --- centers = model.centers # [10, D] field_at_centers = model.field(centers) # [10, D] fp_loss = (field_at_centers ** 2).mean() # penalty for non-zero flow at attractors # --- Attraction loss: gently pull final states toward their class centers --- # (we can optionally compute this, but not strictly needed) # with torch.no_grad(): # z_final = ... we'd have to re-run, so skip for speed # Total loss loss = ce_loss + 0.1 * fp_loss # weight can be tuned loss.backward() optimizer.step() total_loss += loss.item() * x.size(0) pred = logits.argmax(dim=1) correct += (pred == y).sum().item() total += x.size(0) return total_loss / total, correct / total def evaluate(model, loader, device): model.eval() correct = 0 total = 0 with torch.no_grad(): for x, y in loader: x, y = x.to(device), y.to(device) logits = model(x) pred = logits.argmax(dim=1) correct += (pred == y).sum().item() total += x.size(0) return correct / total # ============================================================ # Main: Load MNIST, Train, Test # ============================================================ def main(): device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') print(f"Using device: {device}") # ------ Data ------ transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]) train_ds = datasets.MNIST('../data', train=True, download=True, transform=transform) test_ds = datasets.MNIST('../data', train=False, download=True, transform=transform) train_loader = DataLoader(train_ds, batch_size=128, shuffle=True, num_workers=2, pin_memory=True) test_loader = DataLoader(test_ds, batch_size=1000, shuffle=False, num_workers=2, pin_memory=True) # ------ Model ------ model = XYFlowClassifier(latent_dim=64, int_time=2.0).to(device) optimizer = optim.Adam(model.parameters(), lr=1e-3) scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.5) # ------ Training loop ------ epochs = 15 for epoch in range(1, epochs + 1): train_loss, train_acc = train_epoch(model, train_loader, optimizer, device) scheduler.step() test_acc = evaluate(model, test_loader, device) print(f"Epoch {epoch:2d} | Loss: {train_loss:.4f} | Train Acc: {train_acc:.4f} | Test Acc: {test_acc:.4f}") print(f"\nFinal test accuracy: {test_acc:.4f}") if __name__ == '__main__': main()