gnn_tracking.graph_construction.graph_builder
=============================================

.. py:module:: gnn_tracking.graph_construction.graph_builder


Classes
-------

.. autoapisummary::

   gnn_tracking.graph_construction.graph_builder.GraphBuilder


Functions
---------

.. autoapisummary::

   gnn_tracking.graph_construction.graph_builder.get_two_hop_tuples


Module Contents
---------------

.. py:function:: 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.


.. py:class:: 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.

   :param indir:
   :param outdir:
   :param pixel_only: Only consider pixel detector
   :param redo:
   :param phi_slope_max:
   :param z0_max:
   :param dR_max:
   :param 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
   :param directed: Build directed edges
   :param measurement_mode:
   :param write_output: Save graphs?
   :param log_level:
   :param collect_data: Deprecated: Directly load the data into memory
   :param edge_augmentation: Add more edges (e.g., adding next-neighbor connections).
                             Needs ``remove_intersecting`` to be false


   .. py:attribute:: indir


   .. py:attribute:: outdir


   .. py:attribute:: pixel_only
      :value: True



   .. py:attribute:: redo
      :value: True



   .. py:attribute:: phi_slope_max
      :value: 0.005



   .. py:attribute:: z0_max
      :value: 200



   .. py:attribute:: dR_max
      :value: 1.7



   .. py:attribute:: feature_names
      :value: ('r', 'phi', 'z', 'eta_rz', 'u', 'v', 'charge_frac', 'leta', 'lphi', 'lx', 'ly', 'lz', 'geta', 'gphi')



   .. py:attribute:: feature_scale


   .. py:attribute:: _data_list
      :value: []



   .. py:attribute:: directed
      :value: False



   .. py:attribute:: measurement_mode
      :value: False



   .. py:attribute:: write_output
      :value: True



   .. py:attribute:: measurements
      :value: []



   .. py:attribute:: logger


   .. py:attribute:: _collect_data
      :value: True



   .. py:attribute:: _remove_intersecting
      :value: True



   .. py:attribute:: _edge_augmentation
      :value: None



   .. py:property:: data_list


   .. py:method:: get_measurements()


   .. py:method:: calc_dphi(phi1: numpy.ndarray, phi2: numpy.ndarray) -> numpy.ndarray

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



   .. py:method:: calc_eta(r: numpy.ndarray, z: numpy.ndarray) -> numpy.ndarray

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



   .. py:method:: get_dataframe(evt: torch_geometric.data.Data, evtid: int) -> pandas.DataFrame

      Converts pytorch geometric data object to pandas dataframe

      :param evt: pytorch geometric data object
      :param evtid: event id

      :returns: pandas dataframe



   .. py:method:: select_edges(hits1: pandas.DataFrame, hits2: pandas.DataFrame, layer1: int, layer2: int) -> pandas.DataFrame

      Build edges between two layers

      :param hits1: Information about hit 1
      :param hits2: Information about hit 2
      :param layer1: Layer number for hit 1
      :param layer2: Layer number for hit 2

      :returns: Dictionary containing edge indices and extra information



   .. py:method:: 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



   .. py:method:: build_edges(hits: pandas.DataFrame) -> tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray, numpy.ndarray]

      Build edges between hits

      :param 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)



   .. py:method:: 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

      :param point_cloud: Hit dataframe, see `get_dataframe`
      :param edge_index: See `build_edges`
      :param edge_attr: See `build_edges`
      :param y: See `build_edges`
      :param evtid: Event ID
      :param s: Sector

      :returns: Pytorch geometric data object



   .. py:method:: get_n_truth_edges(df: pandas.DataFrame) -> dict[float, int]


   .. py:method:: get_event_id_sector_from_str(name: str) -> tuple[int, int]
      :staticmethod:


      Parses input file names.

      :param name: Input file name

      :returns: Event id, sector Id



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

      Main processing loop

      :param start:
      :param stop:
      :param only_sector: Only process files for this sector. If < 0 (default): process
                          all sectors.

      Returns: