class_docstr: The method class for GraphEBM algorithm proposed in the paper `GraphEBM: Molecular Graph Generation with Energy-Based Models <https://arxiv.org/abs/2102.00546>`_. This class provides interfaces for running random generation, goal-directed generation (including property
  optimization and constrained optimization), and compositional generation with GraphEBM algorithm. Please refer to the `example codes <https://github.com/divelab/DIG/tree/dig-stable/examples/ggraph/GraphEBM>`_ for usage examples.
  
  Args:
      n_atom (int): Maximum number of atoms.
      n_atom_type (int): Number of possible atom types.
      n_edge_type (int): Number of possible bond types.
      hidden (int): Hidden dimensions.
      device (torch.device, optional): The device where the model is deployed.
class_docstr_tok: ["The","method","class","for","GraphEBM","algorithm","proposed","in","the","paper","`","GraphEBM",":","Molecular","Graph","Generation","with","Energy-Based","Models","<","https",":","\/\/arxiv.org\/abs\/2102.00546",">","`","_",".","This","class","provides","interfaces","for","running","random","generation",",","goal-directed","generation","(","including","property","optimization","and","constrained","optimization",")",",","and","compositional","generation","with","GraphEBM","algorithm",".","Please","refer","to","the","`","example","codes","<","https",":","\/\/github.com\/divelab\/DIG\/tree\/dig-stable\/examples\/ggraph\/GraphEBM",">","`","_","for","usage","examples",".","Args",":","n_atom","(","int",")",":","Maximum","number","of","atoms",".","n_atom_type","(","int",")",":","Number","of","possible","atom","types",".","n_edge_type","(","int",")",":","Number","of","possible","bond","types",".","hidden","(","int",")",":","Hidden","dimensions",".","device","(","torch.device",",","optional",")",":","The","device","where","the","model","is","deployed","."]
docstr: Running graph generation for goal-directed generation task: constrained property optimization.
  
  Args:
      checkpoint_path (str): The path of the trained model, *i.e.*, the .pt file.
      initialization_loader: The data loader for loading samples to initialize the Langevin dynamics. It is supposed to use dig.ggraph.dataset.QM9/ZINC250k as the dataset class, and apply torch_geometric.data.DenseDataLoader to it to form the data loader.
      c (float): The scaling hyperparameter for dequantization.
      ld_step (int): The number of iteration steps of Langevin dynamics.
      ld_noise (float): The standard deviation of the added noise in Langevin dynamics.
      ld_step_size (int): The step size of Langevin dynamics.
      clamp (bool): Whether to use gradient clamp in Langevin dynamics.
      atomic_num_list (list): The list used to indicate atom types.
      train_smiles (list): A list of smiles string corresponding to training samples.
  
  :rtype:
      mols_0_list (list), mols_2_list (list), mols_4_list (list), mols_6_list (list), imp_0_list (list), imp_2_list (list), imp_4_list (list), imp_4_list (list): They are lists of optimized molecules (represented by rdkit Chem.Mol objects) and the corresponding improvements under the threshold 0.0, 0.2, 0.4, 0.6, respectively.   
docstr_tok: ["Running","graph","generation","for","goal-directed","generation","task",":","constrained","property","optimization",".","Args",":","checkpoint_path","(","str",")",":","The","path","of","the","trained","model",",","*","i.e",".","*",",","the",".pt","file",".","initialization_loader",":","The","data","loader","for","loading","samples","to","initialize","the","Langevin","dynamics",".","It","is","supposed","to","use","dig.ggraph.dataset.QM9\/ZINC250k","as","the","dataset","class",",","and","apply","torch_geometric.data.DenseDataLoader","to","it","to","form","the","data","loader",".","c","(","float",")",":","The","scaling","hyperparameter","for","dequantization",".","ld_step","(","int",")",":","The","number","of","iteration","steps","of","Langevin","dynamics",".","ld_noise","(","float",")",":","The","standard","deviation","of","the","added","noise","in","Langevin","dynamics",".","ld_step_size","(","int",")",":","The","step","size","of","Langevin","dynamics",".","clamp","(","bool",")",":","Whether","to","use","gradient","clamp","in","Langevin","dynamics",".","atomic_num_list","(","list",")",":","The","list","used","to","indicate","atom","types",".","train_smiles","(","list",")",":","A","list","of","smiles","string","corresponding","to","training","samples",".",":","rtype",":","mols_0_list","(","list",")",",","mols_2_list","(","list",")",",","mols_4_list","(","list",")",",","mols_6_list","(","list",")",",","imp_0_list","(","list",")",",","imp_2_list","(","list",")",",","imp_4_list","(","list",")",",","imp_4_list","(","list",")",":","They","are","lists","of","optimized","molecules","(","represented","by","rdkit","Chem.Mol","objects",")","and","the","corresponding","improvements","under","the","threshold","0.0",",","0.2",",","0.4",",","0.6",",","respectively","."]
code: def run_const_prop_opt(
      self,
      checkpoint_path,
      initialization_loader,
      c,
      ld_step,
      ld_noise,
      ld_step_size,
      clamp,
      atomic_num_list,
      train_smiles,
  ):
      r"""
      Running graph generation for goal-directed generation task: constrained property optimization.
  
      Args:
          checkpoint_path (str): The path of the trained model, *i.e.*, the .pt file.
          initialization_loader: The data loader for loading samples to initialize the Langevin dynamics. It is supposed to use dig.ggraph.dataset.QM9/ZINC250k as the dataset class, and apply torch_geometric.data.DenseDataLoader to it to form the data loader.
          c (float): The scaling hyperparameter for dequantization.
          ld_step (int): The number of iteration steps of Langevin dynamics.
          ld_noise (float): The standard deviation of the added noise in Langevin dynamics.
          ld_step_size (int): The step size of Langevin dynamics.
          clamp (bool): Whether to use gradient clamp in Langevin dynamics.
          atomic_num_list (list): The list used to indicate atom types.
          train_smiles (list): A list of smiles string corresponding to training samples.
  
      :rtype:
          mols_0_list (list), mols_2_list (list), mols_4_list (list), mols_6_list (list), imp_0_list (list), imp_2_list (list), imp_4_list (list), imp_4_list (list): They are lists of optimized molecules (represented by rdkit Chem.Mol objects) and the corresponding improvements under the threshold 0.0, 0.2, 0.4, 0.6, respectively.
      """
      print("Loading paramaters from {}".format(checkpoint_path))
      self.energy_function.load_state_dict(torch.load(checkpoint_path))
      parameters = self.energy_function.parameters()
  
      mols_0_list = [None] * 800
      mols_2_list = [None] * 800
      mols_4_list = [None] * 800
      mols_6_list = [None] * 800
  
      imp_0_list = [0] * 800
      imp_2_list = [0] * 800
      imp_4_list = [0] * 800
      imp_6_list = [0] * 800
  
      for i, batch in enumerate(tqdm(initialization_loader)):
          ### Initialization
          gen_x = batch.x.to(self.device).to(dtype=torch.float32)
          gen_adj = batch.adj.to(self.device).to(dtype=torch.float32)
  
          ori_mols = gen_mol_from_one_shot_tensor(
              gen_adj, gen_x, atomic_num_list, correct_validity=True
          )
          ori_smiles = [Chem.MolToSmiles(mol) for mol in ori_mols]
  
          gen_x.requires_grad = True
          gen_adj.requires_grad = True
          requires_grad(parameters, False)
          self.energy_function.eval()
  
          noise_x = torch.randn_like(gen_x, device=self.device)
          noise_adj = torch.randn_like(gen_adj, device=self.device)
  
          ### Langevin dynamics
          for k in range(ld_step):
              noise_x.normal_(0, ld_noise)
              noise_adj.normal_(0, ld_noise)
              gen_x.data.add_(noise_x.data)
              gen_adj.data.add_(noise_adj.data)
  
              gen_out = self.energy_function(gen_adj, gen_x)
              gen_out.sum().backward()
              if clamp:
                  gen_x.grad.data.clamp_(-0.1, 0.1)
                  gen_adj.grad.data.clamp_(-0.1, 0.1)
  
              gen_x.data.add_(gen_x.grad.data, alpha=-ld_step_size)
              gen_adj.data.add_(gen_adj.grad.data, alpha=-ld_step_size)
  
              gen_x.grad.detach_()
              gen_x.grad.zero_()
              gen_adj.grad.detach_()
              gen_adj.grad.zero_()
  
              gen_x.data.clamp_(0, 1 + c)
              gen_adj.data.clamp_(0, 1)
  
              gen_x_t = copy.deepcopy(gen_x)
              gen_adj_t = copy.deepcopy(gen_adj)
              gen_adj_t = (
                  gen_adj_t + gen_adj_t.permute(0, 1, 3, 2)
              ) / 2
  
              gen_mols = gen_mol_from_one_shot_tensor(
                  gen_adj_t,
                  gen_x_t,
                  atomic_num_list,
                  correct_validity=True,
              )
              gen_smiles = [Chem.MolToSmiles(mol) for mol in gen_mols]
  
              for mol_idx in range(len(gen_smiles)):
                  if gen_mols[mol_idx] is not None:
                      tmp_mol = gen_mols[mol_idx]
                      ori_mol = ori_mols[mol_idx]
                      imp_p = calculate_min_plogp(
                          tmp_mol
                      ) - calculate_min_plogp(ori_mol)
                      current_sim = reward_target_molecule_similarity(
                          tmp_mol, ori_mol
                      )
                      if current_sim >= 0.0:
                          if imp_p > imp_0_list[mol_idx]:
                              mols_0_list[mol_idx] = tmp_mol
                      if current_sim >= 0.2:
                          if imp_p > imp_2_list[mol_idx]:
                              mols_2_list[mol_idx] = tmp_mol
                      if current_sim >= 0.4:
                          if imp_p > imp_4_list[mol_idx]:
                              mols_4_list[mol_idx] = tmp_mol
                      if current_sim >= 0.6:
                          if imp_p > imp_6_list[mol_idx]:
                              mols_6_list[mol_idx] = tmp_mol
  
      return (
          mols_0_list,
          mols_2_list,
          mols_4_list,
          mols_6_list,
          imp_0_list,
          imp_2_list,
          imp_4_list,
          imp_4_list,
      )
  
imports: import time
  import os
  import copy
  import torch
  from torch.optim import Adam
  from tqdm import tqdm
  from rdkit import Chem
  from dig.ggraph.method import Generator
  from dig.ggraph.utils import gen_mol_from_one_shot_tensor
  from dig.ggraph.utils import qed, calculate_min_plogp, reward_target_molecule_similarity
  from .energy_func import EnergyFunc
  from .util import rescale_adj, requires_grad, clip_grad
label_desc: Use image node_id 6 for calling the GraphEBM obj's underlying member method code with example usage: obj.run_const_prop_opt(checkpoint_path, initialization_loader, c, ld_step, ld_noise, ld_step_size, clamp, atomic_num_list, train_smiles) and returns: mols_0_list,  mols_2_list,  mols_4_list,  mols_6_list,  imp_0_list,  imp_2_list,  imp_4_list,  imp_4_list