Molecule graph environment

config/env/alaninedipeptide.yaml → config/env/conformers/alaninedipeptide.yaml

@@ -1,7 +1,7 @@
 defaults:
   - base
 
-_target_: gflownet.envs.alaninedipeptide.AlanineDipeptide
+_target_: gflownet.envs.conformers.alaninedipeptide.AlanineDipeptide
 
 path_to_dataset: './data/alanine_dipeptide_conformers_1.npy'
 url_to_dataset: 'https://drive.google.com/uc?id=1r1KRGcpBhR3xaS8yt2i64dfMnJGgNj4C'

config/env/conformers/conformer.yaml

@@ -0,0 +1,36 @@
+defaults:
+  - base
+
+_target_: gflownet.envs.conformers.conformer.Conformer
+
+# smiles: 'CC(C(=O)NC)NC(=O)C' # alanine dipeptide
+# smiles: 'CC(C)Cc1ccc(cc1)[C@@H](C)C(=O)O' # ibuprofen
+smiles: 'O=C(c1ccc2n1CCC2C(=O)O)c3ccccc3' # ketorolac
+# smiles: 'CCCCCC1=CC(=C(C(=C1)O)C2C=C(CCC2C(=C)C)C)O' # cannabidiol
+# smiles: 'CN1C2CCC1C(C(C2)OC(=O)C3=CC=CC=C3)C(=O)OC' # cocaine
+n_torsion_angles: 2
+reward_sampling_method: nested
+
+id: conformer
+policy_encoding_dim_per_angle: null
+# Fixed length of trajectories
+length_traj: 10
+vonmises_min_concentration: 1e-3
+# Parameters of the fixed policy output distribution
+n_comp: 3
+fixed_distribution:
+  vonmises_mean: 0.0
+  vonmises_concentration: 0.5
+# Parameters of the random policy output distribution
+random_distribution:
+  vonmises_mean: 0.0
+  vonmises_concentration: 0.001
+# Buffer
+buffer:
+  data_path: null
+  train: null
+  test:
+    type: grid
+    n: 1000
+    output_csv: conformer_test.csv
+    output_pkl: conformer_test.pkl

config/experiments/conformer.yaml

@@ -0,0 +1,52 @@
+# @package _global_
+
+defaults:
+   - override /env: conformers/conformer
+   - override /gflownet: trajectorybalance
+   - override /policy: conformers/mlp
+   - override /proxy: conformers/tblite
+   - override /logger: wandb
+
+# Environment
+env:
+  length_traj: 10
+  policy_encoding_dim_per_angle: 10
+  policy_type: mlp
+  n_comp: 5
+  vonmises_min_concentration: 4
+  reward_func: boltzmann
+  reward_beta: 32
+  reward_sampling_method: nested
+
+# GFlowNet hyperparameters
+gflownet:
+  random_action_prob: 0.1
+  optimizer:
+    batch_size:
+      forward: 100
+      backward_dataset: 0
+      backward_replay: 0
+    lr: 0.00001
+    z_dim: 16
+    lr_z_mult: 1000
+    n_train_steps: 40000
+    lr_decay_period: 1000000
+
+# WandB
+logger:
+  lightweight: True
+  project_name: "gflownet"
+  tags:
+    - gflownet
+    - continuous
+    - molecule
+  test:
+    period: 5000
+    n: 10000
+  checkpoints:
+    period: 5000
+
+# Hydra
+hydra:
+  run:
+    dir: ${user.logdir.root}/molecule/${now:%Y-%m-%d_%H-%M-%S}

config/experiments/icml23/alaninedipeptide.yaml

@@ -1,9 +1,9 @@
 # @package _global_
 
 defaults:
-   - override /env: alaninedipeptide
+   - override /env: conformers/alaninedipeptide
    - override /gflownet: trajectorybalance
-   - override /proxy: molecule
+   - override /proxy: conformers/rf
    - override /logger: wandb
    - override /user: sasha
 

config/policy/conformers/egnn.yaml

@@ -0,0 +1,16 @@
+_target_: gflownet.policy.conformers.egnn.EGNNPolicy
+
+forward:
+  n_gnn_layers: 7
+  n_node_mlp_layers: 2
+  n_pool_mlp_layers: 2
+  egnn_hidden_dim: 128
+  node_mlp_hidden_dim: 128
+  pool_mlp_hidden_dim: 128
+backward:
+  n_gnn_layers: 7
+  n_node_mlp_layers: 2
+  n_pool_mlp_layers: 2
+  egnn_hidden_dim: 128
+  node_mlp_hidden_dim: 128
+  pool_mlp_hidden_dim: 128

config/policy/conformers/mlp.yaml

@@ -0,0 +1,13 @@
+_target_: gflownet.policy.base.Policy
+
+forward:
+  type: mlp
+  n_hid: 512
+  n_layers: 5
+  checkpoint: forward
+backward:
+  type: mlp
+  n_hid: 512
+  n_layers: 5
+  shared_weights: False
+  checkpoint: backward

config/proxy/molecule.yaml → config/proxy/conformers/rf.yaml

@@ -1,4 +1,4 @@
-_target_: gflownet.proxy.molecule.RFMoleculeEnergy
+_target_: gflownet.proxy.conformers.rf.RFMoleculeEnergy
 
 path_to_model: './data/random_forest_reward_100.pkl'
-url_to_model: 'https://drive.google.com/uc?id=1OpQNC8WWIsMh8K4olfSaQRFlj3emYThF'
+url_to_model: 'https://drive.google.com/uc?id=1OpQNC8WWIsMh8K4olfSaQRFlj3emYThF'

config/proxy/conformers/tblite.yaml

@@ -0,0 +1 @@
+_target_: gflownet.proxy.conformers.tblite.TBLiteMoleculeEnergy

config/proxy/conformers/torchani.yaml

@@ -0,0 +1,4 @@
+_target_: gflownet.proxy.conformers.torchani.TorchANIMoleculeEnergy
+
+model: ANI2x
+use_ensemble: True

config/proxy/conformers/xtb.yaml

@@ -0,0 +1,3 @@
+_target_: gflownet.proxy.conformers.xtb.XTBMoleculeEnergy
+
+method: gfnff

gflownet/envs/__init__.py

@@ -1 +1 @@
-__all__ = ["base", "grid", "aptamers"]
+__all__ = ["base", "grid", "aptamers", "alaninedipeptide"]

gflownet/envs/base.py

@@ -686,6 +686,26 @@ def statebatch2policy(
             tfloat(states, float_type=self.float, device=self.device)
         )
 
+    def statebatch2kde(self, states: List[List]) -> npt.NDArray[np.float32]:
+        """
+        Prepares a batch of states in "GFlowNet format" for the proxy. Typically,
+        this will be the same as the statebatch2proxy, but in cases in which proxy
+        input is already processed (e.g., conformers, with list of torsion angles
+        converted to 3D positions of atoms), this can be overwritten to preserve KDE.
+        """
+        return self.statebatch2proxy(states)
+
+    def statetorch2kde(
+        self, states: TensorType["batch_size", "state_dim"]
+    ) -> TensorType["batch_size", "state_proxy_dim"]:
+        """
+        Prepares a batch of states in torch "GFlowNet format" for the KDE. Typically,
+        this will be the same as the statetorch2proxy, but in cases in which proxy
+        input is already processed (e.g., conformers, with list of torsion angles
+        converted to 3D positions of atoms), this can be overwritten to preserve KDE.
+        """
+        return self.statetorch2proxy(states)
+
     def policy2state(self, state_policy: List) -> List:
         """
         Converts the model (e.g. one-hot encoding) version of a state given as

gflownet/envs/alaninedipeptide.py → gflownet/envs/conformers/alaninedipeptide.py

@@ -8,8 +8,8 @@
 
 from gflownet.envs.ctorus import ContinuousTorus
 from gflownet.utils.molecule import constants
-from gflownet.utils.molecule.atom_positions_dataset import AtomPositionsDataset
-from gflownet.utils.molecule.conformer_base import ConformerBase
+from gflownet.utils.molecule.datasets import AtomPositionsDataset
+from gflownet.utils.molecule.rdkit_conformer import RDKitConformer
 
 
 class AlanineDipeptide(ContinuousTorus):
@@ -26,7 +26,7 @@ def __init__(
             path_to_dataset, url_to_dataset
         )
         atom_positions = self.atom_positions_dataset.sample()
-        self.conformer = ConformerBase(
+        self.conformer = RDKitConformer(
             atom_positions, constants.ad_smiles, constants.ad_free_tas
         )
         n_dim = len(self.conformer.freely_rotatable_tas)

gflownet/envs/conformers/conformer.py

@@ -0,0 +1,138 @@
+import copy
+from typing import List, Optional, Tuple
+
+import numpy as np
+import numpy.typing as npt
+from rdkit import Chem
+from rdkit.Chem import AllChem
+from torchtyping import TensorType
+
+from gflownet.envs.ctorus import ContinuousTorus
+from gflownet.utils.molecule.constants import ad_atom_types
+from gflownet.utils.molecule.featurizer import MolDGLFeaturizer
+from gflownet.utils.molecule.rdkit_conformer import RDKitConformer
+from gflownet.utils.molecule.rotatable_bonds import find_rotor_from_smile
+
+
+class Conformer(ContinuousTorus):
+    """
+    Extension of continuous torus to conformer generation. Based on AlanineDipeptide,
+    but accepts any molecule (defined by SMILES and freely rotatable torsion angles).
+    """
+
+    def __init__(
+        self,
+        smiles: str,
+        n_torsion_angles: Optional[int] = 2,
+        torsion_indices: Optional[List[int]] = None,
+        policy_type: str = "mlp",
+        **kwargs,
+    ):
+        if torsion_indices is None:
+            # We hard code default torsion indices for Alanine Dipeptide to preserve
+            # backward compatibility.
+            if smiles == "CC(C(=O)NC)NC(=O)C" and n_torsion_angles == 2:
+                torsion_indices = [2, 1]
+            else:
+                torsion_indices = list(range(n_torsion_angles))
+
+        atom_positions = Conformer._get_positions(smiles)
+        torsion_angles = Conformer._get_torsion_angles(smiles, torsion_indices)
+        self.conformer = RDKitConformer(atom_positions, smiles, torsion_angles)
+
+        # Conversions
+        self.statebatch2oracle = self.statebatch2proxy
+        self.statetorch2oracle = self.statetorch2proxy
+        if policy_type == "gnn":
+            self.statebatch2policy = self.statebatch2policy_gnn
+        elif policy_type != "mlp":
+            raise ValueError(
+                f"Unrecognized policy_type = {policy_type}, expected either 'mlp' or 'gnn'."
+            )
+
+        super().__init__(n_dim=len(self.conformer.freely_rotatable_tas), **kwargs)
+
+        self.sync_conformer_with_state()
+
+        self.graph = MolDGLFeaturizer(ad_atom_types).mol2dgl(self.conformer.rdk_mol)
+
+    @staticmethod
+    def _get_positions(smiles: str) -> npt.NDArray:
+        mol = Chem.MolFromSmiles(smiles)
+        mol = Chem.AddHs(mol)
+        AllChem.EmbedMolecule(mol, randomSeed=0)
+        return mol.GetConformer().GetPositions()
+
+    @staticmethod
+    def _get_torsion_angles(smiles: str, indices: List[int]) -> List[Tuple[int]]:
+        torsion_angles = find_rotor_from_smile(smiles)
+        torsion_angles = [torsion_angles[i] for i in indices]
+        return torsion_angles
+
+    def sync_conformer_with_state(self, state: List = None):
+        if state is None:
+            state = self.state
+        for idx, ta in enumerate(self.conformer.freely_rotatable_tas):
+            self.conformer.set_torsion_angle(ta, state[idx])
+        return self.conformer
+
+    def statebatch2proxy(self, states: List[List]) -> npt.NDArray:
+        """
+        Returns a list of proxy states, each being a numpy array with dimensionality
+        (n_atoms, 4), in which the first column encodes atomic number, and the last
+        three columns encode atom positions.
+        """
+        states_proxy = []
+        for st in states:
+            conf = self.sync_conformer_with_state(st)
+            states_proxy.append(
+                np.concatenate(
+                    [
+                        conf.get_atomic_numbers()[..., np.newaxis],
+                        conf.get_atom_positions(),
+                    ],
+                    axis=1,
+                )
+            )
+        return np.array(states_proxy)
+
+    def statetorch2proxy(self, states: TensorType["batch", "state_dim"]) -> npt.NDArray:
+        return self.statebatch2proxy(states.cpu().numpy())
+
+    def statebatch2policy_gnn(self, states: List[List]) -> npt.NDArray[np.float32]:
+        """
+        Returns an array of GNN-format policy inputs with dimensionality
+        (n_states, n_atoms, 4), in which the first three columns encode atom positions,
+        and the last column encodes current timestep.
+        """
+        policy_input = []
+        for state in states:
+            conformer = self.sync_conformer_with_state(state)
+            positions = conformer.get_atom_positions()
+            policy_input.append(
+                np.concatenate(
+                    [positions, np.full((positions.shape[0], 1), state[-1])],
+                    axis=1,
+                )
+            )
+        return np.array(policy_input)
+
+    def statebatch2kde(self, states: List[List]) -> npt.NDArray[np.float32]:
+        return np.array(states)[:, :-1]
+
+    def statetorch2kde(
+        self, states: TensorType["batch_size", "state_dim"]
+    ) -> TensorType["batch_size", "state_proxy_dim"]:
+        return states.cpu().numpy()[:, :-1]
+
+    def __deepcopy__(self, memo):
+        cls = self.__class__
+        new_instance = cls.__new__(cls)
+
+        for attr_name, attr_value in self.__dict__.items():
+            if attr_name != "conformer":
+                setattr(new_instance, attr_name, copy.copy(attr_value))
+
+        new_instance.conformer = self.conformer
+
+        return new_instance