3phs category

Co-authored-by: Copilot <copilot@github.com>

src/datasets.py

@@ -238,5 +238,39 @@ class doublePhotonInferenceDataset(Dataset):
         dummy_label = np.zeros((8,), dtype=np.float32)  ### dummy label
         return sample, torch.tensor(dummy_label, dtype=torch.float32)
 
+    def __len__(self):
+        return self.length
+
+class triplePhotonDataset(Dataset):
+    def __init__(self, sampleList, sampleRatio, datasetName):
+        self.sampleFileList = sampleList
+        self.sampleRatio = sampleRatio
+        self.datasetName = datasetName
+        all_samples = []
+        all_labels = []
+        for sampleFile in self.sampleFileList:
+            if '.npz' in sampleFile:
+                data = np.load(sampleFile)
+                all_samples.append(data['samples'])
+                all_labels.append(data['labels'])
+            elif '.h5' in sampleFile:
+                import h5py
+                with h5py.File(sampleFile, 'r') as f:
+                    samples = f['clusters'][:]
+                    labels = f['labels'][:]
+                    all_samples.append(samples)
+                    all_labels.append(labels)
+        self.samples = np.concatenate(all_samples, axis=0)
+        self.labels = np.concatenate(all_labels, axis=0)
+        ### total number of samples
+        self.length = int(self.samples.shape[0] * self.sampleRatio)
+        print(f"[{self.datasetName} dataset] \t Total number of samples: {self.length}")
+    def __getitem__(self, index):
+        sample = self.samples[index]
+        label = self.labels[index]
+        label[:, :2] -= self.samples.shape[-1] / 2.  ### adjust labels to be centered at sample center
+        sample = torch.tensor(sample, dtype=torch.float32).unsqueeze(0)
+        label = torch.tensor(label, dtype=torch.float32)
+        return sample, label
     def __len__(self):
         return self.length

src/models.py

@@ -17,6 +17,13 @@ def get_double_photon_model_class(version):
         raise ValueError(f"Model class '{class_name}' not found.")
     return cls
 
+def get_triple_photon_model_class(version):
+    class_name = f'triplePhotonNet_{version}'
+    cls = globals().get(class_name)
+    if cls is None:
+        raise ValueError(f"Model class '{class_name}' not found.")
+    return cls
+
 class singlePhotonNet_250909(nn.Module):
     def weight_init(self):
         for m in self.modules():
@@ -415,4 +422,65 @@ class doublePhotonNet_260507(nn.Module): ## adapted from 251124, removed max poo
         flat_feat = c3.view(c3.size(0), -1)  # [B, 6272]
         
         coords = self.fc(flat_feat)     # [B, 4]
+        return coords
+
+
+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):
+        super().__init__()
+        # Backbone: deeper for 9x9 input containing 3 photons
+        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1) ### 9x9
+        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1) ## 9x9
+        self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1) ## 9x9
+        self.conv4 = nn.Conv2d(128, 128, kernel_size=3) ## 7x7
+        
+        # Spatial Attention Module
+        self.spatial_attn = nn.Sequential(
+            nn.Conv2d(128, 1, kernel_size=1),
+            nn.Sigmoid()
+        )
+        
+        self.fc = nn.Sequential(
+            nn.Linear(128 * 7 * 7, 512),
+            nn.ReLU(),
+            nn.Dropout(0.3),
+            nn.Linear(512, 128),
+            nn.ReLU(),
+            nn.Dropout(0.3),
+            nn.Linear(128, 6)
+        )
+        
+        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
+        nSize = 9  # should match the input size of the model
+        x = np.linspace(-nSize/2. + 0.5, nSize/2. - 0.5, nSize)
+        y = np.linspace(-nSize/2. + 0.5, nSize/2. - 0.5, 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, 9, 9]
+        c1 = F.relu(self.conv1(x))      # [B, 32, 9, 9]
+        c2 = F.relu(self.conv2(c1))     # [B, 64, 9, 9]
+        c3 = F.relu(self.conv3(c2))     # [B, 128, 9, 9]
+        c4 = F.relu(self.conv4(c3))     # [B, 128, 7, 7]
+        
+        attn = self.spatial_attn(c4)    # [B, 1, 7, 7]
+        c4 = c4 * attn                  # [B, 128, 7, 7]
+        
+        flat_feat = c4.view(c4.size(0), -1)  # [B, 6272]
+        
+        coords = self.fc(flat_feat)     # [B, 6]
         return coords