EXESS: Geometry Opt.¶

EXESS can perform geometry optimization using its ability to calculate gradients. It can do this using QM (via its core inbuilt capabilities), MM (via OpenMM), or ML (via AIMNet) techniques, dubbed Q4ML as a shorthand of QM/MM/ML.

Reminder

Don’t forget to set RUSH_TOKEN and RUSH_PROJECT!

QM optimization¶

Note

Optimization requires gradient-capable methods (RHF, RI-HF, or RI-MP2). See the EXESS docs for method support details.

The exess submodule in rush-py provides a dedicated function for doing optimization, which automatically selects EXESS’s Optimization driver. By default, it’s only necessary to provide the QDX Topology file and the number of iterations to run:

optimization_qm.py¶

from rush import exess
from rush.client import RunOpts, save_object

out = exess.optimization(
    "benzene_t.json",
    100,  # Number of optimization iterations 
    standard_orientation="None",
    run_opts=RunOpts(name="Tutorial: Optimization using QM"),
    collect=True,
)

Help, my output conformer is in a totally different position!

Setting standard_orientation="None" prevents EXESS from translating or rotating the input geometry, making direct comparisons easier.

ML optimization¶

To use ML with the AIMNet model for the whole system, set qm_fragments and mm_fragments to the empty list. EXESS will infer that ml_fragments should be the remaining ones, i.e. all of them.

It would also be fine to set ml_fragments to all the fragments as well (and in this case, it would only be necessary to set one of the other two to the empty list so that the third can be inferred), but this requires constructing that complete list which isn’t as convenient.

optimization_ml.py¶

from rush import exess
from rush.client import RunOpts, save_object

out = exess.optimization(
    "benzene_t.json",
    100,
    basis="STO-2G",
    optimization_keywords=exess.OptimizationKeywords(
        coordinate_system="Cartesian",
        algorithm="LBFGS",
        lbfgs_keywords=exess.LBFGSKeywords(),
    ),
    standard_orientation="None",
    qm_fragments=[],
    mm_fragments=[],
    run_opts=RunOpts(name="Tutorial: Optimization using ML"),
    collect=True,
)

Note

Using the smallest basis set available, i.e. STO-2G, reduces memory requirements for non-QM runs and has no effect on the calculation or results.

Optimizing your optimization keywords

The values shown for optimization_keywords are the only supported ones for non-QM runs. Setting the coordinate system to Cartesian is mandatory, and while other algorithms may work, using LBFGS is highly recommended.

Working with the output¶

The first output is a vector of Topologies, one for each step, which effectively provides the optimization trajectory. The second output is a vector of dictionaries, also one for each step, containing the total energy and the max gradient component at that step. The following example shows how these can be accessed and used:

import json
from rush import Topology

out_traj_path, out_info_path = [save_object(obj["path"]) for obj in out]
with open(out_traj_path) as f1, open(out_info_path) as f2:
    out_traj, out_info = [json.load(f) for f in (f1, f2)]

print("Num steps to convergence:", len(out_traj))

out_traj = [Topology.from_json(t) for t in out_traj]
print("First Atom's Coords")
print(f"  First step: {out_traj[0].geometry[:3]}")
print(f"  Final step: {out_traj[-1].geometry[:3]}")

# The below are only provided for QM regions
print("Final Step Info")  
print(f"  Total energy: {out_info[-1]['total_energy']:.5f} Eh")
print(f"  Max gradient component: {out_info[-1]['max_gradient_component']:.2} Å")

Limitation

The second output only provides data for QM regions; if there are none, the dictionaries will be empty.