From e479ad658645140537af5be0abe8b7f77401fe12 Mon Sep 17 00:00:00 2001
From: "xiangyu.xie" <xiangyu.xie@psi.ch>
Date: Wed, 10 Jun 2026 11:11:48 +0200
Subject: [PATCH] Add new 2-ph model without dropout and padding

Co-authored-by: Copilot <copilot@github.com>
---
 src/models.py | 55 +++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 55 insertions(+)

diff --git a/src/models.py b/src/models.py
index 3cc2a1c..65aa189 100644
--- a/src/models.py
+++ b/src/models.py
@@ -424,6 +424,61 @@ class doublePhotonNet_260507(nn.Module): ## adapted from 251124, removed max poo
         coords = self.fc(flat_feat)     # [B, 4]
         return coords
 
+class doublePhotonNet_260608(nn.Module): ## adapted from 260507, removed padding and dropout
+    def __init__(self, nSize=6):
+        super().__init__()
+        self.nSize = nSize
+        # Backbone: 3 CNN layers
+        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)
+        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=0)
+        self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=0)
+        
+        # Spatial Attention Module
+        self.spatial_attn = nn.Sequential(
+            nn.Conv2d(128, 1, kernel_size=1),
+            nn.Sigmoid()
+        )
+        
+        self.fc = nn.Sequential(
+            nn.Linear(128 * (self.nSize-4) * (self.nSize-4), 512),
+            nn.ReLU(),
+            nn.Linear(512, 128),
+            nn.ReLU(),
+            nn.Linear(128, 4)
+        )
+        
+        self._init_weights()
+        self._init_coords()
+    
+    def _init_coords(self):
+        # Create a coordinate grid; moved from dataset generation to model initialization for lower traffic and more flexibility
+        x = np.linspace(-self.nSize/2. + 0.5, self.nSize/2. - 0.5, self.nSize)
+        y = np.linspace(-self.nSize/2. + 0.5, self.nSize/2. - 0.5, self.nSize)
+        x_grid, y_grid = np.meshgrid(x, y, indexing='ij')  # (nSize,nSize), (nSize,nSize)
+        self.x_grid = torch.tensor(np.expand_dims(x_grid, axis=0)).float().contiguous().to('cuda')  # (1, nSize, nSize)
+        self.y_grid = torch.tensor(np.expand_dims(y_grid, axis=0)).float().contiguous().to('cuda')  # (1, nSize, nSize)
+
+    def _init_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+            elif isinstance(m, nn.Linear):
+                nn.init.xavier_uniform_(m.weight)
+                nn.init.zeros_(m.bias)
+    
+    def forward(self, x):
+        x = torch.cat((x, self.x_grid.expand(x.size(0), -1, -1, -1), self.y_grid.expand(x.size(0), -1, -1, -1)), dim=1)  # [B, 3, nSize, nSize]
+        c1 = F.relu(self.conv1(x))      # [B, 32, nSize, nSize]
+        c2 = F.relu(self.conv2(c1))     # [B, 64, nSize-2, nSize-2]
+        c3 = F.relu(self.conv3(c2))     # [B, 128, nSize-4, nSize-4]
+        
+        attn = self.spatial_attn(c3)    # [B, 1, nSize-4, nSize-4]
+        c3 = c3 * attn                  # [B, 128, nSize-4, nSize-4]
+        
+        flat_feat = c3.view(c3.size(0), -1)  # [B, 128 * (nSize-4) * (nSize-4)]
+        
+        preditions = self.fc(flat_feat)     # [B, 4]
+        return preditions
 
 class triplePhotonNet_260529(nn.Module): ## adapted from doublePhotonNet_260507, add one more conv layer and increase capacity of FC layers, for 3-photon pileup with 9x9 input
     def __init__(self):