of_batch

class OFBatch(*args: Any, **kwargs: Any)

A batch object for OFData.

Copied from torch_geometric.data.Batch with the addition of the list_keys attribute, which just appends attributes to a list instead of collating them, which is needed for square matrices of varying size.

classmethod from_data_list(data_list: List[BaseData], follow_batch: List[str] | None = None, exclude_keys: List[str] | None = None, list_keys: Sequence[str] | None = None) → Self

Constructs a Batch object from a list of Data or HeteroData objects.

The assignment vector batch is created on the fly. In addition, creates assignment vectors for each key in follow_batch. Will exclude any keys given in exclude_keys. Keys in list_keys will be appended to a list instead of collated (useful for square matrices of varying size).

class OFCollater(dataset: Dataset | Sequence[BaseData] | DatasetAdapter, follow_batch: List[str] | None = None, exclude_keys: List[str] | None = None, list_keys: Sequence[str] | None = None)

Collater for OF-DFT data.

Copy paste from torch_geometric.data.collate.Collater except for using our custom OFBatch and the list_keys attribute.

__call__(batch: List[Any]) → Any

Collates a list of data objects into a single object.

__init__(dataset: Dataset | Sequence[BaseData] | DatasetAdapter, follow_batch: List[str] | None = None, exclude_keys: List[str] | None = None, list_keys: Sequence[str] | None = None)

Collater for OF-DFT data.

collate(cls: Type[T], data_list: List[BaseData], increment: bool = True, add_batch: bool = True, follow_batch: Iterable[str] | None = None, exclude_keys: Iterable[str] | None = None, list_keys: Sequence[str] | None = None) → Tuple[T, Dict[str, Tensor | Dict[str, Tensor]], Dict[str, Tensor | Dict[str, Tensor]]]

Collates a list of data objects into a single object of type cls.

collate can handle both homogeneous and heterogeneous data objects by individually collating all their stores. In addition, collate can handle nested data structures such as dictionaries and lists.