Fix radial/angular predictors + other changes

analyses/analysis.py

@@ -30,7 +30,7 @@
 from configs import root_dirs
 
 try:
-    from symphony.data import input_pipeline_tf
+    from symphony.data import input_pipeline
     import tensorflow as tf
 
     tf.config.experimental.set_visible_devices([], "GPU")
@@ -253,14 +253,7 @@ def load_from_workdir(
     # Check that the config was loaded correctly.
     assert config is not None
     config = ml_collections.ConfigDict(config)
-    if 'max_targets_per_graph' in config:
-        config.root_dir = root_dirs.get_root_dir(
-            config.dataset, config.get("fragment_logic", "nn"), config.max_targets_per_graph
-        )
-    else:
-        config.root_dir = root_dirs.get_root_dir(
-            config.dataset, config.get("fragment_logic", "nn")
-        )
+    config.root_dir = root_dirs.get_root_dir(config.dataset)
 
     # Mimic what we do in train.py.
     rng = jax.random.PRNGKey(config.rng_seed)
@@ -295,31 +288,35 @@ def load_from_workdir(
     else:
         if init_graphs is None:
             logging.info("Initializing dummy model with init_graphs from dataloader")
-            datasets = input_pipeline_tf.get_datasets(dataset_rng, config)
-            train_iter = datasets["train"].as_numpy_iterator()
+            datasets = input_pipeline.get_datasets(dataset_rng, config)
+            train_iter = datasets["train"]
             init_graphs = next(train_iter)
         else:
             logging.info("Initializing dummy model with provided init_graphs")
 
         params = jax.jit(net.init)(init_rng, init_graphs)
 
     tx = train.create_optimizer(config)
-    dummy_state = train_state.TrainState.create(
-        apply_fn=net.apply, params=params, tx=tx
-    )
+    # dummy_state = train_state.TrainState.create(
+    #     apply_fn=net.apply, params=params, tx=tx
+    # )
 
     # Load the actual values.
     checkpoint_dir = os.path.join(workdir, "checkpoints")
     ckpt = checkpoint.Checkpoint(checkpoint_dir, max_to_keep=5)
-    data = ckpt.restore({"best_state": dummy_state, "metrics_for_best_state": None})
-    best_state = jax.tree_util.tree_map(jnp.asarray, data["best_state"])
+    data = ckpt.restore_dict()["state"]
+    best_state = jax.tree_util.tree_map(
+        jnp.asarray,
+        train_state.TrainState.create(
+            apply_fn=net.apply, params=data["best_params"], tx=tx
+    ))
     best_state_in_eval_mode = best_state.replace(apply_fn=eval_net.apply)
 
     return (
         config,
         best_state,
         best_state_in_eval_mode,
-        cast_keys_as_int(data["metrics_for_best_state"]),
+        cast_keys_as_int(data["metrics_for_best_params"]),
     )
 
 

analyses/conditional_generation.py

@@ -11,23 +11,13 @@
 
 from absl import flags, app
 import analyses.generate_molecules as generate_molecules
-from symphony.data.datasets import qm9
-
+from symphony.data.datasets import tmqm
+from configs.root_dirs import get_root_dir
+from symphony import datatypes
 
 
 workdir = "/home/ameyad/spherical-harmonic-net/workdirs/qm9_bessel_embedding_attempt6_edm_splits/e3schnet_and_nequip/interactions=3/l=5/position_channels=2/channels=64"
 outputdir = "conditional_generation"
-beta_species = 1.0
-beta_position = 1.0
-step = "7530000"
-num_seeds_per_chunk = 25
-max_num_atoms = 35
-visualize = False
-num_mols = 1000
-
-all_mols = qm9.load_qm9("../qm9_data", use_edm_splits=True, check_molecule_sanity=False)
-test_mols = all_mols[-num_mols:]
-train_mols = all_mols[:num_mols]
 
 
 def get_fragment_list(mols: Sequence[ase.Atoms], num_mols: int):
@@ -46,36 +36,42 @@ def get_fragment_list(mols: Sequence[ase.Atoms], num_mols: int):
 
 
 def main(unused_argv: Sequence[str]):
+    radial_cutoff = 5.0
     beta_species = 1.0
     beta_position = 1.0
     step = flags.FLAGS.step
     num_seeds_per_chunk = 1
-    max_num_atoms = 200
-    max_num_steps = 10
-    num_mols = 20
+    max_num_atoms = 50
+    num_mols = 500
+    avg_neighbors_per_atom = 32
+
+    atomic_numbers = np.arange(1, 81)
 
     all_mols = tmqm.load_tmqm("../tmqm_data")
     mols_by_split = {"train": all_mols[:num_mols], "test": all_mols[-num_mols:]}
 
     for split, split_mols in mols_by_split.items():
         # Ensure that the number of molecules is a multiple of num_seeds_per_chunk.
         mol_list = get_fragment_list(split_mols, num_mols)
-        mol_list = split_mols[
-            : num_seeds_per_chunk * (len(split_mols) // num_seeds_per_chunk)
+        mol_list = mol_list[
+            : num_seeds_per_chunk * (len(mol_list) // num_seeds_per_chunk)
         ]
         print(f"Number of fragments for {split}: {len(mol_list)}")
 
-        gen_mol_list = generate_molecules.generate_molecules(
+        gen_mol_list = generate_molecules.generate_molecules_from_workdir(
             flags.FLAGS.workdir,
             os.path.join(flags.FLAGS.outputdir, split),
+            radial_cutoff,
             beta_species,
             beta_position,
             step,
+            flags.FLAGS.steps_for_weight_averaging,
             len(mol_list),
             num_seeds_per_chunk,
             mol_list,
             max_num_atoms,
-            max_num_steps,
+            avg_neighbors_per_atom,
+            atomic_numbers,
             flags.FLAGS.visualize,
         )
 
@@ -101,4 +97,14 @@ def main(unused_argv: Sequence[str]):
         "best",
         "Step number to load model from. The default corresponds to the best model.",
     )
+    flags.DEFINE_list(
+        "steps_for_weight_averaging",
+        None,
+        "Steps to average parameters over. If None, the model at the given step is used.",
+    )
+    flags.DEFINE_bool(
+        "seed_structure",
+        False,
+        "Add initial atom of the missing element to structure"
+    )
     app.run(main)

analyses/elements.ini

@@ -1,5 +1,4 @@
   1  H  0.46  1.20  0.200    1.00000    1.00000    1.00000
-  1  D  0.46  1.20  0.200    0.80000    0.80000    1.00000
   2 He  1.22  1.40  1.220    0.85100    1.00000    1.00000
   3 Li  1.57  1.40  0.590    0.80000    0.50200    1.00000
   4 Be  1.12  1.40  0.270    0.76100    1.00000    0.00000
@@ -94,4 +93,4 @@
  93 Np  1.56  2.16  0.750    0.00000    0.50200    1.00000
  94 Pu  1.64  2.16  0.860    0.00000    0.42000    1.00000
  95 Am  1.73  2.16  0.975    0.32900    0.36100    0.94900
- 96 XX  0.80  1.00  0.800    0.47100    0.36100    0.89000
+ 96 XX  0.80  1.00  0.800    0.47100    0.36100    0.89000