EGNN policy

config/env/conformers/conformer.yaml

@@ -8,6 +8,8 @@ _target_: gflownet.envs.conformers.conformer.Conformer
 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
+# smiles: 'Cc1ccc(cc1Nc2nccc(n2)c3cccnc3)NC(=O)c4ccc(cc4)CN5CCN(CC5)C' # imatinib
+# smiles: 'CC(C)c4nc(CN(C)C(=O)N[C@@H](C(C)C)C(=O)N[C@@H](Cc1ccccc1)C[C@H](O)[C@H](Cc2ccccc2)NC(=O)OCc3cncs3)cs4' # ritonavir
 n_torsion_angles: 2
 reward_sampling_method: nested
 

config/experiments/iclr23/egnn_torchani.yaml

@@ -0,0 +1,58 @@
+# @package _global_
+
+defaults:
+   - override /env: conformers/conformer
+   - override /gflownet: trajectorybalance
+   - override /policy: conformers/egnn
+   - override /proxy: conformers/torchani
+   - override /logger: wandb
+
+# Environment
+env:
+  smiles: CC(C)Cc1ccc(cc1)[C@@H](C)C(=O)O
+  n_torsion_angles: 2
+  remove_hs: True
+  length_traj: 5
+  policy_encoding_dim_per_angle: 10
+  policy_type: gnn
+  n_comp: 5
+  vonmises_min_concentration: 4
+  reward_func: boltzmann
+  reward_beta: 32
+  reward_sampling_method: nested
+  buffer:
+    replay_capacity: 1000
+
+# GFlowNet hyperparameters
+gflownet:
+  random_action_prob: 0.1
+  replay_sampling: weighted
+  optimizer:
+    batch_size:
+      forward: 90
+      backward_dataset: 0
+      backward_replay: 10
+    lr: 0.0001
+    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: 2000
+    n: 10000
+  checkpoints:
+    period: 2000
+
+# Hydra
+hydra:
+  run:
+    dir: ${user.logdir.root}/molecule/${now:%Y-%m-%d_%H-%M-%S}

config/experiments/iclr23/mlp_torchani.yaml

@@ -0,0 +1,57 @@
+# @package _global_
+
+defaults:
+   - override /env: conformers/conformer
+   - override /gflownet: trajectorybalance
+   - override /policy: conformers/mlp
+   - override /proxy: conformers/torchani
+   - override /logger: wandb
+
+# Environment
+env:
+  smiles: Cc1ccc(cc1Nc2nccc(n2)c3cccnc3)NC(=O)c4ccc(cc4)CN5CCN(CC5)C
+  n_torsion_angles: 5
+  length_traj: 5
+  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
+  buffer:
+    replay_capacity: 1000
+
+# GFlowNet hyperparameters
+gflownet:
+  random_action_prob: 0.1
+  replay_sampling: weighted
+  optimizer:
+    batch_size:
+      forward: 80
+      backward_dataset: 0
+      backward_replay: 20
+    lr: 0.0001
+    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: 2000
+    n: 10000
+  checkpoints:
+    period: 2000
+
+# Hydra
+hydra:
+  run:
+    dir: ${user.logdir.root}/molecule/${now:%Y-%m-%d_%H-%M-%S}

config/experiments/tree.yaml

@@ -38,13 +38,12 @@ gflownet:
 
 # MLP policy
 policy:
-  forward:
-    n_hid: 256
-    n_layers: 3
-  backward:
+  shared:
     type: mlp
     n_hid: 256
     n_layers: 3
+  forward: null
+  backward:
     shared_weights: False
 
 # WandB

config/policy/conformers/egnn.yaml

@@ -1,16 +1,12 @@
 _target_: gflownet.policy.conformers.egnn.EGNNPolicy
 
-forward:
+shared:
   n_gnn_layers: 7
-  n_node_mlp_layers: 2
-  n_pool_mlp_layers: 2
+  n_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
+  mlp_hidden_dim: 256
+  separate_mlp_per_torsion: True
+  use_fake_edges: True
+  fake_edge_radius: 2.0
+forward: null
+backward: null

config/policy/conformers/mlp.yaml

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

config/policy/mlp.yaml

@@ -1,5 +1,7 @@
 _target_: gflownet.policy.base.Policy
 
+shared: null
+
 forward:
   type: mlp
   n_hid: 128

gflownet/envs/base.py

@@ -406,8 +406,7 @@ def step_backwards(
             return self.state, action, False
         parents, parents_a = self.get_parents()
         state_next = parents[parents_a.index(action)]
-        self.state = state_next
-        self.done = False
+        self.set_state(state_next, done=False)
         self.n_actions += 1
         return self.state, action, True
 

gflownet/envs/conformers/conformer.py

@@ -1,8 +1,10 @@
 import copy
 from typing import List, Optional, Tuple
 
+import dgl
 import numpy as np
 import numpy.typing as npt
+import torch
 from rdkit import Chem
 from rdkit.Chem import AllChem
 from torchtyping import TensorType
@@ -11,7 +13,7 @@
 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
+from gflownet.utils.molecule.rotatable_bonds import find_rotor_from_smiles
 
 
 class Conformer(ContinuousTorus):
@@ -26,13 +28,16 @@ def __init__(
         n_torsion_angles: Optional[int] = 2,
         torsion_indices: Optional[List[int]] = None,
         policy_type: str = "mlp",
+        remove_hs: bool = True,
         **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]
+            elif n_torsion_angles == -1:
+                torsion_indices = None
             else:
                 torsion_indices = list(range(n_torsion_angles))
 
@@ -50,12 +55,27 @@ def __init__(
                 f"Unrecognized policy_type = {policy_type}, expected either 'mlp' or 'gnn'."
             )
 
+        self.graph = MolDGLFeaturizer(ad_atom_types).mol2dgl(self.conformer.rdk_mol)
+        # TODO: use DGL conformer instead
+        rotatable_edges = [ta[1:3] for ta in torsion_angles]
+        for i in range(self.graph.num_edges()):
+            if (
+                self.graph.edges()[0][i].item(),
+                self.graph.edges()[1][i].item(),
+            ) not in rotatable_edges:
+                self.graph.edata["rotatable_edges"][i] = False
+
+        # Hydrogen removal
+        self.remove_hs = remove_hs
+        self.hs = torch.where(self.graph.ndata["atom_features"][:, 0] == 1)[0]
+        self.non_hs = torch.where(self.graph.ndata["atom_features"][:, 0] != 1)[0]
+        if remove_hs:
+            self.graph = dgl.remove_nodes(self.graph, self.hs)
+
         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)
@@ -64,9 +84,12 @@ def _get_positions(smiles: str) -> npt.NDArray:
         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]
+    def _get_torsion_angles(
+        smiles: str, indices: Optional[List[int]]
+    ) -> List[Tuple[int]]:
+        torsion_angles = find_rotor_from_smiles(smiles)
+        if indices is not None:
+            torsion_angles = [torsion_angles[i] for i in indices]
         return torsion_angles
 
     def sync_conformer_with_state(self, state: List = None):
@@ -109,6 +132,8 @@ def statebatch2policy_gnn(self, states: List[List]) -> npt.NDArray[np.float32]:
         for state in states:
             conformer = self.sync_conformer_with_state(state)
             positions = conformer.get_atom_positions()
+            if self.remove_hs:
+                positions = positions[self.non_hs]
             policy_input.append(
                 np.concatenate(
                     [positions, np.full((positions.shape[0], 1), state[-1])],

gflownet/envs/crystals/crystal.py

@@ -522,6 +522,41 @@ def statetorch2oracle(
             dim=1,
         )
 
+    def set_state(self, state: List, done: Optional[bool] = False):
+        super().set_state(state, done)
+
+        stage = self._get_stage(state)
+
+        composition_done = stage in [Stage.SPACE_GROUP, Stage.LATTICE_PARAMETERS]
+        space_group_done = stage == Stage.LATTICE_PARAMETERS
+        lattice_parameters_done = done
+
+        self.composition.set_state(self._get_composition_state(state), composition_done)
+        self.space_group.set_state(self._get_space_group_state(state), space_group_done)
+        self.lattice_parameters.set_state(
+            self._get_lattice_parameters_state(state), lattice_parameters_done
+        )
+
+        """
+        We synchronize LatticeParameter's lattice system with the one of SpaceGroup
+        (if it was set) or reset it to the default triclinic otherwise. Why this is 
+        needed:
+        1) the first case is necessary for backward sampling, where we start from
+           an arbitrary terminal state, and need to synchronize the LatticeParameter's
+           lattice system to what that state indicates,
+        2) the second case is also necessary in backward sampling, but when we 
+           transition from Stage.LATTICE_PARAMETERS to Stage.SPACE_GROUP. We then need
+           to reset the lattice system to the default triclinic, such that its
+           source is back to the original one, and corresponds to the source of the
+           general Crystal environment.
+        """
+        lattice_system = self.space_group.lattice_system
+        if lattice_system != "None":
+            self.lattice_parameters.lattice_system = lattice_system
+        else:
+            self.lattice_parameters.lattice_system = TRICLINIC
+        self.lattice_parameters._set_source()
+
     def state2readable(self, state: Optional[List[int]] = None) -> str:
         if state is None:
             state = self.state

gflownet/envs/crystals/spacegroup.py

@@ -158,7 +158,8 @@ def get_mask_invalid_actions_forward(
             mask[-1] = False
             return mask
         state_type = self.get_state_type(state)
-        # No constraints if neither crystal-lattice system nor point symmetry selected
+        # If neither crystal-lattice system nor point symmetry selected, apply only
+        # composition-compatibility constraints
         if cls_idx == 0 and ps_idx == 0:
             crystal_lattice_systems = [
                 (self.cls_idx, idx + 1, state_type)