docstr: inverse flow to generate molecule
  Args: 
      temp: temperature of normal distributions, we sample from (0, temp^2 * I)
docstr_tok: ["inverse","flow","to","generate","molecule","Args",":","temp",":","temperature","of","normal","distributions",",","we","sample","from","(","0",",","temp^2","*","I",")"]
code: def generate(
      self, atom_list, temperature=[0.3, 0.3], min_atoms=7, max_atoms=48
  ):
      """
      inverse flow to generate molecule
      Args:
          temp: temperature of normal distributions, we sample from (0, temp^2 * I)
      """
      with torch.no_grad():
          num2bond = {
              0: Chem.rdchem.BondType.SINGLE,
              1: Chem.rdchem.BondType.DOUBLE,
              2: Chem.rdchem.BondType.TRIPLE,
          }
          num2atom = {i: atom_list[i] for i in range(len(atom_list))}
  
          if self.dp:
              cur_node_features = torch.zeros(
                  [1, max_atoms, self.node_dim]
              ).cuda()
              cur_adj_features = torch.zeros(
                  [1, self.bond_dim, max_atoms, max_atoms]
              ).cuda()
          else:
              cur_node_features = torch.zeros(
                  [1, max_atoms, self.node_dim]
              )
              cur_adj_features = torch.zeros(
                  [1, self.bond_dim, max_atoms, max_atoms]
              )
  
          node_features_each_iter_backup = (
              cur_node_features.clone()
          )  # backup of features, updated when newly added node is connected to previous subgraph
          adj_features_each_iter_backup = cur_adj_features.clone()
  
          rw_mol = Chem.RWMol()  # editable mol
          mol = None
  
          is_continue = True
          edge_idx = 0
          total_resample = 0
          each_node_resample = np.zeros([max_atoms])
  
          for i in range(max_atoms):
              if not is_continue:
                  break
              if i < self.edge_unroll:
                  edge_total = i  # edge to sample for current node
                  start = 0
              else:
                  edge_total = self.edge_unroll
                  start = i - self.edge_unroll
              # first generate node
              ## reverse flow
  
              prior_node_dist = torch.distributions.OneHotCategorical(
                  logits=self.node_base_log_probs[i] * temperature[0]
              )
              latent_node = prior_node_dist.sample().view(1, -1)
  
              if self.dp:
                  latent_node = self.flow_core.module.reverse(
                      cur_node_features,
                      cur_adj_features,
                      latent_node,
                      mode=0,
                  ).view(
                      -1
                  )  # (9, )
              else:
                  latent_node = self.flow_core.reverse(
                      cur_node_features,
                      cur_adj_features,
                      latent_node,
                      mode=0,
                  ).view(
                      -1
                  )  # (9, )
  
              feature_id = torch.argmax(latent_node).item()
              cur_node_features[0, i, feature_id] = 1.0
              cur_adj_features[0, :, i, i] = 1.0
              rw_mol.AddAtom(Chem.Atom(num2atom[feature_id]))
  
              if i == 0:
                  is_connect = True
              else:
                  is_connect = False
  
              for j in range(edge_total):
                  valid = False
                  resample_edge = 0
                  edge_dis = self.edge_base_log_probs[edge_idx].clone()
                  invalid_bond_type_set = set()
                  while not valid:
                      if (
                          len(invalid_bond_type_set) < 3
                          and resample_edge <= 50
                      ):  # haven't sampled all possible bond type or is not stuck in the loop
                          prior_edge_dist = (
                              torch.distributions.OneHotCategorical(
                                  logits=edge_dis / temperature[1]
                              )
                          )
                          latent_edge = prior_edge_dist.sample().view(
                              1, -1
                          )
                          latent_id = torch.argmax(latent_edge, dim=1)
  
                          if self.dp:
                              latent_edge = (
                                  self.flow_core.module.reverse(
                                      cur_node_features,
                                      cur_adj_features,
                                      latent_edge,
                                      mode=1,
                                      edge_index=torch.Tensor(
                                          [[j + start, i]]
                                      )
                                      .long()
                                      .cuda(),
                                  ).view(-1)
                              )  # (4, )
                          else:
                              latent_edge = self.flow_core.reverse(
                                  cur_node_features,
                                  cur_adj_features,
                                  latent_edge,
                                  mode=1,
                                  edge_index=torch.Tensor(
                                      [[j + start, i]]
                                  ).long(),
                              ).view(
                                  -1
                              )  # (4, )
                          edge_discrete_id = torch.argmax(
                              latent_edge
                          ).item()
                      else:
                          assert (
                              resample_edge > 50
                              or len(invalid_bond_type_set) == 3
                          )
                          edge_discrete_id = 3  # we have no choice but to choose not to add edge between (i, j+start)
  
                      cur_adj_features[
                          0, edge_discrete_id, i, j + start
                      ] = 1.0
                      cur_adj_features[
                          0, edge_discrete_id, j + start, i
                      ] = 1.0
                      if edge_discrete_id == 3:  # virtual edge
                          valid = True
                      else:  # single/double/triple bond
                          rw_mol.AddBond(
                              i, j + start, num2bond[edge_discrete_id]
                          )
                          valid = check_valency(rw_mol)
                          if valid:
                              is_connect = True
                              # print(num2bond_symbol[edge_discrete_id])
                          else:  # backtrack
                              edge_dis[latent_id] = float("-inf")
                              rw_mol.RemoveBond(i, j + start)
                              cur_adj_features[
                                  0, edge_discrete_id, i, j + start
                              ] = 0.0
                              cur_adj_features[
                                  0, edge_discrete_id, j + start, i
                              ] = 0.0
                              total_resample += 1.0
                              each_node_resample[i] += 1.0
                              resample_edge += 1
  
                              invalid_bond_type_set.add(
                                  edge_discrete_id
                              )
  
                  edge_idx += 1
  
              if (
                  is_connect
              ):  # new generated node has at least one bond with previous node, do not stop generation, backup mol from rw_mol to mol
                  is_continue = True
                  mol = rw_mol.GetMol()
                  node_features_each_iter_backup = (
                      cur_node_features.clone()
                  )  # update node backup since new node is valid
                  adj_features_each_iter_backup = (
                      cur_adj_features.clone()
                  )
              else:
                  is_continue = False
                  cur_mol_size = mol.GetNumAtoms()
                  if cur_mol_size >= min_atoms:
                      rw_mol = Chem.RWMol(mol)
                      cur_node_features = (
                          node_features_each_iter_backup.clone()
                      )
                      cur_adj_features = (
                          adj_features_each_iter_backup.clone()
                      )
                      cur_node_features_tmp = cur_node_features.clone()
                      cur_adj_features_tmp = cur_adj_features.clone()
  
                      mol_demon_edit = Chem.RWMol(rw_mol)
                      added_num = np.random.randint(1, 5)
                      for _ in range(added_num):
                          last_id2 = mol_demon_edit.AddAtom(
                              Chem.Atom(6)
                          )
                          cur_node_features_tmp[0, last_id2, 0] = 1.0
                          cur_adj_features_tmp[
                              0, :, last_id2, last_id2
                          ] = 1.0
  
                          flag_success = False
                          count = 0
                          while True:
                              if count > 100:
                                  break
                              if last_id2 > 12:
                                  j = np.random.randint(1, 13)
                              else:
                                  j = np.random.randint(1, last_id2 + 1)
                              mol_demon_edit.AddBond(
                                  int(last_id2 - j),
                                  int(last_id2),
                                  Chem.rdchem.BondType.SINGLE,
                              )
                              cur_adj_features_tmp[
                                  0, 0, last_id2 - j, last_id2
                              ] = 1.0
                              cur_adj_features_tmp[
                                  0, 0, last_id2, last_id2 - j
                              ] = 1.0
  
                              valid = check_valency(mol_demon_edit)
                              if valid:
                                  flag_success = True
                                  break
                              else:
                                  mol_demon_edit.RemoveBond(
                                      int(last_id2 - j), int(last_id2)
                                  )
                                  cur_adj_features_tmp[
                                      0, 0, last_id2 - j, last_id2
                                  ] = 0.0
                                  cur_adj_features_tmp[
                                      0, 0, last_id2, last_id2 - j
                                  ] = 0.0
                                  count += 1
  
                          if flag_success:
                              rw_mol = Chem.RWMol(mol_demon_edit)
                              cur_node_features = (
                                  cur_node_features_tmp.clone()
                              )
                              cur_adj_features = (
                                  cur_adj_features_tmp.clone()
                              )
                              mol = rw_mol.GetMol()
                              node_features_each_iter_backup = (
                                  cur_node_features.clone()
                              )
                              adj_features_each_iter_backup = (
                                  cur_adj_features.clone()
                              )
                          else:
                              break
  
          # mol = rw_mol.GetMol() # mol backup
          assert mol is not None, "mol is None..."
  
          final_mol = convert_radical_electrons_to_hydrogens(mol)
          num_atoms = final_mol.GetNumAtoms()
  
          pure_valid = 0
          if total_resample == 0:
              pure_valid = 1.0
  
          return final_mol, pure_valid, num_atoms
  
file_path: DIG/dig/ggraph/method/GraphDF/model/graphflow.py
imports: import torch
  import torch.nn
  import numpy
  from rdkit import Chem
  from dig.ggraph.utils import check_valency, convert_radical_electrons_to_hydrogens
  from .disgraphaf import DisGraphAF
label_desc: Use image node_id 3 for calling the GraphFlowModel obj's underlying member method code with example usage: obj.generate(atom_list, temperature, min_atoms, max_atoms) and returns: final_mol,  pure_valid,  num_atoms