|
| 1 | +"""Snowstorm dataset module hosted on the IceCube Collaboration servers.""" |
| 2 | + |
| 3 | +import pandas as pd |
| 4 | +import re |
| 5 | +import os |
| 6 | +from typing import Dict, Any, Optional, List, Tuple, Union |
| 7 | +from glob import glob |
| 8 | +from sklearn.model_selection import train_test_split |
| 9 | + |
| 10 | +from graphnet.data.constants import FEATURES, TRUTH |
| 11 | +from graphnet.data.curated_datamodule import IceCubeHostedDataset |
| 12 | +from graphnet.data.utilities import query_database |
| 13 | +from graphnet.models.graphs import GraphDefinition |
| 14 | + |
| 15 | +AVAILABLE_RUN_IDS = [ |
| 16 | + *list(range(22010, 22019)), |
| 17 | + *list(range(22042, 22051)), |
| 18 | + *list(range(22078, 22087)), |
| 19 | +] |
| 20 | + |
| 21 | + |
| 22 | +class SnowStormDataset(IceCubeHostedDataset): |
| 23 | + """IceCube SnowStorm simulation dataset. |
| 24 | +
|
| 25 | + More information can be found at |
| 26 | + https://wiki.icecube.wisc.edu/index.php/SnowStorm_MC#File_Locations |
| 27 | + This is a IceCube Collaboration simulation dataset. |
| 28 | + Requires a username and password. |
| 29 | + """ |
| 30 | + |
| 31 | + _experiment = "IceCube SnowStorm dataset" |
| 32 | + _creator = "Aske Rosted" |
| 33 | + _citation = "arXiv:1909.01530" |
| 34 | + _available_backends = ["sqlite"] |
| 35 | + |
| 36 | + _pulsemaps = ["SRTInIcePulses"] |
| 37 | + _truth_table = "truth" |
| 38 | + _pulse_truth = None |
| 39 | + _features = FEATURES.SNOWSTORM |
| 40 | + _event_truth = TRUTH.SNOWSTORM |
| 41 | + _data_root_dir = "/data/ana/graphnet/Snowstorm_l2" |
| 42 | + |
| 43 | + def __init__( |
| 44 | + self, |
| 45 | + run_ids: List[int], |
| 46 | + graph_definition: GraphDefinition, |
| 47 | + download_dir: str, |
| 48 | + truth: Optional[List[str]] = None, |
| 49 | + features: Optional[List[str]] = None, |
| 50 | + train_dataloader_kwargs: Optional[Dict[str, Any]] = None, |
| 51 | + validation_dataloader_kwargs: Optional[Dict[str, Any]] = None, |
| 52 | + test_dataloader_kwargs: Optional[Dict[str, Any]] = None, |
| 53 | + ): |
| 54 | + """Construct SnowStormDataset. |
| 55 | +
|
| 56 | + Args: |
| 57 | + run_ids: List of RunIDs to include. |
| 58 | + graph_definition: Method that defines the data representation. |
| 59 | + download_dir: Directory to download dataset to. |
| 60 | + truth (Optional): List of event-level truth to include. Will |
| 61 | + include all available information if not given. |
| 62 | + features (Optional): List of input features from pulsemap to use. |
| 63 | + If not given, all available features will be |
| 64 | + used. |
| 65 | + train_dataloader_kwargs (Optional): Arguments for the training |
| 66 | + DataLoader. Default None. |
| 67 | + validation_dataloader_kwargs (Optional): Arguments for the |
| 68 | + validation DataLoader, Default None. |
| 69 | + test_dataloader_kwargs (Optional): Arguments for the test |
| 70 | + DataLoader. Default None. |
| 71 | + """ |
| 72 | + assert all( |
| 73 | + [i in AVAILABLE_RUN_IDS for i in run_ids] |
| 74 | + ), f"RunIDs must be in {AVAILABLE_RUN_IDS}. You provided {run_ids}" |
| 75 | + self._run_ids = run_ids |
| 76 | + self._zipped_files = [ |
| 77 | + os.path.join(self._data_root_dir, f"{s}.tar.gz") for s in run_ids |
| 78 | + ] |
| 79 | + |
| 80 | + super().__init__( |
| 81 | + graph_definition=graph_definition, |
| 82 | + download_dir=download_dir, |
| 83 | + truth=truth, |
| 84 | + features=features, |
| 85 | + backend="sqlite", |
| 86 | + train_dataloader_kwargs=train_dataloader_kwargs, |
| 87 | + validation_dataloader_kwargs=validation_dataloader_kwargs, |
| 88 | + test_dataloader_kwargs=test_dataloader_kwargs, |
| 89 | + ) |
| 90 | + |
| 91 | + def _prepare_args( |
| 92 | + self, backend: str, features: List[str], truth: List[str] |
| 93 | + ) -> Tuple[Dict[str, Any], Union[List[int], None], Union[List[int], None]]: |
| 94 | + """Prepare arguments for dataset.""" |
| 95 | + assert backend == "sqlite" |
| 96 | + dataset_paths = [] |
| 97 | + for rid in self._run_ids: |
| 98 | + dataset_paths += glob( |
| 99 | + os.path.join(self.dataset_dir, str(rid), "**/*.db"), |
| 100 | + recursive=True, |
| 101 | + ) |
| 102 | + |
| 103 | + # get event numbers from all datasets |
| 104 | + event_no = [] |
| 105 | + |
| 106 | + # get RunID |
| 107 | + pattern = rf"{re.escape(self.dataset_dir)}/(\d+)/.*" |
| 108 | + event_counts: Dict[str, int] = {} |
| 109 | + event_counts = {} |
| 110 | + for path in dataset_paths: |
| 111 | + |
| 112 | + # Extract the ID |
| 113 | + match = re.search(pattern, path) |
| 114 | + assert match |
| 115 | + run_id = match.group(1) |
| 116 | + |
| 117 | + query_df = query_database( |
| 118 | + database=path, |
| 119 | + query=f"SELECT event_no FROM {self._truth_table}", |
| 120 | + ) |
| 121 | + query_df["path"] = path |
| 122 | + event_no.append(query_df) |
| 123 | + |
| 124 | + # save event count for description |
| 125 | + if run_id in event_counts: |
| 126 | + event_counts[run_id] += query_df.shape[0] |
| 127 | + else: |
| 128 | + event_counts[run_id] = query_df.shape[0] |
| 129 | + |
| 130 | + event_no = pd.concat(event_no, axis=0) |
| 131 | + |
| 132 | + # split the non-unique event numbers into train/val and test |
| 133 | + train_val, test = train_test_split( |
| 134 | + event_no, |
| 135 | + test_size=0.10, |
| 136 | + random_state=42, |
| 137 | + shuffle=True, |
| 138 | + ) |
| 139 | + |
| 140 | + train_val = train_val.groupby("path") |
| 141 | + test = test.groupby("path") |
| 142 | + |
| 143 | + # parse into right format for CuratedDataset |
| 144 | + train_val_selection = [] |
| 145 | + test_selection = [] |
| 146 | + for path in dataset_paths: |
| 147 | + train_val_selection.append( |
| 148 | + train_val["event_no"].get_group(path).tolist() |
| 149 | + ) |
| 150 | + test_selection.append(test["event_no"].get_group(path).tolist()) |
| 151 | + |
| 152 | + dataset_args = { |
| 153 | + "truth_table": self._truth_table, |
| 154 | + "pulsemaps": self._pulsemaps, |
| 155 | + "path": dataset_paths, |
| 156 | + "graph_definition": self._graph_definition, |
| 157 | + "features": features, |
| 158 | + "truth": truth, |
| 159 | + } |
| 160 | + |
| 161 | + self._create_comment(event_counts) |
| 162 | + |
| 163 | + return dataset_args, train_val_selection, test_selection |
| 164 | + |
| 165 | + @classmethod |
| 166 | + def _create_comment(cls, event_counts: Dict[str, int] = {}) -> None: |
| 167 | + """Print the number of events in each RunID.""" |
| 168 | + fixed_string = ( |
| 169 | + " Simulation produced by the IceCube Collaboration, " |
| 170 | + + "https://wiki.icecube.wisc.edu/index.php/SnowStorm_MC#File_Locations" # noqa: E501 |
| 171 | + ) |
| 172 | + tot = 0 |
| 173 | + runid_string = "" |
| 174 | + for k, v in event_counts.items(): |
| 175 | + runid_string += f"RunID {k} contains {v:10d} events\n" |
| 176 | + tot += v |
| 177 | + cls._comments = ( |
| 178 | + f"Contains ~{tot/1e6:.1f} million events:\n" |
| 179 | + + runid_string |
| 180 | + + fixed_string |
| 181 | + ) |
| 182 | + |
| 183 | + def _get_dir_name(self, source_file_path: str) -> str: |
| 184 | + file_name = os.path.basename(source_file_path).split(".")[0] |
| 185 | + return str(os.path.join(self.dataset_dir, file_name)) |
0 commit comments