class_docstr: A `Pytorch Geometric <https://pytorch-geometric.readthedocs.io/en/latest/index.html>`_ data interface for datasets used in molecule generation.
  
  .. note::
      When creating data loaders of this dataset class, only dig.ggraph3D.dataset.collate_fn can be used as the collate function.
  
  Args:
      root (string, optional): Root directory where the dataset should be saved. (default: :obj:`./`)
      subset_idxs (list, optional): if it is not None, only the data located 
          at the indexs in subset_idxs of the dataset will be sampled.
          (default: :obj:`None`)
      transform (callable, optional): A function/transform that takes in an
          :obj:`torch_geometric.data.Data` object and returns a transformed
          version. The data object will be transformed before every access.
          (default: :obj:`None`)
      pre_transform (callable, optional): A function/transform that takes in
          an :obj:`torch_geometric.data.Data` object and returns a
          transformed version. The data object will be transformed before
          being saved to disk. (default: :obj:`None`)
      pre_filter (callable, optional): A function that takes in an
          :obj:`torch_geometric.data.Data` object and returns a boolean
          value, indicating whether the data object should be included in the
          final dataset. (default: :obj:`None`)
docstr: Gets the data object at index :idx:.
  
  Args:
      idx: The index of the data that you want to reach.
  :rtype: :class:`dict` a python dict with the following items:
              "atom_type" --- the atom types of previously generated geometries at each generation step;
              "position" --- the atom coordinates of previously generated geometries at each generation step;
              "batch" --- the identity indexes used to discriminate different generation step;
              "focus" --- the index of focus atom at each generation step;
              "c1_focus" --- the index of c1 and focus atom at each generation step;
              "c2_c1_focus" --- the index of c2, c1, and focus atom at each generation step;
              "new_atom_type" --- the atom types to be generated at each generation step;
              "new_dist" --- the distances to be generated at each generation step;
              "new_angle" --- the angles to be generated at each generation step;
              "new_torsion" --- the torsion angles to be generated at each generation step;
              "cannot_focus" --- the labels denoting whether each atom can serve as the focus atom or not in the previously generated geometries at each generation step.
docstr_tok: ["Gets","the","data","object","at","index",":","idx",":",".","Args",":","idx",":","The","index","of","the","data","that","you","want","to","reach",".",":","rtype",":",":","class",":","`","dict","`","a","python","dict","with","the","following","items",":","``","atom_type","''","--","-","the","atom","types","of","previously","generated","geometries","at","each","generation","step",";","``","position","''","--","-","the","atom","coordinates","of","previously","generated","geometries","at","each","generation","step",";","``","batch","''","--","-","the","identity","indexes","used","to","discriminate","different","generation","step",";","``","focus","''","--","-","the","index","of","focus","atom","at","each","generation","step",";","``","c1_focus","''","--","-","the","index","of","c1","and","focus","atom","at","each","generation","step",";","``","c2_c1_focus","''","--","-","the","index","of","c2",",","c1",",","and","focus","atom","at","each","generation","step",";","``","new_atom_type","''","--","-","the","atom","types","to","be","generated","at","each","generation","step",";","``","new_dist","''","--","-","the","distances","to","be","generated","at","each","generation","step",";","``","new_angle","''","--","-","the","angles","to","be","generated","at","each","generation","step",";","``","new_torsion","''","--","-","the","torsion","angles","to","be","generated","at","each","generation","step",";","``","cannot_focus","''","--","-","the","labels","denoting","whether","each","atom","can","serve","as","the","focus","atom","or","not","in","the","previously","generated","geometries","at","each","generation","step","."]
code: def get(self, idx):
      r"""Gets the data object at index :idx:.
  
      Args:
          idx: The index of the data that you want to reach.
      :rtype: :class:`dict` a python dict with the following items:
                  "atom_type" --- the atom types of previously generated geometries at each generation step;
                  "position" --- the atom coordinates of previously generated geometries at each generation step;
                  "batch" --- the identity indexes used to discriminate different generation step;
                  "focus" --- the index of focus atom at each generation step;
                  "c1_focus" --- the index of c1 and focus atom at each generation step;
                  "c2_c1_focus" --- the index of c2, c1, and focus atom at each generation step;
                  "new_atom_type" --- the atom types to be generated at each generation step;
                  "new_dist" --- the distances to be generated at each generation step;
                  "new_angle" --- the angles to be generated at each generation step;
                  "new_torsion" --- the torsion angles to be generated at each generation step;
                  "cannot_focus" --- the labels denoting whether each atom can serve as the focus atom or not in the previously generated geometries at each generation step.
      """
      atom_type, position, con_mat = (
          self.atom_type_list[idx],
          self.position_list[idx],
          self.con_mat_list[idx],
      )
      atom_valency = torch.sum(con_mat, dim=1)
  
      squared_dist = torch.sum(
          torch.square(position[:, None, :] - position[None, :, :]),
          dim=-1,
      )
      nx_graph = nx.from_numpy_matrix(squared_dist.numpy())
      edges = list(
          tree.minimum_spanning_edges(
              nx_graph, algorithm="prim", data=False
          )
      )
  
      focus_node_id, target_node_id = zip(*edges)
      # print(focus_node_id, target_node_id)
  
      node_perm = torch.cat(
          (torch.tensor([0]), torch.tensor(target_node_id))
      )
      position = position[node_perm]
      atom_type = atom_type[node_perm]
      con_mat = con_mat[node_perm][:, node_perm]
      squared_dist = squared_dist[node_perm][:, node_perm]
      atom_valency = atom_valency[node_perm]
      # print(con_mat)s
  
      focus_node_id = torch.tensor(focus_node_id)
      steps_focus = torch.nonzero(
          focus_node_id[:, None] == node_perm[None, :]
      )[:, 1]
      steps_c1_focus, steps_c2_c1_focus = torch.empty(
          [0, 2], dtype=int
      ), torch.empty([0, 3], dtype=int)
      steps_batch, steps_position, steps_atom_type = (
          torch.empty([0, 1], dtype=int),
          torch.empty([0, 3], dtype=position.dtype),
          torch.empty([0, 1], dtype=atom_type.dtype),
      )
      steps_cannot_focus = torch.empty([0, 1], dtype=float)
      steps_dist, steps_angle, steps_torsion = (
          torch.empty([0, 1], dtype=float),
          torch.empty([0, 1], dtype=float),
          torch.empty([0, 1], dtype=float),
      )
      idx_offsets = torch.cumsum(
          torch.arange(len(atom_type) - 1), dim=0
      )
  
      for i in range(len(atom_type) - 1):
          partial_con_mat = con_mat[: i + 1, : i + 1]
          valency_sum = partial_con_mat.sum(dim=1, keepdim=True)
          steps_cannot_focus = torch.cat(
              (
                  steps_cannot_focus,
                  (valency_sum == atom_valency[: i + 1, None]).float(),
              )
          )
  
          one_step_focus = steps_focus[i]
          focus_pos, new_pos = position[one_step_focus], position[i + 1]
          one_step_dis = torch.norm(new_pos - focus_pos)
          steps_dist = torch.cat((steps_dist, one_step_dis.view(1, 1)))
  
          if i > 0:
              mask = torch.ones([i + 1], dtype=torch.bool)
              mask[one_step_focus] = False
              c1_dists = squared_dist[one_step_focus, : i + 1][mask]
              one_step_c1 = torch.argmin(c1_dists)
              if one_step_c1 >= one_step_focus:
                  one_step_c1 += 1
              steps_c1_focus = torch.cat(
                  (
                      steps_c1_focus,
                      torch.tensor([one_step_c1, one_step_focus]).view(
                          1, 2
                      )
                      + idx_offsets[i],
                  )
              )
  
              c1_pos = position[one_step_c1]
              a = ((c1_pos - focus_pos) * (new_pos - focus_pos)).sum(
                  dim=-1
              )
              b = torch.cross(
                  c1_pos - focus_pos, new_pos - focus_pos
              ).norm(dim=-1)
              one_step_angle = torch.atan2(b, a)
              steps_angle = torch.cat(
                  (steps_angle, one_step_angle.view(1, 1))
              )
  
              if i > 1:
                  mask[one_step_c1] = False
                  c2_dists = squared_dist[one_step_c1, : i + 1][mask]
                  one_step_c2 = torch.argmin(c2_dists)
                  if one_step_c2 >= min(one_step_c1, one_step_focus):
                      one_step_c2 += 1
                      if one_step_c2 >= max(
                          one_step_c1, one_step_focus
                      ):
                          one_step_c2 += 1
                  steps_c2_c1_focus = torch.cat(
                      (
                          steps_c2_c1_focus,
                          torch.tensor(
                              [one_step_c2, one_step_c1, one_step_focus]
                          ).view(1, 3)
                          + idx_offsets[i],
                      )
                  )
  
                  c2_pos = position[one_step_c2]
                  plane1 = torch.cross(
                      focus_pos - c1_pos, new_pos - c1_pos
                  )
                  plane2 = torch.cross(
                      focus_pos - c1_pos, c2_pos - c1_pos
                  )
                  a = (plane1 * plane2).sum(
                      dim=-1
                  )  # cos_angle * |plane1| * |plane2|
                  b = (
                      torch.cross(plane1, plane2) * (focus_pos - c1_pos)
                  ).sum(dim=-1) / torch.norm(focus_pos - c1_pos)
                  one_step_torsion = torch.atan2(b, a)
                  steps_torsion = torch.cat(
                      (steps_torsion, one_step_torsion.view(1, 1))
                  )
  
          one_step_position = position[: i + 1]
          steps_position = torch.cat(
              (steps_position, one_step_position), dim=0
          )
          one_step_atom_type = atom_type[: i + 1]
          steps_atom_type = torch.cat(
              (steps_atom_type, one_step_atom_type.view(-1, 1))
          )
          steps_batch = torch.cat(
              (steps_batch, torch.tensor([i]).repeat(i + 1).view(-1, 1))
          )
  
      steps_focus += idx_offsets
      steps_new_atom_type = atom_type[1:]
      steps_torsion[steps_torsion <= 0] += 2 * pi
  
      data_batch = {}
      data_batch["atom_type"] = steps_atom_type.view(-1)
      data_batch["position"] = steps_position
      data_batch["batch"] = steps_batch.view(-1)
      data_batch["focus"] = steps_focus[:, None]
      data_batch["c1_focus"] = steps_c1_focus
      data_batch["c2_c1_focus"] = steps_c2_c1_focus
      data_batch["new_atom_type"] = steps_new_atom_type.view(-1)
      data_batch["new_dist"] = steps_dist
      data_batch["new_angle"] = steps_angle
      data_batch["new_torsion"] = steps_torsion
      data_batch["cannot_focus"] = steps_cannot_focus.view(-1).float()
  
      return data_batch
  
imports: import os
  import os.path
  import numpy
  import torch
  import networkx
  from networkx.algorithms import tree
  from math import pi
  from rdkit import Chem
  from rdkit.Chem.rdchem import BondType
  from torch_geometric.data import Data, InMemoryDataset, download_url, extract_tar
label_desc: Use image node_id 10 for calling the QM93DGEN obj's underlying member method code with example usage: obj.get(idx) and returns: data_batch