<jupyter_start><jupyter_text>GraphNet<jupyter_code>from ariadne.graph_net.graph_utils.graph_prepare_utils import to_pandas_graph_from_df, get_pd_line_graph, \
    apply_nodes_restrictions, apply_edge_restriction, construct_output_graph
from ariadne.transformations import Compose, ConstraintsNormalize, ToCylindrical

from ariadne_v2.inference import IModelLoader

import torch

suff_df = ('_p', '_c')
gin.bind_parameter('get_pd_line_graph.restrictions_0',(-0.07, 0.07))
gin.bind_parameter('get_pd_line_graph.restrictions_1', (-0.32, 0.32))
gin.bind_parameter('get_pd_line_graph.suffix_c', '_c')
gin.bind_parameter('get_pd_line_graph.suffix_p','_p')
gin.bind_parameter('get_pd_line_graph.spec_kwargs', {'suffix_c': '_c', 
                                                     'suffix_p':'_p', 
                                                     'axes':['r', 'phi', 'z']} )
_edge_restriction = 0.15

class GraphModelLoader(IModelLoader):    
    def __call__(self):
        from ariadne.graph_net.model import GraphNet_v1
        import torch
        
        gin.bind_parameter('GraphNet_v1.input_dim', 5)
        gin.bind_parameter('GraphNet_v1.hidden_dim', 128)
        gin.bind_parameter('GraphNet_v1.n_iters', 1)
        
        def weights_update_g(model, checkpoint):
            model_dict = model.state_dict()    
            pretrained_dict =  checkpoint['state_dict']
            real_dict = {}
            for (k,v) in model_dict.items():
                needed_key = None
                for pretr_key in pretrained_dict:
                    if k in pretr_key:
                        needed_key = pretr_key
                        break
                assert needed_key is not None, "key %s not in pretrained_dict %r!" % (k, pretrained_dict.keys())
                real_dict[k] = pretrained_dict[needed_key]
        
            model.load_state_dict(real_dict)
            model.eval()
            return model
        
        path_g = '/zfs/hybrilit.jinr.ru/user/g/gooldan/graph/ariadne/lightning_logs/GraphNet_v1/version_32/epoch=201-step=271689.ckpt'
        #path_g = '/zfs/hybrilit.jinr.ru/user/g/gooldan/bes/ariadne/lightning_logs/version_63115/checkpoints/epoch=49.ckpt'
        #path_g = '/Users/egor/prog/dubna/ariadne/version_32/epoch=201-step=271689.ckpt'
        #path_g = 'C:\\Users\\egor\\dubna\\ariadne\\lightning_logs\\GraphNet_v1\\version_32\\epoch=201-step=271689.ckpt'
        
        checkpoint_g = torch.load(path_g) if torch.cuda.is_available() else torch.load(path_g, map_location=torch.device('cpu'))
        model_g = weights_update_g(model=GraphNet_v1(), 
                           checkpoint=checkpoint_g)
        model_hash = {"path_g":path_g, 'gin':gin.config_str(), 'model': '%r' % model_g, 'edge':_edge_restriction}
        return model_hash, model_g



from collections import namedtuple
GraphWithIndices = namedtuple('Graph', ['X', 'Ri', 'Ro', 'y', 'v1v2v3', 'ev_id' ])

transformer_g = Compose([
    DropSpinningTracks(),
    DropShort(),
    DropEmpty(),
    ToCylindrical(),
    ConstraintsNormalize(
        columns=('r', 'phi', 'z'),
        constraints = {'phi': [-3.15, 3.15], 'r': [80.0, 167.0], 'z': [-423.5, 423.5]},
        use_global_constraints = True
    ),
])


def construct_graph_with_indices(graph, v1v2v3, ev_id):
    return GraphWithIndices(graph.X, graph.Ri, graph.Ro, graph.y, v1v2v3, ev_id)


def get_graph(event):
    event = event[['event','x','y','z','station','track', 'index_old']]
    
    try:
        event = transformer_g(event)
    except AssertionError as err:
        print("ASS error %r" % err)
        return None
    
    event.index = event['index_old'].values
    event =  event[['event','r','phi','z','station','track']]
    
    G = to_pandas_graph_from_df(event, suffixes=suff_df, compute_is_true_track=True)

    
    nodes_t, edges_t = get_pd_line_graph(G, apply_nodes_restrictions)

    edges_filtered = apply_edge_restriction(edges_t, edge_restriction=_edge_restriction)
    graph = construct_output_graph(nodes_t, edges_filtered, ['y_p', 'y_c', 'z_p', 'z_c', 'z'],
                                     [np.pi, np.pi, 1., 1., 1.], 'edge_index_p', 'edge_index_c')
    ev_id = event.event.values[0]
    graph_with_inds = construct_graph_with_indices(graph,
                                                   edges_filtered[['from_ind', 'cur_ind', 'to_ind']].values, ev_id)
    return graph_with_inds

from ariadne.graph_net.dataset import collate_fn
def eval_event(tgt_graph, model_g):
    batch_input, batch_target = collate_fn([tgt_graph])
    with torch.no_grad():
        y_pred = model_g(batch_input['inputs']).numpy().flatten() > 0.15
    
    eval_df = pd.DataFrame(columns=['track_pred', 'hit_id_0', 'hit_id_1', 'hit_id_2'])
    eval_df['track_pred'] = y_pred
    eval_df[['hit_id_0', 'hit_id_1', 'hit_id_2']] = tgt_graph.v1v2v3
    return eval_df

N_STATIONS = 3

evaluator = EventEvaluator(parse_cfg, global_transformer, N_STATIONS)
events = evaluator.prepare(model_loader=GraphModelLoader())[0]
all_results = evaluator.build_all_tracks()
model_results = evaluator.run_model(get_graph, eval_event)
results_graphnet = evaluator.solve_results(model_results, all_results)<jupyter_output>read entry f91f90ab5797b64c7f7a8cbe13bb828c hit
[prepare]: started processing a df 3.txt with 225430 rows:
read entry 5639f1cfcd7cf42c31e8c9182ca5052b hit
[prepare] finished
[prepare] loading your model(s)...
[prepare] finished loading your model(s)...
[build_all_tracks] start
read entry 42ec67a68680c1434caa8b24ccee7444 hit
read entry 6f57b057fa883e3d891f2f5a7e573dd8 hit


processed: 5000: 100%|██████████| 4327/4327 [00:35<00:00, 122.26it/s]
[build_all_tracks] cache miss, finish
[run model] start
read entry d15e0a8b7ac8f4f1d0e0aaea19bf1458 hit
read entry 5b47a873a946fef433e055422a914bec hit


processed: 5000: 100%|██████████| 4327/4327 [04:23<00:00, 16.42it/s]
[run model] cache miss, finish
[solve results] start
[solve results] finish
[solve results] final stats:
==========EVALUATION RESULTS==========
Total events evaluated: 4327
Total tracks evaluated: 19659
Track Efficiency (recall): 0.9644 
Track Purity (precision): 0.6822 
Fully reconstructed event ratio: 0.8597
Mean cpu time per e[...]<jupyter_text>Tracknet<jupyter_code>import torch

from ariadne.graph_net.graph_utils.graph_prepare_utils import to_pandas_graph_from_df

from ariadne.utils.model import get_checkpoint_path, weights_update
from ariadne.tracknet_v2.model import TrackNETv2, TrackNETv2BES
from ariadne.tracknet_v2_1.model import TrackNetClassifier
from ariadne.transformations import Compose, ConstraintsNormalize, ToCylindrical, DropSpinningTracks, DropShort

from torch.nn import functional as F

import faiss


NUM_POINTS_TO_SEARCH = 5
threshold = 0.63

class TrackNetModelLoader(IModelLoader):    
    def __call__(self):
        tracknet_input_features=3
        tracknet_conv_features=32
        
        tracknet_ckpt_path_dict = {'model_dir': '/zfs/hybrilit.jinr.ru/user/p/pgonchar/ariadne/lightning_logs/TrackNETv2BES', 
                               'version': 'version_3', 'checkpoint': 'latest'}
        classifier_ckpt_path_dict = {'model_dir': '/zfs/store5.hydra.local/user/n/nuvard/ariadne/tracknet_clf_ovsmplg_10_grus_2_focal_pos_weight_5_gamma_2.00_alpha_0.50_bz_32/TrackNetClassifier', 
                                 'version': 'version_0', 'checkpoint': 'latest'}
    #     tracknet_ckpt_path_dict = {'model_dir': '/Users/egor/prog/dubna/models2606/zfs 2/hybrilit.jinr.ru/user/n/nuvard/ariadne/lightning_logs/TrackNETv2', 
    #                            'version': 'version_22', 'checkpoint': 'latest'}
    #     classifier_ckpt_path_dict = {'model_dir': '/Users/egor/prog/dubna/models2606/zfs 3/hybrilit.jinr.ru/user/n/nuvard/ariadne/lightning_logs/TrackNetClassifier', 
    #                              'version': 'version_108', 'checkpoint': 'latest'}
    
        path_to_tracknet_ckpt = get_checkpoint_path(**tracknet_ckpt_path_dict)
        path_to_classifier_ckpt = get_checkpoint_path(**classifier_ckpt_path_dict)
    
        model = weights_update(model=TrackNETv2BES(input_features=tracknet_input_features,
                                            conv_features=tracknet_conv_features,
                                            rnn_type='gru',
                                            batch_first=True),
                           checkpoint=torch.load(path_to_tracknet_ckpt, map_location=torch.device(DEVICE)))
        model.eval()
        
        model.to(DEVICE)
        class_model = weights_update(model=TrackNetClassifier(coord_size=3, num_gru_states=2),
                                 checkpoint=torch.load(path_to_classifier_ckpt, map_location=torch.device(DEVICE)))
        class_model.eval()
        class_model.to(DEVICE)
        model_hash = {
            "tracknet_ckpt_path_dict":path_to_tracknet_ckpt,
            "classifier_ckpt_path_dict":classifier_ckpt_path_dict,
            'gin':gin.config_str(), 
            'class_model': '%r' % class_model, 
            'model': '%r' % model,
            'edge':_edge_restriction,
            'NUM_POINTS_TO_SEARCH':NUM_POINTS_TO_SEARCH,
            'threshold':threshold
        }
        
        return model_hash, [model, class_model]

DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
NUM_COMPONENTS = 2


SUFX = ['_p', '_c']
COLS = ['z', 'phi','r']
NO_Z_COLS = ['z', 'phi'] 
_constraints = {'r': [80., 167.], 'phi': [-3.15, 3.15], 'z': [-423.5, 423.5]}
_columns=('r', 'phi', 'z')


tracknet_transformer = Compose([
        DropSpinningTracks(),
        DropShort(num_stations=3),
        ToCylindrical(drop_old=True),
        ConstraintsNormalize(columns=_columns, constraints=_constraints)
    ])

def build_index(target_df):
    cont = np.ascontiguousarray(target_df[NO_Z_COLS ].values)
    return store_in_index(cont, num_components=len(NO_Z_COLS ))

def search(points, index):
    cont = np.ascontiguousarray(points[NO_Z_COLS ].values)
    return search_in_index(cont, index, 10)

def to_cart(df):
    graph=to_pandas_graph_from_df(df, SUFX, save_index=False)
    seeds = graph[(graph.station_p == 0) & (graph.station_c == 1)]
    target = df[df.station == 2]
    return seeds, target.reset_index(drop=True)

def seeds_to_input(seeds_df):
    return seeds_df[['z_p', 'phi_p', 'r_p', 'z_c', 'phi_c', 'r_c']].values.reshape((-1, 2, 3))

def preprocess_one_event(event_df):
    
    event_df = event_df[['event','x','y','z','station','track', 'index_old']]
    try:
        event_df = tracknet_transformer(event_df)
    except AssertionError as err:
        print("ASS error %r" % err)
        return None
    #event_df.rename(columns={'index': 'index_old'}, inplace=True)
    
    seeds_unsorted, target = to_cart(event_df)

    batch = seeds_to_input(seeds_unsorted)
    chunk_data_len = torch.tensor(np.full(len(batch), 2), dtype=torch.int64).to(DEVICE)
    
    return seeds_unsorted, target, batch, chunk_data_len


def faiss_seeds_for_class(seeds, target, pred):
    index = faiss.IndexFlatL2(len(NO_Z_COLS))
    db = np.ascontiguousarray(target[NO_Z_COLS].values).astype('float32')
    index.add(db)
    what_search = np.ascontiguousarray(pred[:,-1,:2].astype('float32'))
    d, idx = index.search(what_search, NUM_POINTS_TO_SEARCH)

    return d, idx

def in_ellipse_check(pnts_x, ellipse_x, ellipse_dist_w, pnts_y, ellipse_y, ellipse_dist_h):
    return (np.square(pnts_x - ellipse_x) / np.square(ellipse_dist_w) + np.square(pnts_y - ellipse_y) / np.square(ellipse_dist_h)) <= 1

def filter_by_ellipse(seeds, target, pred, pred_gru):

    dist, idx = faiss_seeds_for_class(seeds, target, pred)
    idx = idx.reshape(-1)

    seeds[['pred_x', 'pred_y', 'pred_w', 'pred_h']] = pred[:, -1, :]

    selected_pnts = target.iloc[idx]
    new_idx = seeds.index.repeat(NUM_POINTS_TO_SEARCH)
    seeds_for_pnts = seeds.loc[new_idx]

    ellipse_x = seeds_for_pnts['pred_x'].values
    ellipse_y = seeds_for_pnts['pred_y'].values
    ellipse_dist_w = seeds_for_pnts['pred_w'].values
    ellipse_dist_h = seeds_for_pnts['pred_h'].values
    
    pnts_x = selected_pnts[['z']].values.T[0]
    pnts_y = selected_pnts[['phi']].values.T[0]
    pnts_z = selected_pnts[['r']].values.T[0]
    
    selected_pnts['in_ellipse'] = in_ellipse_check(pnts_x, ellipse_x, ellipse_dist_w, pnts_y, ellipse_y, ellipse_dist_h)
    
    seeds_for_pnts['index_next'] = selected_pnts.index_old.values
    seeds_for_pnts['in_ellipse'] = selected_pnts.in_ellipse.values
    seeds_for_pnts['z_next'] = selected_pnts.z.values
    seeds_for_pnts['phi_next'] = selected_pnts.phi.values
    seeds_for_pnts['r_next'] = selected_pnts.r.values
    seeds_for_pnts.track = np.logical_and(seeds_for_pnts.track.values, (seeds_for_pnts.track_c.values == selected_pnts.track.values))
    
    gru_for_pnts = np.repeat(pred_gru, NUM_POINTS_TO_SEARCH, 0)

    return seeds_for_pnts[seeds_for_pnts.in_ellipse.values], gru_for_pnts[seeds_for_pnts.in_ellipse.values]
    
def final_class_pred(class_model, filtered_gru, filtered, threshold, num_gru_states=2):
    changed_gru = filtered_gru[:, -num_gru_states:, :]
    if num_gru_states > 1:
        changed_gru = changed_gru.reshape(-1,changed_gru.shape[-1]*num_gru_states)
    class_pred = F.sigmoid(class_model(torch.tensor(changed_gru).to(DEVICE), 
                                       torch.tensor(filtered[['z_next', 'phi_next','r_next']].values).to(DEVICE)))
    filtered['class_pred_score'] = (class_pred).detach().cpu().numpy()
    filtered['track_pred'] = filtered.class_pred_score > threshold
    return filtered


def run_tracknet_eval(input_, model):    
    seeds, target, batch, chunk_data_len = input_
    model, class_model = model
    
    pred, pred_gru = model(torch.tensor(batch).to(DEVICE), chunk_data_len, return_gru_states=True)
    pred = pred.cpu().detach().numpy()
    pred_gru = pred_gru.cpu().detach().numpy()

    filtered, filtered_gru = filter_by_ellipse(seeds, target, pred, pred_gru)    
    final = final_class_pred(class_model, filtered_gru, filtered, threshold)
    
    result = final[['track_pred']]
    result['hit_id_0'] = final[['index_old_p']]
    result['hit_id_1'] = final[['index_old_c']]
    result['hit_id_2'] = final[['index_next']]
    
    return result
N_STATIONS = 3

evaluator = EventEvaluator(parse_cfg, global_transformer, N_STATIONS)
events = evaluator.prepare(model_loader=TrackNetModelLoader())[0]
all_results = evaluator.build_all_tracks()
model_results = evaluator.run_model(preprocess_one_event, run_tracknet_eval)
results_tracknet = evaluator.solve_results(model_results, all_results)<jupyter_output>read entry f91f90ab5797b64c7f7a8cbe13bb828c hit
[prepare]: started processing a df 3.txt with 225430 rows:
read entry 71a6900aedcc23e4e58dff85357dadf7 hit
[prepare] finished
[prepare] loading your model(s)...
[prepare] finished loading your model(s)...
[build_all_tracks] start
read entry 60ffa467b248517cc41f1d806c774283 hit
read entry 7d32631190fe6e11a2e300e617843753 hit


processed: 5000: 100%|██████████| 4327/4327 [00:35<00:00, 123.54it/s]
[build_all_tracks] cache miss, finish
[run model] start
read entry e3b42d8c8ff5ffe980e36691599eb74c hit
read entry 32ead15b67cd0ef6dcccf72f2f68bb09 hit


processed: 5000: 100%|██████████| 4327/4327 [02:58<00:00, 24.23it/s]
[run model] cache miss, finish
[solve results] start
[solve results] finish
[solve results] final stats:
==========EVALUATION RESULTS==========
Total events evaluated: 4327
Total tracks evaluated: 19983
Track Efficiency (recall): 0.9499 
Track Purity (precision): 0.7092 
Fully reconstructed event ratio: 0.8116
Mean cpu time per e[...]<jupyter_text>Kalman<jupyter_code>def kalman_func_efficiency_pt(ax, col_widths, color_box, color_line):
    draw_for_raw_data(
        ax,
        data_x=[ 0.05, 0.15, 0.25, 0.35, 0.45, 0.55, 0.65, 0.75, 0.85, 0.95, 1.05],
        data_y=[0.439614, 0.916571, 0.943983, 0.966013, 0.969349, 0.990453, 0.972414, 0.979275, 0.976744, 1, 1],
        data_y_err=[0.034498, 0.00934838, 0.00740631, 0.00655114, 0.00754448, 0.00475043, 0.00961773, 0.0102547, 0.016252, 0, 0],
        col_widths=col_widths,
        color_box=color_box, color_line=color_line, model_name="<NAME>", subtitle="efficiency"
    )
def kalman_func_efficiency_aphi(ax, col_widths, color_box, color_line):
    
    draw_for_raw_data(
        ax,
        data_x=[-3, -2.5, -2.2, -1.9, -1.6, -1.3, -0.9, -0.6, -0.266, 0, 0.3, 0.5, 0.9, 1.3, 1.5, 2.0, 2.3, 2.6, 2.9],    
        data_y=[0.941463, 0.941463, 0.920168,0.920168, 0.935872,0.935872, 0.923379,0.923379, 0.947791,0.947791, 0.934354, 0.923387, 0.923387, 0.911523,0.911523, 0.940945,0.940945, 0.934211, 0.934211],    
        data_y_err=[0.016396,0.016396, 0.0124228,0.0124228, 0.0109669,0.0109669, 0.0117897,0.0117897, 0.00996813,0.00996813, 0.0115851, 0.0119427,0.0119427, 0.012882,0.012882, 0.0104587,0.0104587, 0.0164185,0.0164185],
        col_widths=col_widths,
        color_box=color_box, color_line=color_line, model_name="Kalman Filter", subtitle="efficiency"
    )

def kalman_func_efficiency_cost(ax, col_widths, color_box, color_line):    
    draw_for_raw_data(
        ax,
        data_x=[ -0.9, -0.7, -0.5, -0.3, -0.1, 0.1, 0.3, 0.5, 0.7, 0.9],
        data_y=[0.850694, 0.921105, 0.947368, 0.960334, 0.962025, 0.977011, 0.972973, 0.940265, 0.911765, 0.790274],
        data_y_err=[0.0210004, 0.0119723, 0.0102456, 0.00891769, 0.00877912, 0.00718556, 0.00739396, 0.0111473, 0.0134912, 0.0224449],
        col_widths=col_widths,
        color_box=color_box, color_line=color_line, model_name="Kalman Filter", subtitle="efficiency"
    )


def kalman_func_purity_pt(ax, col_widths, color_box, color_line):
    draw_for_raw_data(
        ax,
        data_x=[ 0.05, 0.15, 0.25, 0.35, 0.45, 0.55, 0.65, 0.75, 0.85, 0.95, 1.05],
        data_y=[0.679104, 0.96977, 0.95288, 0.953548, 0.93186, 0.94533, 0.955932, 0.913043, 0.875, 0.875, 0.6],
        data_y_err=[0.0403272, 0.00595386, 0.00685681, 0.00755998, 0.0108138, 0.0108501, 0.0119499, 0.0195845, 0.0337539, 0.0522913, 0.154919],
        col_widths=col_widths,
        color_box=color_box, color_line=color_line, model_name="Kalman Filter", subtitle="purity"
    )
def kalman_func_purity_aphi(ax, col_widths, color_box, color_line):
    
    draw_for_raw_data(
        ax,
        data_x=[-3, -2.5, -2.2, -1.9, -1.6, -1.3, -0.9, -0.6, -0.266, 0, 0.3, 0.5, 0.9, 1.3, 1.5, 2.0, 2.3, 2.6, 2.9],    
        data_y=[0.984694,0.984694, 0.901235,0.901235, 0.945344, 0.945344, 0.941884,  0.941884,0.955466, 0.955466,0.961712, 0.952183, 0.952183, 0.938559,0.938559, 0.969574,0.969574, 0.972603,0.972603],    
        data_y_err=[0.0087691,0.0087691, 0.0135333,0.0135333, 0.010227,0.010227, 0.0104736,0.0104736, 0.00928094,0.00928094, 0.00910676, 0.00972924,0.00972924, 0.0110532,0.0110532, 0.00773551,0.00773551, 0.0110306, 0.0110306],
        col_widths=col_widths,
        color_box=color_box, color_line=color_line, model_name="Kalman Filter", subtitle="purity"
    )

def kalman_func_purity_cost(ax, col_widths, color_box, color_line):    
    draw_for_raw_data(
        ax,
        data_x=[ -0.9, -0.7, -0.5, -0.3, -0.1, 0.1, 0.3, 0.5, 0.7, 0.9],
        data_y=[0.935115, 0.974948, 0.957447, 0.964361, 0.928717, 0.900424, 0.94929, 0.93819, 0.961814, 0.992366],
        data_y_err=[0.0152179, 0.00714078, 0.00931053, 0.0084884, 0.0116117, 0.0137826, 0.00988149, 0.0113143, 0.00936249, 0.00537712],
        col_widths=col_widths,
        color_box=color_box, color_line=color_line, model_name="Kalman Filter", subtitle="purity"
    )<jupyter_output><empty_output><jupyter_text>Graphs<jupyter_code>from scipy.interpolate import make_interp_spline, BSpline
from scipy import stats

MODES = {
    "RECALL":1,
    "PRECISION":2
}

def get_diagram_arr_linspace(all_real_hits, found_hits, start, end, num, col, mode):
    spac = np.linspace(start, end, num=num)
    step = (spac[1] - spac[0]) / 2
    arr = []
    spac_ret = []
    track_count_for_arr = []
    for i in range(len(spac)-1):
        beg = spac[i]
        end = spac[i+1]
        elems_real = all_real_hits[(all_real_hits[col] >= beg) & (all_real_hits[col] < end)]
        elems_pred = found_hits[(found_hits[col] >= beg) & (found_hits[col] < end)]
        #print(beg,end)
        #print(len(elems_pred), len(elems_real))
        if mode == MODES["RECALL"]:
            if elems_real.empty:
                #arr.append(1.)
                continue
            else:
                arr.append(len(elems_pred) / len(elems_real))
                track_count_for_arr.append(len(elems_real))
        elif mode == MODES["PRECISION"]:
            if elems_pred.empty:
                #arr.append(1.)
                continue
            else:
                reco_true = len(elems_pred[elems_pred.track != -1])
                arr.append(reco_true / len(elems_pred))
                track_count_for_arr.append(len(elems_real))
        else:
            raise NotImplemented

        spac_ret.append(spac[i])
    
    return np.array(arr), np.array(spac_ret), np.array(track_count_for_arr)


def get_diagram_for_boxplot(all_real_hits, found_hits, start, end, num, col, bin_array, int_mode, mode):
    y_vals, x_vals, counts_ = get_diagram_arr_linspace(all_real_hits, found_hits, start, end, num, col, mode)

    np_y_vals = np.array(y_vals)
    bin_array = bin_array

    
    res_x_array = []
    res_box_data = []
    mean_box_array = []
    track_counts_ed = []
    for i in range(len(bin_array)-1):
        beg = bin_array[i]
        end = bin_array[i+1]
        y_ind = np.where((x_vals>=beg) & (x_vals<end))
        y_vals_corr  = np_y_vals[y_ind]
        track_counts_corr = counts_[y_ind]
        
        y_vals_corr = y_vals_corr[~np.isnan(y_vals_corr)]
        if len(y_vals_corr) == 0:
             continue
                
        #print(beg, end, i, stats.describe(y_vals_corr))
        #print(y_vals_corr)
#         if len(np.where(y_vals_corr < 1)) > 0:
#             print(y_vals_corr)
#             print(y_vals)
#             print(beg,end)
        
        res_box_data.append(y_vals_corr)
        track_counts_ed.append(np.sum(track_counts_corr))
        
        delta = 0 if int_mode else (end-beg)/2
        res_x_array.append(beg + delta)
        mean_box_array.append(np.mean(y_vals_corr))

    return res_box_data, np.array(res_x_array), np.array(mean_box_array), np.array(track_counts_ed)
    #plt.boxplot(res_box_data, positions=bin_array)

def boxplot_style(bp):
    #for box in bp['boxes']:
            # change outline color
        #box.set( color='#7570b3', linewidth=2)
            # change fill color
        #box.set( facecolor = 'silver' )
    
    ## change color and linewidth of the whiskers
    #for whisker in bp['whiskers']:
    #    whisker.set(color='#7570b3', linewidth=2)
    #
    ### change color and linewidth of the caps
    #for cap in bp['caps']:
    #    cap.set(color='#7570b3', linewidth=2)
    #
    ### change color and linewidth of the medians
    #for median in bp['medians']:
    #    median.set(color='tab:red', linewidth=3, zorder=30)

    for median in bp['means']:
        median.set(color='tab:green', linewidth=2, ls='-', zorder=10)
        
        
##########################
###########################

from numpy import linalg as LA
import matplotlib as mpl
mpl.rcParams['figure.dpi'] = 300

def draw_for_col(ax, tracks_real, tracks_pred_true,
                 col, col_pretty, total_events, n_ticks=150, 
                 n_avg_ticks=-1, x_ticks=8,
                 custom_title=None,
                 ticks_custom=True,
                 with_boxplot=False,
                 int_mode=False, save_disk=True,
                 custom_draw_funcs=[],
                 diagram_func=None,
                 color_ax_ticks=False,
                 subtitle=None,
                 model_name="NO_MODEL_NAME",
                 graph_idx=0, total_indices=-1, spacing=-1,
                 color_line='tab:orange', color_box='red', mode=-1
                ):
    real_xticks_count = x_ticks
    
    color_ax_ticks = len(custom_draw_funcs) > 0 and color_ax_ticks
    n_avg_ticks = n_ticks // 5 if n_avg_ticks < 0 else n_avg_ticks
    
    delta = 1e-4 if not int_mode else 1
    
    start = tracks_real[tracks_real[col] > -np.inf][col].min()
    end = tracks_real[tracks_real[col] < np.inf][col].max()+delta

    initial, spac, counts_ = get_diagram_arr_linspace(tracks_real, tracks_pred_true, start, end, n_ticks, col, mode)

    maxX = int(end) if int_mode else end
    ax.set_xlabel(col_pretty)
    ax.plot(spac, initial, alpha=0.0, lw=1)
    
    if not int_mode and ticks_custom:
        ax.set_xticks(np.linspace(start, maxX, x_ticks))
    else:
        ax.locator_params(axis='x', nbins=x_ticks)

    if diagram_func is None:
        ax.set_yticks(np.round(np.linspace(0, 1, 11), decimals=2))
    x_ticks = ax.get_xticks()
    
    if with_boxplot:
        old_ticks = x_ticks
        delta_x = (x_ticks[1] - x_ticks[0])/2

        diagram_func = get_diagram_for_boxplot if diagram_func is None else diagram_func
        
        box_data, ticks_x, mean_data, counts_ed = diagram_func(tracks_real, tracks_pred_true, start, end, 
                                                               n_ticks, col, x_ticks, int_mode, mode)
        if total_indices > 0:       
            
            width_tuned = delta_x * 2 / total_indices
            pos_tuned = ticks_x - delta_x
            spacings = width_tuned * 0.4
            width_real = width_tuned * 0.6
            
            ticks_x = pos_tuned + width_tuned * graph_idx + spacings/2 + width_real/2
            delta_x = width_real
        box_data_t = []
        ticks_x_t = []
        delta_x_t = []
        mean_data_t = []
        for i, box in enumerate(box_data):
            if len(box) > 4:
                box_data_t.append(box)
                ticks_x_t.append(ticks_x[i])
                mean_data_t.append(mean_data[i])

        box_data = np.array(box_data_t)
        ticks_x = np.array(ticks_x_t)
        mean_data = np.array(mean_data_t)

        bp = ax.boxplot(box_data, positions=ticks_x, 
                    manage_ticks=False, meanline=True, showmeans=True,
                    widths=delta_x,patch_artist=True, sym='',zorder=3,boxprops=dict(facecolor=color_box, alpha=0.3))
        boxplot_style(bp)
        ret_widths = delta_x
        # mean line
        xnew = np.linspace(ticks_x.min(), ticks_x.max(), 500)
        spl = make_interp_spline(ticks_x, mean_data, k=1)
        power_smooth = spl(xnew)
        ax.plot(xnew, power_smooth, ls='--', color=color_line, label=f'{model_name} mean track {subtitle}', lw=3, zorder=15)
        #ax.set_xticks(old_ticks)
        if graph_idx == 0:
            for i in range(len(ticks_x)):
                y_pos = 1.01
                x_pos = ticks_x[i] - delta_x/2
                ax.text(x_pos, y_pos, f"{counts_ed[i]}", zorder=80)
    
    if int_mode or not ticks_custom:
        ax.locator_params(axis='x', nbins=real_xticks_count)
    
    #if diagram_func is None:
    ax.set_ylim((-0.05, 1.05))
    ticks = ax.get_xticks()
    step = ticks[1] - ticks[0]
    if graph_idx == total_indices-1:
        ax.set_xlim((np.min(ticks_x) - width_tuned-step, np.max(ticks_x)+width_real + step/2))
    
    for draw_f in custom_draw_funcs:
        draw_f(ax)
    
    return ret_widths
    
    ####plt.locator_params(axis='y', nbins=16)
   

    #plt.tight_layout()
    #plt.rcParams['savefig.facecolor']='white'
    #os.makedirs('../output', exist_ok=True)
    #plt.savefig('../output/new_img_track_eff_%s_ev%r_t%d.png'%(col, total_events, n_ticks), dpi=300)
    #plt.show()

def draw_for_raw_data(ax, data_x, data_y, data_y_err, color_box, color_line, col_widths, model_name, subtitle):
    data_x = np.array(data_x)
    data_y_init = np.array(data_y)
    dataep = data_y + np.array(data_y_err)
    dataem = data_y - np.array(data_y_err)

    data_y = np.expand_dims(data_y,axis=-1)
    dataep = np.expand_dims(dataep,axis=-1)
    dataem = np.expand_dims(dataem,axis=-1)

    data_y = np.concatenate((data_y, dataep, dataem), axis=1).T
    delta_x = (data_x[1]-data_x[0]) / 2
    
    width_tuned = delta_x * 2 / 3
    pos_tuned = data_x - delta_x
    spacings = width_tuned * 0.4
    width_real = width_tuned * 0.6
    
    ticks_x = pos_tuned + width_tuned * 2 + spacings/2 + width_real/2
    
    bp = ax.boxplot(data_y, positions=ticks_x, 
                    manage_ticks=False, meanline=True, showmeans=True,
                    widths=col_widths, patch_artist=True, sym='',zorder=3,boxprops=dict(facecolor=color_box, alpha=0.3))
    boxplot_style(bp)
    xnew = np.linspace(ticks_x.min(), ticks_x.max(), 500)
    mean_data = data_y_init
    spl = make_interp_spline(ticks_x, mean_data, k=1)  # type: BSpline
    power_smooth = spl(xnew)
    
    ax.plot(xnew, power_smooth, ls='--', color=color_line, label=f'{model_name} mean track {subtitle}', lw=3, zorder=15)
    
    pass
    
def to_recall(all_tracks_df):
    tracks_real = all_tracks_df[all_tracks_df.pred != -1]
    tracks_pred_true = all_tracks_df[all_tracks_df.pred == 1]
    return tracks_real, tracks_pred_true, MODES["RECALL"], "efficiency"

def to_precision(all_tracks_df):
    tracks_real = all_tracks_df[all_tracks_df.track != -1]
    tracks_pred_true = all_tracks_df[(all_tracks_df.pred == 1) | (all_tracks_df.pred == -1)]
    return tracks_real, tracks_pred_true, MODES["PRECISION"], "purity"

def plot_model_results(model_names, model_results_arr, mode_func, funcs_for_kalman=None):
        
    fig = plt.figure(figsize=(14,16), dpi=80)
    ax1, ax2, ax3 = fig.subplots(nrows=3, ncols=1)
#    ax1, ax2 = fig.subplots(nrows=1, ncols=2, sharey='all')
    
    total_models = len(model_names) 
    if funcs_for_kalman:
        total_models+=1

    colormap = plt.cm.tab10 #nipy_spectral, Set1,Paired   
    colors = [colormap(i) for i in np.linspace(0, 1, 8)]
    
    widths_for_ax = []
    tit = ""
    for idx, (model_name, model_results) in enumerate(zip(model_names, model_results_arr)):

        all_tracks_df = model_results[0]
        
        all_tracks_df['pt'] = LA.norm(all_tracks_df[['px','py']].values, axis=1)
        all_tracks_df['cos_t'] = (all_tracks_df[['pz']].values/ LA.norm(all_tracks_df[['px','py','pz']].values, axis=1, keepdims=True))
        all_tracks_df['a_phi'] = np.arctan2(all_tracks_df[['px']].values, all_tracks_df[['py']].values)
        
        tracks_real, tracks_pred_true, mode, subtitle  = mode_func(all_tracks_df)

        color_box = colors[idx]        
        color_line = colors[idx]
        n_events = model_results_arr[0][0].event_id.nunique()
        tit = f"Models track {subtitle} on {n_events} events"
        ax1.set_title(tit)
        ax1.set_ylabel(f'Track {subtitle}', fontsize=12)
        #ax2.set_ylabel(f'Track {subtitle}', fontsize=12)
        #ax3.set_ylabel(f'Track {subtitle}', fontsize=12)
        
        widths_ax = draw_for_col(ax1, tracks_real, tracks_pred_true, 'pt', '$pt$', n_events, 350, 
                 n_avg_ticks=48, x_ticks=20, ticks_custom=False, with_boxplot=True, model_name=model_name,
                    graph_idx=idx, total_indices=total_models, subtitle=subtitle,
                    color_line=color_line, color_box=color_box, mode=mode)
        
        widths_for_ax.append([widths_ax])
        
        
        widths_ax = draw_for_col(ax2, tracks_real, tracks_pred_true, 'a_phi', '$a_\phi$',n_events, 350, 
                     n_avg_ticks=48, x_ticks=20, with_boxplot=True, model_name=model_name, 
                     graph_idx=idx, total_indices=total_models, subtitle=subtitle,
                     color_line=color_line, color_box=color_box, mode=mode)
        widths_for_ax[-1].append(widths_ax)
        #
        ##enable on big dataset
        widths_ax = draw_for_col(ax3, tracks_real, tracks_pred_true,'cos_t', '$cos_t$',n_events, 350, 
                     n_avg_ticks=48,x_ticks=20, ticks_custom=False, with_boxplot=True, model_name=model_name,
                     graph_idx=idx, total_indices=total_models, subtitle=subtitle,
                     color_line=color_line, color_box=color_box, mode=mode)
        widths_for_ax[-1].append(widths_ax)
        
    if funcs_for_kalman:
        for idx, ax in enumerate([ax1,ax2,ax3]):
            color_box = colors[-2]        
            color_line = colors[-2]
            funcs_for_kalman[idx](ax, widths_for_ax[-1][idx], color_box, color_line)
        
    for ax in [ax1,ax2,ax3]:
        ax.grid()
        ax.legend(bbox_to_anchor=(1, -0.05))
    
    fig = plt.gcf()
    handles, labels = ax1.get_legend_handles_labels()
    #fig.legend(handles, labels, loc='lower right', bbox_to_anchor=(1, -0.05))
    #plt.title()
    plt.tight_layout()
    plt.rcParams['savefig.facecolor']='white'
    
    os.makedirs('../output', exist_ok=True)
    plt.savefig('../output/%s.png'%(''.join([i if (ord(i) < 128) and (ord(i)>ord('a')) else '_' for i in tit])), dpi=300)
    
    plt.show()
    pass

plot_model_results(["TrackNetV3", "RDGraphNet"],
                   [results_tracknet, results_graphnet],
                   to_precision, funcs_for_kalman=[kalman_func_purity_pt, kalman_func_purity_aphi, kalman_func_purity_cost])
plot_model_results(["TrackNetV3", "RDGraphNet"],
                   [results_tracknet, results_graphnet],
                   to_recall, funcs_for_kalman=[kalman_func_efficiency_pt, kalman_func_efficiency_aphi, kalman_func_efficiency_cost])
def prepare_data_for_box(df, 
                         column_x, column_y, 
                         n_columns, int_mode):
    if int_mode:
        min_x = df[column_x].min()
        max_x = df[column_x].max()
        assert(pd.api.types.is_integer_dtype(min_x), f'not integer dtype minx={min_x}')
        
        box_data = []
        x_ticks = []
        if n_columns == -1:
            n_columns = max_x - min_x + 1
        else:
            assert(False)

        all_values = df[column_y]
        
        for start, finish in zip(range(min_x, max_x-1), range(min_x+1, max_x)):
            vals = all_values[(df[column_x] >= start) & (df[column_x] < finish)].values
            box_data.append(vals)
            x_ticks.append(start)
        
        vals = all_values[(df[column_x] >= finish)].values
        x_ticks.append(finish)
        box_data.append(vals)
            
    return np.array(box_data), np.array(x_ticks)

def draw_custom_boxplot(ax, draw_data, 
                        graph_idx, graphs_count, 
                        color_line, color_box, model_name, subtitle):    
    box_data, x_ticks = draw_data

    delta_x = (x_ticks[1] - x_ticks[0])/2
    
    if graphs_count > 0:
        width_tuned = delta_x * 2 / graphs_count
        pos_tuned = x_ticks - delta_x
        spacings = width_tuned * 0.7
        width_real = width_tuned * 0.3
        
        ticks_x = pos_tuned + width_tuned * graph_idx + spacings/2 + width_real/2
        delta_x = width_real
            
    bp = ax.boxplot(box_data, positions=ticks_x, 
                    manage_ticks=False, meanline=True, showmeans=True,
                    widths=delta_x,
                    patch_artist=True, sym='',zorder=3,
                    boxprops=dict(facecolor=color_box, alpha=0.3))
    mean_data = [arr.mean() for arr in box_data]
    boxplot_style(bp)
    xnew = np.linspace(ticks_x.min(), ticks_x.max(), 500)
    spl = make_interp_spline(ticks_x, mean_data, k=1)
    power_smooth = spl(xnew)
    ax.plot(xnew, power_smooth, ls='--', color=color_line, label=f'{model_name} mean processing speed', lw=3, zorder=15)
    start, end = ax.get_xlim()
    ax.xaxis.set_ticks(x_ticks)
    

def draw_speed_graphs(model_names, model_results_arr):
    fig = plt.figure(figsize=(14,6), dpi=80)
    ax1, ax2 = fig.subplots(nrows=1, ncols=2)
    
    colormap = plt.cm.tab10 #nipy_spectral, Set1,Paired   
    colors = [colormap(i) for i in np.linspace(0, 1, 8)]
    total_models = len(model_names) 

    for idx, (model_name, model_results) in enumerate(zip(model_names, model_results_arr)):
        all_events_df = model_results[1]
        events_per_sec_cpu = 1. / all_events_df['cpu_time'].values
        events_per_sec_gpu = 1. / all_events_df['gpu_time'].values
        all_events_df['cpu_per_sec'] = events_per_sec_cpu
        all_events_df['gpu_per_sec'] = events_per_sec_gpu
        
        color_box = colors[idx]        
        color_line = colors[idx]
        draw_data = prepare_data_for_box(all_events_df, 
                                         'multiplicity', 'cpu_per_sec', 
                                         n_columns=-1, int_mode=True)
        draw_custom_boxplot(ax1, draw_data, idx, total_models, 
                            color_line=color_line, color_box=color_box, 
                            model_name=model_name, subtitle='CPU')
        
        draw_data = prepare_data_for_box(all_events_df, 
                                         'multiplicity', 'gpu_per_sec', 
                                         n_columns=-1, int_mode=True)
        draw_custom_boxplot(ax2, draw_data, idx, total_models,
                            color_line=color_line, color_box=color_box,
                            model_name=model_name, subtitle='GPU')
        ax1.set_title('CPU')
        ax2.set_title('GPU')
    
    for ax in [ax1, ax2]:
        ax.grid()
    
    ax2.legend(bbox_to_anchor=(1, -0.1))
        
    fig = plt.gcf()
    handles, labels = ax1.get_legend_handles_labels()
    tit = 'processing speed'
    plt.tight_layout()
    plt.rcParams['savefig.facecolor']='white'
    
    os.makedirs('../output', exist_ok=True)
    
    plt.suptitle('Models processing speed (events per second) vs multiplicity')
    ax1.set_ylabel('events/sec')
    ax1.set_xlabel('multiplicity')
    ax2.set_xlabel('multiplicity')
    plt.savefig('../output/%s.png'%(''.join([i if (ord(i) < 128) and (ord(i)>ord('a')) else '_' for i in tit])), dpi=300)
    plt.show()
draw_speed_graphs(["RDGraphNet", 'TrackNetV2'], [results_graphnet, results_tracknet])<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1294731021407694, 'fraction_numerical': 0.05998367432919925, 'mean_word_length': 2.971763229449906, 'pattern_counts': {'":': 5, '<': 30, '<?xml version=': 0, '>': 32, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 29, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '634861', 'license': 'MIT', 'n_tokens_mistral': 15832, 'n_tokens_neox': 14250, 'n_words': 2282, 'path': 'notebooks/evaluate_modern.ipynb', 'repo_name': 't3hseus/ariadne', 'size': '37984'}