Crystal environment with continuous lattice parameters

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

config/env/ccube.yaml

@@ -12,21 +12,17 @@ n_dim: 2
 kappa: 1e-3
 # Policy
 min_incr: 0.1
-n_comp: 1
+n_comp: 2
 epsilon: 1e-6
 beta_params_min: 0.1
 beta_params_max: 100.0
 fixed_distr_params:
-  beta_params_min: ${env.beta_params_min}
-  beta_params_max: ${env.beta_params_max}
   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_params_min: ${env.beta_params_min}
-  beta_params_max: ${env.beta_params_max}
   beta_weights: 1.0
   beta_alpha: 10.0
   beta_beta: 10.0

config/env/crystals/ccrystal.yaml

@@ -0,0 +1,26 @@
+defaults:
+  - base
+
+_target_: gflownet.envs.crystals.ccrystal.CCrystal
+
+# Composition config
+id: ccrystal
+composition_kwargs:
+  elements: 89
+# Lattice parameters config
+lattice_parameters_kwargs:
+  min_length: 1.0
+  max_length: 350.0
+  min_angle: 50.0
+  max_angle: 150.0
+# Space group config
+space_group_kwargs:
+  space_groups_subset: null
+# Stoichiometry <-> space group check
+do_stoichiometry_sg_check: True
+
+# Buffer
+buffer:
+  data_path: null
+  train: null
+  test: nulll

config/env/crystals/clattice_parameters.yaml

@@ -0,0 +1,43 @@
+defaults:
+  - base
+
+_target_: gflownet.envs.crystals.clattice_parameters.CLatticeParameters
+
+id: clattice_parameters
+
+# Lattice system
+lattice_system: triclinic
+# Allowed ranges of size and angles
+min_length: 1.0
+max_length: 350.0
+min_angle: 50.0
+max_angle: 150.0
+
+# Policy
+min_incr: 0.1
+n_comp: 2
+epsilon: 1e-6
+beta_params_min: 0.1
+beta_params_max: 100.0
+fixed_distribution:
+  beta_weights: 1.0
+  beta_alpha: 10.0
+  beta_beta: 10.0
+  bernoulli_bts_prob: 0.1
+  bernoulli_eos_prob: 0.1
+random_distribution:
+  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: clp_test.csv
+    output_pkl: clp_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/experiments/ccube/corners.yaml

@@ -17,16 +17,12 @@ env:
   beta_params_max: 100.0
   min_incr: 0.1
   fixed_distr_params:
-    beta_params_min: ${env.beta_params_min}
-    beta_params_max: ${env.beta_params_max}
     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_params_min: ${env.beta_params_min}
-    beta_params_max: ${env.beta_params_max}
     beta_weights: 1.0
     beta_alpha: 10.0
     beta_beta: 10.0

config/experiments/ccube/hyperparams_search_20230920_batch1.yaml

@@ -65,9 +65,8 @@ jobs:
       env:
         __value__: ccube
         n_comp: 5
-        fixed_distr_params:
-          beta_params_min: 0.01
-          beta_params_max: 100.0
+        beta_params_min: 0.01
+        beta_params_max: 100.0
         random_distr_params:
           beta_weights: 1.0
           beta_alpha: 100.0
@@ -80,9 +79,8 @@ jobs:
       env:
         __value__: ccube
         n_comp: 5
-        fixed_distr_params:
-          beta_params_min: 0.1
-          beta_params_max: 100.0
+        beta_params_min: 0.1
+        beta_params_max: 100.0
         random_distr_params:
           beta_weights: 1.0
           beta_alpha: 100.0
@@ -95,9 +93,8 @@ jobs:
       env:
         __value__: ccube
         n_comp: 5
-        fixed_distr_params:
-          beta_params_min: 1
-          beta_params_max: 100.0
+        beta_params_min: 1
+        beta_params_max: 100.0
         random_distr_params:
           beta_weights: 1.0
           beta_alpha: 100.0
@@ -110,9 +107,8 @@ jobs:
       env:
         __value__: ccube
         n_comp: 5
-        fixed_distr_params:
-          beta_params_min: 0.01
-          beta_params_max: 1000.0
+        beta_params_min: 0.01
+        beta_params_max: 1000.0
         random_distr_params:
           beta_weights: 1.0
           beta_alpha: 100.0
@@ -125,9 +121,8 @@ jobs:
       env:
         __value__: ccube
         n_comp: 5
-        fixed_distr_params:
-          beta_params_min: 0.1
-          beta_params_max: 1000.0
+        beta_params_min: 0.1
+        beta_params_max: 1000.0
         random_distr_params:
           beta_weights: 1.0
           beta_alpha: 100.0
@@ -140,9 +135,8 @@ jobs:
       env:
         __value__: ccube
         n_comp: 5
-        fixed_distr_params:
-          beta_params_min: 1
-          beta_params_max: 1000.0
+        beta_params_min: 1
+        beta_params_max: 1000.0
         random_distr_params:
           beta_weights: 1.0
           beta_alpha: 100.0

config/experiments/ccube/hyperparams_search_20230920_batch2.yaml

@@ -65,9 +65,8 @@ jobs:
       env:
         __value__: ccube
         n_comp: 5
-        fixed_distr_params:
-          beta_params_min: 0.01
-          beta_params_max: 100.0
+        beta_params_min: 0.01
+        beta_params_max: 100.0
         random_distr_params:
           beta_weights: 1.0
           beta_alpha: 100.0
@@ -80,9 +79,8 @@ jobs:
       env:
         __value__: ccube
         n_comp: 5
-        fixed_distr_params:
-          beta_params_min: 0.1
-          beta_params_max: 100.0
+        beta_params_min: 0.1
+        beta_params_max: 100.0
         random_distr_params:
           beta_weights: 1.0
           beta_alpha: 100.0
@@ -95,9 +93,8 @@ jobs:
       env:
         __value__: ccube
         n_comp: 5
-        fixed_distr_params:
-          beta_params_min: 1
-          beta_params_max: 100.0
+        beta_params_min: 1
+        beta_params_max: 100.0
         random_distr_params:
           beta_weights: 1.0
           beta_alpha: 100.0
@@ -110,9 +107,8 @@ jobs:
       env:
         __value__: ccube
         n_comp: 5
-        fixed_distr_params:
-          beta_params_min: 0.01
-          beta_params_max: 1000.0
+        beta_params_min: 0.01
+        beta_params_max: 1000.0
         random_distr_params:
           beta_weights: 1.0
           beta_alpha: 100.0
@@ -125,9 +121,8 @@ jobs:
       env:
         __value__: ccube
         n_comp: 5
-        fixed_distr_params:
-          beta_params_min: 0.1
-          beta_params_max: 1000.0
+        beta_params_min: 0.1
+        beta_params_max: 1000.0
         random_distr_params:
           beta_weights: 1.0
           beta_alpha: 100.0
@@ -140,9 +135,8 @@ jobs:
       env:
         __value__: ccube
         n_comp: 5
-        fixed_distr_params:
-          beta_params_min: 1
-          beta_params_max: 1000.0
+        beta_params_min: 1
+        beta_params_max: 1000.0
         random_distr_params:
           beta_weights: 1.0
           beta_alpha: 100.0

config/experiments/ccube/hyperparams_search_20230920_batch3.yaml

@@ -65,9 +65,8 @@ jobs:
       env:
         __value__: ccube
         n_comp: 2
-        fixed_distr_params:
-          beta_params_min: 0.01
-          beta_params_max: 100.0
+        beta_params_min: 0.01
+        beta_params_max: 100.0
         random_distr_params:
           beta_weights: 1.0
           beta_alpha: 100.0
@@ -80,9 +79,8 @@ jobs:
       env:
         __value__: ccube
         n_comp: 2
-        fixed_distr_params:
-          beta_params_min: 0.1
-          beta_params_max: 100.0
+        beta_params_min: 0.1
+        beta_params_max: 100.0
         random_distr_params:
           beta_weights: 1.0
           beta_alpha: 100.0
@@ -95,9 +93,8 @@ jobs:
       env:
         __value__: ccube
         n_comp: 2
-        fixed_distr_params:
-          beta_params_min: 1
-          beta_params_max: 100.0
+        beta_params_min: 1
+        beta_params_max: 100.0
         random_distr_params:
           beta_weights: 1.0
           beta_alpha: 100.0
@@ -110,9 +107,8 @@ jobs:
       env:
         __value__: ccube
         n_comp: 2
-        fixed_distr_params:
-          beta_params_min: 0.01
-          beta_params_max: 1000.0
+        beta_params_min: 0.01
+        beta_params_max: 1000.0
         random_distr_params:
           beta_weights: 1.0
           beta_alpha: 100.0
@@ -125,9 +121,8 @@ jobs:
       env:
         __value__: ccube
         n_comp: 2
-        fixed_distr_params:
-          beta_params_min: 0.1
-          beta_params_max: 1000.0
+        beta_params_min: 0.1
+        beta_params_max: 1000.0
         random_distr_params:
           beta_weights: 1.0
           beta_alpha: 100.0
@@ -140,9 +135,8 @@ jobs:
       env:
         __value__: ccube
         n_comp: 2
-        fixed_distr_params:
-          beta_params_min: 1
-          beta_params_max: 1000.0
+        beta_params_min: 1
+        beta_params_max: 1000.0
         random_distr_params:
           beta_weights: 1.0
           beta_alpha: 100.0