Transforms

class NormalizeAdj(delta: float = 0.85)

Bases: BaseTransform

Transforms the adjacency matrix \(\mathbf{A}\) by applying the following transformation:

\[\mathbf{A} \to \mathbf{I} - \delta \mathbf{L}\]

where \(\mathbf{L}\) is the normalized Laplacian of the graph and \(\delta\) is a scaling factor.

Parameters:

delta (int, optional) – Scaling factor for the Laplacian. (default: 0.85)

forward(data: Data) → Data

Applies the normalization transform to the graph while preserving edge attributes.

The transform computes the normalized Laplacian and rescales it with \(-\delta\). It also handles self-loops and concatenates additional edge attributes if available. Duplicate entries are coalesced by summing their values.

Parameters:

data (Data) – A Data object containing graph data.

Returns:

The transformed data object with updated edge_index, edge_weight, and (optionally) edge_attr.

Return type:

Data

class SortNodes

Bases: BaseTransform

Sorts the nodes of a graph based on their labels.

forward(data: Data) → Data

Sorts the nodes of the graph according to their labels.

The function sorts the node labels, reassigns node indices accordingly, and updates the associated attributes. If edge attributes exist, they are re-sorted based on the new node ordering.

Parameters:

data (Data) – A Data object containing graph data with attributes edge_index, y, x, and optionally edge_attr.

Returns:

The data object with sorted nodes and updated attributes.

Return type:

Data

class PreCoarsening(poolers: SRCPooling | str | Tuple[str, Dict[str, Any]] | Dict[str, Any] | Sequence[SRCPooling | str | Tuple[str, Dict[str, Any]] | Dict[str, Any]], input_key: str | None = None, output_key: str = 'pooled_data')

Bases: BaseTransform

A transform that precomputes a hierarchy of pooled (coarsened) graphs and attaches them to the input Data object.

Takes one or more pooling operators from SRCPooling to build a multi-level pooling hierarchy. Pre-coarsenable poolers share the rollout contract exposed by multi_level_precoarsening(); by default, this rollout greedily repeats single-level precoarsening().

Some poolers customize single-level behavior (for example, precoarsening() and precoarsening() keep a fixed assignment width), while others customize the full rollout (for example, multi_level_precoarsening()). Poolers must be non-trainable, i.e., they should not have learnable parameters. The graph is recursively coarsened for as many levels as given in poolers. At each level, a coarsened adjacency matrix and, optionally, a pooled batch is computed. The result is stored as a list of intermediate pooled subgraphs in Data, which downstream GNN models can consume.

Parameters:

• poolers –

  Per-level pooler configuration. Can be a single pooler or a sequence of level configs. A single value is treated as one level. Each entry can be one of:

  • a pre-instantiated pooler instance;

  • a pooler alias string, e.g. "ndp";

  • a tuple ("eigen", {"k": 5});

  • a dictionary with keys {"pooler": "<name>", ...kwargs} or {"name": "<name>", ...kwargs}.

  To use the same pooler for multiple levels, pass a sequence (e.g. [pooler, pooler, pooler] or ["ndp", "ndp", "ndp"]).

• input_key (str, optional) – The key in the data object from which to read the graph data. If None, uses the default data object.

• output_key (str, optional) – The key in the data object where the pooled graphs will be stored. Defaults to "pooled_data".

forward(data: Data) → Data

Attach pooled levels to data[self.output_key].

Execution is run-based: adjacent identical pooler configs are collapsed and each run is executed once via multi_level_precoarsening(). Returned levels are still appended one-by-one, preserving the original external contract (len(pooled_data) == number of requested levels).

Example

from tgp.data.transforms import PreCoarsening

transform = PreCoarsening(
    poolers=["ndp", "ndp", "sep", "sep", "graclus"]
)
data = transform(data)

# Internal execution uses collapsed runs:
# [(ndp, 2), (sep, 2), (graclus, 1)]
# but the output still contains five levels:
# data.pooled_data = [lvl1, lvl2, lvl3, lvl4, lvl5]