gnn_tracking.graph_construction.graph_builder

gnn_tracking.graph_construction.graph_builder#

Classes#

GraphBuilder

Build graphs out of the input data.

Functions#

get_two_hop_tuples(→ set[tuple[int, int]])

Given a list of tuples (a, b), returns the set of tuples (x, y)

Module Contents#

gnn_tracking.graph_construction.graph_builder.get_two_hop_tuples(tuples: list[tuple[int, int]]) → set[tuple[int, int]]#

Given a list of tuples (a, b), returns the set of tuples (x, y) where (x, t) and (t, y) are in the input list.

class gnn_tracking.graph_construction.graph_builder.GraphBuilder(indir: str | os.PathLike, outdir: str | os.PathLike, *, pixel_only=True, redo=True, phi_slope_max=0.005, z0_max=200, dR_max=1.7, remove_intersecting=True, directed=False, measurement_mode=False, write_output=True, log_level=0, collect_data=True, edge_augmentation: str | None = None)#

Build graphs out of the input data.

Parameters:

  • indir

  • outdir

  • pixel_only – Only consider pixel detector

  • redo

  • phi_slope_max

  • z0_max

  • dR_max

  • remove_intersecting – Remove “ambiguous” edges, see Fig. 3 in “Charged particle tracking via edge-classifying interaction networks” http://arxiv.org/abs/2103.16701 and mark the remaining ones as incorrect edges

  • directed – Build directed edges

  • measurement_mode

  • write_output – Save graphs?

  • log_level

  • collect_data – Deprecated: Directly load the data into memory

  • edge_augmentation – Add more edges (e.g., adding next-neighbor connections). Needs remove_intersecting to be false

indir#
outdir#
pixel_only#
redo#
phi_slope_max#
z0_max#
dR_max#
feature_names#
feature_scale#
_data_list = []#
directed#
measurement_mode#
write_output#
measurements = []#
level#
logger#
_collect_data#
_remove_intersecting#
_edge_augmentation#
property data_list#
get_measurements()#
calc_dphi(phi1: numpy.ndarray, phi2: numpy.ndarray) → numpy.ndarray#

Computes phi2-phi1 given in range [-pi,pi]

calc_eta(r: numpy.ndarray, z: numpy.ndarray) → numpy.ndarray#

Computes pseudorapidity (https://en.wikipedia.org/wiki/Pseudorapidity)

get_dataframe(evt: torch_geometric.data.Data, evtid: int) → pandas.DataFrame#

Converts pytorch geometric data object to pandas dataframe

Parameters:

  • evt – pytorch geometric data object

  • evtid – event id

Returns:

pandas dataframe

select_edges(hits1: pandas.DataFrame, hits2: pandas.DataFrame, layer1: int, layer2: int) → pandas.DataFrame#

Build edges between two layers

Parameters:

  • hits1 – Information about hit 1

  • hits2 – Information about hit 2

  • layer1 – Layer number for hit 1

  • layer2 – Layer number for hit 2

Returns:

Dictionary containing edge indices and extra information

correct_truth_labels(hits: pandas.DataFrame, edges: pandas.DataFrame, y: numpy.ndarray, particle_ids: numpy.ndarray) → tuple[numpy.ndarray, int]#

Corrects for extra edges surviving the barrel intersection cut, i.e. for each particle counts the number of extra “transition edges” crossing from a barrel layer to an innermost endcap layer; the sum is n_incorrect - [edges] = n_edges x 2 - [y] = n_edges - [particle_ids] = n_edges

Returns:

corrected truth labels, number of incorrect edges

build_edges(hits: pandas.DataFrame) → tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray, numpy.ndarray]#

Build edges between hits

Parameters:

hits – Point cloud dataframe

Returns:

edge_index (2 x num edges), edge_attr (edge features x num edges), y (truth label, shape = num edges), edge_pt (pt of track belong to first hit)

to_pyg_data(point_cloud, edge_index: numpy.ndarray, edge_attr: numpy.ndarray, y: numpy.ndarray, evtid: int = -1, s: int = -1) → torch_geometric.data.Data#

Convert hit dataframe and edges to pytorch geometric data object

Parameters:

  • point_cloud – Hit dataframe, see get_dataframe

  • edge_index – See build_edges

  • edge_attr – See build_edges

  • y – See build_edges

  • evtid – Event ID

  • s – Sector

Returns:

Pytorch geometric data object

get_n_truth_edges(df: pandas.DataFrame) → dict[float, int]#
static get_event_id_sector_from_str(name: str) → tuple[int, int]#

Parses input file names.

Parameters:

name – Input file name

Returns:

Event id, sector Id

process(start=0, stop=1, *, only_sector: int = -1, progressbar=False)#

Main processing loop

Parameters:

  • start

  • stop

  • only_sector – Only process files for this sector. If < 0 (default): process all sectors.

Returns:

Contents