diff --git a/.gitignore b/.gitignore index d8d702060..7f98e22cf 100644 --- a/.gitignore +++ b/.gitignore @@ -154,6 +154,7 @@ data/examples/output/ !/graphnet/src/graphnet/models/pretrained/**/**/**/**/**.pth # Exception to geometry tables !/data/geometry_tables/**/**.parquet +!/data/image_mapping_tables/**/**.parquet !/data/tests/sqlite/upgrade_genie_step4_140028_000998_first_5_frames/upgrade_genie_step4_140028_000998_first_5_frames.db !/data/tests/parquet/oscNext_genie_level7_v02/merged/** !data/tests/parquet/oscNext_genie_level7_v02/oscNext_genie_level7_v02_first_5_frames.parquet diff --git a/data/image_mapping_tables/IC86_CNN_mapping.parquet b/data/image_mapping_tables/IC86_CNN_mapping.parquet new file mode 100644 index 000000000..591cfc785 Binary files /dev/null and b/data/image_mapping_tables/IC86_CNN_mapping.parquet differ diff --git a/data/image_mapping_tables/prometheus_CNN_mapping.parquet b/data/image_mapping_tables/prometheus_CNN_mapping.parquet new file mode 100644 index 000000000..aba42350e Binary files /dev/null and b/data/image_mapping_tables/prometheus_CNN_mapping.parquet differ diff --git a/data/tests/images/IC86lower_deepcore_test.npy b/data/tests/images/IC86lower_deepcore_test.npy new file mode 100644 index 000000000..178a09047 Binary files /dev/null and b/data/tests/images/IC86lower_deepcore_test.npy differ diff --git a/data/tests/images/IC86main_array_test.npy b/data/tests/images/IC86main_array_test.npy new file mode 100644 index 000000000..628cbfd71 Binary files /dev/null and b/data/tests/images/IC86main_array_test.npy differ diff --git a/data/tests/images/IC86upper_deepcore_test.npy b/data/tests/images/IC86upper_deepcore_test.npy new file mode 100644 index 000000000..24a3cc697 Binary files /dev/null and b/data/tests/images/IC86upper_deepcore_test.npy differ diff --git a/examples/04_training/09_train_cnn.py b/examples/04_training/09_train_cnn.py new file mode 100644 index 000000000..7aa837163 --- /dev/null +++ b/examples/04_training/09_train_cnn.py @@ -0,0 +1,343 @@ +"""Example of training a CNN Model.""" + +import os +from typing import Any, Dict, List, Optional + +from pytorch_lightning.loggers import WandbLogger +import torch +from torch.optim.adam import Adam + +from graphnet.constants import EXAMPLE_DATA_DIR, EXAMPLE_OUTPUT_DIR +from graphnet.data.constants import TRUTH +from graphnet.models import StandardModel +from graphnet.models.cnn import LCSC +from graphnet.models.data_representation import PercentileClusters +from graphnet.models.task.reconstruction import EnergyReconstruction +from graphnet.training.callbacks import PiecewiseLinearLR +from graphnet.training.loss_functions import LogCoshLoss +from graphnet.utilities.argparse import ArgumentParser +from graphnet.utilities.logging import Logger +from graphnet.data.dataset import SQLiteDataset +from graphnet.data.dataset import ParquetDataset +from graphnet.models.detector import ORCA150 +from torch_geometric.data import Batch +from graphnet.models.data_representation.images import ExamplePrometheusImage + +# Constants +features = ["sensor_id", "sensor_string_id", "t"] +truth = TRUTH.PROMETHEUS + + +def main( + path: str, + pulsemap: str, + target: str, + truth_table: str, + gpus: Optional[List[int]], + max_epochs: int, + early_stopping_patience: int, + batch_size: int, + num_workers: int, + wandb: bool = False, +) -> None: + """Run example.""" + # Construct Logger + logger = Logger() + + # Initialise Weights & Biases (W&B) run + if wandb: + # Make sure W&B output directory exists + wandb_dir = "./wandb/" + os.makedirs(wandb_dir, exist_ok=True) + wandb_logger = WandbLogger( + project="example-script", + entity="graphnet-team", + save_dir=wandb_dir, + log_model=True, + ) + + logger.info(f"features: {features}") + logger.info(f"truth: {truth}") + + # Configuration + config: Dict[str, Any] = { + "path": path, + "pulsemap": pulsemap, + "batch_size": batch_size, + "num_workers": num_workers, + "target": target, + "early_stopping_patience": early_stopping_patience, + "fit": { + "gpus": gpus, + "max_epochs": max_epochs, + }, + "dataset_reference": ( + SQLiteDataset if path.endswith(".db") else ParquetDataset + ), + } + + archive = os.path.join(EXAMPLE_OUTPUT_DIR, "train_cnn_model") + run_name = "lcsc_{}_example".format(config["target"]) + if wandb: + # Log configuration to W&B + wandb_logger.experiment.config.update(config) + + # First we need to define how the image is constructed. + # This is done using an ImageDefinition. + + # An ImageDefinition combines two components: + + # 1. A pixel definition, which defines how the pixel data is + # represented. Since an image has always fixed dimensions this + # pixel definition is also responsible to represent the data in + # a way such that this fixed dimensions can be achieved. + # Normally, this could mean that light pulses that arrive at + # the same optical module must be aggregated to a + # fixed-dimensional vector. + # A pixel definition works exactly the same as the + # a node definition in the graph scenerio. + + # 2. A pixel mapping, which defines where each pixel is located + # in the final image. This is highly detector specific, as it + # depends on the geometry of the detector. + + # An ImageDefinition can be used to create multiple images of + # a single event. In the example of IceCube, you can e.g + # create three images, one for the so called main array, + # one for the upper deep core and one for the lower deep + # core. Essentially, these are just different areas in + # the detector. + + # Here we use the PercentileClusters pixel definition, which + # aggregates the light pulses that arrive at the same optical + # module with percentiles. + print(features) + pixel_definition = PercentileClusters( + cluster_on=["sensor_id", "sensor_string_id"], + percentiles=[10, 50, 90], + add_counts=True, + input_feature_names=features, + ) + + # The final image definition used here is the ExamplePrometheusImage, + # which is a detector specific pixel mapping. + # It maps optical modules into the image + # using the sensor_string_id and sensor_id + # (number of the optical module). + # The detector class standardizes the input features, + # so that the features are in a ML friendly range. + # For the mapping of the optical modules to the image it is + # essential to not change the value of the sensor_id and + # sensor_string_id. Therefore we need to make sure that + # these features are not standardized, which is done by the + # `replace_with_identity` argument of the detector. + image_definition = ExamplePrometheusImage( + detector=ORCA150( + replace_with_identity=[ + "sensor_id", + "sensor_string_id", + ], + ), + node_definition=pixel_definition, + input_feature_names=features, + string_label="sensor_string_id", + dom_number_label="sensor_id", + ) + + # Use SQLiteDataset to load in data + # The input here depends on the dataset being used, + # in this case the Prometheus dataset. + dataset = SQLiteDataset( + path=config["path"], + pulsemaps=config["pulsemap"], + truth_table=truth_table, + features=features, + truth=truth, + data_representation=image_definition, + ) + + # Create the training and validation dataloaders. + training_dataloader = torch.utils.data.DataLoader( + dataset=dataset, + batch_size=config["batch_size"], + num_workers=config["num_workers"], + collate_fn=Batch.from_data_list, + ) + + validation_dataloader = torch.utils.data.DataLoader( + dataset=dataset, + batch_size=config["batch_size"], + num_workers=config["num_workers"], + collate_fn=Batch.from_data_list, + ) + + # Building model + + # Define architecture of the backbone, in this example + # the LCSC architecture from Alexander Harnisch is used. + backbone = LCSC( + num_input_features=image_definition.nb_outputs, + out_put_dim=2, + input_norm=True, + num_conv_layers=5, + conv_filters=[5, 10, 20, 40, 60], + kernel_size=3, + image_size=(8, 9, 22), # dimensions of the example image + pooling_type=[ + "Avg", + None, + "Avg", + None, + "Avg", + ], + pooling_kernel_size=[ + [1, 1, 2], + None, + [2, 2, 2], + None, + [2, 2, 2], + ], + pooling_stride=[ + [1, 1, 2], + None, + [2, 2, 2], + None, + [2, 2, 2], + ], + num_fc_neurons=50, + norm_list=True, + norm_type="Batch", + ) + # Define the task. + # Here an energy reconstruction, with a LogCoshLoss function. + # The target and prediction are transformed using the log10 function. + # When infering the prediction is transformed back to the + # original scale using 10^x. + task = EnergyReconstruction( + hidden_size=backbone.nb_outputs, + target_labels=config["target"], + loss_function=LogCoshLoss(), + transform_prediction_and_target=lambda x: torch.log10(x), + transform_inference=lambda x: torch.pow(10, x), + ) + # Define the full model, which includes the backbone, task(s), + # along with typical machine learning options such as + # learning rate optimizers and schedulers. + model = StandardModel( + data_representation=image_definition, + backbone=backbone, + tasks=[task], + optimizer_class=Adam, + optimizer_kwargs={"lr": 1e-03, "eps": 1e-03}, + scheduler_class=PiecewiseLinearLR, + scheduler_kwargs={ + "milestones": [ + 0, + len(training_dataloader) / 2, + len(training_dataloader) * config["fit"]["max_epochs"], + ], + "factors": [1e-2, 1, 1e-02], + }, + scheduler_config={ + "interval": "step", + }, + ) + + # Training model + model.fit( + training_dataloader, + validation_dataloader, + early_stopping_patience=config["early_stopping_patience"], + logger=wandb_logger if wandb else None, + **config["fit"], + ) + + # Get predictions + additional_attributes = model.target_labels + assert isinstance(additional_attributes, list) # mypy + + results = model.predict_as_dataframe( + validation_dataloader, + additional_attributes=additional_attributes + ["event_no"], + gpus=config["fit"]["gpus"], + ) + + # Save predictions and model to file + db_name = path.split("/")[-1].split(".")[0] + path = os.path.join(archive, db_name, run_name) + logger.info(f"Writing results to {path}") + os.makedirs(path, exist_ok=True) + + # Save results as .csv + results.to_csv(f"{path}/cnn_results.csv") + + # Save model config and state dict - Version safe save method. + # This method of saving models is the safest way. + model.save_state_dict(f"{path}/cnn_state_dict.pth") + model.save_config(f"{path}/cnn_model_config.yml") + + +if __name__ == "__main__": + + # Parse command-line arguments + parser = ArgumentParser( + description=""" +Train GNN model without the use of config files. +""" + ) + + parser.add_argument( + "--path", + help="Path to dataset file (default: %(default)s)", + default=f"{EXAMPLE_DATA_DIR}/sqlite/prometheus/prometheus-events.db", + ) + + parser.add_argument( + "--pulsemap", + help="Name of pulsemap to use (default: %(default)s)", + default="total", + ) + + parser.add_argument( + "--target", + help=( + "Name of feature to use as regression target (default: " + "%(default)s)" + ), + default="total_energy", + ) + + parser.add_argument( + "--truth-table", + help="Name of truth table to be used (default: %(default)s)", + default="mc_truth", + ) + + parser.with_standard_arguments( + "gpus", + ("max-epochs", 1), + "early-stopping-patience", + ("batch-size", 16), + ("num-workers", 2), + ) + + parser.add_argument( + "--wandb", + action="store_true", + help="If True, Weights & Biases are used to track the experiment.", + ) + + args, unknown = parser.parse_known_args() + + main( + args.path, + args.pulsemap, + args.target, + args.truth_table, + args.gpus, + args.max_epochs, + args.early_stopping_patience, + args.batch_size, + args.num_workers, + args.wandb, + ) diff --git a/src/graphnet/constants.py b/src/graphnet/constants.py index 1b4519bb8..a12d37df1 100644 --- a/src/graphnet/constants.py +++ b/src/graphnet/constants.py @@ -21,6 +21,14 @@ TEST_PARQUET_DATA = os.path.join( TEST_DATA_DIR, "parquet", _test_dataset_name, "merged" ) +TEST_IMAGE_DIR = os.path.join(TEST_DATA_DIR, "images") +TEST_IC86MAIN_IMAGE = os.path.join(TEST_IMAGE_DIR, "IC86main_array_test.npy") +TEST_IC86LOWERDC_IMAGE = os.path.join( + TEST_IMAGE_DIR, "IC86lower_deepcore_test.npy" +) +TEST_IC86UPPERDC_IMAGE = os.path.join( + TEST_IMAGE_DIR, "IC86upper_deepcore_test.npy" +) # Example data EXAMPLE_DATA_DIR = os.path.join(DATA_DIR, "examples") @@ -41,3 +49,12 @@ ICECUBE_GEOMETRY_TABLE_DIR = os.path.join(GEOMETRY_TABLE_DIR, "icecube") PROMETHEUS_GEOMETRY_TABLE_DIR = os.path.join(GEOMETRY_TABLE_DIR, "prometheus") LIQUIDO_GEOMETRY_TABLE_DIR = os.path.join(GEOMETRY_TABLE_DIR, "liquid-o") + +# Image Mapping Tables +IMAGE_MAPPING_TABLE_DIR = os.path.join(DATA_DIR, "image_mapping_tables") +IC86_CNN_MAPPING = os.path.join( + IMAGE_MAPPING_TABLE_DIR, "IC86_CNN_mapping.parquet" +) +PROMETHEUS_CNN_MAPPING = os.path.join( + IMAGE_MAPPING_TABLE_DIR, "prometheus_CNN_mapping.parquet" +) diff --git a/src/graphnet/data/extractors/icecube/i3filtermapextractor.py b/src/graphnet/data/extractors/icecube/i3filtermapextractor.py index 357cab15a..cba10ed46 100644 --- a/src/graphnet/data/extractors/icecube/i3filtermapextractor.py +++ b/src/graphnet/data/extractors/icecube/i3filtermapextractor.py @@ -11,8 +11,8 @@ class I3FilterMapExtractor(I3Extractor): """Class for extracting I3FilterMap properties. - This class extracts the boolean condition of the I3FilterMask from the - I3FilterMap in the frame. + This class extracts the boolean condition of the I3FilterMask from + the I3FilterMap in the frame. """ def __init__( diff --git a/src/graphnet/models/cnn/__init__.py b/src/graphnet/models/cnn/__init__.py new file mode 100644 index 000000000..dfaf35e40 --- /dev/null +++ b/src/graphnet/models/cnn/__init__.py @@ -0,0 +1,5 @@ +"""CNN-specific modules, for performing the main learnable operations.""" + +from .cnn import CNN +from .theos_muonE_upgoing import TheosMuonEUpgoing +from .lcsc import LCSC diff --git a/src/graphnet/models/cnn/cnn.py b/src/graphnet/models/cnn/cnn.py new file mode 100644 index 000000000..2453790e4 --- /dev/null +++ b/src/graphnet/models/cnn/cnn.py @@ -0,0 +1,35 @@ +"""Base CNN-specific `Model` class(es).""" + +from abc import abstractmethod + +from torch import Tensor +from torch_geometric.data import Data + +from graphnet.models import Model + + +class CNN(Model): + """Base class for all core CNN models in graphnet.""" + + def __init__(self, nb_inputs: int, nb_outputs: int) -> None: + """Construct `CNN`.""" + # Base class constructor + super().__init__() + + # Member variables + self._nb_inputs = nb_inputs + self._nb_outputs = nb_outputs + + @property + def nb_inputs(self) -> int: + """Return number of input features.""" + return self._nb_inputs + + @property + def nb_outputs(self) -> int: + """Return number of output features.""" + return self._nb_outputs + + @abstractmethod + def forward(self, data: Data) -> Tensor: + """Apply learnable forward pass in model.""" diff --git a/src/graphnet/models/cnn/lcsc.py b/src/graphnet/models/cnn/lcsc.py new file mode 100644 index 000000000..1500cbc81 --- /dev/null +++ b/src/graphnet/models/cnn/lcsc.py @@ -0,0 +1,465 @@ +"""Module for the Lightning CNN signal classifier (LCSC). + +All credits go to Alexander Harnisch (https://github.com/AlexHarn) +""" + +from .cnn import CNN +import torch +from torch_geometric.data import Data +from typing import List, Union, Tuple + + +class LCSC(CNN): + """Lightning CNN Signal Classifier (LCSC). + + All credits go to Alexander Harnisch ( + https://github.com/AlexHarn) + + Intended to be used with the IceCube 86 image containing + only the Main Array image. + """ + + def __init__( + self, + num_input_features: int, + out_put_dim: int = 2, + input_norm: bool = True, + num_conv_layers: int = 8, + conv_filters: List[int] = [50, 50, 50, 50, 50, 50, 50, 10], + kernel_size: Union[int, List[Union[int, List[int]]]] = 3, + padding: Union[str, int, List[Union[str, int]]] = "Same", + pooling_type: List[Union[None, str]] = [ + None, + "Avg", + None, + "Avg", + None, + "Avg", + None, + "Avg", + ], + pooling_kernel_size: List[Union[None, int, List[int]]] = [ + None, + [1, 1, 2], + None, + [2, 2, 2], + None, + [2, 2, 2], + None, + [2, 2, 2], + ], + pooling_stride: Union[int, List[Union[None, int, List[int]]]] = [ + None, + [1, 1, 2], + None, + [2, 2, 2], + None, + [2, 2, 2], + None, + [2, 2, 2], + ], + num_fc_neurons: int = 50, + norm_list: bool = True, + norm_type: str = "Batch", + image_size: Tuple[int, int, int] = (10, 10, 60), + ) -> None: + """Initialize the Lightning CNN signal classifier (LCSC). + + Args: + num_input_features (int): Number of input features. + out_put_dim (int): Number of output dimensions of final MLP. + Defaults to 2. + input_norm (bool): Whether to apply normalization to the input. + Defaults to True. + num_conv_layers (int): Number of convolutional layers. + Defaults to 8. + conv_filters (List[int]): List of number of convolutional + filters to use in hidden layers. + Defaults to [50, 50, 50, 50, 50, 50, 50, 50, 10]. + NOTE needs to have the length of `num_conv_layers`. + kernel_size (int, List[int], or List[List[int]]): + Size of the convolutional kernels. + Options are: + int: single integer for all dimensions + and all layers, + e.g. 3 would equal [3, 3, 3] for each layer. + list: list of integers specifying the kernel size, + for each layer for all dimensions equally, + e.g. [3, 5, 6] would equal [[3,3,3], [5,5,5], [6,6,6]]. + NOTE: needs to have the length of `num_conv_layers`. + If a list of lists is provided, each list will be used + for the corresponding layer as kernel size. + NOTE: If a list if passed it needs to have the length + of `num_conv_layers`. + padding (str, int, or List[int]]): Padding for the + convolutional layers. + Options are: + 'Same' for same convolutional padding, + int: single integer for all dimensions and all layers, + e.g. 1 would equal [1, 1, 1]. + list: list of integers specifying the padding for each + dimension, for each layer equally, + e.g. [1, 2, 3]. + NOTE: If a list is passed it needs to have the length + of `num_conv_layers`. + Defaults to 'Same'. + pooling_type (List[None,str]): List of pooling types + for layers. + Options are + None : No pooling is used, + 'Avg' : Average pooling is used, + 'Max' : Max pooling is used + Defaults to [ + None, 'Avg', + None, 'Avg', + None, 'Avg', + None, 'Avg' + ]. + NOTE: the length of the list must be equal to + `num_conv_layers`. + pooling_kernel_size (List[Union[int,List[int]]]): + List of pooling kernel sizes for each layer. + If an integer is provided, it will be used for all layers. + In case of a list the options for its elements are: + list: list of integers for each dimension, e.g. [1, 1, 2]. + int: single integer for all dimensions, + e.g. 2 would equal [2, 2, 2]. + If None, no pooling is applied. + NOTE: If a list is passed it needs to have the length + of `num_conv_layers`. + Defaults to [ + None, [1, 1, 2], + None, [2, 2, 2], + None, [2, 2, 2], + None, [2, 2, 2] + ]. + pooling_stride (int or List[Union[None,int]]): + List of pooling strides for each layer. + If an integer is provided, it will be used for all layers. + In case of a list the options for its elements are: + list: list of integers for each dimension, e.g. [1, 1, 2]. + int: single integer for all dimensions, + e.g. 2 would equal [2, 2, 2]. + If None, no pooling is applied. + NOTE: If a list is passed it needs to have the length + of `num_conv_layers`. + Defaults to [ + None, [1, 1, 2], + None, [2, 2, 2], + None, [2, 2, 2], + None, [2, 2, 2] + ]. + num_fc_neurons (int): Number of neurons in the + fully connected layers. + Defaults to 50. + norm_list (bool or List[bool]): Whether to apply normalization + for each convolutional layer. + If a boolean is provided, it will be used for all layers. + Defaults to True. + NOTE: If a list is passed it needs to have the length + of `num_conv_layers`. + norm_type (str): Type of normalization to use. + Options are 'Batch' or 'Instance'. + Defaults to 'Batch'. + image_size (Tuple[int, int, int]): Size of the input image + in the format (height, width, depth). + NOTE: Only needs to be changed if the input image is not + the standard IceCube 86 image size. + """ + super().__init__(nb_inputs=num_input_features, nb_outputs=out_put_dim) + + # Check input parameters + if isinstance(conv_filters, int): + conv_filters = [conv_filters for _ in range(num_conv_layers)] + else: + if not isinstance(conv_filters, list): + raise TypeError( + ( + f"`conv_filters` must be a " + f"list or an integer, not {type(conv_filters)}!" + ) + ) + if len(conv_filters) != num_conv_layers: + raise ValueError( + f"`conv_filters` must have {num_conv_layers} " + f"elements, not {len(conv_filters)}!" + ) + + if isinstance(kernel_size, int): + kernel_size = [ # type: ignore[assignment] + [kernel_size, kernel_size, kernel_size] + for _ in range(num_conv_layers) + ] + else: + if not isinstance(kernel_size, list): + raise TypeError( + ( + "`kernel_size` must be a list or an " + f"integer, not {type(kernel_size)}!" + ) + ) + if len(kernel_size) != num_conv_layers: + raise ValueError( + ( + f"`kernel_size` must have {num_conv_layers} " + f"elements, not {len(kernel_size)}!" + ) + ) + + if isinstance(padding, int): + padding = [padding for _ in range(num_conv_layers)] + elif isinstance(padding, str): + if padding.lower() == "same": + padding = ["same" for i in range(num_conv_layers)] + else: + raise ValueError( + ( + "`padding` must be 'Same' or an integer, " + f"not {padding}!" + ) + ) + else: + if not isinstance(padding, list): + raise TypeError( + ( + f"`padding` must be a list or " + f"an integer, not {type(padding)}!" + ) + ) + if len(padding) != num_conv_layers: + raise ValueError( + f"`padding` must have {num_conv_layers} " + f"elements, not {len(padding)}!" + ) + + if isinstance(pooling_kernel_size, int): + pooling_kernel_size = [ + pooling_kernel_size for i in range(num_conv_layers) + ] + else: + if not isinstance(pooling_kernel_size, list): + raise TypeError( + ( + "`pooling_kernel_size` must be a list or " + f"an integer, not {type(pooling_kernel_size)}!" + ) + ) + if len(pooling_kernel_size) != num_conv_layers: + raise ValueError( + ( + f"`pooling_kernel_size` must have " + f"{num_conv_layers} elements, not " + f"{len(pooling_kernel_size)}!" + ) + ) + + if isinstance(pooling_stride, int): + pooling_stride = [pooling_stride for i in range(num_conv_layers)] + else: + if not isinstance(pooling_stride, list): + raise TypeError( + ( + "`pooling_stride` must be a list or an integer, " + f"not {type(pooling_stride)}!" + ) + ) + if len(pooling_stride) != num_conv_layers: + raise ValueError( + ( + f"`pooling_stride` must have {num_conv_layers} " + f"elements, not {len(pooling_stride)}!" + ) + ) + + if isinstance(norm_list, bool): + self._norm_list = [norm_list for i in range(num_conv_layers)] + else: + if not isinstance(norm_list, list): + raise TypeError( + ( + "`norm_list` must be a list or a boolean, " + f"not {type(norm_list)}!" + ) + ) + if len(norm_list) != num_conv_layers: + raise ValueError( + ( + f"`norm_list` must have {num_conv_layers} " + f"elements, not {len(norm_list)}!" + ) + ) + self._norm_list = norm_list + + if norm_type.lower() == "instance": + norm_class = torch.nn.InstanceNorm3d + if input_norm: + self.input_normal = torch.nn.InstanceNorm3d(num_input_features) + elif norm_type.lower() == "batch": + norm_class = torch.nn.BatchNorm3d + if input_norm: + # No momentum or learnable parameters for input normalization, + # just use the average + self.input_normal = torch.nn.BatchNorm3d( + num_input_features, momentum=None, affine=False + ) + else: + raise ValueError( + ( + "`norm_type` has to be 'instance' or " + f"'batch, not '{norm_type}'!" + ) + ) + + # Initialize layers + self.conv = torch.nn.ModuleList() + self.pool = torch.nn.ModuleList() + self.input_norm = input_norm + + self.normal = torch.nn.ModuleList() + dimensions: List[int] = [ + num_input_features, + *image_size, + ] # (nb_features per pixel, height, width, depth) + for i in range(num_conv_layers): + self.conv.append( + torch.nn.Conv3d( + dimensions[0], + conv_filters[i], + kernel_size=kernel_size[i], + padding=padding[i], + ) + ) + dimensions = self._calc_output_dimension( + dimensions, + conv_filters[i], + kernel_size[i], + padding[i], + ) + if pooling_type[i] is None or pooling_type[i] == "None": + self.pool.append(None) + elif pooling_type[i] == "Avg": + self.pool.append( + torch.nn.AvgPool3d( + kernel_size=pooling_kernel_size[i], + stride=pooling_stride[i], + ) + ) + dimensions = self._calc_output_dimension( + dimensions, + out_channels=dimensions[ + 0 + ], # same out channels as input channels for pooling + kernel_size=pooling_kernel_size[i], + stride=pooling_stride[i], + ) + elif pooling_type[i] == "Max": + self.pool.append( + torch.nn.MaxPool3d( + kernel_size=pooling_kernel_size[i], + stride=pooling_stride[i], + ) + ) + dimensions = self._calc_output_dimension( + dimensions, + out_channels=dimensions[ + 0 + ], # same out channels as input channels for pooling + kernel_size=pooling_kernel_size[i], + stride=pooling_stride[i], + ) + else: + raise ValueError( + "Pooling type must be 'None', 'Avg' or 'Max'!" + ) + if self._norm_list[i]: + self.normal.append(norm_class(dimensions[0])) + else: + self.normal.append(None) + + latent_dim = ( + dimensions[0] * dimensions[1] * dimensions[2] * dimensions[3] + ) + + self.flatten = torch.nn.Flatten() + self.fc1 = torch.nn.Linear(latent_dim, num_fc_neurons) + self.fc2 = torch.nn.Linear(num_fc_neurons, out_put_dim) + + def _calc_output_dimension( + self, + dimensions: List[int], + out_channels: int, + kernel_size: Union[None, int, List[int]], + padding: Union[str, int, List[int]] = 0, + stride: Union[None, int, List[int]] = 1, + ) -> List[int]: + """Calculate the output dimension after a CNN layers. + + Works for Conv3D, MaxPool3D and AvgPool3D layers. + + Args: + dimensions (Tuple[int]): Current dimensions of the input tensor. + (C,H,W,D) where C is the number of channels, + H is the height, W is the width and D is the depth. + out_channels (int): Number of output channels. + kernel_size (Union[int,List[int]]): Size of the kernel. + If an integer is provided, it will be used for all dimensions. + padding (Union[int,List[int]]): Padding size. + If an integer is provided, it will be used for all dimensions. + If 'Same', the padding will be calculated to keep the + output size the same as the input size. + Defaults to 0. + stride (Union[int,List[int]]): Stride size. + If an integer is provided, it will be used for all dimensions. + Defaults to 1. + + Returns: + Tuple[int]: New dimensions after the layer. + + NOTE: For the pooling layers, set out_channels equal to the + input channels. Since they do not change the number of channels. + """ + krnl_sz: int + if isinstance(padding, str): + if not padding.lower() == "same": + raise ValueError( + f"`padding` must be 'Same' or an integer, not {padding}!" + ) + dimensions[0] = out_channels + else: + for i in range(1, 4): + if isinstance(kernel_size, list): + krnl_sz = kernel_size[i - 1] + else: + assert isinstance(kernel_size, int) + krnl_sz = kernel_size + if isinstance(padding, list): + pad = padding[i - 1] + else: + pad = padding + if isinstance(stride, list): + strd = stride[i - 1] + else: + assert isinstance(stride, int) + strd = stride + dimensions[i] = (dimensions[i] + 2 * pad - krnl_sz) // strd + 1 + + return dimensions + + def forward(self, data: Data) -> torch.Tensor: + """Forward pass of the LCSC.""" + assert len(data.x) == 1, "Only Main Array image is supported for LCSC" + x = data.x[0] + if self.input_norm: + x = self.input_normal(x) + for i in range(len(self.conv)): + x = self.conv[i](x) + if self.pool[i] is not None: + x = self.pool[i](x) + x = torch.nn.functional.elu(x) + if self.normal[i] is not None: + x = self.normal[i](x) + + x = self.flatten(x) + x = torch.nn.functional.elu(self.fc1(x)) + x = torch.nn.functional.elu(self.fc2(x)) + return x diff --git a/src/graphnet/models/cnn/theos_muonE_upgoing.py b/src/graphnet/models/cnn/theos_muonE_upgoing.py new file mode 100644 index 000000000..3bbb0ed03 --- /dev/null +++ b/src/graphnet/models/cnn/theos_muonE_upgoing.py @@ -0,0 +1,411 @@ +"""CNN used for muon energy reconstruction in IceCube. + +Copy of `upgoing_muon_energy` model from +https://github.com/IceCubeOpenSource/i3deepice/tree/master + +Class and variable names are kept for +compatibility with the original code. +""" + +from typing import Tuple, Union + +import torch +from torch import nn +from pytorch_lightning import LightningModule +from torch_geometric.data import Data +from .cnn import CNN + + +class Conv3dBN(LightningModule): + """The Conv3dBN module from Theos CNN model.""" + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: Tuple[int, int, int], + padding: Union[str, Tuple[int, int, int]], + bias: bool = False, + ): + """Create a Conv3dBN module. + + Args: + in_channels: Number of input channels. + out_channels: Number of output channels. + kernel_size: Size of the kernel. + padding: Padding of the kernel. + bias: If True, bias is used in the Convolution. + """ + super().__init__() + + self.conv = nn.Conv3d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + padding=padding, + bias=bias, + ) + + self.bn = nn.BatchNorm3d(out_channels) + self.activation = nn.ReLU() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Forward pass of the Conv3dBN.""" + return self.activation(self.bn(self.conv(x))) + + +class InceptionBlock4(LightningModule): + """The inception_block4 module from Theos CNN model.""" + + def __init__( + self, + in_channels: int, + out_channels: int, + t0: int = 2, + t1: int = 4, + t2: int = 5, + n_pool: int = 3, + ): + """Create a InceptionBlock4 module. + + Args: + in_channels: Number of input channels. + out_channels: Number of output channels. + t0: Size of the first kernel sequence. + t1: Size of the second kernel sequence. + t2: Size of the third kernel sequence. + n_pool: Size of the pooling kernel. + """ + super().__init__() + + self.tower0 = nn.Sequential( + Conv3dBN( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=(t0, 1, 1), + padding="same", + ), + Conv3dBN( + in_channels=out_channels, + out_channels=out_channels, + kernel_size=(1, t0, 1), + padding="same", + ), + Conv3dBN( + in_channels=out_channels, + out_channels=out_channels, + kernel_size=(1, 1, t0), + padding="same", + ), + ) + + self.tower1 = nn.Sequential( + Conv3dBN( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=(t1, 1, 1), + padding="same", + ), + Conv3dBN( + in_channels=out_channels, + out_channels=out_channels, + kernel_size=(1, t1, 1), + padding="same", + ), + Conv3dBN( + in_channels=out_channels, + out_channels=out_channels, + kernel_size=(1, 1, t1), + padding="same", + ), + ) + + self.tower4 = nn.Sequential( + Conv3dBN( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=(1, 1, t2), + padding="same", + ), + ) + + self.tower3 = nn.Sequential( + nn.MaxPool3d( + kernel_size=(n_pool, n_pool, n_pool), + stride=(1, 1, 1), + padding=(n_pool // 2, n_pool // 2, n_pool // 2), + ), + Conv3dBN( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=(1, 1, 1), + padding="same", + ), + ) + self.out_channels = out_channels * 4 + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Forward pass of the InceptionBlock4.""" + ret = torch.cat( + [ + self.tower0(x), + self.tower1(x), + self.tower3(x), + self.tower4(x), + ], + dim=1, + ) + return ret + + +class InceptionResnet(LightningModule): + """The inception_resnet module from Theos CNN model.""" + + def __init__( + self, + in_channels: int, + out_channels: int, + t1: int = 2, + t2: int = 4, + n_pool: int = 3, + scale: float = 0.1, + ): + """Create a InceptionResnet module. + + Args: + in_channels: Number of input channels. + out_channels: Number of output channels. + t1: Size of the first kernel sequence. + t2: Size of the second kernel sequence. + n_pool: Size of the pooling kernel. + scale: Scaling factor for the residual connection. + """ + super().__init__() + self._scale = scale + self.tower1 = nn.Sequential( + Conv3dBN( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=(1, 1, 1), + padding="same", + ), + Conv3dBN( + in_channels=out_channels, + out_channels=out_channels, + kernel_size=(t1, 1, 1), + padding="same", + ), + Conv3dBN( + in_channels=out_channels, + out_channels=out_channels, + kernel_size=(1, t1, 1), + padding="same", + ), + Conv3dBN( + in_channels=out_channels, + out_channels=out_channels, + kernel_size=(1, 1, t1), + padding="same", + ), + ) + self.tower2 = nn.Sequential( + Conv3dBN( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=(1, 1, 1), + padding="same", + ), + Conv3dBN( + in_channels=out_channels, + out_channels=out_channels, + kernel_size=(t2, 1, 1), + padding="same", + ), + Conv3dBN( + in_channels=out_channels, + out_channels=out_channels, + kernel_size=(1, t2, 1), + padding="same", + ), + Conv3dBN( + in_channels=out_channels, + out_channels=out_channels, + kernel_size=(1, 1, t2), + padding="same", + ), + ) + self.tower3 = nn.Sequential( + nn.MaxPool3d( + kernel_size=(n_pool, n_pool, n_pool), + stride=(1, 1, 1), + padding=(n_pool // 2, n_pool // 2, n_pool // 2), + ), + Conv3dBN( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=(1, 1, 1), + padding="same", + ), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Forward pass of the Conv".""" + tmp = torch.cat( + [ + self.tower1(x), + self.tower2(x), + self.tower3(x), + ], + dim=1, + ) + return x + self._scale * tmp + + +class TheosMuonEUpgoing(CNN): + """The TheosMuonEUpgoing module.""" + + def __init__(self, nb_inputs: int = 15, nb_outputs: int = 16) -> None: + """Construct `TheosMuonEUpgoing`. + + Args: + nb_inputs: Number of input features. + nb_outputs: Number of output features. + """ + super().__init__(nb_inputs, nb_outputs) + self.inceptionblocks4 = nn.Sequential( + InceptionBlock4( + in_channels=nb_inputs, + out_channels=18, + t0=2, + t1=5, + t2=8, + ), + InceptionBlock4( + in_channels=18 * 4, + out_channels=18, + t0=2, + t1=3, + t2=7, + ), + InceptionBlock4( + in_channels=18 * 4, + out_channels=18, + t0=2, + t1=4, + t2=8, + ), + InceptionBlock4( + in_channels=18 * 4, + out_channels=18, + t0=3, + t1=5, + t2=9, + ), + InceptionBlock4( + in_channels=18 * 4, + out_channels=18, + t0=2, + t1=8, + t2=9, + ), + ) + self.avgpool1 = nn.AvgPool3d((2, 2, 3)) + self.bn1 = nn.BatchNorm3d(18 * 4) + tmp = [ + InceptionResnet( + in_channels=18 * 4, + out_channels=24, + t2=3, + ), + InceptionResnet( + in_channels=24 * 3, + out_channels=24, + t2=4, + ), + InceptionResnet( + in_channels=24 * 3, + out_channels=24, + t2=5, + ), + ] + for _ in range(5): + tmp = tmp + [ + InceptionResnet( + in_channels=24 * 3, + out_channels=24, + t2=3, + ), + InceptionResnet( + in_channels=24 * 3, + out_channels=24, + t2=4, + ), + InceptionResnet( + in_channels=24 * 3, + out_channels=24, + t2=5, + ), + ] + + self.resblocks1 = nn.Sequential(*tmp) + self.avgpool2 = nn.AvgPool3d((1, 1, 2)) + self.bn2 = nn.BatchNorm3d(24 * 3) + tmp = [] + for _ in range(6): + tmp = tmp + [ + InceptionResnet( + in_channels=24 * 3, + out_channels=24, + t2=3, + ), + InceptionResnet( + in_channels=24 * 3, + out_channels=24, + t2=4, + ), + InceptionResnet( + in_channels=24 * 3, + out_channels=24, + t2=5, + ), + ] + self.resblocks2 = nn.Sequential(*tmp) + self.convs111 = nn.Sequential( + nn.Conv3d( + in_channels=24 * 3, + out_channels=64, + kernel_size=(1, 1, 1), + padding=(0, 0, 0), + ), + nn.ReLU(), + nn.Conv3d( + in_channels=64, + out_channels=4, + kernel_size=(1, 1, 1), + padding=(0, 0, 0), + ), + nn.ReLU(), + ) + self.avgpool3 = nn.AvgPool3d((1, 1, 2)) + self.mlps = nn.Sequential( + nn.Linear(500, 120), + nn.Linear(120, 64), + nn.Linear(64, 16), + ) + + def forward(self, data: Data) -> torch.Tensor: + """Apply learnable forward pass in model.""" + assert len(data.x) == 1, "Only one image expected" + x = data.x[0] + x = self.inceptionblocks4(x) + x = self.avgpool1(x) + x = self.bn1(x) + x = self.resblocks1(x) + x = self.avgpool2(x) + x = self.bn2(x) + x = self.resblocks2(x) + x = self.convs111(x) + x = self.avgpool3(x) + x = nn.Flatten()(x) + x = self.mlps(x) + return x diff --git a/src/graphnet/models/data_representation/__init__.py b/src/graphnet/models/data_representation/__init__.py index 84dd64331..0eb8b7898 100644 --- a/src/graphnet/models/data_representation/__init__.py +++ b/src/graphnet/models/data_representation/__init__.py @@ -17,4 +17,9 @@ PercentileClusters, NodeAsDOMTimeSeries, IceMixNodes, + ClusterSummaryFeatures, +) +from .images import ( + ImageDefinition, + IC86Image, ) diff --git a/src/graphnet/models/data_representation/graphs/nodes/nodes.py b/src/graphnet/models/data_representation/graphs/nodes/nodes.py index 064073bdd..1730e79fa 100644 --- a/src/graphnet/models/data_representation/graphs/nodes/nodes.py +++ b/src/graphnet/models/data_representation/graphs/nodes/nodes.py @@ -502,6 +502,8 @@ class ClusterSummaryFeatures(NodeDefinition): For more details on some of the features see Theo Glauchs thesis (chapter 5.3): https://mediatum.ub.tum.de/node?id=1584755 + + NOTE: number of pulses per cluster is not mentioned/used in the thesis """ def __init__( diff --git a/src/graphnet/models/data_representation/images/__init__.py b/src/graphnet/models/data_representation/images/__init__.py new file mode 100644 index 000000000..fe7c1ea54 --- /dev/null +++ b/src/graphnet/models/data_representation/images/__init__.py @@ -0,0 +1,10 @@ +"""Modules for mapping images. + +´ImageDefinition´ defines the nodes and the mapping, and contains +general image-manipulation.´PixelMapping´ defines how raw data is +mapped into the regular sized image. +""" + +from .image_definition import ImageDefinition +from .images import IC86Image, ExamplePrometheusImage +from .mappings import IC86PixelMapping, ExamplePrometheusMapping diff --git a/src/graphnet/models/data_representation/images/image_definition.py b/src/graphnet/models/data_representation/images/image_definition.py new file mode 100644 index 000000000..0fbbe0a10 --- /dev/null +++ b/src/graphnet/models/data_representation/images/image_definition.py @@ -0,0 +1,163 @@ +"""Modules for defining images. + +These are self-contained image definitions that hold all the image- +altering code in graphnet. These modules define what image-based models +sees as input and can be passed to dataloaders during training and +deployment. +""" + +from typing import List, Optional, Dict, Union, Any, Callable +import torch +import numpy as np +from numpy.random import Generator + +from graphnet.models.detector import Detector +from graphnet.models.data_representation import DataRepresentation +from graphnet.models.data_representation.graphs import NodeDefinition +from torch_geometric.data import Data +from .mappings import PixelMapping + + +class ImageDefinition(DataRepresentation): + """An Abstract class to create ImageDefinitions from.""" + + def __init__( + self, + detector: Detector, + node_definition: NodeDefinition, + pixel_mapping: PixelMapping, + input_feature_names: Optional[List[str]] = None, + dtype: Optional[torch.dtype] = torch.float, + perturbation_dict: Optional[Dict[str, float]] = None, + seed: Optional[Union[int, Generator]] = None, + add_inactive_sensors: bool = False, + sensor_mask: Optional[List[int]] = None, + string_mask: Optional[List[int]] = None, + ): + """Construct `ImageDefinition`. + + ´node_definition´ defines the processing of raw data into + what will later be saved in the individual pixels + + ´pixel_mapping´ defines the mapping of the processed + data to the images. + + NOTE: pixel_mappings may require specific node_definitions. + + Args: + detector: The corresponding ´Detector´ representing the data. + node_definition: Definition of nodes. + pixel_mapping: Definition of Mapping form nodes to pixels. + input_feature_names: Names of each column in expected input data + that will be built into a image. If not provided, + it is automatically assumed that all features in `Detector` is + used. + dtype: data type used for node features. e.g. ´torch.float´ + perturbation_dict: Dictionary mapping a feature name to a standard + deviation according to which the values for this + feature should be randomly perturbed. Defaults + to None. + seed: seed or Generator used to randomly sample perturbations. + Defaults to None. + add_inactive_sensors: If True, inactive sensors will be appended + to the graph with padded pulse information. Defaults to False. + sensor_mask: A list of sensor id's to be masked from the graph. Any + sensor listed here will be removed from the graph. + Defaults to None. + string_mask: A list of string id's to be masked from the graph. + Defaults to None. + sort_by: Name of node feature to sort by. Defaults to None. + """ + # Base class constructor + super().__init__( + detector=detector, + input_feature_names=input_feature_names, + dtype=dtype, + perturbation_dict=perturbation_dict, + seed=seed, + add_inactive_sensors=add_inactive_sensors, + sensor_mask=sensor_mask, + string_mask=string_mask, + repeat_labels=False, + ) + + self._node_definition = node_definition + self._pixel_mapping = pixel_mapping + + def _set_output_feature_names( + self, input_feature_names: List[str] + ) -> List[str]: + """Set the final data output feature names.""" + # Set input data column names for pixel definition + self._node_definition.set_output_feature_names(input_feature_names) + + # get output data column names for pixel mapping + self._pixel_mapping._set_image_feature_names( + self._node_definition._output_feature_names + ) + return self._pixel_mapping.image_feature_names + + def forward( # type: ignore + self, + input_features: np.ndarray, + input_feature_names: List[str], + truth_dicts: Optional[List[Dict[str, Any]]] = None, + custom_label_functions: Optional[Dict[str, Callable[..., Any]]] = None, + loss_weight_column: Optional[str] = None, + loss_weight: Optional[float] = None, + loss_weight_default_value: Optional[float] = None, + data_path: Optional[str] = None, + ) -> Data: + """Construct graph as ´Data´ object. + + Args: + input_features: Input features for graph construction. + Shape ´[num_rows, d]´ + input_feature_names: name of each column. Shape ´[,d]´. + truth_dicts: Dictionary containing truth labels. + custom_label_functions: Custom label functions. + loss_weight_column: Name of column that holds loss weight. + Defaults to None. + loss_weight: Loss weight associated with event. Defaults to None. + loss_weight_default_value: default value for loss weight. + Used in instances where some events have + no pre-defined loss weight. Defaults to None. + data_path: Path to dataset data files. Defaults to None. + + Returns: + graph + """ + data = super().forward( + input_features=input_features, + input_feature_names=input_feature_names, + truth_dicts=truth_dicts, + custom_label_functions=custom_label_functions, + loss_weight_column=loss_weight_column, + loss_weight=loss_weight, + loss_weight_default_value=loss_weight_default_value, + data_path=data_path, + ) + + # data processing + data.x = self._node_definition(data.x) + + # set data type + data.x = data.x.type(self.dtype) + + # create image + data = self._pixel_mapping(data, self.output_feature_names) + + if not isinstance(data.x, list): + data.x = [data.x] + + nb_nodes = [] + for i, x in enumerate(data.x): + data.x[i] = x.type(self.dtype) + + # setting number of nodes as product of C*(D*)H*W + nb_nodes.append(np.prod(list(data.x[i].size()[2:]))) + + # set num_nodes equals number of pixels in all imagessurpress warning + data.num_nodes = torch.tensor(np.sum(nb_nodes)) + + return data diff --git a/src/graphnet/models/data_representation/images/images.py b/src/graphnet/models/data_representation/images/images.py new file mode 100644 index 000000000..25cd104d8 --- /dev/null +++ b/src/graphnet/models/data_representation/images/images.py @@ -0,0 +1,133 @@ +"""A module containing different image representations in GraphNeT.""" + +from typing import List, Optional, Any +import torch + +from graphnet.models.data_representation.graphs import NodeDefinition +from graphnet.models.detector import Detector, IceCube86, ORCA150 + +from .image_definition import ImageDefinition +from .mappings import IC86PixelMapping, ExamplePrometheusMapping + + +class IC86Image(ImageDefinition): + """Class creating a image for IC86 DNN data.""" + + def __init__( + self, + node_definition: NodeDefinition, + input_feature_names: List[str], + include_lower_dc: bool = True, + include_upper_dc: bool = True, + string_label: str = "string", + dom_number_label: str = "dom_number", + dtype: Optional[torch.dtype] = torch.float, + detector: Optional[Detector] = None, + **kwargs: Any, + ) -> None: + """Construct `IC86DNNImage`. + + Args: + node_definition: Definition of nodes. + input_feature_names: Names of each column in expected input data + that will be built into a image. + include_lower_dc: If True, the lower DeepCore will be included. + include_upper_dc: If True, the upper DeepCore will be included. + string_label: The label for the string number in the data. + dom_number_label: The label for the DOM number in the data. + dtype: data type used for node features. e.g. ´torch.float´ + detector: The corresponding ´Detector´ representing the data. + """ + # Default detector with unstandardized input features + if detector is None: + detector = IceCube86( + replace_with_identity=input_feature_names, + ) + else: + assert isinstance(detector, IceCube86) + node_definition.set_output_feature_names(input_feature_names) + assert ( + string_label in input_feature_names + ), f"String label '{string_label}' not in input feature names" + assert ( + dom_number_label in input_feature_names + ), f"DOM number label '{dom_number_label}' not in input feature names" + + # Base class constructor + pixel_mapping = IC86PixelMapping( + string_label=string_label, + dom_number_label=dom_number_label, + pixel_feature_names=node_definition._output_feature_names, + include_lower_dc=include_lower_dc, + include_upper_dc=include_upper_dc, + dtype=dtype, + ) + + super().__init__( + detector=detector, + node_definition=node_definition, + pixel_mapping=pixel_mapping, # PixelMapping, + input_feature_names=input_feature_names, + add_inactive_sensors=False, + **kwargs, + ) + + +class ExamplePrometheusImage(ImageDefinition): + """Class creating a image for Prometheus. + + This Image was created to be used in the example scripts. This is + not intended to be used for purposes beyond that. + """ + + def __init__( + self, + node_definition: NodeDefinition, + input_feature_names: List[str], + string_label: str = "sensor_string_id", + dom_number_label: str = "sensor_id", + dtype: Optional[torch.dtype] = torch.float, + detector: Optional[Detector] = None, + **kwargs: Any, + ) -> None: + """Construct `ExamplePrometheusImage`. + + Args: + node_definition: Definition of nodes. + input_feature_names: Names of each column in expected input data + that will be built into a image. + string_label: The label for the string number in the data. + dom_number_label: The label for the DOM number in the data. + dtype: data type used for node features. e.g. ´torch.float´ + detector: The corresponding ´Detector´ representing the data. + """ + # Default detector with unstandardized input features + if detector is None: + detector = ORCA150( + replace_with_identity=input_feature_names, + ) + + node_definition.set_output_feature_names(input_feature_names) + assert ( + string_label in input_feature_names + ), f"String label '{string_label}' not in input feature names" + assert ( + dom_number_label in input_feature_names + ), f"DOM number label '{dom_number_label}' not in input feature names" + + # Base class constructor + pixel_mapping = ExamplePrometheusMapping( + string_label=string_label, + sensor_number_label=dom_number_label, + pixel_feature_names=node_definition._output_feature_names, + dtype=dtype, + ) + + super().__init__( + detector=detector, + node_definition=node_definition, + pixel_mapping=pixel_mapping, # PixelMapping, + input_feature_names=input_feature_names, + add_inactive_sensors=False, + **kwargs, + ) diff --git a/src/graphnet/models/data_representation/images/mappings/__init__.py b/src/graphnet/models/data_representation/images/mappings/__init__.py new file mode 100644 index 000000000..5f246f891 --- /dev/null +++ b/src/graphnet/models/data_representation/images/mappings/__init__.py @@ -0,0 +1,11 @@ +"""Modules for mapping images for different detectors. + +´PixelMapping´ defines how raw data is mapped into the regular sized +image. +""" + +from .pixel_mappings import ( + PixelMapping, + IC86PixelMapping, + ExamplePrometheusMapping, +) diff --git a/src/graphnet/models/data_representation/images/mappings/pixel_mappings.py b/src/graphnet/models/data_representation/images/mappings/pixel_mappings.py new file mode 100644 index 000000000..558b90c04 --- /dev/null +++ b/src/graphnet/models/data_representation/images/mappings/pixel_mappings.py @@ -0,0 +1,398 @@ +"""Classes for mapping pixel data to images.""" + +from abc import abstractmethod +from typing import List +from torch_geometric.data import Data +import torch +import pandas as pd +import numpy as np + +from graphnet.models import Model +from graphnet.constants import IC86_CNN_MAPPING, PROMETHEUS_CNN_MAPPING + + +class PixelMapping(Model): + """Abstract class for mapping pixel data to images.""" + + def __init__( + self, + pixel_feature_names: List[str], + ) -> None: + """Construct `PixelMapping`.""" + super().__init__(name=__name__, class_name=self.__class__.__name__) + self._set_image_feature_names(pixel_feature_names) + + @abstractmethod + def forward(self, data: Data, data_feature_names: List[str]) -> Data: + """Map pixel data to images. + + Args: + data: The input data containing pixel features. + data_feature_names: Names of each column in expected input data + that will be built into a image. + + Returns: + Data: The output data with images as features. + NOTE: The output data.x should be a list of tensors, + where each tensor corresponds to an image. + + Make sure to add a batch dimension to the tensors. E.g a picture + with dimensions CxHxW = 10x64x64 should be returned as + 1x10x64x64. + """ + raise NotImplementedError + + @abstractmethod + def _set_image_feature_names(self, input_feature_names: List[str]) -> None: + """Set the final image feature names.""" + raise NotImplementedError + + @property + @abstractmethod + def shape( + self, + ) -> List[List[int]]: + """Return the shape of the output images as a list of tuples. + + In the dimensions (F,D,H,W) where F is the number of features + per pixel. And D,H,W are the dimension of the image + """ + pass + + +class IC86PixelMapping(PixelMapping): + """Mapping for the IceCube86. + + This mapping is based on the CNN mapping used + in the multiple IceCube86 analysis. + For further details see: https://arxiv.org/abs/2101.11589 + """ + + def __init__( + self, + dtype: torch.dtype, + pixel_feature_names: List[str], + string_label: str = "string", + dom_number_label: str = "dom_number", + include_main_array: bool = True, + include_lower_dc: bool = True, + include_upper_dc: bool = True, + ): + """Construct `IC86PixelMapping`. + + Args: + dtype: data type used for node features. e.g. ´torch.float´ + string_label: Name of the feature corresponding + to the DOM string number. Values Integers betweem 1 - 86 + dom_number_label: Name of the feature corresponding + to the DOM number (1 - 60). Values Integers between 1 - 60 + where 1 is the dom with the highest z coordinate. + pixel_feature_names: Names of each column in expected input data + that will be built into a image. + include_main_array: If True, the main array will be included. + include_lower_dc: If True, the lower DeepCore will be included. + include_upper_dc: If True, the upper DeepCore will be included. + + Raises: + ValueError: If no array type is included. + + NOTE: Expects input data to be DOMs with aggregated features. + """ + if not np.any( + [include_main_array, include_lower_dc, include_upper_dc] + ): + raise ValueError("Include at least one array type.") + + self._dtype = dtype + self._string_label = string_label + self._dom_number_label = dom_number_label + self._pixel_feature_names = pixel_feature_names + + self._set_indeces(pixel_feature_names, dom_number_label, string_label) + + self._nb_cnn_features = ( + len(pixel_feature_names) - 2 + ) # 2 for string and dom_number + + self._include_main_array = include_main_array + self._include_lower_dc = include_lower_dc + self._include_upper_dc = include_upper_dc + + # read mapping from parquet file + df = pd.read_parquet(IC86_CNN_MAPPING) + df.sort_values( + by=["string", "dom_number"], + ascending=[True, True], + inplace=True, + ) + + # Set the index to string and dom_number for faster lookup + df.set_index( + ["string", "dom_number"], + inplace=True, + drop=False, + ) + + self._mapping = df + super().__init__(pixel_feature_names=pixel_feature_names) + + def _set_indeces( + self, + feature_names: List[str], + dom_number_label: str, + string_label: str, + ) -> None: + """Set the indices for the features.""" + self._cnn_features_idx = [] + for feature in feature_names: + if feature == dom_number_label: + self._dom_number_idx = feature_names.index(feature) + elif feature == string_label: + self._string_idx = feature_names.index(feature) + else: + self._cnn_features_idx.append(feature_names.index(feature)) + + def forward(self, data: Data, data_feature_names: List[str]) -> Data: + """Map pixel data to images.""" + # Initialize output arrays + if self._include_main_array: + main_arr = torch.zeros( + (self._nb_cnn_features, 10, 10, 60), + dtype=self._dtype, + ) + if self._include_upper_dc: + upper_dc_arr = torch.zeros( + (self._nb_cnn_features, 8, 10), + dtype=self._dtype, + ) + if self._include_lower_dc: + lower_dc_arr = torch.zeros( + (self._nb_cnn_features, 8, 50), + dtype=self._dtype, + ) + + # data.x is expected to be a tensor with shape (N, F) + # where N is the number of nodes and F is the number of features. + x = data.x + + # Direct coordinate and feature extraction + string_dom_number = x[ + :, [self._string_idx, self._dom_number_idx] + ].int() + batch_row_features = x[:, self._cnn_features_idx] + + # look up the mapping for string and dom_number + match_indices = self._mapping.loc[ + zip(*string_dom_number.t().tolist()) + ][ + ["string", "dom_number", "mat_ax0", "mat_ax1", "mat_ax2"] + ].values.astype( + int + ) + + # Copy CNN features to the appropriate arrays + for i, row in enumerate(match_indices): + # Select appropriate array and indexing + if row[0] < 79: # Main Array + if self._include_main_array: + main_arr[ + :, + row[2], # mat_ax0 + row[3], # mat_ax1 + row[4], # mat_ax2 + ] = batch_row_features[i] + + elif row[1] < 11: # Upper DeepCore + if self._include_upper_dc: + upper_dc_arr[ + :, + row[2], # mat_ax0 + row[3], # mat_ax1 + ] = batch_row_features[i] + + else: # Lower DeepCore + if self._include_lower_dc: + lower_dc_arr[ + :, + row[2], # mat_ax0 + row[3], # mat_ax1 + ] = batch_row_features[i] + + # unqueeze to add dimension for batching + # with collate_fn Batch.from_data_list + ret: List[torch.Tensor] = [] + if self._include_main_array: + ret.append(main_arr.unsqueeze(0)) + if self._include_upper_dc: + ret.append(upper_dc_arr.unsqueeze(0)) + if self._include_lower_dc: + ret.append(lower_dc_arr.unsqueeze(0)) + + # Set list of images as data.x + data.x = ret + return data + + def _set_image_feature_names(self, input_feature_names: List[str]) -> None: + """Set the final output feature names.""" + # string and dom_number are only used for mapping + # and will not be included in the output features. + self.image_feature_names = [ + infeature + for infeature in input_feature_names + if infeature not in [self._string_label, self._dom_number_label] + ] + + @property + def shape( + self, + ) -> List[List[int]]: + """Return the shape of the output images as a list of tuples.""" + ret = [] + if self._include_main_array: + ret.append([self._nb_cnn_features, 10, 10, 60]) + if self._include_upper_dc: + ret.append([self._nb_cnn_features, 1, 8, 10]) + if self._include_lower_dc: + ret.append([self._nb_cnn_features, 1, 8, 50]) + return ret + + +class ExamplePrometheusMapping(PixelMapping): + """Mapping for the Prometheus detector. + + This mapping is made for example purposes and is not optimized for + any specific use case. There is no guarantee that this mapping will + work with all Prometheus data. + """ + + def __init__( + self, + dtype: torch.dtype, + pixel_feature_names: List[str], + string_label: str = "sensor_string_id", + sensor_number_label: str = "sensor_id", + ): + """Construct `ExamplePrometheusMapping`. + + Args: + dtype: data type used for node features. e.g. ´torch.float´ + string_label: Name of the feature corresponding + to the sensor string number. + sensor_number_label: Name of the feature corresponding + to the sensor number + pixel_feature_names: Names of each column in expected input data + that will be built into a image. + + Raises: + ValueError: If no array type is included. + + NOTE: Expects input data to be sensors with aggregated features. + """ + self._dtype = dtype + self._string_label = string_label + self._sensor_number_label = sensor_number_label + self._pixel_feature_names = pixel_feature_names + + self._set_indeces( + pixel_feature_names, sensor_number_label, string_label + ) + + self._nb_cnn_features = ( + len(pixel_feature_names) - 2 + ) # 2 for string and sensor number + + # read mapping from parquet file + df = pd.read_parquet(PROMETHEUS_CNN_MAPPING) + df.sort_values( + by=["sensor_string_id", "sensor_id"], + ascending=[True, True], + inplace=True, + ) + + # Set the index to string and sensor_number for faster lookup + df.set_index( + ["sensor_string_id", "sensor_id"], + inplace=True, + drop=False, + ) + + self._mapping = df + super().__init__(pixel_feature_names=pixel_feature_names) + + def _set_indeces( + self, + feature_names: List[str], + sensor_number_label: str, + string_label: str, + ) -> None: + """Set the indices for the features.""" + self._cnn_features_idx = [] + for feature in feature_names: + if feature == sensor_number_label: + self._sensor_number_idx = feature_names.index(feature) + elif feature == string_label: + self._string_idx = feature_names.index(feature) + else: + self._cnn_features_idx.append(feature_names.index(feature)) + + def forward(self, data: Data, data_feature_names: List[str]) -> Data: + """Map pixel data to images.""" + # Initialize output arrays + image_tensor = torch.zeros( + (self._nb_cnn_features, 8, 9, 22), + dtype=self._dtype, + ) + + # data.x is expected to be a tensor with shape (N, F) + # where N is the number of nodes and F is the number of features. + x = data.x + + # Direct coordinate and feature extraction + string_sensor_number = x[ + :, [self._string_idx, self._sensor_number_idx] + ].int() + batch_row_features = x[:, self._cnn_features_idx] + + # look up the mapping for string and sensor_number + match_indices = self._mapping.loc[ + zip(*string_sensor_number.t().tolist()) + ][ + ["sensor_string_id", "sensor_id", "mat_ax0", "mat_ax1", "mat_ax2"] + ].values.astype( + int + ) + + # Copy CNN features to the appropriate arrays + for i, row in enumerate(match_indices): + # Select appropriate array and indexing + image_tensor[ + :, + row[2], # mat_ax0 + row[3], # mat_ax1 + row[4], # mat_ax2 + ] = batch_row_features[i] + + # unqueeze to add dimension for batching + # with collate_fn Batch.from_data_list + ret: List[torch.Tensor] = [image_tensor.unsqueeze(0)] + + # Set list of images as data.x + data.x = ret + return data + + def _set_image_feature_names(self, input_feature_names: List[str]) -> None: + """Set the final output feature names.""" + # string and sensor_number are only used for mapping + # and will not be included in the output features. + self.image_feature_names = [ + infeature + for infeature in input_feature_names + if infeature not in [self._string_label, self._sensor_number_label] + ] + + @property + def shape( + self, + ) -> List[List[int]]: + """Return the shape of the output images as a list of tuples.""" + return [[self._nb_cnn_features, 8, 9, 22]] diff --git a/tests/models/test_image_definition.py b/tests/models/test_image_definition.py new file mode 100644 index 000000000..0c00b596d --- /dev/null +++ b/tests/models/test_image_definition.py @@ -0,0 +1,91 @@ +from graphnet.models.data_representation import IC86Image +from graphnet.models.data_representation import NodesAsPulses +from graphnet.models.detector import IceCube86 +import torch +from torch_geometric.data import Data +import pandas as pd +import numpy as np +from graphnet.constants import ( + IC86_CNN_MAPPING, + TEST_IC86MAIN_IMAGE, + TEST_IC86UPPERDC_IMAGE, + TEST_IC86LOWERDC_IMAGE, +) + + +def test_image_definition() -> None: + """Test the ImageDefinition class for IC86 DNN data.""" + # Define input feature names + + grid = pd.read_parquet(IC86_CNN_MAPPING) + grid = grid.loc[:, ["string", "dom_number"]] + grid["redundant_string"] = grid["string"].copy() + grid["redundant_dom_number"] = grid["dom_number"].copy() + dtype = torch.float32 + + # Create a NodeDefinition instance + node_def = NodesAsPulses( + input_feature_names=grid.columns.tolist(), + ) + + detector = IceCube86(replace_with_identity=grid.columns.tolist()) + + # Create an instance of TestImageIC86 + image_definition = IC86Image( + node_definition=node_def, + input_feature_names=grid.columns.tolist(), + include_lower_dc=True, + include_upper_dc=True, + string_label="string", + dom_number_label="dom_number", + dtype=dtype, + detector=detector, + ) + + assert ( + image_definition.nb_outputs == 2 + ), "Expected 2 outputs, got {}".format(image_definition.nb_outputs) + + output_feature_names = grid.columns.tolist() + output_feature_names.remove("string") + output_feature_names.remove("dom_number") + + assert image_definition.output_feature_names == output_feature_names, ( + f"Output feature names do not match expected output: " + f"{image_definition.output_feature_names} != {output_feature_names}" + ) + + image = image_definition( + grid.values, + input_feature_names=grid.columns.tolist(), + ) + + assert isinstance( + image, Data + ), "Expected output to be a torch_geometric.data.Data object" + assert isinstance(image.x, list), "Expected image.x to be a list" + assert np.all( + [isinstance(x, torch.Tensor) for x in image.x] + ), "Expected all elements in image.x to be torch.Tensor" + assert ( + len(image.x) == 3 + ), "Expected image.x to have 3 elements, got {}".format(len(image.x)) + assert ( + "num_nodes" in image.keys() + ), "Expected 'num_nodes' in image attributes" + + image_list = [ + TEST_IC86MAIN_IMAGE, + TEST_IC86UPPERDC_IMAGE, + TEST_IC86LOWERDC_IMAGE, + ] + for i, img in enumerate(image_list): + expected_image = torch.tensor(np.load(img), dtype=dtype).unsqueeze(0) + assert image.x[i].size() == expected_image.size(), ( + f"Image at index {i} size mismatch: " + f"expected {torch.tensor(expected_image).size()}," + f"got {image.x[i].size()}" + ) + assert torch.equal( + image.x[i], expected_image + ), f"Image at index {i} does not match expected image" diff --git a/tests/models/test_pixel_mapping.py b/tests/models/test_pixel_mapping.py new file mode 100644 index 000000000..49e5157e5 --- /dev/null +++ b/tests/models/test_pixel_mapping.py @@ -0,0 +1,205 @@ +"""Unit tests for node definitions.""" + +import numpy as np +import pandas as pd +import torch +from torch_geometric.data import Data +from copy import deepcopy +from graphnet.models.data_representation.images import IC86PixelMapping +from graphnet.constants import ( + TEST_IC86MAIN_IMAGE, + IC86_CNN_MAPPING, + TEST_IC86UPPERDC_IMAGE, + TEST_IC86LOWERDC_IMAGE, +) +import pytest + + +def basic_checks_picture(picture: Data, dtype: torch.dtype) -> None: + """Basic checks for the output of pixel mapping.""" + assert isinstance( + picture, Data + ), f"Output should be a Data object got {type(picture)}" + assert isinstance( + picture.x, list + ), f"x should be a list of tensors got {type(picture.x)}" + assert np.all( + [isinstance(picture.x[i], torch.Tensor) for i in range(len(picture.x))] + ), ( + "All tensors in x should be torch.Tensors", + f"got {[type(picture.x[i]) for i in range(len(picture.x))]}", + ) + assert np.all( + [picture.x[i].dtype == dtype for i in range(len(picture.x))] + ), ( + "All tensors in x should have the dtype specified in pixel_mapping", + f"got {[picture.x[i].dtype for i in range(len(picture.x))]}", + ) + + +def test_pixel_mappings() -> None: + """Test pixel mapping for IC86 DNN mapping.""" + # definitions + dtype = torch.float32 + pixel_feature_names = ["string", "dom_number", "data1", "data2"] + string_label = "string" + dom_number_label = "dom_number" + + # Create dummy data + dummy_data = Data( + x=torch.tensor( + [[1, 2, 5.8, 1e-4], [79, 46, 3.7, 1e-18], [84, 9, 6.87, 2e5]], + dtype=dtype, + ), + ) + + # Construct node definition + # This defines each DOM as a cluster, and will summarize pulses seen by + # DOMs using percentiles. + pixel_mapping = IC86PixelMapping( + dtype=dtype, + pixel_feature_names=pixel_feature_names, + string_label=string_label, + dom_number_label=dom_number_label, + include_lower_dc=True, + include_upper_dc=True, + ) + + # Apply node definition to torch tensor with raw pulses + picture = pixel_mapping(dummy_data, pixel_feature_names) + new_features = pixel_mapping.image_feature_names + n_features = len(new_features) + + # Check the output + basic_checks_picture(picture, dtype) + + # More checks + assert ( + len(pixel_mapping.shape) == 3 + ), f"Expected shape to be 3 got {len(pixel_mapping.shape)}" + assert pixel_mapping.shape == [ + (n_features, 10, 10, 60), + (n_features, 1, 8, 10), + (n_features, 1, 8, 50), + ], ( + f"Expected shape to be [({n_features},10,10,60), " + f"({n_features},1,8,10), ({n_features},1,8,50)] got " + f"{pixel_mapping.shape}" + ) + assert isinstance( + new_features, list + ), f"Output should be a list of feature names got {type(new_features)}" + assert new_features == [ + "data1", + "data2", + ], f"Expected feature to be ['data1', 'data2'] names got: {new_features}" + assert len(picture.x) == 3, ( + "There should be three tensors in x ", + f"got list with length {len(picture.x)}" + "(main array, upper DeepCore, lower DeepCore)", + ) + assert picture.x[0].size() == torch.Size( + [1, 2, 10, 10, 60] + ), f"Main array should have shape (1,2,10,10,60) got {picture.x[0].size()}" + assert picture.x[1].size() == torch.Size( + [1, 2, 8, 10] + ), f"upper DeepCore should have shape (1,2,8,10) got {picture.x[1].size()}" + assert picture.x[2].size() == torch.Size( + [1, 2, 8, 50] + ), f"lower DeepCore should have shape (1,2,8,50) got {picture.x[2].size()}" + assert not torch.all( + picture.x[0] == 0 + ), "Main array should not be all zeros, got all zeros." + assert not torch.all( + picture.x[1] == 0 + ), "Upper DeepCore should not be all zeros, got all zeros." + assert not torch.all( + picture.x[2] == 0 + ), "Lower DeepCore should not be all zeros, got all zeros." + + # Try string and dom_number that does not exist + dummy_data = Data( + x=torch.tensor( + [ + [100, 5, 5.8, 1e-4], + [54, 230, 3.7, 1e-18], + [1294, 500, 6.87, 2e5], + ], + dtype=dtype, + ), + ) + + # should raise KeyError since the string and dom_number + # do not exist in the mapping + with pytest.raises(KeyError): + picture = pixel_mapping(dummy_data, pixel_feature_names) + + +def test_segments_mapping() -> None: + """Test pixel mapping for IC86 main array.""" + # definitions + dtype = torch.float32 + string_label = "string" + dom_number_label = "dom_number" + pixel_feature_names = [ + "string", + "dom_number", + "redundant_string", + "redundant_dom_number", + ] + + # Load the grid mapping + # This is a mapping from string and dom_number to the pixel coordinates + # in the main array, upper DeepCore and lower DeepCore. + # Running the grid mapping through the pixel mapping will + # create the full images for the main array, upper DeepCore + # and lower DeepCore. + grid = pd.read_parquet(IC86_CNN_MAPPING) + grid = grid.loc[:, ["string", "dom_number"]] + grid["redundant_string"] = grid["string"].copy() + grid["redundant_dom_number"] = grid["dom_number"].copy() + grid = Data(x=torch.tensor(grid.to_numpy(), dtype=dtype)) + + # Test the pixel mapping for the main array, upper and lower DeepCore + for image, inc_main, inc_upc, inc_lowdc, label in zip( + [TEST_IC86MAIN_IMAGE, TEST_IC86UPPERDC_IMAGE, TEST_IC86LOWERDC_IMAGE], + [True, False, False], + [False, True, False], + [False, False, True], + ["main array", "upper deepcore", "lower deepcore"], + ): + tmp = deepcopy(grid) + pixel_mapping = IC86PixelMapping( + dtype=dtype, + pixel_feature_names=pixel_feature_names, + string_label=string_label, + dom_number_label=dom_number_label, + include_main_array=inc_main, + include_lower_dc=inc_lowdc, + include_upper_dc=inc_upc, + ) + picture = pixel_mapping(tmp, pixel_feature_names) + tensor_image: torch.tensor = torch.tensor( + np.load(image), dtype=dtype + ).unsqueeze(0) + + # Check the output + basic_checks_picture(picture, dtype) + + # More checks + assert len(picture.x) == 1, ( + "There should be one tensor in x ", + f"got list with length {len(picture.x)}", + ) + assert picture.x[0].size() == tensor_image.size(), ( + f"{label} should have shape {tensor_image.size()} " + f"got {picture.x[0].size()}" + ) + assert not torch.all( + picture.x[0] == 0 + ), f"{label} should not be all zeros, got all zeros." + # Check if the tensor matches the expected image + assert torch.equal(tensor_image, picture.x[0]), ( + f"{label} should match the expected" + " main array from IC86 DNN mapping." + )