<jupyter_start><jupyter_text>Predict next track using RNN from Steve<jupyter_code>from __future__ import print_function
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
torch.__version__
from trackml.dataset import load_event
from trackml.randomize import shuffle_hits
from trackml.score import score_event
import os
import numpy as np
import pandas as pd
import glob
import math


import time
from utils import timeSince
from tqdm import tqdm

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import seaborn as sns
%matplotlib inline

class HitGausPredictor(nn.Module):
    """
    A PyTorch module for particle track state estimation and hit prediction.

    This module is an RNN which takes a sequence of hits and produces a
    Gaussian shaped prediction for the location of the next hit.
    """

    def __init__(self, hidden_dim=5, batch_size=64, device=None):
        super(HitGausPredictor, self).__init__()
        input_dim = 3
        output_dim = 2
        self.lstm = nn.LSTM(input_dim, hidden_dim, batch_first=True)
#         out_size = int(output_dim * (output_dim + 3) / 2)
        out_size = 4
        self.fc = nn.Linear(hidden_dim, out_size)
        self.device = device
        self.batch_size = batch_size
        

    def forward(self, x):
        """Might want to accept also the radius of the target layer."""
        input_size = x.size()

        # Initialize the LSTM hidden state
        h = (torch.zeros(self.lstm.num_layers, self.batch_size, self.lstm.hidden_size, device=self.device),
             torch.zeros(self.lstm.num_layers, self.batch_size, self.lstm.hidden_size, device=self.device))
        
        # Apply the LSTM module
        x, h = self.lstm(x, h)
        # Squash layer axis into batch axis
        x = x.contiguous().view(-1, x.size(-1))

        # Apply linear layer
        output = self.fc(x)
        
        # Extract and transform the gaussian parameters
        means = output[:, :2]
        variances = output[:, 2:4] # ensure it's positive
        variances = torch.exp(variances)

        # Expand the layer axis again, just for consistency/interpretability
        means = means.contiguous().view(self.batch_size, -1, 2)
        variances = variances.contiguous().view(self.batch_size, -1, 2)
        return means, variances
def gaus_llh_loss(outputs, targets):
    """Custom gaussian log-likelihood loss function"""
    means, covs = outputs
    
#     means = means_[:, 1:, :]
#     covs = covs_[:, 1:, :]
        
    batches = means.size(0)
    hits = means.size(1)
    

    # Flatten layer axis into batch axis to use batch matrix operations
    means = means.contiguous().view(means.size(0)*means.size(1), means.size(2))
    covs = covs.contiguous().view(covs.size(0)*covs.size(1), covs.size(2))
    targets = targets.contiguous().view(targets.size(0)*targets.size(1),
                                        targets.size(2))
    

#     # Calculate the residual error
    res = targets - means
    prob = 0.5*res*res /covs # + torch.log(covs)
#     logdet = 2 * torch.logdet(covs)

#     prob = prob.contiguous().view(batches, hits, -1)

    return torch.sum(prob)/batches/hits
batch_size = 64
model = HitGausPredictor(batch_size=batch_size)
import pickle
track_list_raw = pickle.load(open('ten_hists.npy', 'rb'))
for event in track_list_raw:
    for tracks in event:
        tracks[:, 0] = (tracks[:, 0] - 913.681763)/692.430542 # r
        tracks[:, 1] = (tracks[:, 1] - 0.009939)/1.823752     # phi
        tracks[:, 2] = (tracks[:, 2] + 2.315056)/1061.912476  # z
track_arrays = np.concatenate([np.array(x) for x in track_list_raw])
print(track_arrays.shape)
total_samples = track_arrays.shape[0]
n_training = int(total_samples/2) - 1
n_testing = n_training
plt.hist(track_arrays[:, :, 0].flatten())
plt.hist(track_arrays[:, :, 1].flatten())
plt.hist(track_arrays[:, :, 2].flatten())
batch_size = 64
from utils import random_choice
import random
def random_input():
    event = random_choice(track_list[0:1000])
    rand_start = random.randint(0, len(event)-batch_size-1)
    input = np.array(event[rand_start:rand_start+batch_size])
    input_t = torch.from_numpy(input)
    target_t = torch.from_numpy(input[:, 1:, 1:])
    return input_t, target_t
model = HitGausPredictor(batch_size=batch_size)
optimizer = optim.Adam(model.parameters())
input_t = torch.from_numpy(track_arrays[0:64, :-1, :])
print(input_t.size())
with torch.no_grad():
    output = model(input_t)
    print(output[0].size())
    print(output[1].size())
output[1]
gaus_llh_loss(output, torch.from_numpy(track_arrays[0:64, 1:, 1:]))
from utils import tunable_parameters
print('total tunable parameters:', tunable_parameters(model))
def train(n_epochs):
    loss_train = []
#     save_every = int(n_epochs/50)
    save_every = 1
    for i in tqdm(range(n_epochs)):
        
        sum_loss = 0.
        
        nbatches = int(n_training/batch_size)
        if i==0:
            print("number of batches:", nbatches)
        for ibatch in range(nbatches):
            start = ibatch*batch_size
            end = start + batch_size

            input = track_arrays[start:end]
            input_t = torch.from_numpy(input[:, :-1, :])
            target_t = torch.from_numpy(input[:, 1:, 1:])
    
        
            model.zero_grad()
            output_t = model(input_t)
            loss = gaus_llh_loss(output_t, target_t)
            loss.backward()
            optimizer.step()
            sum_loss += loss.item()
        
        if (i % save_every == 0):
            loss_train.append(sum_loss/nbatches)
    return loss_train
res = train(50)
plt.plot(np.arange(0, len(res)), res)
res
def cal_res(test_track):
    test_t = torch.from_numpy(test_track)
    target_t = torch.from_numpy(test_track[:, 1:, 1:])
    with torch.no_grad():
        means, covs = model(test_t)
        means = means[:, 1:, :]
        covs = covs[:, 1:, :]
        
        means = means.contiguous().view(means.size(0)*means.size(1), means.size(2))
        covs = covs.contiguous().view(covs.size(0)*covs.size(1), covs.size(2))
        target_t = target_t.contiguous().view(target_t.size(0)*target_t.size(1),
                                        target_t.size(2))
        
        res = means - target_t
        return res, covs
test_t = torch.from_numpy(track_arrays[0:64])
test_t[0]
with torch.no_grad():
    output = model(test_t)
print(output[0][0])
t_t = torch.from_numpy(np.array(track_list[900][0:64])[:, 1:, 1:])
print(gaus_llh_loss(output, t_t))
print(n)
residule, covs = cal_res(np.array(track_list[800][0:64]))
res_phi = residule.numpy()[:, 0].reshape(residule.size(0))
res_z = residule.numpy()[:, 1].reshape(residule.size(0))
plt.hist(res_phi)
plt.hist(res_z)
a = residule.numpy()[:, 1]
b = covs.numpy()[:, 1]
sns.jointplot(a, b)
a = residule.numpy()[:, 0]
b = covs.numpy()[:, 0]
sns.jointplot(a, b)
print(total_samples)
residule, covs = cal_res(track_arrays[894060:894060+64])
res_phi = residule.numpy()[:, 0]
res_z = residule.numpy()[:, 1]
plt.hist(res_phi/covs.numpy()[:, 0])
plt.hist(res_phi)
plt.hist(res_z)
a = residule.numpy()[:, 1]
b = covs.numpy()[:, 1]
sns.jointplot(a, b)
a = residule.numpy()[:, 0]
b = covs.numpy()[:, 0]
sns.jointplot(a, b)
target_t = torch.from_numpy(np.array(track_list[0:64])[:, 1:, :2])
input_t.size()
target_t.size()
loss_func = gaus_llh_loss
model = HitGausPredictor()
output = model(input_t)
print(output[0].size())
print(output[1].size())
gaus_llh_loss(output, target_t)
particles[particles['nhits'] == 10]
hits, cells, particles, truth = load_event("input/train_1/event000001000")
from utils import get_features
hits_features = get_features(hits)
hits_features['phi'].describe()
hits_features['z'].describe()
hits_features['r'].describe()
hits_features['j'] = (hits_features['r'] - 913.681763)/692.430542
hits_features['j'].describe()
torch.zeros<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.13917981072555205, 'fraction_numerical': 0.0294006309148265, 'mean_word_length': 3.2982646420824295, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '307293', 'license': 'MIT', 'n_tokens_mistral': 2980, 'n_tokens_neox': 2832, 'n_words': 603, 'path': 'ipynb/rnn steve.ipynb', 'repo_name': 'xju2/trackingML', 'size': '7925'}