class_docstr: The method class for G-SphereNet algorithm proposed in the paper `An Autoregressive Flow Model for 3D Molecular Geometry Generation from Scratch <https://openreview.net/forum?id=C03Ajc-NS5W>`_. 
  This class provides interfaces for running training and generation with G-SphereNet algorithm. 
  Please refer to the `example codes <https://github.com/divelab/DIG/tree/dig-stable/examples/ggraph3D/G_SphereNet>`_ for usage examples.
docstr: Running graph generation for random generation task.
  
  Args:
      model_conf_dict (dict): The python dict for configuring the model hyperparameters.
      checkpoint_path (str): The path to the saved model checkpoint file.
      n_mols (int, optional): The number of molecular geometries to generate. (default: :obj:`1000`)
      chunk_size (int, optional): The maximum number of molecular geometries that are allowed to be generated in parallel. (default: :obj:`100`)
      num_min_node (int, optional): The minimum number of nodes in the generated molecular geometries. (default: :obj:`7`)
      num_max_node (int, optional): the maximum number of nodes in the generated molecular geometries. (default: :obj:`25`)
      temperature (list, optional): a list of four float numbers, the temperature parameter of prior distribution. (default: :obj:`[1.0, 1.0, 1.0, 1.0]`)
      focus_th (float, optional): The threshold for focus node classification. (default: :obj:`0.5`)
  
  :rtype:
      mol_dicts,
      A python dict where the key is the number of atoms, and the value indexed by that key is another python dict storing the atomic
      number matrix (indexed by the key '_atomic_numbers') and the coordinate tensor (indexed by the key '_positions') of all generated molecular geometries with that atom number. 
docstr_tok: ["Running","graph","generation","for","random","generation","task",".","Args",":","model_conf_dict","(","dict",")",":","The","python","dict","for","configuring","the","model","hyperparameters",".","checkpoint_path","(","str",")",":","The","path","to","the","saved","model","checkpoint","file",".","n_mols","(","int",",","optional",")",":","The","number","of","molecular","geometries","to","generate",".","(","default",":",":","obj",":","`","1000","`",")","chunk_size","(","int",",","optional",")",":","The","maximum","number","of","molecular","geometries","that","are","allowed","to","be","generated","in","parallel",".","(","default",":",":","obj",":","`","100","`",")","num_min_node","(","int",",","optional",")",":","The","minimum","number","of","nodes","in","the","generated","molecular","geometries",".","(","default",":",":","obj",":","`","7","`",")","num_max_node","(","int",",","optional",")",":","the","maximum","number","of","nodes","in","the","generated","molecular","geometries",".","(","default",":",":","obj",":","`","25","`",")","temperature","(","list",",","optional",")",":","a","list","of","four","float","numbers",",","the","temperature","parameter","of","prior","distribution",".","(","default",":",":","obj",":","`","[","1.0",",","1.0",",","1.0",",","1.0","]","`",")","focus_th","(","float",",","optional",")",":","The","threshold","for","focus","node","classification",".","(","default",":",":","obj",":","`","0.5","`",")",":","rtype",":","mol_dicts",",","A","python","dict","where","the","key","is","the","number","of","atoms",",","and","the","value","indexed","by","that","key","is","another","python","dict","storing","the","atomic","number","matrix","(","indexed","by","the","key","'_atomic_numbers","'",")","and","the","coordinate","tensor","(","indexed","by","the","key","'_positions","'",")","of","all","generated","molecular","geometries","with","that","atom","number","."]
code: def generate(
      self,
      model_conf_dict,
      checkpoint_path,
      n_mols=1000,
      chunk_size=100,
      num_min_node=7,
      num_max_node=25,
      temperature=[1.0, 1.0, 1.0, 1.0],
      focus_th=0.5,
  ):
      r"""
      Running graph generation for random generation task.
  
      Args:
          model_conf_dict (dict): The python dict for configuring the model hyperparameters.
          checkpoint_path (str): The path to the saved model checkpoint file.
          n_mols (int, optional): The number of molecular geometries to generate. (default: :obj:`1000`)
          chunk_size (int, optional): The maximum number of molecular geometries that are allowed to be generated in parallel. (default: :obj:`100`)
          num_min_node (int, optional): The minimum number of nodes in the generated molecular geometries. (default: :obj:`7`)
          num_max_node (int, optional): the maximum number of nodes in the generated molecular geometries. (default: :obj:`25`)
          temperature (list, optional): a list of four float numbers, the temperature parameter of prior distribution. (default: :obj:`[1.0, 1.0, 1.0, 1.0]`)
          focus_th (float, optional): The threshold for focus node classification. (default: :obj:`0.5`)
  
      :rtype:
          mol_dicts,
          A python dict where the key is the number of atoms, and the value indexed by that key is another python dict storing the atomic
          number matrix (indexed by the key '_atomic_numbers') and the coordinate tensor (indexed by the key '_positions') of all generated molecular geometries with that atom number.
      """
      self.get_model(model_conf_dict, checkpoint_path)
      self.model.eval()
  
      type_to_atomic_number = np.array([1, 6, 7, 8, 9])
      mol_dicts = {}
      num_remain, one_time_gen = n_mols, chunk_size
  
      while num_remain > 0:
          if num_remain > one_time_gen:
              mols = self.model.generate(
                  type_to_atomic_number,
                  one_time_gen,
                  temperature,
                  num_min_node,
                  num_max_node,
                  focus_th,
              )
          else:
              mols = self.model.generate(
                  type_to_atomic_number,
                  num_remain,
                  temperature,
                  num_min_node,
                  num_max_node,
                  focus_th,
              )
  
          for num_atom in mols:
              if not num_atom in mol_dicts.keys():
                  mol_dicts[num_atom] = mols[num_atom]
              else:
                  mol_dicts[num_atom][
                      "_atomic_numbers"
                  ] = np.concatenate(
                      (
                          mol_dicts[num_atom]["_atomic_numbers"],
                          mols[num_atom]["_atomic_numbers"],
                      ),
                      axis=0,
                  )
                  mol_dicts[num_atom]["_positions"] = np.concatenate(
                      (
                          mol_dicts[num_atom]["_positions"],
                          mols[num_atom]["_positions"],
                      ),
                      axis=0,
                  )
                  mol_dicts[num_atom]["_focus"] = np.concatenate(
                      (
                          mol_dicts[num_atom]["_focus"],
                          mols[num_atom]["_focus"],
                      ),
                      axis=0,
                  )
              num_mol = len(mols[num_atom]["_atomic_numbers"])
              num_remain -= num_mol
  
          print(
              "{} molecules are generated!".format(n_mols - num_remain)
          )
  
      return mol_dicts
  
file_path: DIG/dig/ggraph3D/method/G_SphereNet/gspherenet.py
imports: import os
  import torch
  import numpy
  from .model import SphGen
label_desc: Use image node_id 5 for calling the G_SphereNet obj's underlying member method code with example usage: obj.generate(model_conf_dict, checkpoint_path, n_mols, chunk_size, num_min_node, num_max_node, temperature, focus_th) and returns: mol_dicts