code: def reinforce_constrained_optim_one_mol(
      self,
      x,
      adj,
      mol_size,
      raw_smile,
      bfs_perm_origin,
      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))}
  
      self.eval()
  
      cur_mols = []
      cur_mol_imps = []
      cur_mol_sims = []
  
      with torch.no_grad():
          flag_reconstruct_from_node_adj = True
  
          rand_num = np.random.rand()
          if rand_num <= 0.5:
              cur_modify_size = np.random.randint(
                  low=0, high=self.conf_rl["modify_size"]
              )
          else:
              cur_modify_size = 0
  
          keep_size = mol_size - cur_modify_size
  
          org_bfs_perm_origin = bfs_perm_origin
          org_node_features = x  # (1, 38, 9)
          org_adj_features = adj  # (1, 4, 38, 38)
  
          cur_node_features = x.clone()  # (1, 38, 9)
          cur_adj_features = adj.clone()  # (1, 4, 38, 38)
  
          cur_node_features[:, keep_size:, :] = 0.0
          cur_adj_features[:, :, keep_size:, :] = 0.0
          cur_adj_features[:, :, :, keep_size:] = 0.0
  
          rw_mol = Chem.RWMol()  # editable mol
          mol = None
  
          for i in range(mol_size):
              node_id = torch.argmax(
                  org_node_features[0, i]
              ).item()  # (9, )
              if i < keep_size:
                  rw_mol.AddAtom(Chem.Atom(num2atom[node_id]))
              else:
                  pass
  
              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
  
              for j in range(edge_total):
                  edge_id = torch.argmax(
                      org_adj_features[0, :, i, j + start]
                  ).item()  # (4, )
                  if edge_id == 3:  # no bond to add
                      continue
                  if i < keep_size:
                      rw_mol.AddBond(i, j + start, num2bond[edge_id])
                  else:
                      pass
  
          mol = rw_mol.GetMol()
          s_raw = raw_smile
          org_mol_true_raw = Chem.MolFromSmiles(s_raw)
  
          # calculate property for mol org
          org_mol_plogp = calculate_min_plogp(org_mol_true_raw)
  
          if check_chemical_validity(mol) is False or (
              cur_modify_size == 0 and np.random.rand() <= 0.5
          ):  # do not use subgraph, use original graph
              rw_mol = Chem.RWMol(org_mol_true_raw)
              mol = rw_mol.GetMol()
              flag_reconstruct_from_node_adj = (
                  False  # bfs is not reflected in current mol
              )
              keep_size = mol_size
              cur_node_features = org_node_features.clone()
              cur_adj_features = org_adj_features.clone()
              cur_node_features[:, keep_size:, :] = 0.0
              cur_adj_features[:, :, keep_size:, :] = 0.0
              cur_adj_features[:, :, :, keep_size:] = 0.0
  
          assert (
              check_chemical_validity(org_mol_true_raw) is True
          ), "s_raw is %s" % (s_raw)
          assert check_chemical_validity(mol) is True
  
          assert mol.GetNumAtoms() == keep_size
          assert org_mol_true_raw.GetNumAtoms() == mol_size
  
          if self.dp:
              cur_node_features = cur_node_features.cuda()
              cur_adj_features = cur_adj_features.cuda()
  
          is_continue = True
          if keep_size <= self.edge_unroll:
              edge_idx = int(keep_size * (keep_size - 1) / 2)
          else:
              edge_idx = int(
                  self.edge_unroll * (self.edge_unroll - 1) / 2
                  + (keep_size - self.edge_unroll) * self.edge_unroll
              )
  
          min_action_node = max(1, max_size_rl - keep_size)
          added_num = 0
          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()
          for i in range(keep_size, 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.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]))
  
              # 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
                          if flag_reconstruct_from_node_adj:
                              rw_mol.AddBond(
                                  i,
                                  j + start,
                                  num2bond[edge_discrete_id],
                              )
                          else:
                              # current rw_mol is copy from original_mol, which does not reflect the bfs order
                              # so we use bfs_perm_origin to map j+start to the true node_id in rw_mol
                              rw_mol.AddBond(
                                  i,
                                  int(
                                      org_bfs_perm_origin[
                                          j + start
                                      ].item()
                                  ),
                                  num2bond[edge_discrete_id],
                              )
                          valid = check_valency(rw_mol)
                          if valid:
                              is_connect = True
                          else:  # backtrack
                              edge_dis[latent_id] = float("-inf")
                              if flag_reconstruct_from_node_adj:
                                  rw_mol.RemoveBond(i, j + start)
                              else:
                                  rw_mol.RemoveBond(
                                      i,
                                      int(
                                          org_bfs_perm_origin[
                                              j + start
                                          ].item()
                                      ),
                                  )
                              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()
  
                  if check_chemical_validity(mol) is True:
                      current_smile = Chem.MolToSmiles(
                          mol, isomericSmiles=True
                      )
                      tmp_mol1 = Chem.MolFromSmiles(current_smile)
                      current_imp = (
                          calculate_min_plogp(tmp_mol1) - org_mol_plogp
                      )
                      current_sim = reward_target_molecule_similarity(
                          tmp_mol1, org_mol_true_raw
                      )
                      if current_imp > 0:
                          cur_mols.append(tmp_mol1)
                          cur_mol_imps.append(current_imp)
                          cur_mol_sims.append(current_sim)
  
                      if (
                          flag_reconstruct_from_node_adj is False
                      ):  # not reconstructed from adj, can convert.
                          mol_converted = (
                              convert_radical_electrons_to_hydrogens(
                                  mol
                              )
                          )
                          if (
                              check_chemical_validity(mol_converted)
                              is True
                          ):
                              current_smile2 = Chem.MolToSmiles(
                                  mol_converted, isomericSmiles=True
                              )
                              tmp_mol2 = Chem.MolFromSmiles(
                                  current_smile2
                              )
                              current_imp2 = (
                                  calculate_min_plogp(tmp_mol2)
                                  - org_mol_plogp
                              )
                              current_sim2 = (
                                  reward_target_molecule_similarity(
                                      tmp_mol2, org_mol_true_raw
                                  )
                              )
                              if current_imp2 > 0:
                                  cur_mols.append(tmp_mol2)
                                  cur_mol_imps.append(current_imp2)
                                  cur_mol_sims.append(current_sim2)
  
                  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()
                  )
                  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:
                      is_continue = False  # first set as False
  
                      rw_mol = Chem.RWMol(
                          mol
                      )  # important, recover the mol
                      cur_node_features = (
                          node_features_each_iter_backup.clone()
                      )  # recover the backuped features
                      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()
                      cur_mol_size = rw_mol.GetNumAtoms()
  
                      mol_demon_edit = Chem.RWMol(rw_mol)
                      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
                      if cur_mol_size > keep_size:
                          mol_demon_edit.AddBond(
                              cur_mol_size - 1,
                              last_id2,
                              Chem.rdchem.BondType.SINGLE,
                          )
                          cur_adj_features_tmp[
                              0, 0, cur_mol_size - 1, last_id2
                          ] = 1.0
                          cur_adj_features_tmp[
                              0, 0, last_id2, cur_mol_size - 1
                          ] = 1.0
                          valid = check_valency(mol_demon_edit)
                          if valid:
                              flag_success = True
                          else:
                              mol_demon_edit.RemoveBond(
                                  cur_mol_size - 1, last_id2
                              )
                              cur_adj_features_tmp[
                                  0, 0, cur_mol_size - 1, last_id2
                              ] = 0.0
                              cur_adj_features_tmp[
                                  0, 0, last_id2, cur_mol_size - 1
                              ] = 0.0
  
                              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
                                      )
                                  if flag_reconstruct_from_node_adj:
                                      mol_demon_edit.AddBond(
                                          int(last_id2 - j),
                                          int(last_id2),
                                          Chem.rdchem.BondType.SINGLE,
                                      )
                                  else:
                                      mol_demon_edit.AddBond(
                                          int(
                                              org_bfs_perm_origin[
                                                  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:
                                      if flag_reconstruct_from_node_adj:
                                          mol_demon_edit.RemoveBond(
                                              int(last_id2 - j),
                                              int(last_id2),
                                          )
                                      else:
                                          mol_demon_edit.RemoveBond(
                                              int(
                                                  org_bfs_perm_origin[
                                                      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
                      else:
                          count = 0
                          while True:
                              if count > 100:
                                  break
                              if keep_size > 12:
                                  j = np.random.randint(1, 13)
                              else:
                                  j = np.random.randint(
                                      1, keep_size + 1
                                  )
                              if flag_reconstruct_from_node_adj:
                                  mol_demon_edit.AddBond(
                                      int(keep_size - j),
                                      int(keep_size),
                                      Chem.rdchem.BondType.SINGLE,
                                  )
                              else:
                                  mol_demon_edit.AddBond(
                                      int(
                                          org_bfs_perm_origin[
                                              keep_size - j
                                          ]
                                      ),
                                      int(keep_size),
                                      Chem.rdchem.BondType.SINGLE,
                                  )
                              cur_adj_features_tmp[
                                  0, 0, keep_size - j, keep_size
                              ] = 1.0
                              cur_adj_features_tmp[
                                  0, 0, keep_size, keep_size - j
                              ] = 1.0
  
                              valid = check_valency(mol_demon_edit)
                              if valid:
                                  flag_success = True
                                  break
                              else:
                                  if flag_reconstruct_from_node_adj:
                                      mol_demon_edit.RemoveBond(
                                          int(keep_size - j),
                                          int(keep_size),
                                      )
                                  else:
                                      mol_demon_edit.RemoveBond(
                                          int(
                                              org_bfs_perm_origin[
                                                  keep_size - j
                                              ]
                                          ),
                                          int(keep_size),
                                      )
                                  cur_adj_features_tmp[
                                      0, 0, keep_size - j, keep_size
                                  ] = 0.0
                                  cur_adj_features_tmp[
                                      0, 0, keep_size, keep_size - j
                                  ] = 0.0
                                  count += 1
  
                      if flag_success and check_chemical_validity(
                          mol_demon_edit
                      ):  # successfully take one min action.
                          rw_mol = Chem.RWMol(mol_demon_edit)
                          cur_node_features = (
                              cur_node_features_tmp.clone()
                          )
                          cur_adj_features = (
                              cur_adj_features_tmp.clone()
                          )
                          is_continue = True
                          mol = rw_mol.GetMol()
                          if check_chemical_validity(mol) is True:
                              current_smile = Chem.MolToSmiles(
                                  mol, isomericSmiles=True
                              )
                              tmp_mol1 = Chem.MolFromSmiles(
                                  current_smile
                              )
                              current_imp = (
                                  calculate_min_plogp(tmp_mol1)
                                  - org_mol_plogp
                              )
                              current_sim = (
                                  reward_target_molecule_similarity(
                                      tmp_mol1, org_mol_true_raw
                                  )
                              )
                              if current_imp > 0:
                                  cur_mols.append(tmp_mol1)
                                  cur_mol_imps.append(current_imp)
                                  cur_mol_sims.append(current_sim)
  
                              if (
                                  flag_reconstruct_from_node_adj
                                  is False
                              ):  # not reconstructed from adj, can convert.
                                  mol_converted = convert_radical_electrons_to_hydrogens(
                                      mol
                                  )
                                  if (
                                      check_chemical_validity(
                                          mol_converted
                                      )
                                      is True
                                  ):
                                      current_smile2 = Chem.MolToSmiles(
                                          mol_converted,
                                          isomericSmiles=True,
                                      )
                                      tmp_mol2 = Chem.MolFromSmiles(
                                          current_smile2
                                      )
                                      current_imp2 = (
                                          calculate_min_plogp(tmp_mol2)
                                          - org_mol_plogp
                                      )
                                      current_sim2 = reward_target_molecule_similarity(
                                          tmp_mol2, org_mol_true_raw
                                      )
                                      if current_imp2 > 0:
                                          cur_mols.append(tmp_mol2)
                                          cur_mol_imps.append(
                                              current_imp2
                                          )
                                          cur_mol_sims.append(
                                              current_sim2
                                          )
  
                          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()
                          )  #
                          added_num += 1
  
      return cur_mols, cur_mol_imps, cur_mol_sims
  
file_path: DIG/dig/ggraph/method/GraphDF/model/graphflow_con_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, reward_target_molecule_similarity
  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_con_rl obj's underlying member method code with example usage: obj.reinforce_constrained_optim_one_mol(x, adj, mol_size, raw_smile, bfs_perm_origin, atom_list, temperature, max_size_rl) and returns: cur_mols,  cur_mol_imps,  cur_mol_sims