Struct GraphMae 

pub struct GraphMae { /* private fields */ }

Graph Masked Autoencoder.

Maintains:

• A learnable mask token of shape [feature_dim] used to replace masked node features.
• An encoder weight matrix [feature_dim × encoder_dim].
• A decoder weight matrix [encoder_dim × feature_dim].

Implementations

impl GraphMae

pub fn new(feature_dim: usize, config: GraphMaeConfig, seed: u64) -> Self

Construct a new GraphMAE.

Arguments

• feature_dim – dimension of each node’s input features
• config – MAE hyper-parameters
• seed – RNG seed for reproducible initialisation

pub fn mask_features( &self, features: &Array2<f64>, seed: u64, ) -> (Array2<f64>, Vec<usize>)

Apply random feature masking.

Selects a random subset of nodes (fraction ≈ config.mask_rate) and replaces their feature vectors with the learnable mask token.

Arguments

• features – node feature matrix [n_nodes × feature_dim]
• seed – RNG seed (different from the model seed so each call can produce a different mask)

Returns

(masked_features, mask_indices) where mask_indices contains the row indices of the masked nodes (sorted ascending).

pub fn encode(&self, features: &Array2<f64>) -> Array2<f64>

Encode node features: Z = ReLU(X @ W_enc)

Arguments

• features – (possibly masked) node features [n_nodes × feature_dim]

Returns

Latent representations [n_nodes × encoder_dim]

pub fn decode(&self, encoded: &Array2<f64>) -> Array2<f64>

Decode latent representations: X̂ = Z @ W_dec

Arguments

• encoded – latent representations [n_nodes × encoder_dim]

Returns

Reconstructed features [n_nodes × feature_dim]

pub fn sce_loss( &self, original: &Array2<f64>, reconstructed: &Array2<f64>, mask_indices: &[usize], gamma: f64, ) -> f64

Scaled Cosine Error (SCE) reconstruction loss on masked nodes.

L = (1/|M|) Σ_{i∈M} (1 - cosine_sim(reconstructed_i, original_i))^γ

If mask_indices is empty, returns 0.0.

Arguments

• original – original node features [n_nodes × feature_dim]
• reconstructed – decoder output [n_nodes × feature_dim]
• mask_indices – indices of masked nodes
• gamma – exponent ≥ 1 (typical: 2 or 3)

pub fn forward(&self, features: &Array2<f64>, seed: u64) -> (Array2<f64>, f64)

Full GraphMAE forward pass.

1. Mask features randomly.
2. Encode masked features.
3. Decode back to feature space.
4. Compute SCE loss over masked nodes (γ = 2).

Arguments

• features – original node feature matrix [n_nodes × feature_dim]
• seed – RNG seed for the masking step

Returns

(reconstructed_features, sce_loss)

pub fn mask_token(&self) -> &Array1<f64>

The learnable mask token vector [feature_dim].

pub fn feature_dim(&self) -> usize

Input / output feature dimension.

pub fn encoder_dim(&self) -> usize

Encoder output dimension.