gnn_tracking.graph_construction.graph_builder#

Attributes#

DEFAULT_FEATURES

logger

Classes#

GraphBuilder

Build graphs out of the input data.

Functions#

get_logger([name, level])

Sets up global logger.

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.DEFAULT_FEATURES = ('r', 'phi', 'z', 'eta_rz', 'u', 'v', 'charge_frac', 'leta', 'lphi', 'lx', 'ly', 'lz', 'geta', 'gphi')#
gnn_tracking.graph_construction.graph_builder.get_logger(name='gnn-tracking', level=LOG_DEFAULT_LEVEL)#

Sets up global logger.

gnn_tracking.graph_construction.graph_builder.logger#
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

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: