import torch import torch.nn as nn import torch.nn.functional as F from torchvision import datasets, transforms from torch.utils.data import DataLoader # ============================================================================= # PNS CIRCUIT ELEMENTS # ============================================================================= class PNSCircuitLayer(nn.Module): """ A Layer that treats data as 'Particle Numbers'. Internal state is a multi-field tensor: [batch, channels, fields] Fields: 0: Linear (L), 1: Quadratic (Q), 2: Oscillatory (O), 3: Exponential (E), 4: Rest Mass (rho) """ def __init__(self, in_features, out_features): super().__init__() self.in_features = in_features self.out_features = out_features # Library of params for DIFFERENT function types # [type, out, in, params] # Types: 0=Linear, 1=Oscillatory, 2=Exponential self.func_library = nn.Parameter(torch.randn(3, out_features, in_features, 3) * 0.001) # Routing Vector: The 'Knowledge Map' # Determines which function from the library is used for each (out, in) connection # [out, in] - values are soft-maxed to pick a function type self.routing_map = nn.Parameter(torch.randn(out_features, in_features, 3) * 0.001) self.bias_particles = nn.Parameter(torch.randn(out_features, 5) * 0.001) def forward(self, x): # x shape: [batch, in_features, 5] psi_x = x.sum(dim=-1) # [batch, in_features] # 1. THE ROUTING (The "Self-Questioning" Process) # Softmax over the 3 function types to get a weighted combination of subroutines # routing_weights: [out, in, 3] routing_weights = torch.softmax(self.routing_map, dim=-1) # 2. THE SUBROUTINE LIBRARY # we compute the result of all 3 possible function types first all_results = [] # Type 0: Linear (a * psi_x) a0 = self.func_library[0, ..., 0] res_lin = torch.einsum('bn,on->bo', psi_x, a0) all_results.append(res_lin) # Type 1: Oscillatory (a * sin(w * psi_x + p)) a1 = self.func_library[1, ..., 0] w1 = self.func_library[1, ..., 1] p1 = self.func_library[1, ..., 2] wave_input = (w1.unsqueeze(0) * psi_x.unsqueeze(1)) + p1.unsqueeze(0) res_osc = torch.sum(a1.unsqueeze(0) * torch.sin(wave_input), dim=-1) all_results.append(res_osc) # Type 2: Exponential (a * exp(-|w * psi_x|)) a2 = self.func_library[2, ..., 0] w2 = self.func_library[2, ..., 1] res_exp = torch.sum(a2.unsqueeze(0) * torch.exp(-torch.abs(w2.unsqueeze(0) * psi_x.unsqueeze(1))), dim=-1) all_results.append(res_exp) # 3. DYNAMIC PATCHING # The model 'selects' the answer by summing the results based on the routing map. # result = sum_type(routing_weight[type] * result[type]) final_interaction = torch.zeros_like(res_lin) for t in range(3): # Weighted sum for each type across the routing map # route_sum: [out, in] -> combined weighting route_weight_sum = routing_weights[..., t].sum(dim=-1) # Simplified routing final_interaction += route_weight_sum.unsqueeze(0) * all_results[t] # 4. map back to multi-field PNS state fields = [ final_interaction, #Linear final_interaction**2, #Quadratic final_interaction * torch.sin(final_interaction), #Oscillatory torch.exp(-torch.abs(final_interaction)), #Exponential torch.mean(final_interaction, dim=1, keepdim=True).expand_as(final_interaction) #rho ] out_particles = torch.stack(fields, dim=-1) out_particles += self.bias_particles return out_particles class PNSDiode(nn.Module): """ ReLU equivalent: Only allows 'positive charge' particles to flow. """ def forward(self, x): # Collapses the particle to its scalar value to determine the charge sign # psi = L + Q + O + E + rho psi = x.sum(dim=-1, keepdim=True) # Diode gating: if psi < 0, the particle is annihilated to vacuum (0) gate = (psi > 0).float() return x * gate class PNSImpedanceMatcher(nn.Module): """ LayerNorm equivalent: Normalizes the energy (Quadratic field) to 1. """ def forward(self, x): # Normalize using the energy field (index 1) energy = torch.sqrt(x[..., 1]**2 + 1e-6) return x / energy.unsqueeze(-1) class PNSCollapse(nn.Module): """ The Measurement Operator: Final collapse of all fields into a scalar. """ def forward(self, x): # Psi(P) = sum of all fields return x.sum(dim=-1) # ============================================================================= # FULL CIRCUIT MODEL # ============================================================================= class PNSClassifier(nn.Module): def __init__(self): super().__init__() # Input: 28*28 = 784. We map each pixel to a Particle (5 fields) self.layer1 = PNSCircuitLayer(784, 128) self.diode1 = PNSDiode() self.matcher1 = PNSImpedanceMatcher() self.layer2 = PNSCircuitLayer(128, 64) self.diode2 = PNSDiode() self.matcher2 = PNSImpedanceMatcher() self.layer3 = PNSCircuitLayer(64, 10) self.collapse = PNSCollapse() def forward(self, x): # 1. Injection: Convert scalar pixels into Particle Numbers # Map pixel value to the Linear field, others start at zero # x: [batch, 784] -> [batch, 784, 5] batch_size = x.shape[0] p_in = torch.zeros((batch_size, 784, 5), device=x.device) p_in[..., 0] = x # Linear field excitation p_in[..., 4] = 1.0 # Every pixel starts with unit rest mass rho=1 # 2. Circuit Flow x = self.layer1(p_in) x = self.diode1(x) x = self.matcher1(x) x = self.layer2(x) x = self.diode2(x) x = self.matcher2(x) x = self.layer3(x) # 3. Measurement (Collapse) return self.collapse(x) # ============================================================================= # TRAINING LOOP # ============================================================================= def train_mnist(): # Data transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]) train_set = datasets.MNIST('./data', train=True, download=True, transform=transform) test_set = datasets.MNIST('./data', train=False, transform=transform) train_loader = DataLoader(train_set, batch_size=64, shuffle=True) test_loader = DataLoader(test_set, batch_size=1000, shuffle=False) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = PNSClassifier().to(device) optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) criterion = nn.CrossEntropyLoss() model.train() for epoch in range(2): # Small epoch count for demonstration total_loss = 0 for batch_idx, (data, target) in enumerate(train_loader): data, target = data.view(-1, 784).to(device), target.to(device) optimizer.zero_grad() output = model(data) loss = criterion(output, target) loss.backward() optimizer.step() total_loss += loss.item() if batch_idx % 200 == 0: print(f"Epoch {epoch} [{batch_idx*64}/{len(train_loader)*64}] Loss: {loss.item():.4f}") # Evaluation model.eval() correct = 0 with torch.no_grad(): for data, target in test_loader: data, target = data.view(-1, 784).to(device), target.to(device) output = model(data) pred = output.argmax(dim=1, keepdim=True) correct += pred.eq(target.view_as(pred)).sum().item() print(f"\nTest Accuracy: {100. * correct / len(test_loader.dataset):.2f}%") if __name__ == "__main__": train_mnist()