graphnet.datasets.prometheus_datasets module

Public datasets from Prometheus Simulation.

class graphnet.datasets.prometheus_datasets.PublicPrometheusDataset(graph_definition, download_dir, truth, features, backend, train_dataloader_kwargs, validation_dataloader_kwargs, test_dataloader_kwargs)[source]

Bases: ERDAHostedDataset

A generic class for public Prometheus Datasets hosted using ERDA.

Construct CuratedDataset.

Parameters:

  • graph_definition (GraphDefinition) – Method that defines the data representation.

  • download_dir (str) – Directory to download dataset to.

  • truth (Optional) – List of event-level truth to include. Will include all available information if not given.

  • features (Optional) – List of input features from pulsemap to use. If not given, all available features will be used.

  • backend (Optional) – data backend to use. Either “parquet” or “sqlite”. Defaults to “parquet”.

  • train_dataloader_kwargs (Optional) – Arguments for the training DataLoader. Default None.

  • validation_dataloader_kwargs (Optional) – Arguments for the validation DataLoader, Default None.

  • test_dataloader_kwargs (Optional) – Arguments for the test DataLoader. Default None.

class graphnet.datasets.prometheus_datasets.TRIDENTSmall(graph_definition, download_dir, truth, features, backend, train_dataloader_kwargs, validation_dataloader_kwargs, test_dataloader_kwargs)[source]

Bases: PublicPrometheusDataset

Public Dataset for Prometheus simulation of a TRIDENT geometry.

Contains ~ 1 million track events between 10 GeV - 10 TeV.

Construct CuratedDataset.

Parameters:

  • graph_definition (GraphDefinition) – Method that defines the data representation.

  • download_dir (str) – Directory to download dataset to.

  • truth (Optional) – List of event-level truth to include. Will include all available information if not given.

  • features (Optional) – List of input features from pulsemap to use. If not given, all available features will be used.

  • backend (Optional) – data backend to use. Either “parquet” or “sqlite”. Defaults to “parquet”.

  • train_dataloader_kwargs (Optional) – Arguments for the training DataLoader. Default None.

  • validation_dataloader_kwargs (Optional) – Arguments for the validation DataLoader, Default None.

  • test_dataloader_kwargs (Optional) – Arguments for the test DataLoader. Default None.

class graphnet.datasets.prometheus_datasets.PONESmall(graph_definition, download_dir, truth, features, backend, train_dataloader_kwargs, validation_dataloader_kwargs, test_dataloader_kwargs)[source]

Bases: PublicPrometheusDataset

Public Dataset for Prometheus simulation of a P-ONE geometry.

Contains ~ 1 million track events between 10 GeV - 10 TeV.

Construct CuratedDataset.

Parameters:

  • graph_definition (GraphDefinition) – Method that defines the data representation.

  • download_dir (str) – Directory to download dataset to.

  • truth (Optional) – List of event-level truth to include. Will include all available information if not given.

  • features (Optional) – List of input features from pulsemap to use. If not given, all available features will be used.

  • backend (Optional) – data backend to use. Either “parquet” or “sqlite”. Defaults to “parquet”.

  • train_dataloader_kwargs (Optional) – Arguments for the training DataLoader. Default None.

  • validation_dataloader_kwargs (Optional) – Arguments for the validation DataLoader, Default None.

  • test_dataloader_kwargs (Optional) – Arguments for the test DataLoader. Default None.

class graphnet.datasets.prometheus_datasets.BaikalGVDSmall(graph_definition, download_dir, truth, features, backend, train_dataloader_kwargs, validation_dataloader_kwargs, test_dataloader_kwargs)[source]

Bases: PublicPrometheusDataset

Public Dataset for Prometheus simulation of a Baikal-GVD geometry.

Contains ~ 1 million track events between 10 GeV - 10 TeV.

Construct CuratedDataset.

Parameters:

  • graph_definition (GraphDefinition) – Method that defines the data representation.

  • download_dir (str) – Directory to download dataset to.

  • truth (Optional) – List of event-level truth to include. Will include all available information if not given.

  • features (Optional) – List of input features from pulsemap to use. If not given, all available features will be used.

  • backend (Optional) – data backend to use. Either “parquet” or “sqlite”. Defaults to “parquet”.

  • train_dataloader_kwargs (Optional) – Arguments for the training DataLoader. Default None.

  • validation_dataloader_kwargs (Optional) – Arguments for the validation DataLoader, Default None.

  • test_dataloader_kwargs (Optional) – Arguments for the test DataLoader. Default None.