Method __getitem__

python/dgl/data/qm9.py:180–219

r"""Get graph and label by index Parameters ---------- idx : int Item index Returns ------- dgl.DGLGraph The graph contains: - ``ndata['R']``: the coordinates of each atom - ``ndata['Z']``: the atomic

(self, idx)

        return self.label.shape[1]

    def __getitem__(self, idx):
        r"""Get graph and label by index

        Parameters
        ----------
        idx : int
            Item index

        Returns
        -------
        dgl.DGLGraph
            The graph contains:

            - ``ndata['R']``: the coordinates of each atom
            - ``ndata['Z']``: the atomic number

        Tensor
            Property values of molecular graphs
        """
        label = F.tensor(self.label[idx], dtype=F.data_type_dict["float32"])
        n_atoms = self.N[idx]
        R = self.R[self.N_cumsum[idx] : self.N_cumsum[idx + 1]]
        dist = np.linalg.norm(R[:, None, :] - R[None, :, :], axis=-1)
        adj = sp.csr_matrix(dist <= self.cutoff) - sp.eye(
            n_atoms, dtype=np.bool_
        )
        adj = adj.tocoo()
        u, v = F.tensor(adj.row), F.tensor(adj.col)
        g = dgl_graph((u, v))
        g = to_bidirected(g)
        g.ndata["R"] = F.tensor(R, dtype=F.data_type_dict["float32"])
        g.ndata["Z"] = F.tensor(
            self.Z[self.N_cumsum[idx] : self.N_cumsum[idx + 1]],
            dtype=F.data_type_dict["int64"],
        )

        if self._transform is not None:
            g = self._transform(g)

        return g, label

    def __len__(self):
        r"""Number of graphs in the dataset.