xanesnet.datasets.torchgeometric.e3ee module

E3EE PyTorch Geometric dataset implementation.

class xanesnet.datasets.torchgeometric.e3ee.E3EEBatch(*args, **kwargs)

Bases: Protocol

Protocol for batches emitted by E3EEDataset.collate_fn.

Node fields are padded to (batch, max_nodes, ...). Flat edge and path indices are offset into the padded batch * max_nodes layout.

x: Tensor

mask: Tensor

target_site_index: Tensor

edge_src: Tensor

edge_dst: Tensor

edge_weight: Tensor

edge_vec: Tensor

att_dst: Tensor

att_dist: Tensor

att_vec: Tensor

path_j: Tensor

path_k: Tensor

path_r0j: Tensor

path_r0k: Tensor

path_rjk: Tensor

path_cosangle: Tensor

path_batch: Tensor

energies: Tensor

intensities: Tensor

sample_id: list[str]

__init__(*args, **kwargs)

class xanesnet.datasets.torchgeometric.e3ee.E3EEDataset(dataset_type, datasource, root, preload, skip_prepare, split_ratios, split_indexfile, graph_builder, att_graph_builder, use_path_branch, max_paths_per_structure)

Bases: TorchGeometricDataset

E3EE dataset that emits one graph sample per target site.

The dataset preserves target-site ordering and supports both periodic Structures and non-periodic Molecules. All edges and target-site- centered triplet path scalars are precomputed at prepare() time using the shared graph utilities, so the model does not need to rebuild geometry at forward time.

Both the main and attention graphs are constructed by GraphBuilder instances resolved from the graph_builder and att_graph_builder nested configs.

Parameters:

• dataset_type (str) – Registered dataset type name.

• datasource (DataSource) – Raw datasource of pymatgen structures or molecules.

• root (str) – Directory that stores processed .pth files.

• preload (bool) – Whether to preload processed samples.

• skip_prepare (bool) – Whether to reuse existing processed files.

• split_ratios (list[float] | None) – Optional split ratios.

• split_indexfile (str | None) – Optional path to split indices.

• graph_builder (Config) – Main graph builder configuration.

• att_graph_builder (Config) – Attention graph builder configuration.

• use_path_branch (bool) – Whether to precompute target-site-centered paths. The path cutoff is taken from graph_builder.cutoff.

• max_paths_per_structure (int) – Maximum target-site paths saved per structure.

__init__(dataset_type, datasource, root, preload, skip_prepare, split_ratios, split_indexfile, graph_builder, att_graph_builder, use_path_branch, max_paths_per_structure)

Initialize the E3EE dataset.

collate_fn(batch)

Collate target-site-centered graph samples into one padded batch.

Node tensors are padded to (batch, max_nodes, ...) and flat edge/path indices are offset by batch_index * max_nodes.

Parameters:

batch (list[BaseData]) – PyG graph samples loaded by __getitem__.

Return type:

Batch

Returns:

PyG batch with E3EE-specific padded tensors attached.

property signature: Config

Dataset configuration signature.

Returns:

Configuration values that identify this E3EE dataset.