Beamline neural network reconstruction training

.gitignore

@@ -36,6 +36,13 @@ __pycache__/
 *$py.class
 .ipynb_checkpoints
 
+# results and simulation output
+results/
+sim_output/
+.snakemake/
+calibrations/
+fieldmaps/
+
 # test for calorimeter
 calorimeters/test/
 *.d

.gitlab-ci.yml

@@ -128,6 +128,7 @@ include:
   - local: 'benchmarks/calo_pid/config.yml'
   - local: 'benchmarks/campaign/config.yml'
   - local: 'benchmarks/ecal_gaps/config.yml'
+  - local: 'benchmarks/lowq2_reconstruction/config.yml'
   - local: 'benchmarks/tracking_detectors/config.yml'
   - local: 'benchmarks/tracking_performances/config.yml'
   - local: 'benchmarks/tracking_performances_dis/config.yml'
@@ -165,6 +166,7 @@ deploy_results:
     - "collect_results:campaign"
     - "collect_results:ecal_gaps"
     - "collect_results:lfhcal"
+    - "collect_results:lowq2_reconstruction"
     - "collect_results:material_scan"
     - "collect_results:pid"
     - "collect_results:rich"

Snakefile

@@ -37,6 +37,7 @@ include: "benchmarks/beamline/Snakefile"
 include: "benchmarks/calo_pid/Snakefile"
 include: "benchmarks/campaign/Snakefile"
 include: "benchmarks/ecal_gaps/Snakefile"
+include: "benchmarks/lowq2_reconstruction/Snakefile"
 include: "benchmarks/material_scan/Snakefile"
 include: "benchmarks/tracking_performances/Snakefile"
 include: "benchmarks/tracking_performances_dis/Snakefile"

benchmarks/beamline/Snakefile

@@ -1,12 +1,17 @@
 SIMOUTDIR="sim_output/beamline/"
 ANALYSISDIR=SIMOUTDIR+"analysis/"
 
+##########################################################################################
+### Rules for checking the steering of the electron beam through the magnets
+##########################################################################################
+
 rule beamline_steering_sim:
     input:
         warmup="warmup/epic_ip6_extended.edm4hep.root",
         macro=workflow.source_path("beamlineGPS.mac"),
     output:
         SIMOUTDIR+"beamlineTest{CAMPAIGN}.edm4hep.root",
+    cache: True
     shell:
         """
             exec npsim \
@@ -20,12 +25,41 @@ rule beamline_steering_sim:
             --physics.rangecut 100*m
         """
 
+rule beamline_steering_analysis:
+    input:
+        warmup="warmup/epic_ip6_extended.edm4hep.root",
+        script=workflow.source_path("beamlineAnalysis.C"),
+        header=workflow.source_path("shared_functions.h"),
+        data=SIMOUTDIR+"beamlineTest{CAMPAIGN}.edm4hep.root",
+    output:
+        rootfile=ANALYSISDIR+"electron_beamline_beamlineTestAnalysis{CAMPAIGN}.root",
+        beamspot_canvas=ANALYSISDIR+"electron_beamline_beamspot_{CAMPAIGN}.png",
+        x_px_canvas=ANALYSISDIR+"electron_beamline_x_px_{CAMPAIGN}.png",
+        y_py_canvas=ANALYSISDIR+"electron_beamline_y_py_{CAMPAIGN}.png",
+        fitted_position_means_stdevs_canvas=ANALYSISDIR+"electron_beamline_fitted_position_means_stdevs_{CAMPAIGN}.png",
+        fitted_momentum_means_stdevs_canvas=ANALYSISDIR+"electron_beamline_fitted_momentum_means_stdevs_{CAMPAIGN}.png",
+        pipe_parameter_canvas=ANALYSISDIR+"electron_beamline_pipe_parameter_{CAMPAIGN}.png",
+    params:
+        xml=os.getenv("DETECTOR_PATH")+"/epic_ip6_extended.xml",
+    shell:
+        """
+            root -l -b -q '{input.script}("{input.data}", "{output.rootfile}", "{params.xml}",
+            "{output.beamspot_canvas}", "{output.x_px_canvas}", "{output.y_py_canvas}",
+            "{output.fitted_position_means_stdevs_canvas}", "{output.fitted_momentum_means_stdevs_canvas}",
+            "{output.pipe_parameter_canvas}")'
+        """
+
+##########################################################################################
+### Rules for checking the acceptance of electrons at each stage of the beamline
+##########################################################################################
+
 rule beamline_acceptance_sim:
     input:
         warmup="warmup/epic_ip6_extended.edm4hep.root",
         macro=workflow.source_path("acceptanceGPS.mac"),
     output:
         SIMOUTDIR+"acceptanceTest{CAMPAIGN}.edm4hep.root",
+    cache: True
     shell:
         """
             exec npsim \
@@ -34,47 +68,22 @@ rule beamline_acceptance_sim:
             --enableG4GPS \
             --macroFile {input.macro} \
             --compactFile $DETECTOR_PATH/epic_ip6_extended.xml \
-            --printLevel WARNING \
             --outputFile {output} \
             --physics.rangecut 100*m
         """
 
-rule beamline_steering_analysis:
-    input:
-        warmup="warmup/epic_ip6_extended.edm4hep.root",
-        script=workflow.source_path("beamlineAnalysis.C"),
-        header=workflow.source_path("shared_functions.h"),
-        data=SIMOUTDIR+"beamlineTest{CAMPAIGN}.edm4hep.root",
-    output:
-        rootfile=ANALYSISDIR+"beamlineTestAnalysis{CAMPAIGN}.root",
-        beamspot_canvas=ANALYSISDIR+"beamspot_{CAMPAIGN}.png",
-        x_px_canvas=ANALYSISDIR+"x_px_{CAMPAIGN}.png",
-        y_py_canvas=ANALYSISDIR+"y_py_{CAMPAIGN}.png",
-        fitted_position_means_stdevs_canvas=ANALYSISDIR+"fitted_position_means_stdevs_{CAMPAIGN}.png",
-        fitted_momentum_means_stdevs_canvas=ANALYSISDIR+"fitted_momentum_means_stdevs_{CAMPAIGN}.png",
-        pipe_parameter_canvas=ANALYSISDIR+"pipe_parameter_{CAMPAIGN}.png",
-    params:
-        xml=os.getenv("DETECTOR_PATH")+"/epic_ip6_extended.xml",
-    shell:
-        """
-            root -l -b -q '{input.script}("{input.data}", "{output.rootfile}", "{params.xml}",
-            "{output.beamspot_canvas}", "{output.x_px_canvas}", "{output.y_py_canvas}",
-            "{output.fitted_position_means_stdevs_canvas}", "{output.fitted_momentum_means_stdevs_canvas}",
-            "{output.pipe_parameter_canvas}")'
-        """
-
 rule beamline_acceptance_analysis:
     input:
         warmup="warmup/epic_ip6_extended.edm4hep.root",
         script=workflow.source_path("acceptanceAnalysis.C"),
         header=workflow.source_path("shared_functions.h"),
         data=SIMOUTDIR+"acceptanceTest{CAMPAIGN}.edm4hep.root",
     output:
-        rootfile=ANALYSISDIR+"acceptanceTestAnalysis{CAMPAIGN}.root",
-        beampipe_canvas=ANALYSISDIR+"acceptance_in_beampipe_{CAMPAIGN}.png",
-        etheta_canvas=ANALYSISDIR+"acceptance_energy_theta_{CAMPAIGN}.png",
-        etheta_acceptance_canvas=ANALYSISDIR+"acceptance_energy_theta_acceptance_{CAMPAIGN}.png",
-        entries_canvas=ANALYSISDIR+"acceptance_entries_{CAMPAIGN}.png",
+        rootfile=ANALYSISDIR+"electron_beamline_acceptanceTestAnalysis{CAMPAIGN}.root",
+        beampipe_canvas=ANALYSISDIR+"electron_beamline_acceptance_in_beampipe_{CAMPAIGN}.png",
+        etheta_canvas=ANALYSISDIR+"electron_beamline_acceptance_energy_theta_{CAMPAIGN}.png",
+        etheta_acceptance_canvas=ANALYSISDIR+"electron_beamline_acceptance_energy_theta_acceptance_{CAMPAIGN}.png",
+        entries_canvas=ANALYSISDIR+"electron_beamline_acceptance_entries_{CAMPAIGN}.png",
     params:
         xml=os.getenv("DETECTOR_PATH")+"/epic_ip6_extended.xml",
     shell:
@@ -83,28 +92,34 @@ rule beamline_acceptance_analysis:
             "{output.entries_canvas}")'
         """
 
+##########################################################################################
+# Combine results
+##########################################################################################
 rule beamline:
     input:
-        ANALYSISDIR+"beamlineTestAnalysis{CAMPAIGN}.root",
-        ANALYSISDIR+"beamspot_{CAMPAIGN}.png",
-        ANALYSISDIR+"x_px_{CAMPAIGN}.png",
-        ANALYSISDIR+"y_py_{CAMPAIGN}.png",
-        ANALYSISDIR+"fitted_position_means_stdevs_{CAMPAIGN}.png",
-        ANALYSISDIR+"fitted_momentum_means_stdevs_{CAMPAIGN}.png",
-        ANALYSISDIR+"pipe_parameter_{CAMPAIGN}.png",  
-        ANALYSISDIR+"acceptanceTestAnalysis{CAMPAIGN}.root",
-        ANALYSISDIR+"acceptance_in_beampipe_{CAMPAIGN}.png",
-        ANALYSISDIR+"acceptance_energy_theta_{CAMPAIGN}.png",
-        ANALYSISDIR+"acceptance_energy_theta_acceptance_{CAMPAIGN}.png",     
-        ANALYSISDIR+"acceptance_entries_{CAMPAIGN}.png", 
+        ANALYSISDIR+"electron_beamline_beamlineTestAnalysis{CAMPAIGN}.root",
+        ANALYSISDIR+"electron_beamline_beamspot_{CAMPAIGN}.png",
+        ANALYSISDIR+"electron_beamline_x_px_{CAMPAIGN}.png",
+        ANALYSISDIR+"electron_beamline_y_py_{CAMPAIGN}.png",
+        ANALYSISDIR+"electron_beamline_fitted_position_means_stdevs_{CAMPAIGN}.png",
+        ANALYSISDIR+"electron_beamline_fitted_momentum_means_stdevs_{CAMPAIGN}.png",
+        ANALYSISDIR+"electron_beamline_pipe_parameter_{CAMPAIGN}.png",  
+        ANALYSISDIR+"electron_beamline_acceptanceTestAnalysis{CAMPAIGN}.root",
+        ANALYSISDIR+"electron_beamline_acceptance_in_beampipe_{CAMPAIGN}.png",
+        ANALYSISDIR+"electron_beamline_acceptance_energy_theta_{CAMPAIGN}.png",
+        ANALYSISDIR+"electron_beamline_acceptance_energy_theta_acceptance_{CAMPAIGN}.png",
+        ANALYSISDIR+"electron_beamline_acceptance_entries_{CAMPAIGN}.png"
     output:
-        directory("results/beamline/{CAMPAIGN}/")
+        directory("results/beamline/steering_{CAMPAIGN}/")
     shell:
         """
             mkdir {output}        
-            cp {input} {output}
+            cp -r {input} {output}
         """
 
+##########################################################################################
+# Defualt running
+##########################################################################################
 rule beamline_local:
     input:
-        "results/beamline/local/"
+        "results/beamline/steering_local/"

benchmarks/beamline/beamlineAnalysis.C

@@ -316,6 +316,8 @@ int beamlineAnalysis(   TString inFile          = "/scratch/EIC/G4out/beamline/b
             std::cout << "Warning: Only " << h->GetEntries()/nEntries << " of particles contributing to histogram " <<  name
                       << " , which is below the accepted threshold of " << acceptableEntries/nEntries << std::endl;
             pass = 1;
+        } else{
+            std::cout << "Histogram " << name << " has " << h->GetEntries() << " entries." << std::endl;
         }
 
         // Get the pipe radius for this histogram

benchmarks/lowq2_reconstruction/LoadData.py

@@ -0,0 +1,15 @@
+import uproot
+import awkward as ak
+
+def create_arrays(dataFiles,featureName="_TaggerTrackerFeatureTensor_floatData",targetName="_TaggerTrackerTargetTensor_floatData", entries=None, treeName="events"):
+
+    # List of branches to load
+    branches = [featureName,targetName]
+
+    # Load data from concatenated list of files
+    data = uproot.concatenate([f"{file}:{treeName}" for file in dataFiles], branches, entry_stop=entries, library="ak")
+
+    input_data  = data[featureName]
+    target_data = data[targetName]
+
+    return input_data, target_data

benchmarks/lowq2_reconstruction/RegressionModel.py

@@ -0,0 +1,155 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import numpy as np
+
+class ProjectToX0Plane(nn.Module):
+    def forward(self, x):
+        # x shape: (batch, 6) -> [x, y, z, px, py, pz]
+        x0, y0, z0, px, py, pz = x.unbind(dim=1)
+
+        # Normalize momentum components
+        momentum = torch.sqrt(px**2 + py**2 + pz**2)
+        px_norm = px / momentum
+        py_norm = py / momentum
+        pz_norm = pz / momentum
+
+        # Avoid division by zero for px
+        # eps = 1e-8
+        # px_safe = torch.where(px_norm.abs() < eps, eps * torch.sign(px_norm) + eps, px_norm)
+        t = -x0 / px_norm
+
+        y_proj = y0 + py_norm * t
+        z_proj = z0 + pz_norm * t
+
+        # Output: [y_proj, z_proj, px_norm, py_norm]
+        return torch.stack([y_proj, z_proj, px_norm, py_norm], dim=1)
+
+    def project_numpy(self, arr):
+        """
+        Projects a numpy array of shape (N, 6) using the forward method,
+        returns a numpy array of shape (N, 4).
+        """
+        device = next(self.parameters()).device if any(p.device.type != 'cpu' for p in self.parameters()) else 'cpu'
+        x = torch.from_numpy(arr).float().to(device)
+        with torch.no_grad():
+            projected = self.forward(x)
+        return projected.cpu().numpy()
+
+class RegressionModel(nn.Module):
+    def __init__(self):
+        super(RegressionModel, self).__init__()
+        self.project_to_x0 = ProjectToX0Plane()
+        self.fc1  = nn.Linear(4, 512)
+        self.fc2  = nn.Linear(512, 64)
+        self.fc3  = nn.Linear(64, 3)  # Output layer for
+
+        # Normalization parameters
+        self.input_mean = nn.Parameter(torch.zeros(4), requires_grad=False)
+        self.input_std = nn.Parameter(torch.ones(4), requires_grad=False)
+        self.output_mean = nn.Parameter(torch.zeros(3), requires_grad=False)
+        self.output_std = nn.Parameter(torch.ones(3), requires_grad=False)
+
+
+    def forward(self, x):
+        # Conditionally apply projection
+        x = self.project_to_x0(x)
+        # Normalize inputs
+        x = (x - self.input_mean) / self.input_std
+
+        # Pass through the fully connected layers
+        x = self._core_forward(x)
+
+        # Denormalize outputs
+        x = x * self.output_std + self.output_mean
+        return x
+
+    def _core_forward(self, x):
+        # Core fully connected layers
+        x = torch.relu(self.fc1(x))
+        x = torch.relu(self.fc2(x))
+        x = self.fc3(x)
+        return x
+
+    def adapt(self, input_data, output_data):
+        # Normalization
+        self.input_mean.data = input_data.mean(dim=0)
+        self.input_std.data = input_data.std(dim=0)
+        self.output_mean.data = output_data.mean(dim=0)
+        self.output_std.data = output_data.std(dim=0)
+
+def preprocess_data(model, data_loader, adapt=True):
+    inputs  = data_loader.dataset.tensors[0]
+    targets = data_loader.dataset.tensors[1]
+
+
+    projected_inputs = model.project_to_x0(inputs)
+
+    # Compute normalization parameters
+    if adapt:
+        model.adapt(projected_inputs, targets)
+
+    # Normalize inputs and targets
+    normalized_inputs  = (projected_inputs - model.input_mean ) / model.input_std
+    normalized_targets = (targets          - model.output_mean) / model.output_std
+
+    # Replace the dataset with preprocessed data
+    data_loader.dataset.tensors = (normalized_inputs, normalized_targets)
+
+def makeModel():
+    # Create the model
+    model = RegressionModel()
+    # Define the optimizer
+    optimizer = optim.Adam(model.parameters(), lr=0.0004)
+    # Define the loss function
+    criterion = nn.HuberLoss(delta=0.2)  # Huber loss for regression tasks
+
+    return model, optimizer, criterion
+
+def trainModel(epochs, train_loader, val_loader, device):
+
+    model, optimizer, criterion = makeModel()
+
+    model.to(device)
+
+    # Preprocess training and validation data
+    preprocess_data(model, train_loader, adapt=True)
+
+    # Preprocess validation data without adapting
+    preprocess_data(model, val_loader, adapt=False)
+
+    # Move data to the GPU
+    train_loader.dataset.tensors = (train_loader.dataset.tensors[0].to(device), train_loader.dataset.tensors[1].to(device))
+    val_loader.dataset.tensors = (val_loader.dataset.tensors[0].to(device), val_loader.dataset.tensors[1].to(device))
+
+    # Verify that the model parameters are on the GPU
+    for name, param in model.named_parameters():
+        print(f"{name} is on {param.device}")
+
+    for epoch in range(epochs):
+        model.train()
+        running_loss = 0.0
+        for inputs, targets in train_loader:
+            optimizer.zero_grad()
+            outputs = model._core_forward(inputs)
+            loss = criterion(outputs, targets)
+            loss.backward()
+            optimizer.step()
+            running_loss += loss.item() * inputs.size(0)
+
+        epoch_loss = running_loss / len(train_loader.dataset)
+
+
+        # Validation step
+        model.eval()
+        val_loss = 0.0
+        with torch.no_grad():
+            for val_inputs, val_targets in val_loader:
+                val_outputs = model._core_forward(val_inputs)
+                val_loss += criterion(val_outputs, val_targets).item() * val_inputs.size(0)
+
+        val_loss /= len(val_loader.dataset)
+
+        print(f"Epoch [{epoch+1}/{epochs}], Loss: {epoch_loss}, Val Loss: {val_loss}")
+
+    return model