Class TFGWPooling

ot/gnn/_layers.py:20–179

r""" Template Fused Gromov-Wasserstein (TFGW) layer. This layer is a pooling layer for graph neural networks. Computes the fused Gromov-Wasserstein distances between the graph and a set of templates. .. math:: TFGW_{ \overline{ \mathcal{G} }, \alpha }(C,F,h)=[ FGW_{\alpha}(

class TFGWPooling(nn.Module):
    r"""
    Template Fused Gromov-Wasserstein (TFGW) layer. This layer is a pooling layer for graph neural networks.
        Computes the fused Gromov-Wasserstein distances between the graph and a set of templates.


    .. math::
        TFGW_{ \overline{ \mathcal{G} }, \alpha }(C,F,h)=[ FGW_{\alpha}(C,F,h,\overline{C}_k,\overline{F}_k,\overline{h}_k)]_{k=1}^{K}

    where :

    - :math:`\mathcal{G}=\{(\overline{C}_k,\overline{F}_k,\overline{h}_k) \}_{k \in \{1,...,K \}} \}` is the set of :math:`K` templates characterized by their adjacency matrices :math:`\overline{C}_k`, their feature matrices :math:`\overline{F}_k` and their node weights :math:`\overline{h}_k`.
    - :math:`C`, :math:`F` and :math:`h` are respectively the adjacency matrix, the feature matrix and the node weights of the graph.
    - :math:`\alpha` is the trade-off parameter between features and structure for the Fused Gromov-Wasserstein distance.


    Parameters
    ----------
    n_features : int
        Feature dimension of the nodes.
    n_tplt : int
         Number of graph templates.
    n_tplt_nodes : int
        Number of nodes in each template.
    alpha : float, optional
        FGW trade-off parameter (0 < alpha < 1). If None alpha is trained, else it is fixed at the given value.
        Weights features (alpha=0) and structure (alpha=1).
    train_node_weights : bool, optional
        If True, the templates node weights are learned.
        Else, they are uniform.
    multi_alpha: bool, optional
        If True, the alpha parameter is a vector of size n_tplt.
    feature_init_mean: float, optional
        Mean of the random normal law to initialize the template features.
    feature_init_std: float, optional
        Standard deviation of the random normal law to initialize the template features.



    References
    ----------
    .. [53]  Cédric Vincent-Cuaz, Rémi Flamary, Marco Corneli, Titouan Vayer, Nicolas Courty.
            "Template based graph neural network with optimal transport distances"
    """

    def __init__(
        self,
        n_features,
        n_tplt=2,
        n_tplt_nodes=2,
        alpha=None,
        train_node_weights=True,
        multi_alpha=False,
        feature_init_mean=0.0,
        feature_init_std=1.0,
    ):
        r"""
        Template Fused Gromov-Wasserstein (TFGW) layer. This layer is a pooling layer for graph neural networks.