add loss

loss.py

@@ -0,0 +1,158 @@
+
+import torch
+
+def discriminator_loss(generator, discriminator, mol_graph, adj, annot, batch_size, device, grad_pen, lambda_gp,z_edge,z_node):
+    
+    # Compute loss with real molecules.
+    
+    logits_real_disc = discriminator(mol_graph)              
+                        
+    prediction_real =  - torch.mean(logits_real_disc)
+
+    # Compute loss with fake molecules.
+
+    node, edge, node_sample, edge_sample  = generator(z_edge,  z_node)
+
+    graph = torch.cat((node_sample.view(batch_size, -1), edge_sample.view(batch_size, -1)), dim=-1)
+    
+    logits_fake_disc = discriminator(graph.detach())
+        
+    prediction_fake = torch.mean(logits_fake_disc)
+    
+    # Compute gradient loss.
+    
+    eps = torch.rand(mol_graph.size(0),1).to(device)
+    x_int0 = (eps * mol_graph + (1. - eps) * graph).requires_grad_(True)
+
+    grad0 = discriminator(x_int0)
+    d_loss_gp = grad_pen(grad0, x_int0) 
+    
+    # Calculate total loss
+    
+    d_loss = prediction_fake + prediction_real +  d_loss_gp * lambda_gp
+    
+    return node, edge,d_loss
+
+    
+def generator_loss(generator, discriminator, v, adj, annot, batch_size, penalty, matrices2mol, fps_r,submodel):
+    
+    # Compute loss with fake molecules.
+   
+    node, edge, node_sample, edge_sample  = generator(adj,  annot)
+   
+    
+    graph = torch.cat((node_sample.view(batch_size, -1), edge_sample.view(batch_size, -1)), dim=-1)
+   
+ 
+    logits_fake_disc = discriminator(graph)
+    
+    prediction_fake = - torch.mean(logits_fake_disc)
+    
+    # Produce molecules.
+
+    g_edges_hat_sample = torch.max(edge_sample, -1)[1] 
+    g_nodes_hat_sample = torch.max(node_sample , -1)[1]   
+                
+    fake_mol = [matrices2mol(n_.data.cpu().numpy(), e_.data.cpu().numpy(), strict=True) 
+            for e_, n_ in zip(g_edges_hat_sample, g_nodes_hat_sample)]        
+    g_loss = prediction_fake
+    # Compute penalty loss.
+    if submodel == "RL":
+        reward = penalty(fake_mol, fps_r)
+    
+        # Reinforcement Loss
+   
+        rew_fake = v(graph)
+
+        reward_loss =  torch.mean(rew_fake) ** 2 + reward    
+    
+        # Calculate total loss
+    
+        g_loss = prediction_fake + reward_loss * 1
+    
+    
+    return g_loss, fake_mol, g_edges_hat_sample, g_nodes_hat_sample, node, edge
+
+def discriminator2_loss(generator, discriminator, mol_graph, adj, annot, batch_size, device, grad_pen, lambda_gp,akt1_adj,akt1_annot):
+
+    # Generate molecules.
+    
+    dr_edges, dr_nodes = generator(adj, 
+                                annot, 
+                                akt1_adj,
+                                akt1_annot)
+    
+    
+    dr_edges_hat = dr_edges.view(batch_size, -1)
+    
+    dr_nodes_hat = dr_nodes.view(batch_size, -1)
+
+    dr_graph = torch.cat((dr_nodes_hat, dr_edges_hat), dim=-1)                    
+
+    # Compute loss with fake molecules.
+    
+    dr_logits_fake = discriminator(dr_graph.detach())
+
+    d2_loss_fake = torch.mean(dr_logits_fake)
+    
+    # Compute loss with real molecules.
+    
+    dr_logits_real2 = discriminator(mol_graph)
+
+    d2_loss_real = - torch.mean(dr_logits_real2)
+    
+    # Compute gradient loss.  
+    
+    eps_dr = torch.rand(mol_graph.size(0),1).to(device)
+    x_int0_dr = (eps_dr * mol_graph + (1. - eps_dr) * dr_graph).requires_grad_(True)
+
+
+    grad0_dr = discriminator(x_int0_dr)
+    d2_loss_gp = grad_pen(grad0_dr, x_int0_dr) 
+    
+    # Compute total loss.
+    
+    d2_loss = d2_loss_fake + d2_loss_real + d2_loss_gp * lambda_gp
+    
+    return d2_loss
+
+def generator2_loss(generator, discriminator, v, adj, annot, batch_size, penalty, matrices2mol, fps_r,ak1_adj,akt1_annot, submodel):
+    
+    # Generate molecules.
+    
+    dr_edges_g, dr_nodes_g = generator(adj, 
+                                annot, 
+                                ak1_adj,
+                                akt1_annot)
+    
+    dr_edges_hat_g = dr_edges_g.view(batch_size, -1)
+    
+    dr_nodes_hat_g = dr_nodes_g.view(batch_size, -1)
+
+    dr_graph_g = torch.cat((dr_nodes_hat_g, dr_edges_hat_g), dim=-1)      
+    
+    # Compute loss with fake molecules.
+
+    dr_g_edges_hat_sample, dr_g_nodes_hat_sample = torch.max(dr_edges_g, -1)[1], torch.max(dr_nodes_g, -1)[1]
+         
+    g_tra_logits_fake2 = discriminator(dr_graph_g)      
+   
+    g2_loss_fake = - torch.mean(g_tra_logits_fake2)                                                            
+
+    # Reward
+    fake_mol_g = [matrices2mol(n_.data.cpu().numpy(), e_.data.cpu().numpy(), strict=True) 
+                for e_, n_ in zip(dr_g_edges_hat_sample, dr_g_nodes_hat_sample)]       
+    g2_loss =  g2_loss_fake    
+    if submodel == "RL":
+        reward2 = penalty(fake_mol_g, fps_r)
+    
+        # Reinforcement Loss
+        
+        rew_fake2 = v(dr_graph_g)
+        reward_loss2 =  torch.mean(rew_fake2) ** 2 + reward2
+        
+        # Calculate total loss
+        
+        g2_loss =  g2_loss_fake + reward_loss2 * 10    
+    
+    return g2_loss, fake_mol_g, dr_g_edges_hat_sample, dr_g_nodes_hat_sample#, reward2