Merge pull request #16 from shenoynikhil/conformer_fix

Conformer creation fix

README.md

@@ -3,20 +3,12 @@ Implementation of [Equivariant Flow Matching for Molecule Conformer Generation](
 
 ET-Flow is a state-of-the-art generative model for generating small molecule conformations using equivariant transformers and flow matching.
 
-### Install Etflow
+### Install ET-flow
 We are now available on PyPI. Easily install the package using the following command:
 ```bash
 pip install etflow
 ```
 
-### Setup dev Environment
-Run the following commands to setup the environment:
-```bash
-conda env create -n etflow -f env.yml
-conda activate etflow
-# to install the etflow package
-python3 -m pip install -e .
-```
 
 ### Generating Conformations for Custom Smiles
 We have a sample notebook ([generate_confs.ipynb](generate_confs.ipynb)) to generate conformations for custom smiles input. One needs to pass the config and corresponding checkpoint path in order as additional inputs.
@@ -34,6 +26,15 @@ We currently support the following configurations and checkpoint:
 - `qm9-o3`
 - `drugs-so3`
 
+### Setup Dev Environment
+Run the following commands to setup the environment:
+```bash
+conda env create -n etflow -f env.yml
+conda activate etflow
+# to install the etflow package
+python3 -m pip install -e .
+```
+
 ### Preprocessing Data
 To pre-process the data, perform the following steps,
 1. Download the raw GEOM data and unzip the raw data using the following commands,

etflow/commons/covmat.py

@@ -29,6 +29,19 @@ def build_conformer(pos):
     return conformer
 
 
+def set_multiple_rdmol_positions(rdkit_mol, pos):
+    """
+    Args:
+        rdkit_mol:  An `rdkit.Chem.rdchem.Mol` object.
+        pos: (n, N_atoms, 3)
+    """
+    mol = deepcopy(rdkit_mol)
+    for conf_pos in pos:
+        conformer = build_conformer(conf_pos)
+        mol.AddConformer(conformer)
+    return mol
+
+
 def set_rdmol_positions(rdkit_mol, pos):
     """
     Args:

etflow/models/model.py

@@ -1,4 +1,3 @@
-from copy import deepcopy
 from typing import Any, Dict, List, Optional, TypeVar
 
 import numpy as np
@@ -8,7 +7,7 @@
 from torch_geometric.data import Batch
 
 from etflow.commons.configs import CONFIG_DICT
-from etflow.commons.covmat import set_rdmol_positions
+from etflow.commons.covmat import set_multiple_rdmol_positions
 from etflow.commons.featurization import MoleculeFeaturizer, get_mol_from_smiles
 from etflow.commons.utils import signed_volume
 from etflow.models.base import BaseModel
@@ -706,13 +705,9 @@ def sample(
             )
             if as_mol:
                 mol = get_mol_from_smiles(smile)
-                set_rdmol_positions(mol, pos[0])
-                mols = []
-                for i in range(num_samples):
-                    copied_mol = deepcopy(mol)
-                    set_rdmol_positions(copied_mol, pos[i])
-                    mols.append(copied_mol)
-                data[smile] = mols
+                data[smile] = set_multiple_rdmol_positions(mol, pos)
+            else:
+                data[smile] = pos
         return data