gnn_tracking.utils.loading
==========================

.. py:module:: gnn_tracking.utils.loading


Attributes
----------

.. autoapisummary::

   gnn_tracking.utils.loading.DEFAULT_FEATURES


Classes
-------

.. autoapisummary::

   gnn_tracking.utils.loading.TrackingDataset
   gnn_tracking.utils.loading.TrackingDataModule
   gnn_tracking.utils.loading.TestTrackingDataModule


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

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


.. py:class:: TrackingDataset(in_dir: str | os.PathLike | list[str] | list[os.PathLike], *, start=0, stop=None, sector: int | None = None, point_cloud_builder: gnn_tracking.preprocessing.point_cloud_builder.TrackMLPointCloudBuilder | gnn_tracking.preprocessing.point_cloud_builder.CMSPointCloudBuilder | gnn_tracking.preprocessing.point_cloud_builder.MDPointCloudBuilder | None, feature_subset_names: list[str] | None = None, pt_cut: float | None = None)

   Bases: :py:obj:`torch_geometric.data.Dataset`


   Dataset for tracking applications

   :param in_dir: Directory or list of directories containing the data files
   :param start: Index of the first file to be considered (with files from the
                 in dirs considered in order)
   :param stop: Index of the last file to be considered
   :param sector: If not None, only files with this sector number will be considered


   .. py:attribute:: point_cloud_builder


   .. py:attribute:: _processed_paths
      :value: []



   .. py:attribute:: file_number
      :value: 0



   .. py:attribute:: prev_file_number
      :value: -1



   .. py:attribute:: sector_results
      :value: []



   .. py:attribute:: pt_cut
      :value: None



   .. py:method:: _get_paths(in_dir: str | os.PathLike | list[str] | list[os.PathLike], *, start=0, stop: int | None = None, sector: int | None = None) -> list[pathlib.Path]

      Collect all paths that should be in this dataset.



   .. py:method:: len() -> int


   .. py:method:: get(idx: int) -> torch_geometric.data.Data


   .. py:method:: _make_pt_cut(data: torch_geometric.data.Data) -> torch_geometric.data.Data


   .. py:method:: _cut_to_pixel_cms(data: torch_geometric.data.Data) -> torch_geometric.data.Data


.. py:class:: TrackingDataModule(*, identifier: str, train: dict | None = None, val: dict | None = None, test: dict | None = None, predict: dict | None = None, cpus: int = 1, builder_params: dict | None = None, feature_subset_names: list[str] | None = None, pt_cut: float | None = None)

   Bases: :py:obj:`pytorch_lightning.LightningDataModule`


   This subclass of `LightningDataModule` configures all data for the
   ML pipeline.

   :param identifier: Identifier of the dataset (e.g., `graph_v5`)
   :param train: Config dictionary for training data (see below)
   :param val: Config dictionary for validation data (see below)
   :param test: Config dictionary for test data (see below)
   :param cpus: Number of CPUs to use for loading data.

   The following keys are available for each config dictionary:

   - `dirs`: List of dirs to load from (required)
   - `start=0`: Index of first file to load
   - `stop=None`: Index of last file to load
   - `sector=None`: Sector to load from (if None, load all sectors)
   - `batch_size=1`: Batch size

   Training has the following additional keys:

   - `sample_size=None`: Number of samples to load for each epoch
       (if None, load all samples)


   .. py:attribute:: _configs


   .. py:attribute:: _datasets


   .. py:attribute:: _cpus
      :value: 1



   .. py:attribute:: builder_params
      :value: None



   .. py:attribute:: feature_subset_names
      :value: None



   .. py:attribute:: pt_cut
      :value: None



   .. py:property:: datasets
      :type: dict[str, TrackingDataset]



   .. py:method:: _fix_datatypes(dct: dict[str, Any] | None) -> dict[str, Any] | None
      :staticmethod:


      Fix datatypes of config dictionary.
      This is necessary because when configuring values from the command line,
      all values might be strings.



   .. py:method:: _get_dataset(key: str) -> TrackingDataset


   .. py:method:: setup(stage: str) -> None


   .. py:method:: _get_dataloader(key: str) -> torch_geometric.loader.DataLoader


   .. py:method:: train_dataloader()


   .. py:method:: val_dataloader()


   .. py:method:: test_dataloader()


   .. py:method:: predict_dataloader()


.. py:class:: TestTrackingDataModule(graphs: list[torch_geometric.data.Data])

   Bases: :py:obj:`pytorch_lightning.LightningDataModule`


   Version of `TrackingDataLoader` only used for testing purposes.


   .. py:attribute:: graphs


   .. py:attribute:: datasets


   .. py:method:: setup(stage: str) -> None


   .. py:method:: train_dataloader()


   .. py:method:: val_dataloader()


   .. py:method:: test_dataloader()