GraphDataSet Class Reference

Inheritance diagram for GraphDataSet:

Collaboration diagram for GraphDataSet:

Public Member Functions
def	__init__ (self, root, n_files=None, samples=None, features=[], edge_features=[], global_features=[], normalize=None, **kwargs)
def	processed_file_names (self)
def	process (self)

Public Attributes
	root
	Root path.
	normalize
	Normalize.
	n_files
	Number of files.
	node_features
	Node features.
	edge_features
	Edge features.
	global_features
	Global features.
	samples
	Samples.
	slices
	Data and Slices.

Detailed Description

Dataset handler for converting Belle II data to PyTorch geometric InMemoryDataset.

The ROOT format expects the tree in every file to be named ``Tree``,
and all node features to have the format ``feat_FEATNAME``.

.. note:: This expects the files under root to have the structure ``root/**/<file_name>.root``
    where the root path is different for train and val.
    The ``**/`` is to handle subdirectories, e.g. ``sub00``.

Args:
    root (str): Path to ROOT files.
    n_files (int): Load only ``n_files`` files.
    samples (int): Load only ``samples`` events.
    features (list): List of node features names.
    edge_features (list): List of edge features names.
    global_features (list): List of global features names.
    normalize (bool): Whether to normalize input features.

Definition at line 258 of file geometric_datasets.py.

Constructor & Destructor Documentation

__init__()

def __init__	(		self,
			root,
			n_files = None,
			samples = None,
			features = [],
			edge_features = [],
			global_features = [],
			normalize = None,
		**	kwargs
	)

Initialization.

Definition at line 279 of file geometric_datasets.py.

    ):
        """
        Initialization.
        """
        assert isinstance(
            features, list
        ), f'Argument "features" must be a list and not {type(features)}'
        assert len(features) > 0, "You need to use at least one node feature"
 
        
        self.root = Path(root)
 
        
        self.normalize = normalize
 
        
        self.n_files = n_files
        
        self.node_features = features
        
        self.edge_features = edge_features
        
        self.global_features = global_features
        
        self.samples = samples
 
        # Delete processed files, in case
        file_path = Path(self.root, "processed")
        files = list(file_path.glob("*.pt"))
        for f in files:
            f.unlink(missing_ok=True)
 
        # Needs to be called after having assigned all attributes
        super().__init__(root, None, None, None)
 
        
        self.data, self.slices = torch.load(self.processed_paths[0])
 

Member Function Documentation

process()

def process

(

self

)

Processes the data to create graph objects and stores them in ``root/processed/processed_data.pt``
where the root path is different for train and val.

Called internally by PyTorch.

Definition at line 334 of file geometric_datasets.py.

    def process(self):
        """
        Processes the data to create graph objects and stores them in ``root/processed/processed_data.pt``
        where the root path is different for train and val.
 
        Called internally by PyTorch.
        """
        num_samples = _preload(self)
        data_list = [_process_graph(self, i) for i in range(num_samples)]
        data, slices = self.collate(data_list)
        torch.save((data, slices), self.processed_paths[0])
 
        del self.x, self.y, self.avail_samples, data_list, data, slices

processed_file_names()

def processed_file_names

(

self

)

Name of processed file.

Definition at line 328 of file geometric_datasets.py.

    def processed_file_names(self):
        """
        Name of processed file.
        """
        return ["processed_data.pt"]
 

Member Data Documentation

edge_features

edge_features

Edge features.

Definition at line 309 of file geometric_datasets.py.

global_features

global_features

Global features.

Definition at line 311 of file geometric_datasets.py.

n_files

n_files

Number of files.

Definition at line 305 of file geometric_datasets.py.

node_features

node_features

Node features.

Definition at line 307 of file geometric_datasets.py.

normalize

normalize

Normalize.

Definition at line 302 of file geometric_datasets.py.

root

root

Root path.

Definition at line 299 of file geometric_datasets.py.

samples

samples

Samples.

Definition at line 313 of file geometric_datasets.py.

slices

slices

Data and Slices.

Definition at line 325 of file geometric_datasets.py.

---

The documentation for this class was generated from the following file:

• analysis/scripts/grafei/model/geometric_datasets.py