Calibrator.py

###################################################################################################
#
# Calibrator.py
#
# Copyright (C) by Andreas Zoglauer & Shreya Sareen
# All rights reserved.
#  
###################################################################################################


  
###################################################################################################


#from mpl_toolkits.mplot3d import Axes3D
#import matplotlib.pyplot as plt

import random 

import signal
import sys
import time
import math
import csv
import os
import argparse
from datetime import datetime
from functools import reduce

import torch
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt


print("\nA machine-learning based COSI calibrator")
print("========================+++++===========\n")



###################################################################################################
# Step 1: Input parameters
###################################################################################################


# Default parameters

NumberOfComptonEvents = 10000


OutputDirectory = "Output"


# Parse command line:

print("\nParsing the command line (if there is any)\n")

parser = argparse.ArgumentParser(description='Perform training and/or testing for gamma-ray burst localization')
parser.add_argument('-m', '--mode', default='toymodel', help='Choose an input data more: toymodel or simulations')
parser.add_argument('-t', '--toymodeloptions', default='10000:510.99:511.00', help='The toy-model options: source_events:energy_min:energy_max"')
parser.add_argument('-s', '--simulationoptions', default='', help='')
parser.add_argument('-b', '--batchsize', default='256', help='The number of GRBs in one training batch (default: 256 corresponsing to 5 degree grid resolution (64 for 3 degrees))')
parser.add_argument('-o', '--outputdirectory', default='Output', help='Name of the output directory. If it exists, the current data and time will be appended.')
parser

args = parser.parse_args()

  
Mode = (args.mode).lower()
if Mode != 'toymodel' and Mode != 'simulation':
  print("Error: The mode must be either \'toymodel\' or \'simulation\'")
  sys.exit(0)


if Mode == 'toymodel':
  print("CMD-Line: Using toy model")

  ToyModelOptions = args.toymodeloptions.split(":")
  if len(ToyModelOptions) != 3:
    print("Error: You need to give 3 toy model options. You gave {}. Options: {}".format(len(ToyModelOptions), ToyModelOptions))
    sys.exit(0)
  
  NumberOfComptonEvents = int(ToyModelOptions[0])
  if NumberOfComptonEvents <= 10:
    print("Error: You need at least 10 source events and not {}".format(NumberOfComptonEvents))
    sys.exit(0)       
  print("CMD-Line: Toy model: Using {} source events per GRB".format(NumberOfComptonEvents))

  MinimumEnergy = float(ToyModelOptions[1])
  if MinimumEnergy < 0:
    print("Error: You need a non-negative number for the energy, and not {}".format(MinimumEnergy))
    sys.exit(0)
  print("CMD-Line: Toy model: Using {} keV as minimum energy".format(MinimumEnergy))

  MaximumEnergy = float(ToyModelOptions[2])
  if MaximumEnergy < 0:
    print("Error: You need a non-negative number for the energy, and not {}".format(MinimumEnergy))
    sys.exit(0)
  print("CMD-Line: Toy model: Using {} keV as minimum energy".format(MinimumEnergy))

  '''
  NumberOfTrainingBatches = int(ToyModelOptions[3])
  if NumberOfTrainingBatches < 1:
    print("Error: You need a positive number for the number of traing batches and not {}".format(NumberOfTrainingBatches))
    sys.exit(0)
  print("CMD-Line: Toy model: Using {} training batches".format(NumberOfTrainingBatches))

  NumberOfTestingBatches = int(ToyModelOptions[4])
  if NumberOfTestingBatches < 1:
    print("Error: You need a positive number for the number of testing batches and not {}".format(NumberOfTestingBatches))
    sys.exit(0)
  print("CMD-Line: Toy model: Using {} testing batches".format(NumberOfTestingBatches))
  '''

elif Mode == 'simulation':
  print("Error: The simulation mode has not yet implemented")
  sys.exit(0)  

  
MaxBatchSize = int(args.batchsize)
if MaxBatchSize < 1 or MaxBatchSize > 1024:
  print("Error: The batch size must be between 1 && 1024")
  sys.exit(0)  
print("CMD-Line: Using {} as batch size".format(MaxBatchSize))
  
OutputDirectory = args.outputdirectory
# TODO: Add checks
print("CMD-Line: Using \"{}\" as output directory".format(OutputDirectory))

print("\n\n")


# Determine derived parameters


NumberOfTrainingEvents = int(0.8*NumberOfComptonEvents)
NumberOfTestingEvents = NumberOfComptonEvents - NumberOfTrainingEvents


'''
if os.path.exists(OutputDirectory):
  Now = datetime.now()
  OutputDirectory += Now.strftime("_%Y%m%d_%H%M%S")
    
os.makedirs(OutputDirectory)
'''


###################################################################################################
# Step 2: Global functions
###################################################################################################


# Take care of Ctrl-C
Interrupted = False
NInterrupts = 0
def signal_handler(signal, frame):
  global Interrupted
  Interrupted = True        
  global NInterrupts
  NInterrupts += 1
  if NInterrupts >= 2:
    print("Aborting!")
    sys.exit(0)
  print("You pressed Ctrl+C - waiting for graceful abort, or press  Ctrl-C again, for quick exit.")
signal.signal(signal.SIGINT, signal_handler)


# Everything ROOT related can only be loaded here otherwise it interferes with the argparse
#from CalibrationCreatorToyModel import CalibrationCreatorToyModel

# Load MEGAlib into ROOT so that it is usable
import ROOT as M
M.gSystem.Load("$(MEGALIB)/lib/libMEGAlib.so")
M.PyConfig.IgnoreCommandLineOptions = True



###################################################################################################
# Step 3: Create some training, test & verification data sets
###################################################################################################


print("Info: Creating {} Compton events".format(NumberOfComptonEvents))

from CalibrationData import CalibrationData

def generateOneDataSet(_):
  DataSet = CalibrationData()
  DataSet.create()
  return DataSet
  

# Create data sets
TimerCreation = time.time()

TrainingDataSets = []
for i in range(NumberOfTrainingEvents):
    result = generateOneDataSet(i)
    TrainingDataSets.append(generateOneDataSet(i))
print("Info: Created {:,} training data sets. ".format(NumberOfTrainingEvents))

TestingDataSets = []
for i in range(NumberOfTestingEvents):
    result = generateOneDataSet(i)
    TestingDataSets.append(generateOneDataSet(i))
print("Info: Created {:,} testing data sets. ".format(NumberOfTestingEvents))



TimeCreation = time.time() - TimerCreation
print("Info: Total time to create data sets: {:.1f} seconds (= {:,.0f} events/second)".format(TimeCreation, (NumberOfTrainingEvents + NumberOfTestingEvents) / TimeCreation))



###################################################################################################
# Step 4: Setting up the neural network
###################################################################################################
def dataset_to_tensors(datasets):

    X = []
    y = []

    for data in datasets:

        sh1 = data.StripHits[0]
        sh2 = data.StripHits[1]

        features = [
            sh1.DetectorID,
            sh1.StripID,
            sh1.IsHV,
            sh1.ADC,
            sh1.TAC,
            sh2.DetectorID,
            sh2.StripID,
            sh2.IsHV,
            sh2.ADC,
            sh2.TAC
        ]

        hit = data.Hits[0]

        target = [
            hit.X,
            hit.Y,
            hit.Z,
            hit.Energy
        ]

        X.append(features)
        y.append(target)

    X = torch.tensor(X, dtype=torch.float32)
    y = torch.tensor(y, dtype=torch.float32)

    return X, y


# Neural neutral

# -------------------------------
# Converting datasets to tensors
# -------------------------------
train_X, train_y = dataset_to_tensors(TrainingDataSets)
test_X, test_y   = dataset_to_tensors(TestingDataSets)

# -------------------------------
# Normalize inputs
# -------------------------------
X_mean = train_X.mean(0)
X_std  = train_X.std(0) + 1e-8

train_X = (train_X - X_mean) / X_std
test_X  = (test_X  - X_mean) / X_std


# -------------------------------
# Normalize outputs
# -------------------------------
y_mean = train_y.mean(0)
y_std  = train_y.std(0) + 1e-8

train_y_norm = (train_y - y_mean) / y_std
test_y_norm  = (test_y  - y_mean) / y_std


# -------------------------------
# Linear Model
# -------------------------------
model = nn.Sequential(
    nn.Linear(train_X.shape[1], 64),
    nn.ReLU(),
    nn.Linear(64, 64),
    nn.ReLU(),
    nn.Linear(64, 4)
)

# -------------------------------
# Loss + optimizer
# -------------------------------
loss_fn = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.01)

# -------------------------------
# Training
# -------------------------------
for epoch in range(2000):

    pred = model(train_X)  
    loss = loss_fn(pred, train_y_norm)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()


# -------------------------------
# Evaluate
# -------------------------------
with torch.no_grad():

    pred_norm = model(test_X)  

    # unnormalize predictions
    pred = pred_norm * y_std + y_mean


# -------------------------------
# Root Mean Squared Error
# -------------------------------
rmse = torch.sqrt(((pred - test_y) ** 2).mean(dim=0))

print("RMSE per output:")
print("X =", rmse[0].item())
print("Y =", rmse[1].item())
print("Z =", rmse[2].item())
print("Energy =", rmse[3].item())

# -------------------------------
# Plot
# -------------------------------
true_vals = test_y.detach().numpy()
pred_vals = pred.detach().numpy()

labels = ["X", "Y", "Z", "Energy"]

residuals = true_vals - pred_vals

plt.figure(figsize=(12,10))
# add units 
units = ["cm", "cm", "cm", "keV"]

plt.figure(figsize=(12,10))

for i in range(4):
    plt.subplot(2,2,i+1)

    plt.hist(residuals[:,i], bins=50, alpha=0.8)
    plt.axvline(0, color='red', linestyle='--')

    plt.xlabel(f"Measured - Reconstructed ({labels[i]} [{units[i]}])")
    plt.ylabel("Counts")
    plt.title(f"Residuals: {labels[i]} ({units[i]})")

plt.tight_layout()
plt.show()