class_docstr: An implementation of PGExplainer in
  `Parameterized Explainer for Graph Neural Network <https://arxiv.org/abs/2011.04573>`_.
  
  Args:
      model (:class:`torch.nn.Module`): The target model prepared to explain
      in_channels (:obj:`int`): Number of input channels for the explanation network
      explain_graph (:obj:`bool`): Whether to explain graph classification model (default: :obj:`True`)
      epochs (:obj:`int`): Number of epochs to train the explanation network
      lr (:obj:`float`): Learning rate to train the explanation network
      coff_size (:obj:`float`): Size regularization to constrain the explanation size
      coff_ent (:obj:`float`): Entropy regularization to constrain the connectivity of explanation
      t0 (:obj:`float`): The temperature at the first epoch
      t1(:obj:`float`): The temperature at the final epoch
      num_hops (:obj:`int`, :obj:`None`): The number of hops to extract neighborhood of target node
      (default: :obj:`None`)
  
  .. note: For node classification model, the :attr:`explain_graph` flag is False.
    If :attr:`num_hops` is set to :obj:`None`, it will be automatically calculated by calculating the
    :class:`torch_geometric.nn.MessagePassing` layers in the :attr:`model`.
class_docstr_tok: ["An","implementation","of","PGExplainer","in","`","Parameterized","Explainer","for","Graph","Neural","Network","<","https",":","\/\/arxiv.org\/abs\/2011.04573",">","`","_",".","Args",":","model","(",":","class",":","`","torch.nn.Module","`",")",":","The","target","model","prepared","to","explain","in_channels","(",":","obj",":","`","int","`",")",":","Number","of","input","channels","for","the","explanation","network","explain_graph","(",":","obj",":","`","bool","`",")",":","Whether","to","explain","graph","classification","model","(","default",":",":","obj",":","`","True","`",")","epochs","(",":","obj",":","`","int","`",")",":","Number","of","epochs","to","train","the","explanation","network","lr","(",":","obj",":","`","float","`",")",":","Learning","rate","to","train","the","explanation","network","coff_size","(",":","obj",":","`","float","`",")",":","Size","regularization","to","constrain","the","explanation","size","coff_ent","(",":","obj",":","`","float","`",")",":","Entropy","regularization","to","constrain","the","connectivity","of","explanation","t0","(",":","obj",":","`","float","`",")",":","The","temperature","at","the","first","epoch","t1","(",":","obj",":","`","float","`",")",":","The","temperature","at","the","final","epoch","num_hops","(",":","obj",":","`","int","`",",",":","obj",":","`","None","`",")",":","The","number","of","hops","to","extract","neighborhood","of","target","node","(","default",":",":","obj",":","`","None","`",")","..","note",":","For","node","classification","model",",","the",":","attr",":","`","explain_graph","`","flag","is","False",".","If",":","attr",":","`","num_hops","`","is","set","to",":","obj",":","`","None","`",",","it","will","be","automatically","calculated","by","calculating","the",":","class",":","`","torch_geometric.nn.MessagePassing","`","layers","in","the",":","attr",":","`","model","`","."]
docstr: explain the GNN behavior for graph and calculate the metric values.
  The interface for the :class:`dig.evaluation.XCollector`.
  
  Args:
      x (:obj:`torch.Tensor`): Node feature matrix with shape
        :obj:`[num_nodes, dim_node_feature]`
      edge_index (:obj:`torch.Tensor`): Graph connectivity in COO format
        with shape :obj:`[2, num_edges]`
      kwargs(:obj:`Dict`):
        The additional parameters
          - top_k (:obj:`int`): The number of edges in the final explanation results
  
  :rtype: (:obj:`None`, List[torch.Tensor], List[Dict])
docstr_tok: ["explain","the","GNN","behavior","for","graph","and","calculate","the","metric","values",".","The","interface","for","the",":","class",":","`","dig.evaluation.XCollector","`",".","Args",":","x","(",":","obj",":","`","torch.Tensor","`",")",":","Node","feature","matrix","with","shape",":","obj",":","`","[","num_nodes",",","dim_node_feature","]","`","edge_index","(",":","obj",":","`","torch.Tensor","`",")",":","Graph","connectivity","in","COO","format","with","shape",":","obj",":","`","[","2",",","num_edges","]","`","kwargs","(",":","obj",":","`","Dict","`",")",":","The","additional","parameters","-","top_k","(",":","obj",":","`","int","`",")",":","The","number","of","edges","in","the","final","explanation","results",":","rtype",":","(",":","obj",":","`","None","`",",","List","[","torch.Tensor","]",",","List","[","Dict","]",")"]
code: def forward(
      self, x: Tensor, edge_index: Tensor, **kwargs
  ) -> Tuple[None, List, List[Dict]]:
      r"""explain the GNN behavior for graph and calculate the metric values.
      The interface for the :class:`dig.evaluation.XCollector`.
  
      Args:
          x (:obj:`torch.Tensor`): Node feature matrix with shape
            :obj:`[num_nodes, dim_node_feature]`
          edge_index (:obj:`torch.Tensor`): Graph connectivity in COO format
            with shape :obj:`[2, num_edges]`
          kwargs(:obj:`Dict`):
            The additional parameters
              - top_k (:obj:`int`): The number of edges in the final explanation results
  
      :rtype: (:obj:`None`, List[torch.Tensor], List[Dict])
      """
      # set default subgraph with 10 edges
      top_k = (
          kwargs.get("top_k") if kwargs.get("top_k") is not None else 10
      )
      x = x.to(self.device)
      edge_index = edge_index.to(self.device)
  
      self.__clear_masks__()
      logits = self.model(x, edge_index)
      probs = F.softmax(logits, dim=-1)
      pred_labels = probs.argmax(dim=-1)
      embed = self.model.get_emb(x, edge_index)
  
      if self.explain_graph:
          # original value
          probs = probs.squeeze()
          label = pred_labels
          # masked value
          _, edge_mask = self.explain(
              x, edge_index, embed=embed, tmp=1.0, training=False
          )
          data = Data(x=x, edge_index=edge_index)
          selected_nodes = calculate_selected_nodes(
              data, edge_mask, top_k
          )
          masked_node_list = [
              node
              for node in range(data.x.shape[0])
              if node in selected_nodes
          ]
          maskout_nodes_list = [
              node
              for node in range(data.x.shape[0])
              if node not in selected_nodes
          ]
          value_func = GnnNetsGC2valueFunc(
              self.model, target_class=label
          )
  
          masked_pred = gnn_score(
              masked_node_list,
              data,
              value_func=value_func,
              subgraph_building_method="zero_filling",
          )
  
          maskout_pred = gnn_score(
              maskout_nodes_list,
              data,
              value_func,
              subgraph_building_method="zero_filling",
          )
  
          sparsity_score = 1 - len(selected_nodes) / data.x.shape[0]
  
      else:
          node_idx = kwargs.get("node_idx")
          assert (
              kwargs.get("node_idx") is not None
          ), "please input the node_idx"
          # original value
          probs = probs.squeeze()[node_idx]
          label = pred_labels[node_idx]
          # masked value
          x, edge_index, _, subset, _ = self.get_subgraph(
              node_idx, x, edge_index
          )
          new_node_idx = torch.where(subset == node_idx)[0]
          embed = self.model.get_emb(x, edge_index)
          _, edge_mask = self.explain(
              x,
              edge_index,
              embed,
              tmp=1.0,
              training=False,
              node_idx=new_node_idx,
          )
  
          data = Data(x=x, edge_index=edge_index)
          selected_nodes = calculate_selected_nodes(
              data, edge_mask, top_k
          )
          masked_node_list = [
              node
              for node in range(data.x.shape[0])
              if node in selected_nodes
          ]
          maskout_nodes_list = [
              node
              for node in range(data.x.shape[0])
              if node not in selected_nodes
          ]
          value_func = GnnNetsNC2valueFunc(
              self.model, node_idx=new_node_idx, target_class=label
          )
  
          masked_pred = gnn_score(
              masked_node_list,
              data,
              value_func=value_func,
              subgraph_building_method="zero_filling",
          )
  
          maskout_pred = gnn_score(
              maskout_nodes_list,
              data,
              value_func=value_func,
              subgraph_building_method="zero_filling",
          )
  
          sparsity_score = sparsity(
              masked_node_list,
              data,
              subgraph_building_method="zero_filling",
          )
  
      # return variables
      pred_mask = [edge_mask]
      related_preds = [
          {
              "masked": masked_pred,
              "maskout": maskout_pred,
              "origin": probs[label],
              "sparsity": sparsity_score,
          }
      ]
      return None, pred_mask, related_preds
  
imports: import tqdm
  import time
  import torch
  import numpy
  import torch.nn
  import networkx
  from math import sqrt
  from torch import Tensor
  from textwrap import wrap
  from torch.optim import Adam
  import matplotlib.pyplot
  import torch.nn.functional
  from torch_geometric.data import Data
  from torch_geometric.nn.conv import MessagePassing
  from torch_geometric.utils import to_networkx
  from torch_geometric.utils.num_nodes import maybe_num_nodes
  from typing import Tuple, List, Dict, Optional
  from .shapley import gnn_score, GnnNetsNC2valueFunc, GnnNetsGC2valueFunc, sparsity
  from torch_geometric.datasets import MoleculeNet
  from rdkit import Chem
label_desc: Use image node_id 11 for calling the PGExplainer obj's underlying member method code with example usage: obj.forward(x, edge_index) and returns: None,  pred_mask,  related_preds