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 training for goal-directed generation task.
  
  Args:
      loader: The data loader for loading training samples. 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.
      lr (float): The learning rate for training.
      wd (float): The weight decay factor for training.
      max_epochs (int): The maximum number of training epochs.
      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.
      alpha (float): The weight coefficient for loss function.
      save_interval (int): The frequency to save the model parameters to .pt files,
          *e.g.*, if save_interval=2, the model parameters will be saved for every 2 training epochs.
      save_dir (str): the directory to save the model parameters.
docstr_tok: ["Running","training","for","goal-directed","generation","task",".","Args",":","loader",":","The","data","loader","for","loading","training","samples",".","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",".","lr","(","float",")",":","The","learning","rate","for","training",".","wd","(","float",")",":","The","weight","decay","factor","for","training",".","max_epochs","(","int",")",":","The","maximum","number","of","training","epochs",".","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",".","alpha","(","float",")",":","The","weight","coefficient","for","loss","function",".","save_interval","(","int",")",":","The","frequency","to","save","the","model","parameters","to",".pt","files",",","*","e.g",".","*",",","if","save_interval=2",",","the","model","parameters","will","be","saved","for","every","2","training","epochs",".","save_dir","(","str",")",":","the","directory","to","save","the","model","parameters","."]
code: def train_goal_directed(
      self,
      loader,
      lr,
      wd,
      max_epochs,
      c,
      ld_step,
      ld_noise,
      ld_step_size,
      clamp,
      alpha,
      save_interval,
      save_dir,
  ):
      r"""
      Running training for goal-directed generation task.
  
      Args:
          loader: The data loader for loading training samples. 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.
          lr (float): The learning rate for training.
          wd (float): The weight decay factor for training.
          max_epochs (int): The maximum number of training epochs.
          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.
          alpha (float): The weight coefficient for loss function.
          save_interval (int): The frequency to save the model parameters to .pt files,
              *e.g.*, if save_interval=2, the model parameters will be saved for every 2 training epochs.
          save_dir (str): the directory to save the model parameters.
      """
      parameters = self.energy_function.parameters()
      optimizer = Adam(
          parameters, lr=lr, betas=(0.0, 0.999), weight_decay=wd
      )
  
      if not os.path.exists(save_dir):
          os.makedirs(save_dir)
  
      for epoch in range(max_epochs):
          t_start = time.time()
          losses_reg = []
          losses_en = []
          losses = []
          for _, batch in enumerate(tqdm(loader)):
              ### Dequantization
              pos_x = batch.x.to(self.device).to(dtype=torch.float32)
              pos_x += c * torch.rand_like(pos_x, device=self.device)
              pos_adj = batch.adj.to(self.device).to(
                  dtype=torch.float32
              )
              pos_adj += c * torch.rand_like(
                  pos_adj, device=self.device
              )
  
              pos_y = batch.y.to(self.device)
  
              ### Langevin dynamics
              neg_x = torch.rand_like(pos_x, device=self.device) * (
                  1 + c
              )
              neg_adj = torch.rand_like(pos_adj, device=self.device)
  
              pos_adj = rescale_adj(pos_adj)
              neg_x.requires_grad = True
              neg_adj.requires_grad = True
  
              requires_grad(parameters, False)
              self.energy_function.eval()
  
              noise_x = torch.randn_like(neg_x, device=self.device)
              noise_adj = torch.randn_like(neg_adj, device=self.device)
              for _ in range(ld_step):
                  noise_x.normal_(0, ld_noise)
                  noise_adj.normal_(0, ld_noise)
                  neg_x.data.add_(noise_x.data)
                  neg_adj.data.add_(noise_adj.data)
  
                  neg_out = self.energy_function(neg_adj, neg_x)
                  neg_out.sum().backward()
                  if clamp:
                      neg_x.grad.data.clamp_(-0.01, 0.01)
                      neg_adj.grad.data.clamp_(-0.01, 0.01)
  
                  neg_x.data.add_(neg_x.grad.data, alpha=-ld_step_size)
                  neg_adj.data.add_(
                      neg_adj.grad.data, alpha=-ld_step_size
                  )
  
                  neg_x.grad.detach_()
                  neg_x.grad.zero_()
                  neg_adj.grad.detach_()
                  neg_adj.grad.zero_()
  
                  neg_x.data.clamp_(0, 1 + c)
                  neg_adj.data.clamp_(0, 1)
  
              ### Training by backprop
              neg_x = neg_x.detach()
              neg_adj = neg_adj.detach()
              requires_grad(parameters, True)
              self.energy_function.train()
  
              self.energy_function.zero_grad()
  
              pos_out = self.energy_function(pos_adj, pos_x)
              neg_out = self.energy_function(neg_adj, neg_x)
  
              loss_reg = (
                  pos_out**2 + neg_out**2
              )  # energy magnitudes regularizer
              loss_en = (
                  1 + torch.exp(pos_y)
              ) * pos_out - neg_out  # loss for shaping energy function
              loss = loss_en + alpha * loss_reg
              loss = loss.mean()
              loss.backward()
              clip_grad(optimizer)
              optimizer.step()
  
              losses_reg.append(loss_reg.mean())
              losses_en.append(loss_en.mean())
              losses.append(loss)
  
          t_end = time.time()
  
          ### Save checkpoints
          if (epoch + 1) % save_interval == 0:
              torch.save(
                  self.energy_function.state_dict(),
                  os.path.join(
                      save_dir, "epoch_{}.pt".format(epoch + 1)
                  ),
              )
              print("Saving checkpoint at epoch ", epoch + 1)
              print("==========================================")
          print(
              "Epoch: {:03d}, Loss: {:.6f}, Energy Loss: {:.6f}, Regularizer Loss: {:.6f}, Sec/Epoch: {:.2f}".format(
                  epoch + 1,
                  (sum(losses) / len(losses)).item(),
                  (sum(losses_en) / len(losses_en)).item(),
                  (sum(losses_reg) / len(losses_reg)).item(),
                  t_end - t_start,
              )
          )
          print("==========================================")
  
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 4 for calling the GraphEBM obj's underlying member method code with example usage: obj.train_goal_directed(loader, lr, wd, max_epochs, c, ld_step, ld_noise, ld_step_size, clamp, alpha, save_interval, save_dir) without return types