Graph Data Structures ===================== At the core of TorchDrug, we provides several data structures to enable common operations in graph representation learning. Create a Graph -------------- To begin with, let's create a graph. .. code:: python import torch from torchdrug import data edge_list = [[0, 1], [1, 2], [2, 3], [3, 4], [4, 5], [5, 0]] graph = data.Graph(edge_list, num_node=6) graph.visualize() This will plot a ring graph like the following. .. image:: ../../../asset/graph/graph.png :align: center :width: 33% Internally, the graph is stored as a sparse edge list to save memory footprint. For an intuitive comparison, a `scale-free graph`_ may have 1 million nodes and 10 million edges. The dense version takes about 4TB, while the sparse version only requires 120MB. .. _scale-free graph: https://en.wikipedia.org/wiki/Scale-free_network Here are some commonly used properties of the graph. .. code:: python print(graph.num_node) print(graph.num_edge) print(graph.edge_list) print(graph.edge_weight) In some scenarios, the graph may also have type information on its edges. For example, molecules have bond types like ``single bound``, while knowledge graphs have relations like ``consists of``. To construct such a relational graph, we can pass the edge type as a third variable in the edge list. .. code:: python triplet_list = [[0, 1, 0], [1, 2, 1], [2, 3, 0], [3, 4, 1], [4, 5, 0], [5, 0, 1]] graph = data.Graph(triplet_list, num_node=6, num_relation=2) graph.visualize() .. image:: ../../../asset/graph/relational_graph.png :align: center :width: 33% Alternatively, we can also use adjacency matrices to create the above graphs. The normal graph uses a 2D adjacency matrix :math:`A`, where non-zero :math:`A_{i,j}` corresponds to an edge from node :math:`i` to node :math:`j`. The relational graph uses a 3D adjacency matrix :math:`A`, where non-zero :math:`A_{i,j,k}` denotes an edge from node :math:`i` to node :math:`j` with edge type :math:`k`. .. code:: python adjacency = torch.zeros(6, 6) adjacency[edge_list] = 1 graph = data.Graph.from_dense(adjacency) adjacency = torch.zeros(6, 6, 2) adjacency[triplet_list] = 1 graph = data.Graph.from_dense(adjacency) For molecule graphs, TorchDrug supports creating instances from `SMILES`_ strings. For example, the following code creates a benzene molecule. .. _SMILES: https://en.wikipedia.org/wiki/Simplified_molecular-input_line-entry_system .. code:: python mol = data.Molecule.from_smiles("C1=CC=CC=C1") mol.visualize() .. image:: ../../../asset/graph/benzene.png :align: center :width: 33% Once the graph is created, we can transfer it between CPU and GPUs, just like :class:`torch.Tensor`. .. code:: python graph = graph.cuda() print(graph.device) graph = graph.cpu() print(graph.device) Graph Attributes ---------------- A common practice in graph representation learning is to add some graph features as the input of neural networks. Typically, there are three types of features, node-level, edge-level and graph-level features. In TorchDrug, these features are stored as node/edge/graph attributes in the data structure, and are automatically processed during any graph operation. Here we specify some features during the construction of the molecule graph. .. code:: python mol = data.Molecule.from_smiles("C1=CC=CC=C1", atom_feature="default", bond_feature="default", mol_feature="ecfp") print(mol.node_feature.shape) print(mol.edge_feature.shape) print(mol.graph_feature.shape) There are a bunch of popular feature functions provided in :mod:`torchdrug.data.feature`. We may also want to define our own attributes. This only requires to wrap the assignment lines with a context manager. The following example defines edge importance as the reciprocal of node degrees. .. code:: python node_in, node_out = mol.edge_list.t()[:2] with mol.edge(): mol.edge_importance = 1 / graph.degree_in[node_in] + 1 / graph.degree_out[node_out] We can use ``mol.node()`` and ``mol.graph()`` for node- and graph-level attributes respectively. Attributes may also be a reference to node/edge/graph indexes. See :doc:`reference` for more details. Note in order to support batching and masking, attributes should always have the same length as their corresponding components. This means the size of the first dimension of the tensor should be either ``num_node``, ``num_edge`` or ``1``. Batch Graph ----------- Modern deep learning frameworks employs batched operations to accelerate computation. In TorchDrug, we can easily batch same kind of graphs with **arbitary sizes**. Here is an example of creating a batch of 4 graphs. .. code:: python graphs = [graph, graph, graph, graph] batch = data.Graph.pack(graphs) batch.visualize(num_row=1) .. image:: ../../../asset/graph/batch.png This returns a :class:`PackedGraph ` instance with all attributes automatically batched. The essential trick behind this operation is based on a property of graphs. A batch of :math:`n` graphs is equivalent to a large graph with :math:`n` connected components. The equivalent adjacency matrix for a batch is .. math:: A = \begin{bmatrix} A_1 & \cdots & 0 \\ \vdots & \ddots & \vdots \\ 0 & \cdots & A_n \end{bmatrix} where :math:`A_i` is the adjacency of :math:`i`-th graph. To get a single graph from the batch, use the conventional index or :meth:`PackedGraph.unpack `. .. code:: python graph = batch[1] graphs = batch.unpack() One advantage of such batching mechanism is that it does not distinguish a single graph and a batch. In other words, we only need to implement single graph operations, and they can be directly applied as batched operations. This reduces the pain of writing batched operations. Subgraph and Masking -------------------- The graph data structure also provides a bunch of slicing operations to create subgraphs or masked graphs in a sparse manner. Some typical operations include .. code:: python g1 = graph.subgraph([1, 2, 3, 4]) g1.visualize() g2 = graph.node_mask([1, 2, 3, 4]) g2.visualize() g3 = graph.edge_mask([0, 1, 5]) g3.visualize() g4 = g3.compact() g4.visualize() .. image:: ../../../asset/graph/subgraph.png :width: 24% .. image:: ../../../asset/graph/node_mask.png :width: 24% .. image:: ../../../asset/graph/edge_mask.png :width: 24% .. image:: ../../../asset/graph/compact.png :width: 24% All the above operations accept either integer node indexes or binary node masks. :meth:`subgraph() ` extracts a subgraph based on the given nodes. The node ids are re-mapped to produce a compact index. :meth:`node_mask() ` keeps edges among the given nodes. :meth:`edge_mask() ` keeps edges of the given edge indexes. :meth:`compact() ` removes all isolated nodes. The same operations can also be applied to batches. In this case, we need to convert the index of a single graph into the index in a batch. .. code:: python graph_ids = torch.tensor([0, 0, 0, 0, 1, 1, 1, 1, 1, 1]) node_ids = torch.tensor([1, 2, 3, 4, 0, 1, 2, 3, 4, 5]) node_ids += batch.num_cum_nodes[graph_ids] - batch.num_nodes[graph_ids] batch = batch.node_mask(node_ids) batch.visualize(num_row=1) .. image:: ../../../asset/graph/batch_node_mask.png We can also pick a subset of graphs in a batch. .. code:: python batch = batch[[0, 1]] batch.visualize() .. image:: ../../../asset/graph/subbatch.png :align: center :width: 66%