code: def reinforce_optim_one_mol(
      self, atom_list, temperature=[0.3, 0.3], max_size_rl=38
  ):
      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_size_rl, self.node_dim]
          ).cuda()
          cur_adj_features = torch.zeros(
              [1, self.bond_dim, max_size_rl, max_size_rl]
          ).cuda()
      else:
          cur_node_features = torch.zeros(
              [1, max_size_rl, self.node_dim]
          )
          cur_adj_features = torch.zeros(
              [1, self.bond_dim, max_size_rl, max_size_rl]
          )
      self.eval()
  
      with torch.no_grad():
          min_action_node = 1
          rw_mol = Chem.RWMol()  # editable mol
          mol = None
  
          is_continue = True
          edge_idx = 0
  
          added_num = 0
          for i in range(max_size_rl):
              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_old.module.reverse(cur_node_features, cur_adj_features, latent_node, mode=0).view(-1)  # (9, )
                  latent_node = self.flow_core.module.reverse(
                      cur_node_features,
                      cur_adj_features,
                      latent_node,
                      mode=0,
                  ).view(-1)
              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]))
  
              # then generate edges
              if i == 0:
                  is_connect = True
              else:
                  is_connect = False
              # cur_mol_size = mol.GetNumAtoms
              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()
                                  .cuda(),
                              ).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  # virtual edge is alway valid
                      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
                          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
                              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()
                  added_num += 1
  
              else:
                  #! we may need to satisfy min action here
                  if added_num >= min_action_node:
                      # if we have already add 'min_action_node', ignore and let continue be false.
                      # the data added in last iter will be pop afterwards
                      is_continue = False
                  else:
                      pass
  
      if check_chemical_validity(mol) is True:
          current_smile = Chem.MolToSmiles(mol, isomericSmiles=True)
          final_mol = Chem.MolFromSmiles(current_smile)
          num_atoms = final_mol.GetNumAtoms()
          return final_mol, num_atoms
      else:
          return None, None
  
file_path: DIG/dig/ggraph/method/GraphDF/model/graphflow_rl.py
imports: import torch
  import torch.nn
  import numpy
  from rdkit import Chem
  from dig.ggraph.utils import check_chemical_validity, check_valency, calculate_min_plogp, qed
  from dig.ggraph.utils import convert_radical_electrons_to_hydrogens, steric_strain_filter, zinc_molecule_filter
  from .disgraphaf import DisGraphAF
label_desc: Use image node_id 2 for calling the GraphFlowModel_rl obj's underlying member method code with example usage: obj.reinforce_optim_one_mol(atom_list, temperature, max_size_rl) and returns: final_mol,  num_atoms, None,  None