Source code for pipeline.Embedding.process_custom

"""Custom functions for filtering and alterning an event.
"""

import torch
from torch_geometric.data import Data
from utils.tools import tarray
from utils.graphutils import batch2df


def edges_at_least_3_hits(batch: Data) -> Data:
    # at least 3 hits to be classified as valid edges
    # not_reconstructible_mask = batch.n_unique_planes < 3
    df_hits_particles = batch2df.get_df_hits_particles(
        batch=batch, particle_columns=["nhits_velo"]
    )
    df_edges = batch2df.get_df_edges(batch=batch)
    df_edges_particles = batch2df.merge_df_hits_particles_to_edges(
        df_edges=df_edges, df_hits_particles=df_hits_particles
    )
    df_edges_particles["nhits_velo_left"] = df_edges_particles[
        "nhits_velo_left"
    ].fillna(0)
    df_edges_particles["nhits_velo_right"] = df_edges_particles[
        "nhits_velo_right"
    ].fillna(0)

    df_edges_particles["y"] = (
        df_edges_particles["y"]
        & (
            df_edges_particles["nhits_velo_left"]
            == df_edges_particles["nhits_velo_right"]
        )
        & (df_edges_particles["nhits_velo_left"] >= 3)
    )

    df_combined_edges = (
        df_edges_particles.groupby(["hit_idx_left", "hit_idx_right"])["y"]
        .max()
        .rename("combined_y")
        .reset_index()
    )
    df_edges = df_edges.merge(
        df_combined_edges,
        on=["hit_idx_left", "hit_idx_right"],
        how="left",
    ).sort_values("edge_idx")

    batch["y"] = tarray.series_to_tensor(df_edges["combined_y"])
    return batch


def edges_at_least_3_planes(batch: Data) -> Data:
    # at least 3 hits to be classified as valid edges
    # not_reconstructible_mask = batch.n_unique_planes < 3
    df_hits_particles = batch2df.get_df_hits_particles(
        batch=batch, particle_columns=["n_unique_planes"]
    )
    df_edges = batch2df.get_df_edges(batch=batch)
    df_edges_particles = batch2df.merge_df_hits_particles_to_edges(
        df_edges=df_edges, df_hits_particles=df_hits_particles
    )
    df_edges_particles["n_unique_planes_left"] = df_edges_particles[
        "n_unique_planes_left"
    ].fillna(0)
    df_edges_particles["n_unique_planes_right"] = df_edges_particles[
        "n_unique_planes_right"
    ].fillna(0)

    df_edges_particles["y"] = (
        df_edges_particles["y"]
        & (
            df_edges_particles["n_unique_planes_left"]
            == df_edges_particles["n_unique_planes_right"]
        )
        & (df_edges_particles["n_unique_planes_left"] >= 3)
    )

    df_combined_edges = (
        df_edges_particles.groupby(["hit_idx_left", "hit_idx_right"])["y"]
        .max()
        .rename("combined_y")
        .reset_index()
    )
    df_edges = df_edges.merge(
        df_combined_edges,
        on=["hit_idx_left", "hit_idx_right"],
        how="left",
    ).sort_values("edge_idx")

    batch["y"] = tarray.series_to_tensor(df_edges["combined_y"])
    return batch


def remove_edges_in_same_plane(batch: Data) -> Data:
    edge_index_plane = batch.plane[batch.edge_index]
    no_self_edge_mask = edge_index_plane[0] != edge_index_plane[1]
    batch["edge_index"] = batch.edge_index[:, no_self_edge_mask]
    batch["y"] = batch.y[no_self_edge_mask]
    return batch


def edge_features_as_slope(batch: Data) -> Data:
    """Build edge features that correspond to the slope."""
    norm_x = batch.un_x / 14.5
    norm_y = batch.un_y / 14.5
    xe = batch.un_x[batch.edge_index]
    ye = batch.un_y[batch.edge_index]
    ze = batch.un_z[batch.edge_index]

    slopes_yz = (ye[1] - ye[0]) / (ze[1] - ze[0]) / 0.17
    slopes_xz = (xe[1] - xe[0]) / (ze[1] - ze[0]) / 0.17

    assert not torch.isnan(slopes_yz).any()
    assert not torch.isnan(slopes_xz).any()
    # Modify batch definition
    batch["x"] = torch.stack((norm_x, norm_y, batch["x"][:, 2]), dim=1).float()
    batch["edge_features"] = torch.stack((slopes_xz, slopes_yz), dim=1).float()

    return batch


def weights_inversely_proportional_to_nhits(batch: Data) -> Data:
    """Define edge weights that are inversely proportional to the number of hits."""
    node_weights = 7.0 / batch.n_unique_planes
    node_weights = torch.nan_to_num(node_weights, nan=1.0)
    edge_weights = torch.mean(node_weights[batch.edge_index], dim=0)
    batch.edge_weights = edge_weights * batch.edge_index.shape[1] / edge_weights.sum()
    assert not torch.isnan(batch.edge_weights).any(), str(batch.edge_weights)
    return batch