handle nan values if linker size is small

app.py

@@ -40,19 +40,6 @@ if not os.path.exists(diffusion_path):
 ddpm = DDPM.load_from_checkpoint('models/geom_difflinker.ckpt', map_location=device).eval().to(device)
 print('Loaded diffusion model')
 
-
-def sample_fn(_data):
-    output, _ = size_nn.forward(_data, return_loss=False)
-    probabilities = torch.softmax(output, dim=1)
-    distribution = torch.distributions.Categorical(probs=probabilities)
-    samples = distribution.sample()
-    sizes = []
-    for label in samples.detach().cpu().numpy():
-        sizes.append(size_nn.linker_id2size[label])
-    sizes = torch.tensor(sizes, device=samples.device, dtype=torch.long)
-    return sizes
-
-
 def read_molecule_content(path):
     with open(path, "r") as f:
         return "".join(f.readlines())
@@ -72,7 +59,7 @@ def read_molecule(path):
 
 def show_input(input_file):
     if input_file is None:
-        return ''
+        return ['', gr.Radio.update(visible=False, value='Sample 1')]
     if isinstance(input_file, str):
         path = input_file
     else:
@@ -80,15 +67,24 @@ def show_input(input_file):
     extension = path.split('.')[-1]
     if extension not in ['sdf', 'pdb', 'mol', 'mol2']:
         msg = output.INVALID_FORMAT_MSG.format(extension=extension)
-        return output.IFRAME_TEMPLATE.format(html=msg)
+        return [
+            output.IFRAME_TEMPLATE.format(html=msg),
+            gr.Radio.update(visible=False)
+        ]
 
     try:
         molecule = read_molecule_content(path)
     except Exception as e:
-        return f'Could not read the molecule: {e}'
+        return [
+            f'Could not read the molecule: {e}',
+            gr.Radio.update(visible=False)
+        ]
 
     html = output.INITIAL_RENDERING_TEMPLATE.format(molecule=molecule, fmt=extension)
-    return output.IFRAME_TEMPLATE.format(html=html)
+    return [
+        output.IFRAME_TEMPLATE.format(html=html),
+        gr.Radio.update(visible=False)
+    ]
 
 
 def draw_sample(idx, out_files):
@@ -109,7 +105,7 @@ def draw_sample(idx, out_files):
     return output.IFRAME_TEMPLATE.format(html=html)
 
 
-def generate(input_file, n_steps):
+def generate(input_file, n_steps, n_atoms):
     if input_file is None:
         return ''
 
@@ -156,6 +152,21 @@ def generate(input_file, n_steps):
     ddpm.edm.T = n_steps
     assert ddpm.center_of_mass == 'fragments'
 
+    if n_atoms == 0:
+        def sample_fn(_data):
+            out, _ = size_nn.forward(_data, return_loss=False)
+            probabilities = torch.softmax(out, dim=1)
+            distribution = torch.distributions.Categorical(probs=probabilities)
+            samples = distribution.sample()
+            sizes = []
+            for label in samples.detach().cpu().numpy():
+                sizes.append(size_nn.linker_id2size[label])
+            sizes = torch.tensor(sizes, device=samples.device, dtype=torch.long)
+            return sizes
+    else:
+        def sample_fn(_data):
+            return torch.ones(_data['positions'].shape[0], device=device, dtype=torch.long) * n_atoms
+
     for data in dataloader:
         chain, node_mask = ddpm.sample_chain(data, sample_fn=sample_fn, keep_frames=1)
         print('Generated linker')
@@ -208,6 +219,11 @@ with demo:
                 gr.Markdown('Upload the file with 3D-coordinates of the input fragments in .pdb, .mol2 or .sdf format:')
                 input_file = gr.File(file_count='single', label='Input Fragments')
                 n_steps = gr.Slider(minimum=10, maximum=500, label="Number of Denoising Steps", step=10)
+                n_atoms = gr.Slider(
+                    minimum=0, maximum=20,
+                    label="Linker Size: DiffLinker will predict it if set to 0",
+                    step=1
+                )
                 examples = gr.Dataset(
                     components=[gr.File(visible=False)],
                     samples=[['examples/example_1.sdf'], ['examples/example_2.sdf']],
@@ -235,21 +251,21 @@ with demo:
     input_file.change(
         fn=show_input,
         inputs=[input_file],
-        outputs=[visualization],
+        outputs=[visualization, samples],
+    )
+    input_file.clear(
+        fn=lambda: [None, '', gr.Radio.update(visible=False)],
+        inputs=[],
+        outputs=[input_file, visualization, samples],
     )
     examples.click(
-        fn=lambda idx: [
-            f'examples/example_{idx+1}.sdf',
-            10,
-            show_input(f'examples/example_{idx+1}.sdf'),
-            gr.Radio(value='Sample 1', visible=False)
-        ],
+        fn=lambda idx: [f'examples/example_{idx+1}.sdf', 10, 0] + show_input(f'examples/example_{idx+1}.sdf'),
         inputs=[examples],
-        outputs=[input_file, n_steps, visualization, samples]
+        outputs=[input_file, n_steps, n_atoms, visualization, samples]
     )
     button.click(
         fn=generate,
-        inputs=[input_file, n_steps],
+        inputs=[input_file, n_steps, n_atoms],
         outputs=[visualization, output_files, samples],
     )
     samples.change(
@@ -257,10 +273,5 @@ with demo:
         inputs=[samples, output_files],
         outputs=[visualization],
     )
-    input_file.clear(
-        fn=lambda: ['', gr.Radio(value='Sample 1', visible=False)],
-        inputs=[],
-        outputs=[visualization, samples],
-    )
 
 demo.launch(server_name=args.ip)

src/egnn.py

@@ -421,13 +421,20 @@ class Dynamics(nn.Module):
         if self.condition_time:
             h_final = h_final[:, :-1]
 
+        if torch.any(torch.isnan(vel)):
+            print('Found NaN values in velocities')
+            nan_mask = torch.isnan(vel).float()
+            vel = x * nan_mask + torch.nan_to_num(vel) * (1 - nan_mask)
+
+        if torch.any(torch.isnan(h_final)):
+            print('Found NaN values in features')
+            nan_mask = torch.isnan(h_final).float()
+            h_final = h[:, :h_final.shape[1]] * nan_mask + torch.nan_to_num(h_final) * (1 - nan_mask)
+
         vel = vel.view(bs, n_nodes, -1)  # (B, N, 3)
         h_final = h_final.view(bs, n_nodes, -1)  # (B, N, D)
         node_mask = node_mask.view(bs, n_nodes, 1)  # (B, N, 1)
 
-        if torch.any(torch.isnan(vel)) or torch.any(torch.isnan(h_final)):
-            raise utils.FoundNaNException(vel, h_final)
-
         if self.centering:
             vel = utils.remove_mean_with_mask(vel, node_mask)