Geometry Optimization¶
Set up environment (optional)¶
These steps are required for Google Colab, but may work on other systems too:
[1]:
# import locale
# locale.getpreferredencoding = lambda: "UTF-8"
# ! pip uninstall torch torchaudio torchvision numpy -y
# ! uv pip install janus-core[all] data-tutorials torch==2.5.1 --system
# get_ipython().kernel.do_shutdown(restart=True)
[2]:
from ase.io import read
from ase.optimize import FIRE
from weas_widget import WeasWidget
from data_tutorials.data import get_data
from janus_core.calculations.single_point import SinglePoint
from janus_core.calculations.geom_opt import GeomOpt
Use data_tutorials to get the data required for this tutorial:
[3]:
get_data(
url="https://raw.githubusercontent.com/stfc/janus-tutorials/main/data/",
filename=["NaCl-deformed.cif"],
folder="data",
)
try to download NaCl-deformed.cif from https://raw.githubusercontent.com/stfc/janus-tutorials/main/data/ and save it in data/NaCl-deformed.cif
saved in data/NaCl-deformed.cif
Prepare for optimization of a deformed salt structure¶
[4]:
NaCl = read("data/NaCl-deformed.cif")
NaCl.wrap()
v=WeasWidget()
v.from_ase(NaCl)
v.avr.model_style = 1
v.avr.show_hydrogen_bonds = True
v
[4]:
[5]:
sp_mace = SinglePoint(
struct=NaCl.copy(),
arch="mace_mp",
device="cpu",
model_path="small",
calc_kwargs={"default_dtype": "float64"},
properties="energy",
)
init_energy = sp_mace.run()["energy"]
/home/runner/work/janus-core/janus-core/.venv/lib/python3.12/site-packages/e3nn/o3/_wigner.py:10: UserWarning: Environment variable TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD detected, since the`weights_only` argument was not explicitly passed to `torch.load`, forcing weights_only=False.
_Jd, _W3j_flat, _W3j_indices = torch.load(os.path.join(os.path.dirname(__file__), 'constants.pt'))
cuequivariance or cuequivariance_torch is not available. Cuequivariance acceleration will be disabled.
Using Materials Project MACE for MACECalculator with /home/runner/.cache/mace/20231210mace128L0_energy_epoch249model
Using float64 for MACECalculator, which is slower but more accurate. Recommended for geometry optimization.
/home/runner/work/janus-core/janus-core/.venv/lib/python3.12/site-packages/mace/calculators/mace.py:139: UserWarning: Environment variable TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD detected, since the`weights_only` argument was not explicitly passed to `torch.load`, forcing weights_only=False.
torch.load(f=model_path, map_location=device)
To optimize only the atomic positions and not the cell, set filter_func = None:
[6]:
optimized_NaCl = GeomOpt(
struct=sp_mace.struct,
fmax=0.001,
filter_func=None,
)
optimized_NaCl.run()
v=WeasWidget()
v.from_ase(optimized_NaCl.struct)
v.avr.model_style = 1
v.avr.show_hydrogen_bonds = True
v
Step Time Energy fmax
LBFGS: 0 17:11:31 -27.020331 0.022349
LBFGS: 1 17:11:31 -27.020366 0.021806
LBFGS: 2 17:11:31 -27.021393 0.020842
LBFGS: 3 17:11:32 -27.021424 0.020436
LBFGS: 4 17:11:32 -27.021890 0.006714
LBFGS: 5 17:11:32 -27.021928 0.002554
LBFGS: 6 17:11:32 -27.021932 0.001476
LBFGS: 7 17:11:32 -27.021932 0.001378
LBFGS: 8 17:11:32 -27.021932 0.001022
LBFGS: 9 17:11:32 -27.021933 0.000560
[6]:
Check energy has been lowered, and cell is unchanged:
[7]:
print(f"Initial cell: {NaCl.cell.cellpar()}")
print(f"Initial energy: {init_energy}")
print(f"Final cell: {optimized_NaCl.struct.cell.cellpar()}")
print(f"Final energy: {optimized_NaCl.struct.get_potential_energy()}")
Initial cell: [ 5.62 5.62 5.62 89. 90. 90. ]
Initial energy: -27.020331256158105
Final cell: [ 5.62 5.62 5.62 89. 90. 90. ]
Final energy: -27.02193251719944
Optimizing cell vectors and atomic positions¶
Setting filter_kwargs = {"hydrostatic_strain": True} allows the cell lengths to be changed, in addition to atomic positions, but cell angles remain fixed:
[8]:
optimized_NaCl_lengths = GeomOpt(
struct=NaCl.copy(),
arch="mace_mp",
device="cpu",
model_path="small",
calc_kwargs={"default_dtype": "float64"},
fmax=0.001,
filter_kwargs={"hydrostatic_strain": True},
)
optimized_NaCl_lengths.run()
Using Materials Project MACE for MACECalculator with /home/runner/.cache/mace/20231210mace128L0_energy_epoch249model
Using float64 for MACECalculator, which is slower but more accurate. Recommended for geometry optimization.
/home/runner/work/janus-core/janus-core/.venv/lib/python3.12/site-packages/mace/calculators/mace.py:139: UserWarning: Environment variable TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD detected, since the`weights_only` argument was not explicitly passed to `torch.load`, forcing weights_only=False.
torch.load(f=model_path, map_location=device)
Step Time Energy fmax
LBFGS: 0 17:11:32 -27.020331 0.152574
LBFGS: 1 17:11:32 -27.021352 0.149055
LBFGS: 2 17:11:33 -27.043661 0.016830
LBFGS: 3 17:11:33 -27.043729 0.016632
LBFGS: 4 17:11:33 -27.043750 0.016155
LBFGS: 5 17:11:33 -27.043851 0.012843
LBFGS: 6 17:11:33 -27.043969 0.009577
LBFGS: 7 17:11:34 -27.044084 0.009987
LBFGS: 8 17:11:34 -27.044129 0.011949
LBFGS: 9 17:11:34 -27.044146 0.011951
LBFGS: 10 17:11:34 -27.044164 0.011168
LBFGS: 11 17:11:35 -27.044202 0.008924
LBFGS: 12 17:11:35 -27.044258 0.005022
LBFGS: 13 17:11:35 -27.044305 0.003603
LBFGS: 14 17:11:35 -27.044321 0.001500
LBFGS: 15 17:11:35 -27.044322 0.001258
LBFGS: 16 17:11:35 -27.044323 0.001136
LBFGS: 17 17:11:36 -27.044323 0.000967
Check energy has been lowered, and cell lengths have been updated, but angles remain unchanged:
[9]:
print(f"Initial cell: {NaCl.cell.cellpar()}")
print(f"Initial energy: {init_energy}")
print(f"Final cell: {optimized_NaCl_lengths.struct.cell.cellpar()}")
print(f"Final energy: {optimized_NaCl_lengths.struct.get_potential_energy()}")
Initial cell: [ 5.62 5.62 5.62 89. 90. 90. ]
Initial energy: -27.020331256158105
Final cell: [ 5.69023326 5.69023326 5.69023326 89. 90. 90. ]
Final energy: -27.044322814633645
Optimizing at constant pressure and volume¶
Calculations can also be run at a fixed pressure and volume, by setting filter_kwargs = {"scalar_pressure": x, "constant_volume": True}
By default, both the cell lengths and angles will be optimized, in addition to the atomic positions.
We can also set the optimizer function and filter function used, either by passing the function itself (e.g. FIRE) or passing the name of the ASE function (e.g. "ExpCellFilter"):
[10]:
optimized_NaCl_pressure = GeomOpt(
struct=NaCl.copy(),
arch="mace_mp",
device="cpu",
model_path="small",
calc_kwargs={"default_dtype": "float64"},
fmax=0.01,
filter_kwargs={"scalar_pressure": 0.05, "constant_volume": True},
optimizer=FIRE,
filter_func="ExpCellFilter",
)
optimized_NaCl_pressure.run()
Using Materials Project MACE for MACECalculator with /home/runner/.cache/mace/20231210mace128L0_energy_epoch249model
Using float64 for MACECalculator, which is slower but more accurate. Recommended for geometry optimization.
/home/runner/work/janus-core/janus-core/.venv/lib/python3.12/site-packages/mace/calculators/mace.py:139: UserWarning: Environment variable TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD detected, since the`weights_only` argument was not explicitly passed to `torch.load`, forcing weights_only=False.
torch.load(f=model_path, map_location=device)
Step Time Energy fmax
FIRE: 0 17:11:36 -26.964945 0.223903
FIRE: 1 17:11:36 -26.965836 0.166226
FIRE: 2 17:11:36 -26.966752 0.066309
FIRE: 3 17:11:37 -26.966814 0.049491
FIRE: 4 17:11:37 -26.966834 0.046327
FIRE: 5 17:11:37 -26.966871 0.040207
FIRE: 6 17:11:37 -26.966919 0.031551
FIRE: 7 17:11:37 -26.966970 0.020992
FIRE: 8 17:11:37 -26.967020 0.018673
FIRE: 9 17:11:38 -26.967064 0.018201
FIRE: 10 17:11:38 -26.967102 0.017653
FIRE: 11 17:11:38 -26.967143 0.020467
FIRE: 12 17:11:38 -26.967191 0.025421
FIRE: 13 17:11:38 -26.967254 0.025704
FIRE: 14 17:11:39 -26.967332 0.020370
FIRE: 15 17:11:39 -26.967420 0.012820
FIRE: 16 17:11:39 -26.967508 0.012436
FIRE: 17 17:11:39 -26.967596 0.016894
FIRE: 18 17:11:39 -26.967708 0.020040
FIRE: 19 17:11:40 -26.967853 0.011632
FIRE: 20 17:11:40 -26.968001 0.010725
FIRE: 21 17:11:40 -26.968130 0.017548
FIRE: 22 17:11:40 -26.968250 0.007994
Check cell lengths and angles have both been updated:
[11]:
print(f"Initial cell: {NaCl.cell.cellpar()}")
print(f"Initial energy: {init_energy}")
print(f"Final cell: {optimized_NaCl_pressure.struct.cell.cellpar()}")
print(f"Final energy: {optimized_NaCl_pressure.struct.get_potential_energy()}")
Initial cell: [ 5.62 5.62 5.62 89. 90. 90. ]
Initial energy: -27.020331256158105
Final cell: [ 5.61945859 5.61965116 5.61963753 89.96406884 90. 90. ]
Final energy: -27.023632508132863
Comparing MACE to CHGNet and SevenNet¶
[12]:
optimized_NaCl_mace = GeomOpt(
struct=NaCl.copy(),
arch="mace_mp",
device="cpu",
model_path="small",
calc_kwargs={"default_dtype": "float64"},
fmax=0.01,
)
optimized_NaCl_mace.run()
optimized_NaCl_chgnet = GeomOpt(
struct=NaCl.copy(),
arch="chgnet",
device="cpu",
calc_kwargs={"default_dtype": "float64"},
fmax=0.01,
)
optimized_NaCl_chgnet.run()
optimized_NaCl_sevennet = GeomOpt(
struct=NaCl.copy(),
arch="sevennet",
device="cpu",
calc_kwargs={"default_dtype": "float64"},
fmax=0.01,
)
optimized_NaCl_sevennet.run()
Using Materials Project MACE for MACECalculator with /home/runner/.cache/mace/20231210mace128L0_energy_epoch249model
Using float64 for MACECalculator, which is slower but more accurate. Recommended for geometry optimization.
/home/runner/work/janus-core/janus-core/.venv/lib/python3.12/site-packages/mace/calculators/mace.py:139: UserWarning: Environment variable TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD detected, since the`weights_only` argument was not explicitly passed to `torch.load`, forcing weights_only=False.
torch.load(f=model_path, map_location=device)
Step Time Energy fmax
LBFGS: 0 17:11:41 -27.020331 0.156191
LBFGS: 1 17:11:41 -27.021374 0.152692
LBFGS: 2 17:11:41 -27.044614 0.021866
LBFGS: 3 17:11:41 -27.044701 0.021533
LBFGS: 4 17:11:41 -27.044744 0.021790
LBFGS: 5 17:11:42 -27.044926 0.023049
LBFGS: 6 17:11:42 -27.045194 0.023866
LBFGS: 7 17:11:42 -27.045540 0.020973
LBFGS: 8 17:11:42 -27.045741 0.016076
LBFGS: 9 17:11:42 -27.045813 0.016967
LBFGS: 10 17:11:43 -27.045858 0.016120
LBFGS: 11 17:11:43 -27.045941 0.013383
LBFGS: 12 17:11:43 -27.046078 0.011105
LBFGS: 13 17:11:43 -27.046224 0.009142
CHGNet v0.3.0 initialized with 412,525 parameters
CHGNet will run on cpu
Step Time Energy fmax
LBFGS: 0 17:11:47 -29.318993 0.138351
LBFGS: 1 17:11:47 -29.319818 0.136138
LBFGS: 2 17:11:47 -29.342087 0.035508
LBFGS: 3 17:11:48 -29.342720 0.014143
LBFGS: 4 17:11:48 -29.342743 0.014978
LBFGS: 5 17:11:48 -29.342773 0.016334
LBFGS: 6 17:11:48 -29.342850 0.018683
LBFGS: 7 17:11:48 -29.343006 0.021501
LBFGS: 8 17:11:48 -29.343246 0.021310
LBFGS: 9 17:11:49 -29.343517 0.014323
LBFGS: 10 17:11:49 -29.343639 0.008049
Step Time Energy fmax
LBFGS: 0 17:11:52 -27.047630 0.144920
LBFGS: 1 17:11:52 -27.048546 0.141888
LBFGS: 2 17:11:52 -27.070152 0.020650
LBFGS: 3 17:11:52 -27.070215 0.020472
LBFGS: 4 17:11:53 -27.070261 0.020717
LBFGS: 5 17:11:53 -27.070412 0.021701
LBFGS: 6 17:11:53 -27.070656 0.022366
LBFGS: 7 17:11:53 -27.070944 0.019760
LBFGS: 8 17:11:54 -27.071106 0.014357
LBFGS: 9 17:11:54 -27.071163 0.014336
LBFGS: 10 17:11:54 -27.071192 0.013783
LBFGS: 11 17:11:54 -27.071259 0.011845
LBFGS: 12 17:11:55 -27.071373 0.010935
LBFGS: 13 17:11:55 -27.071510 0.009534
[13]:
print(f"Initial energy: {init_energy}")
print(f"Final energy (MACE): {optimized_NaCl_mace.struct.get_potential_energy()}")
print(f"Final energy (CHGNet): {optimized_NaCl_chgnet.struct.get_potential_energy()}")
print(f"Final energy (SevenNet): {optimized_NaCl_sevennet.struct.get_potential_energy()}")
Initial energy: -27.020331256158105
Final energy (MACE): -27.046223650391827
Final energy (CHGNet): -29.343639373779297
Final energy (SevenNet): -27.071510314941406