code: def main():
      writer = SummaryWriter()
      args = arg_parse()
  
      edge_train_mode = args.mode
      print("edge train mode: {}".format(edge_train_mode))
  
      ppi_graph = read_ppi_data(args.ppi_path)
  
      mode = "mixed"
      if mode == "ppi":
          message_passing_graph = ppi_graph
          cmap_graph, knockout_nodes = read_cmap_data(args.data_path)
      elif mode == "mixed":
          message_passing_graph, knockout_nodes = read_cmap_data(
              args.data_path, ppi_graph
          )
  
      print(
          "Each node has gene ID. Example: ",
          message_passing_graph.nodes["ADPGK"],
      )
      print(
          "Each edge has de direction. Example",
          message_passing_graph["ADPGK"]["IL1B"],
      )
      print(
          "Total num edges: ", message_passing_graph.number_of_edges()
      )
  
      # disjoint edge label
      disjoint_split_ratio = 0.1
      val_ratio = 0.1
      disjoint_edge_label_index = []
      val_edges = []
  
      # newly edited
      train_edges = []
      for u in knockout_nodes:
          rand_num = np.random.rand()
          if rand_num < disjoint_split_ratio:
              # add all edges (cmap only) into edge label index
              # cmap is not a multigraph
              disjoint_edge_label_index.extend(
                  [
                      (u, v, edge_key)
                      for v in message_passing_graph.successors(u)
                      for edge_key in message_passing_graph[u][v]
                      if message_passing_graph[u][v][edge_key][
                          "edge_type"
                      ]
                      == 1
                  ]
              )
  
              train_edges.extend(
                  [
                      (u, v, edge_key)
                      for v in message_passing_graph.successors(u)
                      for edge_key in message_passing_graph[u][v]
                      if message_passing_graph[u][v][edge_key][
                          "edge_type"
                      ]
                      == 1
                  ]
              )
          elif rand_num < disjoint_split_ratio + val_ratio:
              val_edges.extend(
                  [
                      (u, v, edge_key)
                      for v in message_passing_graph.successors(u)
                      for edge_key in message_passing_graph[u][v]
                      if message_passing_graph[u][v][edge_key][
                          "edge_type"
                      ]
                      == 1
                  ]
              )
          else:
              train_edges.extend(
                  [
                      (u, v, edge_key)
                      for v in message_passing_graph.successors(u)
                      for edge_key in message_passing_graph[u][v]
                      if message_passing_graph[u][v][edge_key][
                          "edge_type"
                      ]
                      == 1
                  ]
              )
      # add default node types for message_passing_graph
      for node in message_passing_graph.nodes:
          message_passing_graph.nodes[node]["node_type"] = 0
  
      print("Num edges to predict: ", len(disjoint_edge_label_index))
      print("Num edges in val: ", len(val_edges))
      print("Num edges in train: ", len(train_edges))
  
      graph = HeteroGraph(
          message_passing_graph,
          custom={
              "general_splits": [train_edges, val_edges],
              "disjoint_split": disjoint_edge_label_index,
              "task": "link_pred",
          },
      )
  
      graphs = [graph]
      graphDataset = GraphDataset(
          graphs, task="link_pred", edge_train_mode="disjoint"
      )
  
      # Transform dataset
      # de direction (currently using homogeneous graph)
      num_edge_types = 2
  
      graphDataset = graphDataset.apply_transform(
          cmap_transform, num_edge_types=num_edge_types, deep_copy=False
      )
      print(
          "Number of node features: ", graphDataset.num_node_features()
      )
  
      # split dataset
      dataset = {}
      dataset["train"], dataset["val"] = graphDataset.split(
          transductive=True
      )
  
      # sanity check
      print(
          f"dataset['train'][0].edge_label_index.keys(): {dataset['train'][0].edge_label_index.keys()}"
      )
      print(
          f"dataset['train'][0].edge_label_index[(0, 1, 0)].shape[1]: {dataset['train'][0].edge_label_index[(0, 1, 0)].shape[1]}"
      )
      print(
          f"dataset['val'][0].edge_label_index.keys(): {dataset['val'][0].edge_label_index.keys()}"
      )
      print(
          f"dataset['val'][0].edge_label_index[(0, 1, 0)].shape[1]: {dataset['val'][0].edge_label_index[(0, 1, 0)].shape[1]}"
      )
      print(
          f"len(list(dataset['train'][0].G.edges)): {len(list(dataset['train'][0].G.edges))}"
      )
      print(
          f"len(list(dataset['val'][0].G.edges)): {len(list(dataset['val'][0].G.edges))}"
      )
      print(
          f"list(dataset['train'][0].G.edges)[:10]: {list(dataset['train'][0].G.edges)[:10]}"
      )
      print(
          f"list(dataset['val'][0].G.edges)[:10]: {list(dataset['val'][0].G.edges)[:10]}"
      )
  
      # node feature dimension
      input_dim = dataset["train"].num_node_features()
      edge_feat_dim = dataset["train"].num_edge_features()
      num_classes = dataset["train"].num_edge_labels()
      print(
          "Node feature dim: {}; edge feature dim: {}; num classes: {}.".format(
              input_dim, edge_feat_dim, num_classes
          )
      )
      exit()
  
      # relation type is both used for edge features and edge labels
      model = Net(input_dim, edge_feat_dim, num_classes, args).to(
          args.device
      )
      optimizer = torch.optim.Adam(
          model.parameters(), lr=0.001, weight_decay=5e-3
      )
      follow_batch = []  # e.g., follow_batch = ['edge_index']
  
      dataloaders = {
          split: DataLoader(
              ds,
              collate_fn=Batch.collate(follow_batch),
              batch_size=1,
              shuffle=(split == "train"),
          )
          for split, ds in dataset.items()
      }
      print("Graphs after split: ")
      for key, dataloader in dataloaders.items():
          for batch in dataloader:
              print(key, ": ", batch)
  
      train(model, dataloaders, optimizer, args, writer=writer)
  
imports: import argparse
  import copy
  import csv
  import numpy
  import time
  import torch
  import torch.nn
  import torch.nn.functional
  from torch.utils.tensorboard import SummaryWriter
  import networkx
  import numpy
  import sklearn.metrics
  import torch_geometric.transforms
  import torch_geometric.nn
  from torch.utils.data import DataLoader
  from deepsnap.hetero_graph import HeteroGraph
  from deepsnap.dataset import GraphDataset
  from deepsnap.batch import Batch
label_desc: Use image node_id 8 for calling a global function with example usage: main() without return types