Merge branch 'conformer_ai4mat23' of github.com:alexhernandezgarcia/g…

…flownet into rdkit_cluster_baseline

.gitignore

@@ -11,4 +11,6 @@ sh/
 .vscode/
 external/
 !requirements.txt
+!docs/requirements-docs.txt
 .DS_Store
+docs/_build/

.readthedocs.yaml

@@ -0,0 +1,22 @@
+# .readthedocs.yaml
+# Read the Docs configuration file
+# See https://docs.readthedocs.io/en/stable/config-file/v2.html for details
+
+# Required
+version: 2
+
+# Set the version of Python and other tools you might need
+build:
+  os: ubuntu-22.04
+  tools:
+    python: "3.11"
+
+# Build documentation in the docs/ directory with Sphinx
+sphinx:
+  configuration: docs/conf.py
+
+# We recommend specifying your dependencies to enable reproducible builds:
+# https://docs.readthedocs.io/en/stable/guides/reproducible-builds.html
+python:
+  install:
+  - requirements: docs/requirements-docs.txt

LAUNCH.md

@@ -9,8 +9,8 @@ usage: launch.py [-h] [--help-md] [--job_name JOB_NAME] [--outdir OUTDIR]
                  [--cpus_per_task CPUS_PER_TASK] [--mem MEM] [--gres GRES]
                  [--partition PARTITION] [--modules MODULES]
                  [--conda_env CONDA_ENV] [--venv VENV] [--template TEMPLATE]
-                 [--code_dir CODE_DIR] [--jobs JOBS] [--dry-run] [--verbose]
-                 [--force]
+                 [--code_dir CODE_DIR] [--git_checkout GIT_CHECKOUT]
+                 [--jobs JOBS] [--dry-run] [--verbose] [--force]
 
 optional arguments:
   -h, --help            show this help message and exit
@@ -35,6 +35,11 @@ optional arguments:
                         $root/mila/sbatch/template-conda.sh
   --code_dir CODE_DIR   cd before running main.py (defaults to here). Defaults
                         to $root
+  --git_checkout GIT_CHECKOUT
+                        Branch or commit to checkout before running the code.
+                        This is only used if --code_dir='$SLURM_TMPDIR'. If
+                        not specified, the current branch is used. Defaults to
+                        None
   --jobs JOBS           jobs (nested) file name in external/jobs (with or
                         without .yaml). Or an absolute path to a yaml file
                         anywhere Defaults to None
@@ -54,6 +59,7 @@ conda_env     : gflownet
 cpus_per_task : 2
 dry-run       : False
 force         : False
+git_checkout  : None
 gres          : gpu:1
 job_name      : gflownet
 jobs          : None
@@ -74,7 +80,7 @@ a single job from the command-line, or a list of jobs from a `yaml` file.
 
 Examples:
 
-```sh
+```bash
 # using default job configuration, with script args from the command-line:
 $ python mila/launch.py user=$USER logger.do.online=False
 

config/env/ccube.yaml

@@ -0,0 +1,39 @@
+defaults:
+  - base
+
+_target_: gflownet.envs.cube.ContinuousCube
+
+id: ccube
+continuous: True
+func: corners
+# Dimensions of hypercube
+n_dim: 2
+# Constant to restrict interval of test sets
+kappa: 1e-3
+# Policy
+min_incr: 0.1
+n_comp: 1
+epsilon: 1e-6
+beta_params_min: 0.1
+beta_params_max: 100.0
+fixed_distr_params:
+  beta_weights: 1.0
+  beta_alpha: 10.0
+  beta_beta: 10.0
+  bernoulli_bts_prob: 0.1
+  bernoulli_eos_prob: 0.1
+random_distr_params:
+  beta_weights: 1.0
+  beta_alpha: 10.0
+  beta_beta: 10.0
+  bernoulli_bts_prob: 0.1
+  bernoulli_eos_prob: 0.1
+# Buffer
+buffer:
+  data_path: null
+  train: null
+  test:
+    type: grid
+    n: 900
+    output_csv: ccube_test.csv
+    output_pkl: ccube_test.pkl

config/env/conformers/alaninedipeptide.yaml

@@ -12,11 +12,11 @@ length_traj: 10
 vonmises_min_concentration: 1e-3
 # Parameters of the fixed policy output distribution
 n_comp: 3
-fixed_distribution:
+fixed_distr_params:
   vonmises_mean: 0.0
   vonmises_concentration: 0.5
 # Parameters of the random policy output distribution
-random_distribution:
+random_distr_params:
   vonmises_mean: 0.0
   vonmises_concentration: 0.001
 # Buffer

config/env/crystals/lattice_parameters.yaml

@@ -0,0 +1,28 @@
+defaults:
+  - base
+
+_target_: gflownet.envs.crystals.lattice_parameters.LatticeParameters
+
+id: lattice_parameters
+
+lattice_system: triclinic
+min_length: 1.0
+max_length: 350.0
+min_angle: 50.0
+max_angle: 150.0
+grid_size: 10
+max_increment: 1
+
+# Buffer
+buffer:
+  data_path: null
+  train:
+    type: random
+    n: 2000
+    output_csv: lattice_parameters_train.csv
+    output_pkl: lattice_parameters_train.pkl
+  test:
+    type: random
+    n: 2000
+    output_csv: lattice_parameters_test.csv
+    output_pkl: lattice_parameters_test.pkl

config/env/crystals/spacegroup.yaml

@@ -4,6 +4,9 @@ defaults:
 _target_: gflownet.envs.crystals.spacegroup.SpaceGroup
 
 id: spacegroup
+
+# Subset of space groups
+space_groups_subset: null
 # Stoichiometry
 n_atoms: null
 

config/env/ctorus.yaml

@@ -13,11 +13,11 @@ length_traj: 3
 vonmises_min_concentration: 1e-3
 # Parameters of the fixed policy output distribution
 n_comp: 3
-fixed_distribution:
+fixed_distr_params:
   vonmises_mean: 0.0
   vonmises_concentration: 1.0
 # Parameters of the random policy output distribution
-random_distribution:
+random_distr_params:
   vonmises_mean: 0.0
   vonmises_concentration: 0.01
 # Buffer

config/env/htorus.yaml

@@ -13,11 +13,11 @@ policy_encoding_dim_per_angle: null
 length_traj: 3
 vonmises_min_concentration: 1e-3
 # Parameters of the fixed policy output distribution
-fixed_distribution:
+fixed_distr_params:
   vonmises_mean: 0.0
   vonmises_concentration: 0.5
 # Parameters of the random policy output distribution
-random_distribution:
+random_distr_params:
   vonmises_mean: 0.0
   vonmises_concentration: 0.001
 # Buffer

config/env/tree.yaml

@@ -22,10 +22,10 @@ test_args:
 threshold_components: 3
 beta_params_min: 0.1
 beta_params_max: 100.0
-fixed_distribution:
+fixed_distr_params:
   beta_alpha: 2.0
   beta_beta: 5.0
-random_distribution:
+random_distr_params:
   beta_alpha: 1.0
   beta_beta: 1.0
 # Buffer

config/experiments/iclr23/mlp_torchani.yaml → config/experiments/ai4mat23/mlp_torchani.yaml

@@ -21,6 +21,11 @@ env:
   reward_sampling_method: nested
   buffer:
     replay_capacity: 1000
+    # To enable KDE plots, change this to `type: grid` (`n: 1000` was used in the original experiments).
+    test:
+      type: random
+      seed: 0
+      n: 4096
 
 # GFlowNet hyperparameters
 gflownet:

config/experiments/ccube/corners.yaml

@@ -0,0 +1,73 @@
+# @package _global_
+# A configuration that works well with the corners proxy.
+# wandb: https://wandb.ai/alexhg/cube/runs/9u2d3zzh
+
+defaults:
+   - override /env: ccube
+   - override /gflownet: trajectorybalance
+   - override /proxy: corners
+   - override /logger: wandb
+   - override /user: alex
+
+# Environment
+env:
+  n_comp: 5
+  n_dim: 2
+  beta_params_min: 0.1
+  beta_params_max: 100.0
+  min_incr: 0.1
+  fixed_distr_params:
+    beta_weights: 1.0
+    beta_alpha: 10.0
+    beta_beta: 10.0
+    bernoulli_eos_prob: 0.1
+    bernoulli_bts_prob: 0.1
+  random_distr_params:
+    beta_weights: 1.0
+    beta_alpha: 10.0
+    beta_beta: 10.0
+    bernoulli_eos_prob: 0.1
+    bernoulli_bts_prob: 0.1
+  reward_func: identity
+
+# GFlowNet hyperparameters
+gflownet:
+  random_action_prob: 0.1
+  optimizer:
+    batch_size:
+      forward: 100
+    lr: 0.0001
+    z_dim: 16
+    lr_z_mult: 100
+    n_train_steps: 10000
+  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
+
+# WandB
+logger:
+  lightweight: True
+  project_name: "cube"
+  tags: 
+    - gflownet
+    - continuous
+    - ccube
+  test:
+    period: 500
+    n: 1000
+  checkpoints:
+    period: 500
+
+# Hydra
+hydra:
+  run:
+    dir: ${user.logdir.root}/debug/ccube/${now:%Y-%m-%d_%H-%M-%S}

config/experiments/ccube/uniform.yaml

@@ -0,0 +1,73 @@
+# @package _global_
+# A configuration that works well with the uniform proxy.
+# wandb: https://wandb.ai/alexhg/cube/runs/1du9iyr5
+
+defaults:
+   - override /env: ccube
+   - override /gflownet: trajectorybalance
+   - override /proxy: uniform
+   - override /logger: wandb
+   - override /user: alex
+
+# Environment
+env:
+  n_comp: 2
+  n_dim: 2
+  beta_params_min: 0.1
+  beta_params_max: 100.0
+  min_incr: 0.1
+  fixed_distr_params:
+    beta_weights: 1.0
+    beta_alpha: 10.0
+    beta_beta: 10.0
+    bernoulli_eos_prob: 0.1
+    bernoulli_bts_prob: 0.1
+  random_distr_params:
+    beta_weights: 1.0
+    beta_alpha: 10.0
+    beta_beta: 10.0
+    bernoulli_eos_prob: 0.1
+    bernoulli_bts_prob: 0.1
+  reward_func: identity
+
+# GFlowNet hyperparameters
+gflownet:
+  random_action_prob: 0.1
+  optimizer:
+    batch_size:
+      forward: 100
+    lr: 0.0001
+    z_dim: 16
+    lr_z_mult: 100
+    n_train_steps: 10000
+  policy:
+    forward:
+      type: mlp
+      n_hid: 256
+      n_layers: 3
+      checkpoint: forward
+    backward:
+      type: mlp
+      n_hid: 256
+      n_layers: 3
+      shared_weights: False
+      checkpoint: backward
+
+# WandB
+logger:
+  lightweight: True
+  project_name: "cube"
+  tags: 
+    - gflownet
+    - continuous
+    - ccube
+  test:
+    period: 500
+    n: 1000
+  checkpoints:
+    period: 500
+
+# Hydra
+hydra:
+  run:
+    dir: ${user.logdir.root}/debug/ccube/${now:%Y-%m-%d_%H-%M-%S}