Finish todos and update readme

Author: ben rhodes

README.md

@@ -66,14 +66,17 @@ from orb_models.forcefield import atomic_system, pretrained
 from orb_models.forcefield.base import batch_graphs
 
 device = "cpu"  # or device="cuda"
-orbff = pretrained.orb_v2(device=device)
+orbff, system_config = pretrained.orb_v3_conservative_inf_omat(
+  device=device
+  precision="float32-high",   # or "float32-highest" / "float64
+)
 atoms = bulk('Cu', 'fcc', a=3.58, cubic=True)
-graph = atomic_system.ase_atoms_to_atom_graphs(atoms, device=device)
+graph = atomic_system.ase_atoms_to_atom_graphs(atoms, system_config, device=device)
 
 # Optionally, batch graphs for faster inference
 # graph = batch_graphs([graph, graph, ...])
 
-result = orbff.predict(graph)
+result = orbff.predict(graph, split=False)
 
 # Convert to ASE atoms (unbatches the results and transfers to cpu if necessary)
 atoms = atomic_system.atom_graphs_to_ase_atoms(
@@ -94,8 +97,12 @@ from orb_models.forcefield import pretrained
 from orb_models.forcefield.calculator import ORBCalculator
 
 device="cpu" # or device="cuda"
-orbff = pretrained.orb_v2(device=device) # or choose another model using ORB_PRETRAINED_MODELS[model_name]()
-calc = ORBCalculator(orbff, device=device)
+# or choose another model using ORB_PRETRAINED_MODELS[model_name]()
+orbff, system_config = pretrained.orb_v3_conservative_inf_omat(
+  device=device
+  precision="float32-high",   # or "float32-highest" / "float64
+)
+calc = ORBCalculator(orbff, system_config, device=device)
 atoms = bulk('Cu', 'fcc', a=3.58, cubic=True)
 
 atoms.calc = calc
@@ -111,7 +118,7 @@ from ase.optimize import BFGS
 atoms.rattle(0.5)
 print("Rattled Energy:", atoms.get_potential_energy())
 
-calc = ORBCalculator(orbff, device="cpu") # or device="cuda"
+calc = ORBCalculator(orbff, system_config, device="cpu") # or device="cuda"
 dyn = BFGS(atoms)
 dyn.run(fmax=0.01)
 print("Optimized Energy:", atoms.get_potential_energy())
@@ -120,24 +127,43 @@ print("Optimized Energy:", atoms.get_potential_energy())
 Or you can use it to run MD simulations. The script, an example input xyz file and a Colab notebook demonstration are available in the [examples directory.](./examples) This should work with any input, simply modify the input_file and cell_size parameters. We recommend using constant volume simulations.
 
 
+### Floating Point Precision
+
+As shown in usage snippets above, we support 3 floating point precision types: `"float32-high"`, `"float32-highest"` and `"float64"`.
+
+The default value of `"float32-high"` is recommended for maximal acceleration when using A100 / H100 Nvidia GPUs. However, we have observed some performance loss for high-precision calculations involving second and third order properties of the PES. In these cases, we recommend `"float32-highest"`. 
+
+In stark constrast to other universal forcefields, we have not found any benefit to using `"float64"`.
+
 ### Finetuning
 You can finetune the model using your custom dataset.
 The dataset should be an [ASE sqlite database](https://wiki.fysik.dtu.dk/ase/ase/db/db.html#module-ase.db.core).
 ```python
-python finetune.py --dataset=<dataset_name> --data_path=<your_data_path>
+python finetune.py --dataset=<dataset_name> --data_path=<your_data_path> --base_model=<base_model>
 ```
-After the model is finetuned, checkpoints will, by default, be saved to the ckpts folder in the directory you ran the finetuning script from. 
+Where base_model is one of:
+- "orb_v3_conservative_inf_omat"
+- "orb_v3_conservative_20_omat"
+- "orb_v3_direct_inf_omat"
+- "orb_v3_direct_20_omat"
+- "orb_v2"
+
+After the model is finetuned, checkpoints will, by default, be saved to the ckpts folder in the directory you ran the finetuning script from.
 
 You can use the new model and load the checkpoint by:
 ```python
 from orb_models.forcefield import pretrained
 
-model = pretrained.orb_v2(weights_path=<path_to_ckpt>)
+model, system_config = getattr(pretrained, <base_model>)(
+  weights_path=<path_to_ckpt>, 
+  device="cpu",               # or device="cuda"
+  precision="float32-high",   # or precision="float32-highest"
+)
 ```
 
 > ⚠ **Caveats**
 >
-> Our finetuning script is designed for simplicity and advanced users may wish to develop it further. Please be aware that:
+> Our finetuning script is designed for simplicity. We strongly advise users to customise it further for their use-case to get the best performance. Please be aware that:
 > - The script assumes that your ASE database rows contain **energy, forces, and stress** data. To train on molecular data without stress, you will need to edit the code.
 > - **Early stopping** is not implemented. However, you can use the command line argument `save_every_x_epochs` (default is 5), so "retrospective" early stopping can be applied by selecting a suitable checkpoint.
 > - The **learning rate schedule is hardcoded** to be `torch.optim.lr_scheduler.OneCycleLR` with `pct_start=0.05`. The `max_lr`/`min_lr` will be 10x greater/smaller than the `lr` specified via the command line. To get the best performance, you may wish to try other schedulers.
@@ -147,7 +173,6 @@ model = pretrained.orb_v2(weights_path=<path_to_ckpt>)
 
 
 
-
 ### Citing
 
 A preprint describing the model in more detail can be found here: https://arxiv.org/abs/2410.22570

examples/NaClWaterMD.py

@@ -52,7 +52,7 @@ def run_md_simulation(
     atoms.set_pbc([True] * 3)
 
     # Set the calculator
-    atoms.calc = ORBCalculator(model=pretrained.orb_d3_v2(), device=device)
+    atoms.calc = ORBCalculator(*pretrained.orb_d3_v2(), device=device)
 
     # Set the initial velocities
     MaxwellBoltzmannDistribution(atoms, temperature_K=temperature_K)

finetune.py

@@ -231,13 +231,13 @@ def run(args):
     device = utils.init_device(device_id=args.device_id)
     utils.seed_everything(args.random_seed)
 
-    # Make sure to use this flag for matmuls on A100 and H100 GPUs.
+    # Setting this is 2x faster on A100 and H100 
+    # GPUs and does not appear to hurt training
     precision = "float32-high"
 
     # Instantiate model
-
-    # TODO (BEN): make base model configurable!
-    model, system_config = pretrained.orb_v2(device=device, precision=precision)
+    base_model = args.base_model
+    model, system_config = getattr(pretrained, base_model)(device=device, precision=precision)
 
     for param in model.parameters():
         param.requires_grad = True
@@ -385,6 +385,19 @@ def main():
         type=float,
         help="Learning rate. 3e-4 is purely a sensible default; you may want to tune this for your problem.",
     )
+    parser.add_argument(
+        "--base_model",
+        default="orb_v3_conservative_inf_omat",
+        type=str,
+        help="Base model to finetune.",
+        choices=[
+            "orb_v3_conservative_inf_omat",
+            "orb_v3_conservative_20_omat",
+            "orb_v3_direct_inf_omat",
+            "orb_v3_direct_20_omat",
+            "orb_v2",
+        ],
+    )
     args = parser.parse_args()
     run(args)
 

internal/check.py

@@ -31,7 +31,7 @@ def main(model: str, core_model: str):
     )
 
     graph_orig = core_atomic_system.ase_atoms_to_atom_graphs(atoms, sys_config)
-    graph = atomic_system.ase_atoms_to_atom_graphs(atoms)
+    graph = atomic_system.ase_atoms_to_atom_graphs(atoms, sys_config)
 
     pred_orig = original_orbff.predict(graph_orig)
 

orb_models/forcefield/atomic_system.py

@@ -152,13 +152,12 @@ def atom_graphs_to_ase_atoms(
 
     return atoms_list
 
-
 def ase_atoms_to_atom_graphs(
     atoms: ase.Atoms,
+    system_config: SystemConfig,
     *,
     wrap: bool = True,
     edge_method: Optional[EdgeCreationMethod] = None,
-    system_config: Optional[SystemConfig] = None,
     max_num_neighbors: Optional[int] = None,
     system_id: Optional[int] = None,
     half_supercell: bool = False,
@@ -195,8 +194,6 @@ def ase_atoms_to_atom_graphs(
     Returns:
         AtomGraphs object
     """
-    if system_config is None:
-        system_config = SystemConfig(radius=6.0, max_num_neighbors=20)
     if isinstance(atoms.pbc, Iterable) and any(atoms.pbc) and not all(atoms.pbc):
         raise NotImplementedError(
             "We do not support periodicity along a subset of axes. Please ensure atoms.pbc is "

orb_models/forcefield/pretrained.py

@@ -398,7 +398,6 @@ def orb_v3_direct_inf_mpa(
     return model, SystemConfig(radius=6.0, max_num_neighbors=120)
 
 
-
 def orb_v2(
     weights_path: str = "https://orbitalmaterials-public-models.s3.us-west-1.amazonaws.com/forcefields/orb-v2-20241011.ckpt",  # noqa: E501
     device: Union[torch.device, str, None] = None,
@@ -410,8 +409,6 @@ def orb_v2(
     model = load_model_for_inference(
         model, weights_path, device, precision=precision, compile=compile
     )
-    # TODO (BEN): update all functions to return SystemConfig
-    # TODO (BEN): search repo for max_num_neighbors and avoid any hardcoding
 
     return model, SystemConfig(radius=6.0, max_num_neighbors=20)
 

orb_models/forcefield/segment_ops.py

@@ -159,7 +159,7 @@ def scatter_sum(
     Returns:
         torch.Tensor: The output tensor with values scattered and summed.
     """
-    assert reduce == "sum"  # for now, TODO
+    assert reduce == "sum"
     index = _broadcast(index, src, dim)
     if out is None:
         size = list(src.size())