update COM

app.py

@@ -160,6 +160,12 @@ def generate(input_file, n_steps):
         print('Generated linker')
         x = chain[0][:, :, :ddpm.n_dims]
         h = chain[0][:, :, ddpm.n_dims:]
+
+        pos_masked = data['positions'] * data['fragment_mask']
+        N = data['fragment_mask'].sum(1, keepdims=True)
+        mean = torch.sum(pos_masked, dim=1, keepdim=True) / N
+        x = x + mean * node_mask
+
         names = [f'output_{i+1}_{name}' for i in range(N_SAMPLES)]
         save_xyz_file('results', h, x, node_mask, names=names, is_geom=True, suffix='')
         print('Saved XYZ files')