docstr: Args:
      max node unroll: maximal number of nodes in molecules to be generated (default: 38)
      max edge unroll: maximal number of edges to predict for each generated nodes (default: 12, calculated from zink250K data)
  Returns:
      node_masks: node mask for each step
      adj_masks: adjacency mask for each step
      is_node_update_mask: 1 indicate this step is for updating node features
      flow_core_edge_mask: get the distributions we want to model in adjacency matrix
docstr_tok: ["Args",":","max","node","unroll",":","maximal","number","of","nodes","in","molecules","to","be","generated","(","default",":","38",")","max","edge","unroll",":","maximal","number","of","edges","to","predict","for","each","generated","nodes","(","default",":","12",",","calculated","from","zink250K","data",")","Returns",":","node_masks",":","node","mask","for","each","step","adj_masks",":","adjacency","mask","for","each","step","is_node_update_mask",":","1","indicate","this","step","is","for","updating","node","features","flow_core_edge_mask",":","get","the","distributions","we","want","to","model","in","adjacency","matrix"]
code: def initialize_masks(self, max_node_unroll=38, max_edge_unroll=12):
      """
      Args:
          max node unroll: maximal number of nodes in molecules to be generated (default: 38)
          max edge unroll: maximal number of edges to predict for each generated nodes (default: 12, calculated from zink250K data)
      Returns:
          node_masks: node mask for each step
          adj_masks: adjacency mask for each step
          is_node_update_mask: 1 indicate this step is for updating node features
          flow_core_edge_mask: get the distributions we want to model in adjacency matrix
      """
      num_masks = int(
          max_node_unroll
          + (max_edge_unroll - 1) * max_edge_unroll / 2
          + (max_node_unroll - max_edge_unroll) * (max_edge_unroll)
      )
      num_mask_edge = int(num_masks - max_node_unroll)
  
      node_masks1 = torch.zeros(
          [max_node_unroll, max_node_unroll]
      ).byte()
      adj_masks1 = torch.zeros(
          [max_node_unroll, max_node_unroll, max_node_unroll]
      ).byte()
      node_masks2 = torch.zeros([num_mask_edge, max_node_unroll]).byte()
      adj_masks2 = torch.zeros(
          [num_mask_edge, max_node_unroll, max_node_unroll]
      ).byte()
      link_prediction_index = torch.zeros([num_mask_edge, 2]).long()
      flow_core_edge_masks = torch.zeros(
          [max_node_unroll, max_node_unroll]
      ).byte()
  
      cnt = 0
      cnt_node = 0
      cnt_edge = 0
      for i in range(max_node_unroll):
          node_masks1[cnt_node][:i] = 1
          adj_masks1[cnt_node][:i, :i] = 1
          cnt += 1
          cnt_node += 1
  
          edge_total = 0
          if i < max_edge_unroll:
              start = 0
              edge_total = i
          else:
              start = i - max_edge_unroll
              edge_total = max_edge_unroll
          for j in range(edge_total):
              if j == 0:
                  node_masks2[cnt_edge][: i + 1] = 1
                  adj_masks2[cnt_edge] = adj_masks1[
                      cnt_node - 1
                  ].clone()
                  adj_masks2[cnt_edge][i, i] = 1
              else:
                  node_masks2[cnt_edge][: i + 1] = 1
                  adj_masks2[cnt_edge] = adj_masks2[
                      cnt_edge - 1
                  ].clone()
                  adj_masks2[cnt_edge][i, start + j - 1] = 1
                  adj_masks2[cnt_edge][start + j - 1, i] = 1
              cnt += 1
              cnt_edge += 1
      assert cnt == num_masks, "masks cnt wrong"
      assert cnt_node == max_node_unroll, "node masks cnt wrong"
      assert cnt_edge == num_mask_edge, "edge masks cnt wrong"
  
      cnt = 0
      for i in range(max_node_unroll):
          if i < max_edge_unroll:
              start = 0
              edge_total = i
          else:
              start = i - max_edge_unroll
              edge_total = max_edge_unroll
  
          for j in range(edge_total):
              link_prediction_index[cnt][0] = start + j
              link_prediction_index[cnt][1] = i
              cnt += 1
      assert cnt == num_mask_edge, "edge mask initialize fail"
  
      for i in range(max_node_unroll):
          if i == 0:
              continue
          if i < max_edge_unroll:
              start = 0
              end = i
          else:
              start = i - max_edge_unroll
              end = i
          flow_core_edge_masks[i][start:end] = 1
  
      node_masks = torch.cat((node_masks1, node_masks2), dim=0)
      adj_masks = torch.cat((adj_masks1, adj_masks2), dim=0)
  
      node_masks = nn.Parameter(node_masks, requires_grad=False)
      adj_masks = nn.Parameter(adj_masks, requires_grad=False)
      link_prediction_index = nn.Parameter(
          link_prediction_index, requires_grad=False
      )
      flow_core_edge_masks = nn.Parameter(
          flow_core_edge_masks, requires_grad=False
      )
  
      return (
          node_masks,
          adj_masks,
          link_prediction_index,
          flow_core_edge_masks,
      )
  
file_path: DIG/dig/ggraph/method/GraphDF/model/graphflow_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, qed
  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 4 for calling the GraphFlowModel_rl obj's underlying member method code with example usage: obj.initialize_masks(max_node_unroll, max_edge_unroll) and returns: node_masks,  adj_masks,  link_prediction_index,  flow_core_edge_masks