159 lines
5.0 KiB
Python
159 lines
5.0 KiB
Python
"""Module for the PINA dataset classes."""
|
|
|
|
import torch
|
|
from torch.utils.data import Dataset
|
|
from torch_geometric.data import Data
|
|
from ..graph import Graph, LabelBatch
|
|
from ..label_tensor import LabelTensor
|
|
|
|
|
|
class PinaDatasetFactory:
|
|
"""
|
|
TODO: Update docstring
|
|
"""
|
|
|
|
def __new__(cls, conditions_dict, **kwargs):
|
|
"""
|
|
TODO: Update docstring
|
|
"""
|
|
|
|
# Check if conditions_dict is empty
|
|
if len(conditions_dict) == 0:
|
|
raise ValueError("No conditions provided")
|
|
|
|
dataset_dict = {}
|
|
|
|
# Check is a Graph is present in the conditions
|
|
for name, data in conditions_dict.items():
|
|
if not isinstance(data, dict):
|
|
raise ValueError(
|
|
f"Condition '{name}' data must be a dictionary"
|
|
)
|
|
dataset_dict[name] = PinaDataset(data, **kwargs)
|
|
return dataset_dict
|
|
|
|
|
|
class PinaDataset(Dataset):
|
|
"""
|
|
Dataset class for the PINA dataset with :class:`torch.Tensor` and
|
|
:class:`~pina.label_tensor.LabelTensor` data.
|
|
"""
|
|
|
|
def __init__(self, data_dict, automatic_batching=None):
|
|
"""
|
|
Initialize the instance by storing the conditions dictionary.
|
|
|
|
:param dict conditions_dict: A dictionary mapping condition names to
|
|
their respective data. Each key represents a condition name, and the
|
|
corresponding value is a dictionary containing the associated data.
|
|
"""
|
|
|
|
# Store the conditions dictionary
|
|
self.data = data_dict
|
|
self.automatic_batching = (
|
|
automatic_batching if automatic_batching is not None else True
|
|
)
|
|
self.stack_fn = {}
|
|
# Determine stacking functions for each data type (used in collate_fn)
|
|
for k, v in data_dict.items():
|
|
if isinstance(v, LabelTensor):
|
|
self.stack_fn[k] = LabelTensor.stack
|
|
elif isinstance(v, torch.Tensor):
|
|
self.stack_fn[k] = torch.stack
|
|
elif isinstance(v, list) and all(
|
|
isinstance(item, (Data, Graph)) for item in v
|
|
):
|
|
self.stack_fn[k] = LabelBatch.from_data_list
|
|
else:
|
|
raise ValueError(
|
|
f"Unsupported data type for stacking: {type(v)}"
|
|
)
|
|
|
|
def __len__(self):
|
|
return len(next(iter(self.data.values())))
|
|
|
|
def __getitem__(self, idx):
|
|
"""
|
|
Return the data at the given index in the dataset.
|
|
|
|
:param int idx: Index.
|
|
:return: A dictionary containing the data at the given index.
|
|
:rtype: dict
|
|
"""
|
|
|
|
if self.automatic_batching:
|
|
# Return the data at the given index
|
|
return {
|
|
field_name: data[idx] for field_name, data in self.data.items()
|
|
}
|
|
return idx
|
|
|
|
def _getitem_from_list(self, idx_list):
|
|
"""
|
|
Return data from the dataset given a list of indices.
|
|
|
|
:param list[int] idx_list: List of indices.
|
|
:return: A dictionary containing the data at the given indices.
|
|
:rtype: dict
|
|
"""
|
|
|
|
to_return = {}
|
|
for field_name, data in self.data.items():
|
|
if self.stack_fn[field_name] == LabelBatch.from_data_list:
|
|
to_return[field_name] = self.stack_fn[field_name](
|
|
[data[i] for i in idx_list]
|
|
)
|
|
else:
|
|
to_return[field_name] = data[idx_list]
|
|
return to_return
|
|
|
|
|
|
class PinaGraphDataset(Dataset):
|
|
def __init__(self, data_dict, automatic_batching=None):
|
|
"""
|
|
Initialize the instance by storing the conditions dictionary.
|
|
|
|
:param dict conditions_dict: A dictionary mapping condition names to
|
|
their respective data. Each key represents a condition name, and the
|
|
corresponding value is a dictionary containing the associated data.
|
|
"""
|
|
|
|
# Store the conditions dictionary
|
|
self.data = data_dict
|
|
self.automatic_batching = (
|
|
automatic_batching if automatic_batching is not None else True
|
|
)
|
|
|
|
def __len__(self):
|
|
return len(next(iter(self.data.values())))
|
|
|
|
def __getitem__(self, idx):
|
|
"""
|
|
Return the data at the given index in the dataset.
|
|
|
|
:param int idx: Index.
|
|
:return: A dictionary containing the data at the given index.
|
|
:rtype: dict
|
|
"""
|
|
|
|
if self.automatic_batching:
|
|
# Return the data at the given index
|
|
return {
|
|
field_name: data[idx] for field_name, data in self.data.items()
|
|
}
|
|
return idx
|
|
|
|
def _getitem_from_list(self, idx_list):
|
|
"""
|
|
Return data from the dataset given a list of indices.
|
|
|
|
:param list[int] idx_list: List of indices.
|
|
:return: A dictionary containing the data at the given indices.
|
|
:rtype: dict
|
|
"""
|
|
|
|
return {
|
|
field_name: [data[i] for i in idx_list]
|
|
for field_name, data in self.data.items()
|
|
}
|