def test_call(self, node_size, edge_size, globals_size, y_size, model_type,
               n_layers):
   """Check that shapes for call are correct."""
   graphs, n_graphs = self._setup_graphs()
   model = experiments.GNN(node_size, edge_size, globals_size, y_size,
                           model_type, 'relu', templates.TargetType.globals,
                           n_layers)
   y_out = model(graphs).numpy()
   self.assertEqual((n_graphs, y_size), y_out.shape)
 def test_predict(self, node_size, edge_size, globals_size, y_size, model_type,
                  n_layers):
   """Check that predict is 1-D."""
   graphs, n_graphs = self._setup_graphs()
   model = experiments.GNN(node_size, edge_size, globals_size, y_size,
                           model_type, 'relu', templates.TargetType.globals,
                           n_layers)
   model(graphs)
   y_out = model.predict(graphs).numpy()
   self.assertEqual((n_graphs,), y_out.shape)
 def test_get_graph_embedding(self, node_size, edge_size, globals_size, y_size,
                              model_type, n_layers):
   """Check that shapes for get_embedding are correct."""
   graphs, n_graphs = self._setup_graphs()
   model = experiments.GNN(node_size, edge_size, globals_size, y_size,
                           model_type, 'relu', templates.TargetType.globals,
                           n_layers)
   model(graphs)
   emb = model.get_graph_embedding(graphs).numpy()
   self.assertEqual((n_graphs, globals_size), emb.shape)
 def test_get_prediction_weights(self, node_size, edge_size, globals_size,
                                 y_size, model_type, n_layers):
   """Check that we have enough weights as globals vector dim."""
   graphs, _ = self._setup_graphs()
   model = experiments.GNN(node_size, edge_size, globals_size, y_size,
                           model_type, 'relu', templates.TargetType.globals,
                           n_layers)
   model(graphs)
   w = model.get_prediction_weights()
   self.assertEqual(w.shape, globals_size)
 def test_get_gradient(self, node_size, edge_size, globals_size, y_size,
                       model_type, n_layers):
   """Check that shapes for gradients are correct."""
   graphs, _ = self._setup_graphs()
   model = experiments.GNN(node_size, edge_size, globals_size, y_size,
                           model_type, 'relu', templates.TargetType.globals,
                           n_layers)
   model(graphs)
   node_grads, edge_grads = model.get_gradient(graphs)
   self.assertEqual(graphs.nodes.shape, node_grads.shape)
   self.assertEqual(graphs.edges.shape, edge_grads.shape)
 def _setup_graphs_model(self):
     """Setup graphs and smiles if needed."""
     tensorizer = featurization.MolTensorizer()
     smiles = ['CO', 'CCC', 'CN1C=NC2=C1C(=O)N(C(=O)N2C)C']
     graphs = graph_utils.smiles_to_graphs_tuple(smiles, tensorizer)
     # Fix seed so that initialization is deterministic.
     tf.random.set_seed(0)
     model = experiments.GNN(5, 3, 10, 1, models.BlockType('gcn'), 'relu',
                             templates.TargetType.globals, 3)
     model(graphs)
     return graphs, model, tensorizer
 def test_get_gap_activations(self, node_size, edge_size, globals_size, y_size,
                              model_type, n_layers):
   """Check that shapes for gap activation are like global vectors."""
   graphs, _ = self._setup_graphs()
   n_nodes, n_edges = graphs.nodes.shape[0], graphs.edges.shape[0]
   model = experiments.GNN(node_size, edge_size, globals_size, y_size,
                           model_type, 'relu', templates.TargetType.globals,
                           n_layers)
   model(graphs)
   node_act, edge_act = model.get_gap_activations(graphs)
   self.assertEqual((n_nodes, globals_size), node_act.shape)
   self.assertEqual((n_edges, globals_size), edge_act.shape)
  def _setup_experiment(self):
    """Setup graphs and smiles if needed."""
    smiles = ['CO', 'CCC', 'CN1C=NC2=C1C(=O)N(C(=O)N2C)C']
    n = len(smiles)
    smiles_to_mol = functools.partial(
        featurization.smiles_to_mol, infer_hydrogens=True)
    tensorizer = featurization.MolTensorizer(preprocess_fn=smiles_to_mol)
    train_index, test_index = np.arange(n - 1), np.arange(n - 1, n)
    mol_list = [smiles_to_mol(smi) for smi in smiles]
    x = graph_utils.smiles_to_graphs_tuple(smiles, tensorizer)
    task = tasks.get_task(tasks.Task.crippen)
    y = task.get_true_predictions(mol_list)
    atts = task.get_true_attributions(mol_list)
    exp = experiments.ExperimentData.from_data_and_splits(
        x, y, atts, train_index, test_index)
    model = experiments.GNN(5, 3, 10, 1, models.BlockType.gcn, 'relu',
                            templates.TargetType.globals, 2)
    model(x)
    method = techniques.CAM()

    return exp, model, task, method
  def test_get_intermediate_activations_gradients(self, node_size, edge_size,
                                                  globals_size, y_size,
                                                  model_type, n_layers):
    """Check shapes for actvations/gradients, check that predictions match."""
    graphs, _ = self._setup_graphs()
    n_nodes, n_edges = graphs.nodes.shape[0], graphs.edges.shape[0]
    model = experiments.GNN(node_size, edge_size, globals_size, y_size,
                            model_type, 'relu', templates.TargetType.globals,
                            n_layers)
    model(graphs)
    acts, grads, y = model.get_intermediate_activations_gradients(graphs)
    y_expected = model.predict(graphs)

    self.assertTrue(np.allclose(y, y_expected))
    self.assertLen(acts, n_layers)
    self.assertLen(grads, n_layers)
    for (act_nodes, act_edges), (grad_nodes, grad_edges) in zip(acts, grads):
      self.assertEqual((n_nodes, node_size), act_nodes.shape)
      self.assertEqual((n_edges, edge_size), act_edges.shape)
      self.assertEqual((n_nodes, node_size), grad_nodes.shape)
      self.assertEqual((n_edges, edge_size), grad_edges.shape)