Geometry Optimization¶
Set up environment (optional)¶
These steps are required for Google Colab, but may work on other systems too:
[22]:
# 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)
Command-line help and options¶
As with janus singlepoint, we can check the options for geometry optimisation:
[1]:
! janus geomopt --help
Usage: janus geomopt [OPTIONS]
Perform geometry optimization and save optimized structure to file.
╭─ Options ────────────────────────────────────────────────────────────────────╮
│ * --struct PATH Path of │
│ structure to │
│ simulate. │
│ [default: None] │
│ [required] │
│ --optimizer TEXT Name of ASE │
│ optimizer │
│ function to use. │
│ [default: LBFGS] │
│ --fmax FLOAT Maximum force │
│ for convergence, │
│ in eV/Å. │
│ [default: 0.1] │
│ --steps INTEGER Maximum number │
│ of optimization │
│ steps. │
│ [default: 1000] │
│ --arch [mace|mace_mp|ma MLIP │
│ ce_off|m3gnet|ch architecture to │
│ gnet|alignn|seve use for │
│ nnet|nequip|dpa3 calculations. │
│ |orb|mattersim] [default: │
│ mace_mp] │
│ --device [cpu|cuda|mps|xp Device to run │
│ u] calculations on. │
│ [default: cpu] │
│ --model-path TEXT Path to MLIP │
│ model. │
│ [default: None] │
│ --opt-cell-leng… --no-opt-cell-… Optimize cell │
│ vectors, as well │
│ as atomic │
│ positions. │
│ [default: │
│ no-opt-cell-len… │
│ --opt-cell-fully --no-opt-cell-… Fully optimize │
│ the cell │
│ vectors, angles, │
│ and atomic │
│ positions. │
│ [default: │
│ no-opt-cell-ful… │
│ --filter-func TEXT Name of ASE │
│ filter/constrai… │
│ function to use. │
│ If using │
│ --opt-cell-leng… │
│ or │
│ --opt-cell-full… │
│ defaults to │
│ `FrechetCellFil… │
│ if available, │
│ otherwise │
│ `ExpCellFilter`. │
│ [default: None] │
│ --pressure FLOAT Scalar pressure │
│ when optimizing │
│ cell geometry, │
│ in GPa. │
│ [default: 0.0] │
│ --symmetrize --no-symmetrize Whether to │
│ refine symmetry │
│ after geometry │
│ optimization. │
│ [default: │
│ no-symmetrize] │
│ --symmetry-tole… FLOAT Atom │
│ displacement │
│ tolerance for │
│ spglib symmetry │
│ determination, │
│ in Å. │
│ [default: 0.001] │
│ --file-prefix PATH Prefix for │
│ output files, │
│ including │
│ directories. │
│ Default │
│ directory is │
│ ./janus_results, │
│ and default │
│ filename prefix │
│ is inferred from │
│ the input │
│ stucture │
│ filename. │
│ [default: None] │
│ --out PATH Path to save │
│ optimized │
│ structure. │
│ Default is │
│ inferred │
│ `file_prefix`. │
│ [default: None] │
│ --write-traj --no-write-traj Whether to save │
│ a trajectory │
│ file of │
│ optimization │
│ frames. If │
│ traj_kwargs["fi… │
│ is not │
│ specified, it is │
│ inferred from │
│ `file_prefix`. │
│ [default: │
│ no-write-traj] │
│ --read-kwargs DICT Keyword │
│ arguments to │
│ pass to │
│ ase.io.read. │
│ Must be passed │
│ as a dictionary │
│ wrapped in │
│ quotes, e.g. │
│ "{'key': │
│ value}". By │
│ default, │
│ read_kwargs['in… │
│ = -1, so only │
│ the last │
│ structure is │
│ read. │
│ [default: None] │
│ --calc-kwargs DICT Keyword │
│ arguments to │
│ pass to selected │
│ calculator. Must │
│ be passed as a │
│ dictionary │
│ wrapped in │
│ quotes, e.g. │
│ "{'key': │
│ value}". For the │
│ default │
│ architecture │
│ ('mace_mp'), │
│ "{'model': │
│ 'small'}" is set │
│ unless │
│ overwritten. │
│ [default: None] │
│ --minimize-kwar… DICT Keyword │
│ arguments to │
│ pass to │
│ optimizer. Must │
│ be passed as a │
│ dictionary │
│ wrapped in │
│ quotes, e.g. │
│ "{'key': │
│ value}". │
│ [default: None] │
│ --write-kwargs DICT Keyword │
│ arguments to │
│ pass to │
│ ase.io.write │
│ when saving │
│ results. Must be │
│ passed as a │
│ dictionary │
│ wrapped in │
│ quotes, e.g. │
│ "{'key': │
│ value}". │
│ [default: None] │
│ --log PATH Path to save │
│ logs to. Default │
│ is inferred from │
│ `file_prefix` │
│ [default: None] │
│ --tracker --no-tracker Whether to save │
│ carbon emissions │
│ of calculation │
│ [default: │
│ tracker] │
│ --summary PATH Path to save │
│ summary of │
│ inputs, │
│ start/end time, │
│ and carbon │
│ emissions. │
│ Default is │
│ inferred from │
│ `file_prefix`. │
│ [default: None] │
│ --config TEXT Configuration │
│ file. │
│ --help Show this │
│ message and │
│ exit. │
╰──────────────────────────────────────────────────────────────────────────────╯
Running geometry optimisation calculation¶
First, we’ll build a structure to optimise, as we did for single point calculations, but add in a deformation:
[2]:
from pathlib import Path
from ase.build import bulk
from ase.io import write
from weas_widget import WeasWidget
Path("data").mkdir(exist_ok=True)
NaCl = bulk("NaCl", "rocksalt", a=5.63, cubic=True)
NaCl[0].position = [1.5, 1.5, 1.5]
write("../data/NaCl-deformed.xyz", NaCl)
v=WeasWidget()
v.from_ase(NaCl)
v.avr.model_style = 1
v.avr.show_hydrogen_bonds = True
v
[2]:
Now we can optimise the geometry of this structure in a similar manner to running janus singlepoint:
[3]:
! janus geomopt --struct ../data/NaCl-deformed.xyz --no-tracker
/Users/elliottkasoar/Documents/PSDI/janus-core/.venv/lib/python3.12/site-packages/e3nn/o3/_wigner.py:10: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
_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 /Users/elliottkasoar/.cache/mace/20231210mace128L0_energy_epoch249model
Using float64 for MACECalculator, which is slower but more accurate. Recommended for geometry optimization.
/Users/elliottkasoar/Documents/PSDI/janus-core/.venv/lib/python3.12/site-packages/mace/calculators/mace.py:139: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
torch.load(f=model_path, map_location=device)
Step Time Energy fmax
LBFGS: 0 15:51:20 -24.377063 2.732337
LBFGS: 1 15:51:21 -24.576368 2.152513
LBFGS: 2 15:51:21 -25.093690 0.949845
LBFGS: 3 15:51:21 -25.290212 0.685740
LBFGS: 4 15:51:21 -25.403788 0.951695
LBFGS: 5 15:51:21 -25.465213 1.066671
LBFGS: 6 15:51:21 -25.566446 1.095941
LBFGS: 7 15:51:21 -25.695050 0.931041
LBFGS: 8 15:51:21 -25.856683 0.857034
LBFGS: 9 15:51:21 -26.109079 0.811033
LBFGS: 10 15:51:21 -26.298823 0.545193
LBFGS: 11 15:51:21 -26.353032 0.656387
LBFGS: 12 15:51:21 -26.383904 0.710488
LBFGS: 13 15:51:21 -26.485793 0.790107
LBFGS: 14 15:51:21 -26.643853 0.768081
LBFGS: 15 15:51:21 -26.751857 0.477356
LBFGS: 16 15:51:21 -26.792940 0.329784
LBFGS: 17 15:51:21 -26.815123 0.343673
LBFGS: 18 15:51:21 -26.829194 0.334430
LBFGS: 19 15:51:21 -26.845943 0.317718
LBFGS: 20 15:51:21 -26.912307 0.249075
LBFGS: 21 15:51:21 -26.975682 0.229256
LBFGS: 22 15:51:21 -27.018210 0.150833
LBFGS: 23 15:51:21 -27.023510 0.093603
We can also change the optimisation function used, and specify an even lower force convergence criteria, fmax (the maximum force on all individual atoms):
Tip: The optimizer must be a class defined in ASE.
[4]:
! janus geomopt --struct ../data/NaCl-deformed.xyz --optimizer FIRE --fmax 0.005 --no-tracker
/Users/elliottkasoar/Documents/PSDI/janus-core/.venv/lib/python3.12/site-packages/e3nn/o3/_wigner.py:10: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
_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 /Users/elliottkasoar/.cache/mace/20231210mace128L0_energy_epoch249model
Using float64 for MACECalculator, which is slower but more accurate. Recommended for geometry optimization.
/Users/elliottkasoar/Documents/PSDI/janus-core/.venv/lib/python3.12/site-packages/mace/calculators/mace.py:139: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
torch.load(f=model_path, map_location=device)
Step Time Energy fmax
FIRE: 0 15:51:24 -24.377063 2.732337
FIRE: 1 15:51:24 -24.521143 2.311307
FIRE: 2 15:51:24 -24.737816 1.719928
FIRE: 3 15:51:24 -24.954552 1.216180
FIRE: 4 15:51:24 -25.138250 0.893003
FIRE: 5 15:51:24 -25.283693 0.712183
FIRE: 6 15:51:24 -25.398292 0.725931
FIRE: 7 15:51:24 -25.495898 0.776106
FIRE: 8 15:51:24 -25.600795 0.738520
FIRE: 9 15:51:24 -25.727081 0.624992
FIRE: 10 15:51:24 -25.886706 0.633155
FIRE: 11 15:51:24 -26.063737 0.609225
FIRE: 12 15:51:24 -26.237721 0.564647
FIRE: 13 15:51:24 -26.390734 0.525758
FIRE: 14 15:51:24 -26.511572 0.531927
FIRE: 15 15:51:25 -26.597920 0.696601
FIRE: 16 15:51:25 -26.653461 0.790166
FIRE: 17 15:51:25 -26.686575 0.773073
FIRE: 18 15:51:25 -26.705939 0.735325
FIRE: 19 15:51:25 -26.741131 0.662023
FIRE: 20 15:51:25 -26.785861 0.557699
FIRE: 21 15:51:25 -26.832526 0.429577
FIRE: 22 15:51:25 -26.874093 0.288174
FIRE: 23 15:51:25 -26.906094 0.205971
FIRE: 24 15:51:25 -26.928149 0.249613
FIRE: 25 15:51:25 -26.945064 0.282363
FIRE: 26 15:51:25 -26.961246 0.296619
FIRE: 27 15:51:25 -26.981117 0.283688
FIRE: 28 15:51:25 -27.004699 0.232056
FIRE: 29 15:51:25 -27.023875 0.129160
FIRE: 30 15:51:25 -27.023665 0.065418
FIRE: 31 15:51:25 -27.023932 0.063243
FIRE: 32 15:51:25 -27.024426 0.058997
FIRE: 33 15:51:25 -27.025071 0.052890
FIRE: 34 15:51:25 -27.025778 0.045230
FIRE: 35 15:51:25 -27.026455 0.036435
FIRE: 36 15:51:25 -27.027033 0.030872
FIRE: 37 15:51:25 -27.027485 0.033525
FIRE: 38 15:51:25 -27.027863 0.046407
FIRE: 39 15:51:25 -27.028209 0.055288
FIRE: 40 15:51:25 -27.028598 0.057928
FIRE: 41 15:51:25 -27.029081 0.052454
FIRE: 42 15:51:25 -27.029611 0.037632
FIRE: 43 15:51:25 -27.029996 0.013477
FIRE: 44 15:51:25 -27.029974 0.017129
FIRE: 45 15:51:25 -27.029982 0.016447
FIRE: 46 15:51:25 -27.029995 0.015113
FIRE: 47 15:51:25 -27.030013 0.013191
FIRE: 48 15:51:25 -27.030031 0.010771
FIRE: 49 15:51:25 -27.030046 0.007969
FIRE: 50 15:51:25 -27.030057 0.004932
As with single point calculations, this saves a results file, corresponding to the optimised structure, as well as a summary and log:
[5]:
! ls janus_results/NaCl-deformed*
janus_results/NaCl-deformed-geomopt-log.yml
janus_results/NaCl-deformed-geomopt-summary.yml
janus_results/NaCl-deformed-opt.extxyz
We can now see the optimised structure:
[6]:
from ase.io import read
from weas_widget import WeasWidget
traj = read("janus_results/NaCl-deformed-opt.extxyz")
v=WeasWidget()
v.from_ase(traj)
v.avr.model_style = 1
v.avr.show_hydrogen_bonds = True
v
[6]:
Trajectories¶
We can also save the trajectory during optimisation:
[7]:
! janus geomopt --struct ../data/NaCl-deformed.xyz --write-traj --no-tracker
/Users/elliottkasoar/Documents/PSDI/janus-core/.venv/lib/python3.12/site-packages/e3nn/o3/_wigner.py:10: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
_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 /Users/elliottkasoar/.cache/mace/20231210mace128L0_energy_epoch249model
Using float64 for MACECalculator, which is slower but more accurate. Recommended for geometry optimization.
/Users/elliottkasoar/Documents/PSDI/janus-core/.venv/lib/python3.12/site-packages/mace/calculators/mace.py:139: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
torch.load(f=model_path, map_location=device)
Step Time Energy fmax
LBFGS: 0 15:51:28 -24.377063 2.732337
LBFGS: 1 15:51:29 -24.576368 2.152513
LBFGS: 2 15:51:29 -25.093690 0.949845
LBFGS: 3 15:51:29 -25.290212 0.685740
LBFGS: 4 15:51:29 -25.403788 0.951695
LBFGS: 5 15:51:29 -25.465213 1.066671
LBFGS: 6 15:51:29 -25.566446 1.095941
LBFGS: 7 15:51:29 -25.695050 0.931041
LBFGS: 8 15:51:29 -25.856683 0.857034
LBFGS: 9 15:51:29 -26.109079 0.811033
LBFGS: 10 15:51:29 -26.298823 0.545193
LBFGS: 11 15:51:29 -26.353032 0.656387
LBFGS: 12 15:51:29 -26.383904 0.710488
LBFGS: 13 15:51:29 -26.485793 0.790107
LBFGS: 14 15:51:29 -26.643853 0.768081
LBFGS: 15 15:51:29 -26.751857 0.477356
LBFGS: 16 15:51:29 -26.792940 0.329784
LBFGS: 17 15:51:29 -26.815123 0.343673
LBFGS: 18 15:51:29 -26.829194 0.334430
LBFGS: 19 15:51:29 -26.845943 0.317718
LBFGS: 20 15:51:29 -26.912307 0.249075
LBFGS: 21 15:51:29 -26.975682 0.229256
LBFGS: 22 15:51:29 -27.018210 0.150833
LBFGS: 23 15:51:29 -27.023510 0.093603
This creates an additional file, janus_results/NaCl-deformed-traj.extxyz:
[8]:
! ls janus_results/NaCl-deformed*
janus_results/NaCl-deformed-geomopt-log.yml
janus_results/NaCl-deformed-geomopt-summary.yml
janus_results/NaCl-deformed-opt.extxyz
janus_results/NaCl-deformed-traj.extxyz
This allows us to visualise the optimisation:
[9]:
from ase.io import read
from weas_widget import WeasWidget
traj = read("janus_results/NaCl-deformed-traj.extxyz", index=":")
v=WeasWidget()
v.from_ase(traj)
v.avr.model_style = 1
v.avr.show_hydrogen_bonds = True
v
[9]:
Cell optimisation¶
We can also choose to modify the cell vectors during the optimisation. To allow only the cell lengths to change, we can run:
[10]:
! janus geomopt --struct ../data/NaCl-deformed.xyz --write-traj --opt-cell-lengths --no-tracker
/Users/elliottkasoar/Documents/PSDI/janus-core/.venv/lib/python3.12/site-packages/e3nn/o3/_wigner.py:10: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
_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 /Users/elliottkasoar/.cache/mace/20231210mace128L0_energy_epoch249model
Using float64 for MACECalculator, which is slower but more accurate. Recommended for geometry optimization.
/Users/elliottkasoar/Documents/PSDI/janus-core/.venv/lib/python3.12/site-packages/mace/calculators/mace.py:139: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
torch.load(f=model_path, map_location=device)
Step Time Energy fmax
LBFGS: 0 15:51:32 -24.377063 2.732337
LBFGS: 1 15:51:32 -24.593123 2.143240
LBFGS: 2 15:51:32 -25.178180 0.955170
LBFGS: 3 15:51:32 -25.445502 0.670912
LBFGS: 4 15:51:33 -25.641907 0.919610
LBFGS: 5 15:51:33 -25.735611 1.036999
LBFGS: 6 15:51:33 -25.816918 1.015167
LBFGS: 7 15:51:33 -25.937123 0.774473
LBFGS: 8 15:51:33 -26.075553 0.561448
LBFGS: 9 15:51:33 -26.194652 0.423668
LBFGS: 10 15:51:33 -26.290628 0.391005
LBFGS: 11 15:51:33 -26.323319 0.439570
LBFGS: 12 15:51:33 -26.337274 0.452690
LBFGS: 13 15:51:33 -26.372755 0.457902
LBFGS: 14 15:51:33 -26.436676 0.640085
LBFGS: 15 15:51:33 -26.522182 0.786499
LBFGS: 16 15:51:33 -26.636096 0.839171
LBFGS: 17 15:51:33 -26.785166 0.746960
LBFGS: 18 15:51:33 -26.953112 0.390060
LBFGS: 19 15:51:33 -27.019414 0.155454
LBFGS: 20 15:51:34 -27.024708 0.104189
LBFGS: 21 15:51:34 -27.025735 0.100499
LBFGS: 22 15:51:34 -27.040218 0.057352
As before, we can visualise this trajectory:
[11]:
from ase.io import read
from weas_widget import WeasWidget
traj = read("janus_results/NaCl-deformed-traj.extxyz", index=":")
v=WeasWidget()
v.from_ase(traj)
v.avr.model_style = 1
v.avr.show_hydrogen_bonds = True
v
[11]:
We can also allow the cell angles to change:
[12]:
! janus geomopt --struct ../data/NaCl-deformed.xyz --write-traj --opt-cell-fully --no-tracker
/Users/elliottkasoar/Documents/PSDI/janus-core/.venv/lib/python3.12/site-packages/e3nn/o3/_wigner.py:10: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
_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 /Users/elliottkasoar/.cache/mace/20231210mace128L0_energy_epoch249model
Using float64 for MACECalculator, which is slower but more accurate. Recommended for geometry optimization.
/Users/elliottkasoar/Documents/PSDI/janus-core/.venv/lib/python3.12/site-packages/mace/calculators/mace.py:139: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
torch.load(f=model_path, map_location=device)
Step Time Energy fmax
LBFGS: 0 15:51:36 -24.377063 2.732337
LBFGS: 1 15:51:36 -24.593440 2.145517
LBFGS: 2 15:51:37 -25.187807 0.956555
LBFGS: 3 15:51:37 -25.462335 0.675018
LBFGS: 4 15:51:37 -25.663853 0.932226
LBFGS: 5 15:51:37 -25.757785 1.056063
LBFGS: 6 15:51:37 -25.841303 1.038396
LBFGS: 7 15:51:37 -25.968290 0.806821
LBFGS: 8 15:51:37 -26.115858 0.539264
LBFGS: 9 15:51:37 -26.230485 0.404401
LBFGS: 10 15:51:37 -26.323082 0.351277
LBFGS: 11 15:51:37 -26.354901 0.404205
LBFGS: 12 15:51:37 -26.366593 0.418519
LBFGS: 13 15:51:37 -26.402172 0.433457
LBFGS: 14 15:51:37 -26.468589 0.596828
LBFGS: 15 15:51:37 -26.555830 0.733795
LBFGS: 16 15:51:37 -26.664503 0.782841
LBFGS: 17 15:51:38 -26.800773 0.695400
LBFGS: 18 15:51:38 -26.945555 0.318072
LBFGS: 19 15:51:38 -26.988824 0.103903
LBFGS: 20 15:51:38 -26.991370 0.096582
Visualising this:
[13]:
from ase.io import read
from weas_widget import WeasWidget
traj = read("janus_results/NaCl-deformed-traj.extxyz", index=":")
v=WeasWidget()
v.from_ase(traj)
v.avr.model_style = 1
v.avr.show_hydrogen_bonds = True
v
[13]:
Setting the unit cell filter¶
Cell optimisation is carried out by appling an ASE filter to the structure. By default, this is the FrechetCellFilter, but you may wish to apply others, such as the ExpCellFilter. This is passed as a string, and must correspond to a class defined in ASE.
Key word arguments can also be passed to these filters. For example, we can maintain constant volume using the following configuration file:
[14]:
%%writefile geomopt_config_1.yml
struct: ../data/NaCl-deformed.xyz
opt_cell_fully: True
filter_func: ExpCellFilter
minimize_kwargs:
filter_kwargs:
constant_volume: True
tracker: False
Writing geomopt_config_1.yml
Tip: This is equivalent to:
–struct ../data/NaCl-deformed.xyz –opt-cell-fully –filter-func ExpCellFilter –minimize-kwargs “{‘filter_kwargs’: {‘constant_volume’ : True}” –no-tracker
[15]:
! janus geomopt --config geomopt_config_1.yml
/Users/elliottkasoar/Documents/PSDI/janus-core/.venv/lib/python3.12/site-packages/e3nn/o3/_wigner.py:10: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
_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 /Users/elliottkasoar/.cache/mace/20231210mace128L0_energy_epoch249model
Using float64 for MACECalculator, which is slower but more accurate. Recommended for geometry optimization.
/Users/elliottkasoar/Documents/PSDI/janus-core/.venv/lib/python3.12/site-packages/mace/calculators/mace.py:139: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
torch.load(f=model_path, map_location=device)
Step Time Energy fmax
LBFGS: 0 15:51:40 -24.377063 2.732337
LBFGS: 1 15:51:41 -24.592435 2.306591
LBFGS: 2 15:51:41 -25.267905 1.323804
LBFGS: 3 15:51:41 -25.449977 1.003470
LBFGS: 4 15:51:41 -25.521347 0.823847
LBFGS: 5 15:51:41 -25.573640 0.866332
LBFGS: 6 15:51:41 -25.673315 1.004378
LBFGS: 7 15:51:41 -25.811015 1.004781
LBFGS: 8 15:51:41 -25.975638 0.889642
LBFGS: 9 15:51:41 -26.150934 0.689834
LBFGS: 10 15:51:41 -26.319557 0.690018
LBFGS: 11 15:51:41 -26.467568 0.708275
LBFGS: 12 15:51:41 -26.574768 0.722124
LBFGS: 13 15:51:41 -26.637370 0.637261
LBFGS: 14 15:51:41 -26.708598 0.646450
LBFGS: 15 15:51:42 -26.781501 0.681819
LBFGS: 16 15:51:42 -26.844998 0.708926
LBFGS: 17 15:51:42 -26.874141 0.642122
LBFGS: 18 15:51:42 -26.894946 0.530522
LBFGS: 19 15:51:42 -26.919861 0.583009
LBFGS: 20 15:51:42 -26.954422 0.529294
LBFGS: 21 15:51:42 -26.977868 0.252808
LBFGS: 22 15:51:42 -26.985391 0.153658
LBFGS: 23 15:51:42 -26.987735 0.137245
LBFGS: 24 15:51:42 -26.991006 0.146929
LBFGS: 25 15:51:42 -26.998079 0.227195
LBFGS: 26 15:51:42 -27.010168 0.266027
LBFGS: 27 15:51:42 -27.022858 0.207394
LBFGS: 28 15:51:42 -27.029032 0.073577
Visualising this:
[16]:
from ase.io import read
from weas_widget import WeasWidget
traj = read("janus_results/NaCl-deformed-traj.extxyz", index=":")
v=WeasWidget()
v.from_ase(traj)
v.avr.model_style = 1
v.avr.show_hydrogen_bonds = True
v
[16]:
Constant pressure and symmetry refinement¶
We can also choose to optimise at a fixed pressure (in GPa), and refine the symmetry of the final structure:
[17]:
! janus geomopt --struct ../data/NaCl-deformed.xyz --write-traj --pressure 10 --opt-cell-fully --symmetrize --no-tracker
/Users/elliottkasoar/Documents/PSDI/janus-core/.venv/lib/python3.12/site-packages/e3nn/o3/_wigner.py:10: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
_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 /Users/elliottkasoar/.cache/mace/20231210mace128L0_energy_epoch249model
Using float64 for MACECalculator, which is slower but more accurate. Recommended for geometry optimization.
/Users/elliottkasoar/Documents/PSDI/janus-core/.venv/lib/python3.12/site-packages/mace/calculators/mace.py:139: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
torch.load(f=model_path, map_location=device)
Step Time Energy fmax
LBFGS: 0 15:51:45 -13.238868 2.732337
LBFGS: 1 15:51:45 -13.462923 2.165368
LBFGS: 2 15:51:45 -14.202196 0.816332
LBFGS: 3 15:51:45 -14.482717 0.974135
LBFGS: 4 15:51:45 -14.615201 1.357127
LBFGS: 5 15:51:45 -14.725192 1.523082
LBFGS: 6 15:51:46 -14.965894 1.582837
LBFGS: 7 15:51:46 -15.224801 1.365773
LBFGS: 8 15:51:46 -15.565080 1.478877
LBFGS: 9 15:51:46 -16.012663 1.304065
LBFGS: 10 15:51:46 -16.228606 0.830748
LBFGS: 11 15:51:46 -16.309868 0.857244
LBFGS: 12 15:51:46 -16.420749 0.921891
LBFGS: 13 15:51:46 -16.573410 0.762856
LBFGS: 14 15:51:46 -16.657135 0.352620
LBFGS: 15 15:51:46 -16.675171 0.233915
LBFGS: 16 15:51:46 -16.683179 0.220636
LBFGS: 17 15:51:46 -16.689688 0.228274
LBFGS: 18 15:51:46 -16.704695 0.244224
LBFGS: 19 15:51:46 -16.735181 0.260771
LBFGS: 20 15:51:46 -16.783558 0.280588
LBFGS: 21 15:51:47 -16.825748 0.241423
LBFGS: 22 15:51:47 -16.838071 0.184039
LBFGS: 23 15:51:47 -16.842487 0.164277
LBFGS: 24 15:51:47 -16.845307 0.148875
LBFGS: 25 15:51:47 -16.851921 0.108184
LBFGS: 26 15:51:47 -16.859492 0.071406
Visualising this:
[18]:
from ase.io import read
from weas_widget import WeasWidget
traj = read("janus_results/NaCl-deformed-traj.extxyz", index=":")
v=WeasWidget()
v.from_ase(traj)
v.avr.model_style = 1
v.avr.show_hydrogen_bonds = True
v
[18]:
Comparing MACE to SevenNet¶
Finally, let’s compare to the structure optimised by SevenNet:
[19]:
%%writefile geomopt_config_mace.yml
struct: ../data/NaCl-deformed.xyz
arch: mace_mp
opt_cell_fully: True
minimize_kwargs:
filter_kwargs:
constant_volume: True
pressure: 10
file_prefix: janus_results/NaCl-mace
tracker: False
Writing geomopt_config_mace.yml
[20]:
! janus geomopt --config geomopt_config_mace.yml
! janus geomopt --config geomopt_config_mace.yml --arch sevennet --file-prefix janus_results/NaCl-sevennet
/Users/elliottkasoar/Documents/PSDI/janus-core/.venv/lib/python3.12/site-packages/e3nn/o3/_wigner.py:10: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
_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 /Users/elliottkasoar/.cache/mace/20231210mace128L0_energy_epoch249model
Using float64 for MACECalculator, which is slower but more accurate. Recommended for geometry optimization.
/Users/elliottkasoar/Documents/PSDI/janus-core/.venv/lib/python3.12/site-packages/mace/calculators/mace.py:139: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
torch.load(f=model_path, map_location=device)
Step Time Energy fmax
LBFGS: 0 15:51:50 -13.238868 2.732337
LBFGS: 1 15:51:50 -13.438489 2.154819
LBFGS: 2 15:51:50 -13.964822 0.951229
LBFGS: 3 15:51:50 -14.168259 0.691490
LBFGS: 4 15:51:50 -14.287776 0.967388
LBFGS: 5 15:51:50 -14.353338 1.090650
LBFGS: 6 15:51:50 -14.462727 1.132334
LBFGS: 7 15:51:50 -14.599603 0.979719
LBFGS: 8 15:51:50 -14.777078 0.839389
LBFGS: 9 15:51:50 -15.047300 0.768000
LBFGS: 10 15:51:50 -15.234272 0.506532
LBFGS: 11 15:51:51 -15.289339 0.637933
LBFGS: 12 15:51:51 -15.315515 0.687128
LBFGS: 13 15:51:51 -15.381258 0.722061
LBFGS: 14 15:51:51 -15.501767 0.656301
LBFGS: 15 15:51:51 -15.613928 0.438204
LBFGS: 16 15:51:51 -15.657094 0.267788
LBFGS: 17 15:51:51 -15.672232 0.278590
LBFGS: 18 15:51:51 -15.685886 0.274355
LBFGS: 19 15:51:51 -15.695658 0.276833
LBFGS: 20 15:51:51 -15.737869 0.260571
LBFGS: 21 15:51:51 -15.795216 0.210303
LBFGS: 22 15:51:51 -15.844755 0.192020
LBFGS: 23 15:51:51 -15.855027 0.128712
LBFGS: 24 15:51:51 -15.857020 0.105543
LBFGS: 25 15:51:51 -15.859034 0.082504
/Users/elliottkasoar/Documents/PSDI/janus-core/.venv/lib/python3.12/site-packages/e3nn/o3/_wigner.py:10: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
_Jd, _W3j_flat, _W3j_indices = torch.load(os.path.join(os.path.dirname(__file__), 'constants.pt'))
Step Time Energy fmax
LBFGS: 0 15:51:55 -13.224955 2.480491
LBFGS: 1 15:51:55 -13.409037 2.025704
LBFGS: 2 15:51:55 -13.958830 0.882636
LBFGS: 3 15:51:55 -14.118237 0.913999
LBFGS: 4 15:51:56 -14.197076 1.139245
LBFGS: 5 15:51:56 -14.274430 1.231919
LBFGS: 6 15:51:56 -14.405305 1.139788
LBFGS: 7 15:51:56 -14.545731 0.873386
LBFGS: 8 15:51:56 -14.699454 0.873543
LBFGS: 9 15:51:56 -14.915076 0.853270
LBFGS: 10 15:51:56 -15.109297 0.626114
LBFGS: 11 15:51:56 -15.172990 0.807070
LBFGS: 12 15:51:56 -15.206856 0.850213
LBFGS: 13 15:51:56 -15.321530 0.874762
LBFGS: 14 15:51:56 -15.490636 0.744572
LBFGS: 15 15:51:56 -15.617524 0.395816
LBFGS: 16 15:51:56 -15.664673 0.237203
LBFGS: 17 15:51:57 -15.678963 0.249874
LBFGS: 18 15:51:57 -15.693007 0.237574
LBFGS: 19 15:51:57 -15.705279 0.237635
LBFGS: 20 15:51:57 -15.743565 0.255036
LBFGS: 21 15:51:57 -15.800006 0.253323
LBFGS: 22 15:51:57 -15.855382 0.177574
LBFGS: 23 15:51:57 -15.872975 0.098585
Visualising the final structures:
[21]:
from ase.io import read
from weas_widget import WeasWidget
mace_opt = read("janus_results/NaCl-mace-opt.extxyz")
sevennet_opt = read("janus_results/NaCl-sevennet-opt.extxyz")
opt_comparison = [mace_opt, sevennet_opt]
v=WeasWidget()
v.from_ase(opt_comparison)
v.avr.model_style = 1
v.avr.show_hydrogen_bonds = True
v
[21]: