graphnet-team / graphnet

A Deep learning library for neutrino telescopes
https://graphnet-team.github.io/graphnet/
Apache License 2.0
90 stars 92 forks source link

`CuratedDataset` - A method for sharing pre-converted datasets #689

Closed RasmusOrsoe closed 5 months ago

RasmusOrsoe commented 5 months ago

This PR adds functionality for sharing pre-converted and configured datasets as importable DataModules with possibly fixed train, val and test selections. Closes #545. Intended usage is shown below:

from graphnet.models.detector.prometheus import ORCA150
from graphnet.models.graphs import KNNGraph
from graphnet.datasets import TestDataset

# Pick GraphDefinition
graph_definition = KNNGraph(detector = ORCA150())

# Instantiate CuratedDataset - It will download and unzip, build dataloaders
data_module = TestDataset(download_dir='',
                    graph_definition=graph_definition,
                    train_dataloader_kwargs = {'batch_size': 2,
                                                "num_workers": 1})

# Get description of Dataset - this is an auto-generated text using the class properties:
data_module.description()

# Get Dataloaders
train_dataloader = data_module.train_dataloader
val_dataloder = data_module.val_dataloader

# Go Train

The resulting terminal output is:

TestDataset contains data from ARCA Prometheus Simulation and was added to GraphNeT by Rasmus F. Ørsøe. 

 COMMENTS ON USAGE: 
 Rasmus F. Ørsøe: This Dataset should be used for unit tests only. Simulation produced by Stephan Meighen-Berger, U. Melbourne. 

 DATASET DETAILS: 
 pulsemaps: ['photons'] 
 truth table: mc_truth 
 input features: ['sensor_pos_x', 'sensor_pos_y', 'sensor_pos_z', 't']
 pulse truth: None 
 event truth: ['interaction', 'initial_state_energy', 'initial_state_type', 'initial_state_zenith', 'initial_state_azimuth', 'initial_state_x', 'initial_state_y', 'initial_state_z'] 
 Number of training events: 9 
 Number of validation events: 1 
 Number of test events: 0 

In this PR, the following functionality is added:

RasmusOrsoe commented 5 months ago

@Aske-Rosted thank you very much for giving this a read! I've updated the selection and commented on the default value of the backend argument. Please let me know if anything is unclear still.

Aske-Rosted commented 5 months ago

@Aske-Rosted thank you very much for giving this a read! I've updated the selection and commented on the default value of the backend argument. Please let me know if anything is unclear still.

All good