IceCube Neutrino Reconstruction Using PCA and RNNs

sensor_geometry = pd.read_csv('/Volumes/public/Python/icecube/sensor_geometry.csv')
sensor_geometry.head()

plotly.offline.init_notebook_mode()

trace = go.Scatter3d(
    x=sensor_geometry['x'],
    y=sensor_geometry['y'],
    z=sensor_geometry['z'],
    mode='markers',
    marker={'color': sensor_geometry['sensor_id'], 'colorscale': 'Viridis', 'size': 2, 'colorbar': dict(thickness=20)}
)

layout = go.Layout(margin={'l': 0, 'r': 0, 'b': 0, 't': 0})

plot_figure = go.Figure(data=[trace], layout=layout)

plotly.offline.iplot(plot_figure)

train_meta = pd.read_parquet('/Volumes/public/Python/icecube/train_meta.parquet')
train_meta.head()

train_batch_7 = pd.read_parquet('/Volumes/public/Python/icecube/train/batch_7.parquet')
train_batch_7.head()

plt.hist(train_batch_7['charge'], bins=30, range=(0, 6))
plt.xlabel('Pulse Charge [photoelectrons]')
plt.ylabel('Count')
plt.show()

def plot_event(event_id):

    # Get pulse-detecting sensors
    sample_event = train_batch_7.loc[event_id, 'sensor_id']
    sample_sensors = pd.merge(sensor_geometry, sample_event, on='sensor_id', how='inner')

    # Convert polar coordinates (azimuth, zenith) to cartesian coordinates (x, y, z)
    azimuth = train_meta[train_meta['event_id'] == event_id]['azimuth'].squeeze()
    zenith = train_meta[train_meta['event_id'] == event_id]['zenith'].squeeze()
    r = 500
    x_line = r * np.cos(azimuth) * np.sin(zenith)
    y_line = r * np.sin(azimuth) * np.sin(zenith)
    z_line = r * np.cos(zenith)

    # Set up plotly to display in notebook mode
    plotly.offline.init_notebook_mode()

    # Create traces for each scatterplot
    trace = go.Scatter3d(
        x=sensor_geometry['x'],
        y=sensor_geometry['y'],
        z=sensor_geometry['z'],
        mode='markers',
        marker={'size': 2, 'opacity': 0.8, 'color': 'gray'},
        showlegend=False
    )
    sample_trace = go.Scatter3d(
        x=sample_sensors['x'],
        y=sample_sensors['y'],
        z=sample_sensors['z'],
        mode='markers',
        marker={'color': train_batch_7.loc[event_id, 'time'], 'colorscale': 'Viridis', 'size': 5, 'colorbar': dict(thickness=20)}
    )
    label_trace = go.Scatter3d(
        x=[0, x_line],
        y=[0, y_line],
        z=[0, z_line],
        line=dict(color='orange', width=10),
        showlegend=False
    )

    # Define layout margins
    layout = go.Layout(margin={'l': 0, 'r': 0, 'b': 0, 't': 0})

    # Plot the scatterplots
    plot_figure = go.Figure(data=[trace, sample_trace, label_trace], layout=layout)
    plotly.offline.iplot(plot_figure)

plot_event(event_id=20936501)

# Read in parquet data for a single batch, and reset the index
batch_7 = pd.read_parquet('/Volumes/public/Python/icecube/train/batch_7.parquet')
batch_7 = batch_7.reset_index()

# Only keep higher-quality data
batch_7 = batch_7[batch_7['auxiliary'] == False]

# Ignore the charge columns since we won't be using them, and drop duplicates
batch_7 = batch_7[['event_id', 'sensor_id', 'time']]
batch_7 = batch_7.drop_duplicates(subset=['event_id', 'sensor_id'], keep='first')

# Read in sensor_geometry data, and merge with parquet data
sensor_geometry = pd.read_csv('/Volumes/public/Python/icecube/sensor_geometry.csv', low_memory=False)
batch_geo = pd.merge(batch_7, sensor_geometry, on='sensor_id', how='inner')
display(batch_geo.sample(5))
print(batch_geo.shape)

from sklearn.preprocessing import StandardScaler
from sklearn.decomposition import PCA

# Define a DataFrame to hold the results
results = pd.DataFrame(columns=['event_id', 'pca_x', 'pca_y', 'pca_z', 'pca_var', 
                                            'pca_s_x', 'pca_s_y', 'pca_s_z', 'pca_s_var', 
])
results['event_id'] = results['event_id'].astype(int)

# Go through each unique event
for event_id in batch_geo['event_id'].unique():

    # Define the training data
    train = batch_geo[batch_geo['event_id'] == event_id][['x', 'y', 'z', 'time']].to_numpy()

    # Define the scaler, and scale the training data
    scaler = StandardScaler()
    scaled_train = scaler.fit_transform(train)

    # Fit PCA to both scaled and unscaled training data
    pca = PCA(n_components=1, svd_solver='full').fit(train)
    scaled_pca = PCA(n_components=1, svd_solver='full').fit(scaled_train)

    # If x and y are equal to 0, make x a small number so that the azimuth and zenith angles aren't undefined
    if (pca.components_[0][0] == 0) and (pca.components_[0][1] == 0):
        pca.components_[0][0] = 1e-06
    if (scaled_pca.components_[0][0] == 0) and (scaled_pca.components_[0][1] == 0):
        scaled_pca.components_[0][0] = 1e-06

    # Flip the direction of the PCA vectors if time is positive
    if pca.components_[0][3] > 0:
        pca.components_[0] = - pca.components_[0]
    if scaled_pca.components_[0][3] > 0:
        scaled_pca.components_[0] = - scaled_pca.components_[0]
    
    # Save all of this information to a DataFrame
    results.loc[len(results.index)] = [event_id, pca.components_[0][0], pca.components_[0][1], pca.components_[0][2], pca.explained_variance_ratio_[0],
                                                 scaled_pca.components_[0][0], scaled_pca.components_[0][1], scaled_pca.components_[0][2], scaled_pca.explained_variance_ratio_[0],
    ]

display(results.head(5))
print(results.shape)

# Convert from cartesian to spherical coordinates to get the direction of the event — designed to work on DataFrame columns
def cart_to_spherical(df_x, df_y, df_z):
    x2_and_y2 = (df_x ** 2) + (df_y ** 2)
    zenith = np.arccos(df_z / np.sqrt(x2_and_y2 + (df_z ** 2)))
    azimuth = np.sign(df_y) * np.arccos(df_x / np.sqrt(x2_and_y2))
    return azimuth, zenith

# Make conversions
results['pca_azi'], results['pca_zen'] = cart_to_spherical(results['pca_x'], results['pca_y'], results['pca_z'])
results['pca_s_azi'], results['pca_s_zen'] = cart_to_spherical(results['pca_s_x'], results['pca_s_y'], results['pca_s_z'])

# Convert from (-pi, +pi) to (0, 2pi) — this is for single values, not DataFrame columns
def angle_to_2pi(angle):
    angle_2pi = angle % (2 * np.pi)
    if angle_2pi < 0:
        angle_2pi += 2 * np.pi
    return angle_2pi

# Fix angles
results['pca_azi'] = results['pca_azi'].apply(angle_to_2pi)
results['pca_s_azi'] = results['pca_s_azi'].apply(angle_to_2pi)

display(results.head(5))
print(results.shape)

# Get ground truth labels
train_meta = pd.read_parquet('/Volumes/public/Python/icecube/train_meta.parquet')
train_meta = train_meta[['event_id', 'azimuth', 'zenith']]

# Add ground truth labels to the DataFrame
results['event_id'] = results['event_id'].astype(int)
results = pd.merge(results, train_meta, on='event_id', how='left')

def angular_dist(az_true, zen_true, az_pred, zen_pred):
    
    if not (np.all(np.isfinite(az_true)) and
            np.all(np.isfinite(zen_true)) and
            np.all(np.isfinite(az_pred)) and
            np.all(np.isfinite(zen_pred))):
        raise ValueError("All arguments must be finite")
    
    # Pre-compute all sine and cosine values
    sa1 = np.sin(az_true)
    ca1 = np.cos(az_true)
    sz1 = np.sin(zen_true)
    cz1 = np.cos(zen_true)
    
    sa2 = np.sin(az_pred)
    ca2 = np.cos(az_pred)
    sz2 = np.sin(zen_pred)
    cz2 = np.cos(zen_pred)
    
    # Scalar product of the two cartesian vectors (x = sz*ca, y = sz*sa, z = cz)
    scalar_prod = sz1*sz2*(ca1*ca2 + sa1*sa2) + (cz1*cz2)
    
    # Scalar product of two unit vectors is always between -1 and 1, this is against nummerical instability
    scalar_prod =  np.clip(scalar_prod, -1, 1)
    
    # Convert back to an angle (in radian)
    return np.abs(np.arccos(scalar_prod))

# Calculate the angular error
pca_error = angular_dist(results['azimuth'],
                         results['zenith'],
                         results['pca_azi'],
                         results['pca_zen'])
pca_s_error = angular_dist(results['azimuth'],
                           results['zenith'],
                           results['pca_s_azi'],
                           results['pca_s_zen'])
print(f'PCA: {np.average(pca_error):.2f} rad')
print(f'PCA-scaled: {np.average(pca_s_error):.2f} rad')
results['pca_error'] = pca_error

# Make histogram plots
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(13, 5))
ax1.hist(pca_error, bins=20, label='PCA Model', fc=(1, 0.5, 0.2, 0.5), range=(0, np.pi))
ax1.set_title('Angular Error')
ax1.set_xlabel('Ground Truth Difference [rad]')
ax1.set_ylabel('Count')
ax1.legend()
ax2.hist(pca_s_error, bins=20, label='Scaled PCA Model', fc=(0, 0.5, 1, 0.5), range=(0, np.pi))
ax2.set_title('Angular Error')
ax2.set_xlabel('Ground Truth Difference [rad]')
ax2.set_ylabel('Count')
ax2.legend()
plt.show()

# -----------------------------------------------------------------------------
# Imports

import argparse
import glob
import pandas as pd
import numpy as np
import csv


# -----------------------------------------------------------------------------
# Process parquet files

def process_data(input_file, split_name):

    # Read the input file, and reset the index
    batch = pd.read_parquet(input_file)
    batch = batch.reset_index()

    # Only keep higher-quality data
    batch = batch[batch['auxiliary'] == False]

    # Ignore the charge columns since we won't be using them, and re-order sensor_id in order of time
    batch = batch[['event_id', 'sensor_id', 'time']]
    batch = batch.sort_values(by=['event_id', 'time'], axis=0, ascending=True)

    # Add sensor_id by 1: 0-5159 --> 1-5160 because the 0 token is reserved for padding
    batch['sensor_id'] = batch['sensor_id'] + 1

    # Create a nested list of all events [ [event #1's sensors in order], [event #2's sensors in order], ...] and corresponding IDs
    g = batch.groupby('event_id')
    events = g['sensor_id'].apply(list).tolist()
    event_ids = list(g.groups.keys())

    # Truncate events longer than the maximum sequence length
    events = [e[:args.sequencelen] for e in events]

    # Pad events with 0, if shorter than the maximum sequence length
    events = [e + [0] * (args.sequencelen - len(e)) for e in events]

    # Get target values (labels)
    labels = meta.reindex(event_ids).values.tolist()

    # Combine the events and labels into a single list
    events_labels = [e + l for e, l in zip(events, labels)]

    # Save as CSV file
    with open(args.outputdir+split_name+'_'+input_file.split('/')[-1].split('.')[0][6:]+'.csv', 'w') as f:
        write = csv.writer(f)
        write.writerows(events_labels)


# -----------------------------------------------------------------------------

if __name__ == "__main__":

    # Define command line arguments
    parser = argparse.ArgumentParser(description='IceCube data preprocessing')
    parser.add_argument('-metafile', type=str, default='/Volumes/public/Python/icecube/train_meta.parquet', help='metadata filepath')
    parser.add_argument('-parquetdir', type=str, default='/Volumes/public/Python/icecube/train/', help='filepath for the directory containing parquet files')
    parser.add_argument('-trainfrac', type=float, default=0.9, help='fraction of the dataset to be used for the training set')
    parser.add_argument('-sequencelen', type=int, default=300, help='max length of each example in the dataset')
    parser.add_argument('-outputdir', type=str, default='/Volumes/public/Python/icecube/train_2023-02-27/', help='filepath for the output directory')

    # Parse arguments
    args = parser.parse_args()

    # Read the metadata file
    print(f'Reading metadata file...')
    meta = pd.read_parquet(args.metafile)
    meta = meta[['event_id', 'azimuth', 'zenith']]
    meta = meta.set_index('event_id')

    # Grab all of the parquet files, shuffle, and split into training/validation sets
    parquets = glob.glob(args.parquetdir + '*.parquet')
    print(f'Found {len(parquets)} parquet files')
    parquets = np.random.permutation(parquets)
    train_parquets = parquets[:int(len(parquets) * args.trainfrac)]
    valid_parquets = parquets[int(len(parquets) * args.trainfrac):]
    print(f'Split into {len(train_parquets)} training parquets and {len(valid_parquets)} validation parquets')

    # For each parquet file, pre-process the dataset and save to file
    for parquet in train_parquets:
        print(f'Processing train {parquet}')
        process_data(parquet, split_name='train')
    for parquet in valid_parquets:
        print(f'Processing valid {parquet}')
        process_data(parquet, split_name='valid')

class IceCubeDataset(Dataset):
    def __init__(self, data, targ):
        self.data = data
        self.targ = targ

    def __len__(self):
        return len(self.targ)

    def __getitem__(self, idx):
        return self.data[idx], self.targ[idx]

valid_df = pd.read_csv(valid_filepath, names=range(302))
valid_data = torch.tensor(valid_df.iloc[:, :args.trunc].values, dtype=torch.long)
valid_targ = torch.tensor(valid_df.iloc[:, -2:].values, dtype=torch.float32)
valid_dataset = IceCubeDataset(valid_data, valid_targ)
valid_dl = DataLoader(valid_dataset, batch_size=args.bs, shuffle=False, num_workers=args.workers, pin_memory=pinning_mem)

class GRUModel(nn.Module):
    def __init__(self, vocab_size, emb_dim, hidden_dim, n_layers, fc1_dim, output_size):
        super().__init__()
        self.vocab_size = vocab_size
        self.emb_dim = emb_dim
        self.hidden_dim = hidden_dim
        self.n_layers = n_layers
        self.fc1_dim = fc1_dim
        self.output_size = output_size

        self.embedding = nn.Embedding(num_embeddings=self.vocab_size, embedding_dim=self.emb_dim, padding_idx=0)

        self.gru = nn.GRU(input_size=self.emb_dim, hidden_size=self.hidden_dim, num_layers=self.n_layers, 
                          bias=True, batch_first=True, bidirectional=False)
        for name, param in self.gru.named_parameters():
            if 'bias' in name:
                nn.init.constant_(param, 0.0)
            elif 'weight_ih' in name:
                nn.init.xavier_uniform_(param, nn.init.calculate_gain('tanh'))
            elif 'weight_hh' in name:
                nn.init.orthogonal_(param, nn.init.calculate_gain('tanh'))

        self.fc1 = nn.Linear(self.hidden_dim, self.fc1_dim)
        for name, param in self.fc1.named_parameters():
            if 'bias' in name:
                nn.init.constant_(param, 0.0)
            elif 'weight' in name:
                nn.init.xavier_uniform_(param, gain=nn.init.calculate_gain('relu'))

        self.relu = nn.ReLU()

        self.fc2 = nn.Linear(self.fc1_dim, self.output_size)
        for name, param in self.fc2.named_parameters():
            if 'bias' in name:
                nn.init.constant_(param, 0.0)
            elif 'weight' in name:
                nn.init.xavier_uniform_(param, gain=1)

    def forward(self, x):
        seq_lengths = torch.count_nonzero(x, dim=1)
        gru_output, _ = self.gru(self.embedding(x))
        output_last = torch.stack([gru_output[i, l - 1, :] for i, l in enumerate(seq_lengths)], dim=0)
        final_output = self.fc2(self.relu(self.fc1(output_last)))

        return final_output