xanesnet.models.e3ee.layers.branch_path module

3-body path aggregation branch for E3EE.

class xanesnet.models.e3ee.layers.branch_path.PairElementEnergyScattering(max_z, z_emb_dim, e_dim, hidden_dim, out_dim)

Bases: Module

Energy-conditioned element-pair scattering features.

Operates on flat per-path tensors (P, ...) rather than padded (B, n_paths, ...).

Parameters:

• max_z (int) – Maximum atomic number supported by the element embedding.

• z_emb_dim (int) – Embedding dimension for atomic numbers.

• e_dim (int) – Dimension of the energy RBF embedding.

• hidden_dim (int) – Hidden dimension of the output MLP.

• out_dim (int) – Output feature dimension per path.

__init__(max_z, z_emb_dim, e_dim, hidden_dim, out_dim)

Initialize PairElementEnergyScattering.

forward(z_j, z_k, e_feat)

Compute per-path element-pair energy features.

Parameters:

• z_j (Tensor) – Atomic numbers of the j leg, shape (P,).

• z_k (Tensor) – Atomic numbers of the k leg, shape (P,).

• e_feat (Tensor) – Energy RBF features, shape (nE, e_dim).

Return type:

Tensor

Returns:

Per-path energy features of shape (P, nE, out_dim).

class xanesnet.models.e3ee.layers.branch_path.TargetSitePathAggregator(atom_dim, rbf_dim, geom_hidden_dim, scatter_dim, out_dim, cutoff)

Bases: Module

3-body target-site-centered path aggregator for paths (target site, j, k).

Consumes precomputed flat triplet scalars (r0j, r0k, rjk, cos(angle)) and per-path flat atom indices into the padded B * N node layout, together with path_batch for scatter-aggregating into the batch dimension.

Parameters:

• atom_dim (int) – Dimension of invariant per-atom features.

• rbf_dim (int) – Number of Gaussian RBF bases for distance encoding.

• geom_hidden_dim (int) – Hidden dimension of the geometry MLP.

• scatter_dim (int) – Intermediate scatter feature dimension.

• out_dim (int) – Output (latent) dimension.

• cutoff (float) – Radial cutoff used for the cosine envelope weights.

__init__(atom_dim, rbf_dim, geom_hidden_dim, scatter_dim, out_dim, cutoff)

Initialize TargetSitePathAggregator.

forward(h_flat, z_flat, pair_elem_energy, e_feat, path_j, path_k, path_r0j, path_r0k, path_rjk, path_cosangle, path_batch, bsz)

Aggregate path contributions into per-(sample, energy) latent vectors.

Parameters:

• h_flat (Tensor) – Invariant atom features (flattened), shape (B*N, atom_dim).

• z_flat (Tensor) – Atomic numbers (flattened), shape (B*N,).

• pair_elem_energy (PairElementEnergyScattering) – Pre-built element-pair energy module.

• e_feat (Tensor) – Energy RBF features, shape (nE, e_dim).

• path_j (Tensor) – Flat atom indices for leg j (into B*N), shape (P,).

• path_k (Tensor) – Flat atom indices for leg k (into B*N), shape (P,).

• path_r0j (Tensor) – Target-site-to-j distance in Angstrom, shape (P,).

• path_r0k (Tensor) – Target-site-to-k distance in Angstrom, shape (P,).

• path_rjk (Tensor) – j-to-k distance in Angstrom, shape (P,).

• path_cosangle (Tensor) – Cosine of the j-target-site-k angle, shape (P,).

• path_batch (Tensor) – Batch index per path (0..B-1), shape (P,).

• bsz (int) – Batch size.

Return type:

Tensor

Returns:

Latent tensor of shape (B, nE, out_dim).