code: def generate(
      self,
      type_to_atomic_number,
      num_gen=100,
      temperature=[1.0, 1.0, 1.0, 1.0],
      min_atoms=2,
      max_atoms=35,
      focus_th=0.5,
  ):
      with torch.no_grad():
          if self.use_gpu:
              prior_node = torch.distributions.normal.Normal(
                  torch.zeros([self.num_node_types]).cuda(),
                  temperature[0]
                  * torch.ones([self.num_node_types]).cuda(),
              )
              prior_dist = torch.distributions.normal.Normal(
                  torch.zeros([1]).cuda(),
                  temperature[1] * torch.ones([1]).cuda(),
              )
              prior_angle = torch.distributions.normal.Normal(
                  torch.zeros([1]).cuda(),
                  temperature[2] * torch.ones([1]).cuda(),
              )
              prior_torsion = torch.distributions.normal.Normal(
                  torch.zeros([1]).cuda(),
                  temperature[3] * torch.ones([1]).cuda(),
              )
          else:
              prior_node = torch.distributions.normal.Normal(
                  torch.zeros([self.num_node_types]),
                  temperature[0] * torch.ones([self.num_node_types]),
              )
              prior_dist = torch.distributions.normal.Normal(
                  torch.zeros([1]), temperature[1] * torch.ones([1])
              )
              prior_angle = torch.distributions.normal.Normal(
                  torch.zeros([1]), temperature[2] * torch.ones([1])
              )
              prior_torsion = torch.distributions.normal.Normal(
                  torch.zeros([1]), temperature[3] * torch.ones([1])
              )
  
          node_type_emb_block = self.feat_net.init_e.emb
          z = torch.ones([num_gen, 1], dtype=int)
          pos = torch.zeros([num_gen, 1, 3], dtype=torch.float32)
          focuses = torch.zeros([num_gen, 0], dtype=int)
          if self.use_gpu:
              z, pos, focuses = z.cuda(), pos.cuda(), focuses.cuda()
          out_dict = {}
  
          mask_index = lambda mask, p: p[mask].view(num_gen, -1, 3)
          feat_index = lambda node_id, f: f[
              torch.arange(num_gen), node_id
          ]
          pos_index = lambda node_id, p: p[
              torch.arange(num_gen), node_id
          ].view(num_gen, 1, 3)
  
          for i in range(max_atoms):
              batch = (
                  torch.arange(num_gen, device=z.device)
                  .view(num_gen, 1)
                  .repeat(1, i + 1)
              )
              if i == 0:
                  node_feat = node_type_emb_block(z.view(-1))
              elif i == 1:
                  node_feat = self.feat_net.dist_only_forward(
                      z.view(-1), pos.view(-1, 3), batch.view(-1)
                  )
              else:
                  node_feat = self.feat_net(
                      z.view(-1), pos.view(-1, 3), batch.view(-1)
                  )
  
              focus_score = self.focus_mlp(node_feat).view(
                  num_gen, i + 1
              )
              can_focus = torch.logical_and(
                  focus_score < focus_th, z > 0
              )
              complete_mask = can_focus.sum(dim=-1) == 0
  
              if (
                  i > max(0, min_atoms - 2)
                  and torch.sum(complete_mask) > 0
              ):
                  out_dict[i + 1] = {}
                  out_node_types = (
                      z[complete_mask].view(-1, i + 1).cpu().numpy()
                  )
                  out_dict[i + 1][
                      "_atomic_numbers"
                  ] = type_to_atomic_number[out_node_types]
                  out_dict[i + 1]["_positions"] = (
                      pos[complete_mask]
                      .view(-1, i + 1, 3)
                      .cpu()
                      .numpy()
                  )
                  out_dict[i + 1]["_focus"] = (
                      focuses[complete_mask].view(-1, i).cpu().numpy()
                  )
  
              continue_mask = torch.logical_not(complete_mask)
              dirty_mask = torch.nonzero(
                  torch.isnan(focus_score).sum(dim=-1)
              )[:, 0]
              if len(dirty_mask) > 0:
                  continue_mask[dirty_mask] = False
              dirty_mask = torch.nonzero(
                  torch.isinf(focus_score).sum(dim=-1)
              )[:, 0]
              if len(dirty_mask) > 0:
                  continue_mask[dirty_mask] = False
  
              if torch.sum(continue_mask) == 0:
                  break
  
              node_feat = node_feat.view(num_gen, i + 1, -1)
              num_gen = torch.sum(continue_mask).cpu().item()
              z, pos, can_focus, focuses = (
                  z[continue_mask],
                  pos[continue_mask],
                  can_focus[continue_mask],
                  focuses[continue_mask],
              )
              focus_node_id = torch.multinomial(
                  can_focus.float(), 1
              ).view(num_gen)
              node_feat = node_feat[continue_mask]
  
              latent_node = prior_node.sample([num_gen])
              local_node_type_feat, query_batch = feat_index(
                  focus_node_id, node_feat
              ), torch.arange(num_gen, device=node_feat.device)
              key_value_batch = (
                  torch.arange(num_gen, device=node_feat.device)
                  .view(num_gen, 1)
                  .repeat(1, i + 1)
                  .view(-1)
              )
              global_node_type_feat = self.node_att(
                  local_node_type_feat,
                  node_feat.view(num_gen * (i + 1), -1),
                  node_feat.view(num_gen * (i + 1), -1),
                  query_batch,
                  key_value_batch,
              )
              node_type_feat = torch.cat(
                  (local_node_type_feat, global_node_type_feat), dim=-1
              )
  
              latent_node = flow_reverse(
                  self.node_flow_layers, latent_node, node_type_feat
              )
              node_type_id = torch.argmax(latent_node, dim=1)
              node_type_emb = node_type_emb_block(node_type_id)
              node_emb = node_feat * node_type_emb.view(num_gen, 1, -1)
  
              latent_dist = prior_dist.sample([num_gen])
              local_dist_feat = feat_index(focus_node_id, node_emb)
              global_dist_feat = self.dist_att(
                  local_dist_feat,
                  node_emb.view(num_gen * (i + 1), -1),
                  node_emb.view(num_gen * (i + 1), -1),
                  query_batch,
                  key_value_batch,
              )
              dist_feat = torch.cat(
                  (local_dist_feat, global_dist_feat), dim=-1
              )
  
              dist = flow_reverse(
                  self.dist_flow_layers, latent_dist, dist_feat
              )
              # dist = dist.abs()
              if i == 0:
                  new_pos = torch.cat(
                      (
                          dist,
                          torch.zeros_like(dist, device=dist.device),
                          torch.zeros_like(dist, device=dist.device),
                      ),
                      dim=-1,
                  )
              if i > 0:
                  mask = torch.ones([num_gen, i + 1], dtype=torch.bool)
                  mask[torch.arange(num_gen), focus_node_id] = False
                  c1_dists = torch.sum(
                      torch.square(
                          mask_index(mask, pos)
                          - pos_index(focus_node_id, pos)
                      ),
                      dim=-1,
                  )
                  c1_node_id = torch.argmin(c1_dists, dim=-1)
                  c1_node_id[c1_node_id >= focus_node_id] += 1
  
                  latent_angle = prior_angle.sample([num_gen])
                  local_angle_feat = torch.cat(
                      (
                          feat_index(focus_node_id, node_emb),
                          feat_index(c1_node_id, node_emb),
                      ),
                      dim=1,
                  )
                  global_angle_feat = self.angle_att(
                      local_angle_feat,
                      node_emb.view(num_gen * (i + 1), -1),
                      node_emb.view(num_gen * (i + 1), -1),
                      query_batch,
                      key_value_batch,
                  )
                  angle_feat = torch.cat(
                      (local_angle_feat, global_angle_feat), dim=-1
                  )
  
                  angle = flow_reverse(
                      self.angle_flow_layers, latent_angle, angle_feat
                  )
                  # angle = angle.abs()
                  if i == 1:
                      fc1 = feat_index(c1_node_id, pos) - feat_index(
                          focus_node_id, pos
                      )
                      new_pos_x = (
                          torch.cos(angle)
                          * torch.sign(fc1[:, 0:1])
                          * dist
                      )
                      new_pos_y = (
                          torch.sin(angle)
                          * torch.sign(fc1[:, 0:1])
                          * dist
                      )
                      new_pos = torch.cat(
                          (
                              new_pos_x,
                              new_pos_y,
                              torch.zeros_like(
                                  dist, device=dist.device
                              ),
                          ),
                          dim=-1,
                      )
                      new_pos += feat_index(focus_node_id, pos)
                  else:
                      mask[torch.arange(num_gen), c1_node_id] = False
                      c2_dists = torch.sum(
                          torch.square(
                              mask_index(mask, pos)
                              - pos_index(c1_node_id, pos)
                          ),
                          dim=-1,
                      )
                      c2_node_id = torch.argmin(c2_dists, dim=-1)
                      c2_node_id[
                          c2_node_id
                          >= torch.minimum(focus_node_id, c1_node_id)
                      ] += 1
                      c2_node_id[
                          c2_node_id
                          >= torch.maximum(focus_node_id, c1_node_id)
                      ] += 1
  
                      latent_torsion = prior_torsion.sample([num_gen])
                      local_torsion_feat = torch.cat(
                          (
                              feat_index(focus_node_id, node_emb),
                              feat_index(c1_node_id, node_emb),
                              feat_index(c2_node_id, node_emb),
                          ),
                          dim=1,
                      )
                      global_torsion_feat = self.torsion_att(
                          local_torsion_feat,
                          node_emb.view(num_gen * (i + 1), -1),
                          node_emb.view(num_gen * (i + 1), -1),
                          query_batch,
                          key_value_batch,
                      )
                      torsion_feat = torch.cat(
                          (local_torsion_feat, global_torsion_feat),
                          dim=-1,
                      )
  
                      torsion = flow_reverse(
                          self.torsion_flow_layers,
                          latent_torsion,
                          torsion_feat,
                      )
                      new_pos = dattoxyz(
                          pos_index(focus_node_id, pos),
                          pos_index(c1_node_id, pos),
                          pos_index(c2_node_id, pos),
                          dist,
                          angle,
                          torsion,
                      )
  
              z = torch.cat((z, node_type_id[:, None]), dim=1)
              pos = torch.cat((pos, new_pos.view(num_gen, 1, 3)), dim=1)
              focuses = torch.cat(
                  (focuses, focus_node_id[:, None]), dim=1
              )
  
          return out_dict
  
file_path: DIG/dig/ggraph3D/method/G_SphereNet/model/sphgen.py
imports: import torch
  import torch.nn
  import torch.nn.functional
  from .spherenet import SphereNet
  from .net_utils import *
  from .geometric_computing import *
  from .att import MH_ATT
label_desc: Use image node_id 3 for calling the SphGen obj's underlying member method code with example usage: obj.generate(type_to_atomic_number, num_gen, temperature, min_atoms, max_atoms, focus_th) and returns: out_dict