Geometry Optimization¶
Set up environment (optional)¶
These steps are required to run this tutorial with Google Colab. To do so, uncomment and run the cell below.
This will replace pre-installed versions of numpy and torch in Colab with versions that are known to be compatible with janus-core.
It may be possible to skip the steps that uninstall and reinstall torch, which will save a considerable amount of time.
These instructions but may work for other systems too, but it is typically preferable to prepare a virtual environment separately before running this notebook if possible.
[1]:
# import locale
# locale.getpreferredencoding = lambda: "UTF-8"
# ! pip uninstall numpy -y # Uninstall pre-installed numpy
# ! pip uninstall torch torchaudio torchvision transformers -y # Uninstall pre-installed torch
# ! uv pip install torch==2.5.1 # Install pinned version of torch
# ! uv pip install janus-core[mace,orb,visualise] data-tutorials --system # Install janus-core with MACE, Orb, and WeasWidget, and data-tutorials
# get_ipython().kernel.do_shutdown(restart=True) # Restart kernel to update libraries. This may warn that your session has crashed.
To ensure you have the latest version of janus-core installed, compare the output of the following cell to the latest version available at https://pypi.org/project/janus-core/
[2]:
from janus_core import __version__
print(__version__)
0.8.5
Command-line help and options¶
As with janus singlepoint, we can check the options for geometry optimisation:
[3]:
! janus geomopt --help
Usage: janus geomopt [OPTIONS]
Perform geometry optimization and save optimized structure to file.
╭─ Options ────────────────────────────────────────────────────────────────────╮
│ --config TEXT Path to configuration file. │
│ --help Show this message and exit. │
╰──────────────────────────────────────────────────────────────────────────────╯
╭─ MLIP calculator ────────────────────────────────────────────────────────────╮
│ * --arch [mace|mace_mp|mace_off|m MLIP architecture to use │
│ 3gnet|chgnet|alignn|seve for calculations. │
│ nnet|nequip|dpa3|orb|mat [required] │
│ tersim|grace|esen|equifo │
│ rmer|pet_mad|uma|mace_om │
│ ol] │
│ --device [cpu|cuda|mps|xpu] Device to run calculations │
│ on. │
│ [default: cpu] │
│ --model TEXT MLIP model name, or path │
│ to model. │
│ --calc-kwargs DICT Keyword arguments to pass │
│ to selected calculator. │
│ Must be passed as a │
│ dictionary wrapped in │
│ quotes, e.g. "{'key': │
│ value}". │
╰──────────────────────────────────────────────────────────────────────────────╯
╭─ Calculation ────────────────────────────────────────────────────────────────╮
│ * --struct PATH Path of structure to │
│ simulate. │
│ [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] │
│ --opt-cell-lengths --no-opt-cell-len… Optimize cell │
│ vectors, as well as │
│ atomic positions. │
│ [default: │
│ no-opt-cell-lengths] │
│ --opt-cell-fully --no-opt-cell-ful… Fully optimize the │
│ cell vectors, │
│ angles, and atomic │
│ positions. │
│ [default: │
│ no-opt-cell-fully] │
│ --filter TEXT Name of ASE filter │
│ to wrap around │
│ atoms. If using │
│ --opt-cell-lengths │
│ or --opt-cell-fully, │
│ defaults to │
│ `FrechetCellFilter`. │
│ --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-tolera… FLOAT Atom displacement │
│ tolerance for spglib │
│ symmetry │
│ determination, in Å. │
│ [default: 0.001] │
│ --out PATH Path to save │
│ optimized structure. │
│ Default is inferred │
│ `file_prefix`. │
│ --write-traj --no-write-traj Whether to save a │
│ trajectory file of │
│ optimization frames. │
│ If │
│ traj_kwargs["filena… │
│ is not specified, it │
│ is inferred from │
│ `file_prefix`. │
│ [default: │
│ no-write-traj] │
│ --minimize-kwargs DICT Keyword arguments to │
│ pass to geometry │
│ optimizer, including │
│ "opt_kwargs", │
│ "filter_kwargs", and │
│ "traj_kwargs". Must │
│ be passed as a │
│ dictionary wrapped │
│ in quotes, e.g. │
│ "{'key': value}". │
╰──────────────────────────────────────────────────────────────────────────────╯
╭─ Structure I/O ──────────────────────────────────────────────────────────────╮
│ --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. │
│ --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['index'] = -1, so only the last │
│ structure is read. │
│ --write-kwargs DICT Keyword arguments to pass to ase.io.write when │
│ saving any structures. Must be passed as a │
│ dictionary wrapped in quotes, e.g. "{'key': │
│ value}". │
╰──────────────────────────────────────────────────────────────────────────────╯
╭─ Logging/summary ────────────────────────────────────────────────────────────╮
│ --log PATH Path to save logs to. Default is inferred │
│ from `file_prefix` │
│ --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`. │
╰──────────────────────────────────────────────────────────────────────────────╯
Running geometry optimisation calculation¶
First, we’ll build a structure to optimise, as we did for single point calculations, but add in a deformation:
[4]:
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
[4]:
Now we can optimise the geometry of this structure in a similar manner to running janus singlepoint:
[5]:
! janus geomopt --arch mace_mp --struct ../data/NaCl-deformed.xyz --no-tracker
/home/runner/work/janus-core/janus-core/.venv/lib/python3.12/site-packages/codecarbon/core/gpu.py:4: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you.
import pynvml
/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.
Downloading MACE model from 'https://github.com/ACEsuit/mace-mp/releases/download/mace_mp_0/2023-12-10-mace-128-L0_energy_epoch-249.model'
Cached MACE model to /home/runner/.cache/mace/20231210mace128L0_energy_epoch249model
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:197: 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)
Using head Default out of ['Default']
Step Time Energy fmax
LBFGS: 0 15:56:50 -24.377063 2.732337
LBFGS: 1 15:56:50 -24.576368 2.152513
LBFGS: 2 15:56:50 -25.093690 0.949845
LBFGS: 3 15:56:50 -25.290212 0.685740
LBFGS: 4 15:56:50 -25.403788 0.951695
LBFGS: 5 15:56:50 -25.465213 1.066671
LBFGS: 6 15:56:50 -25.566446 1.095941
LBFGS: 7 15:56:51 -25.695050 0.931041
LBFGS: 8 15:56:51 -25.856683 0.857034
LBFGS: 9 15:56:51 -26.109079 0.811033
LBFGS: 10 15:56:51 -26.298823 0.545193
LBFGS: 11 15:56:51 -26.353032 0.656387
LBFGS: 12 15:56:51 -26.383904 0.710488
LBFGS: 13 15:56:51 -26.485793 0.790107
LBFGS: 14 15:56:51 -26.643853 0.768081
LBFGS: 15 15:56:51 -26.751857 0.477356
LBFGS: 16 15:56:51 -26.792940 0.329784
LBFGS: 17 15:56:51 -26.815123 0.343673
LBFGS: 18 15:56:51 -26.829194 0.334430
LBFGS: 19 15:56:51 -26.845943 0.317718
LBFGS: 20 15:56:51 -26.912307 0.249075
LBFGS: 21 15:56:51 -26.975682 0.229256
LBFGS: 22 15:56:51 -27.018210 0.150833
LBFGS: 23 15:56:51 -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.
[6]:
! janus geomopt --arch mace_mp --struct ../data/NaCl-deformed.xyz --optimizer FIRE --fmax 0.005 --no-tracker
/home/runner/work/janus-core/janus-core/.venv/lib/python3.12/site-packages/codecarbon/core/gpu.py:4: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you.
import pynvml
/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:197: 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)
Using head Default out of ['Default']
Step Time Energy fmax
FIRE: 0 15:56:56 -24.377063 2.732337
FIRE: 1 15:56:56 -24.521143 2.311307
FIRE: 2 15:56:57 -24.737816 1.719928
FIRE: 3 15:56:57 -24.954552 1.216180
FIRE: 4 15:56:57 -25.138250 0.893003
FIRE: 5 15:56:57 -25.283693 0.712183
FIRE: 6 15:56:57 -25.398292 0.725931
FIRE: 7 15:56:57 -25.495898 0.776106
FIRE: 8 15:56:57 -25.600795 0.738520
FIRE: 9 15:56:57 -25.727081 0.624992
FIRE: 10 15:56:57 -25.886706 0.633155
FIRE: 11 15:56:57 -26.063737 0.609225
FIRE: 12 15:56:57 -26.237721 0.564647
FIRE: 13 15:56:57 -26.390734 0.525758
FIRE: 14 15:56:57 -26.511572 0.531927
FIRE: 15 15:56:57 -26.597920 0.696601
FIRE: 16 15:56:57 -26.653461 0.790166
FIRE: 17 15:56:57 -26.686575 0.773073
FIRE: 18 15:56:57 -26.705939 0.735325
FIRE: 19 15:56:57 -26.741131 0.662023
FIRE: 20 15:56:57 -26.785861 0.557699
FIRE: 21 15:56:57 -26.832526 0.429577
FIRE: 22 15:56:57 -26.874093 0.288174
FIRE: 23 15:56:57 -26.906094 0.205971
FIRE: 24 15:56:57 -26.928149 0.249613
FIRE: 25 15:56:57 -26.945064 0.282363
FIRE: 26 15:56:58 -26.961246 0.296619
FIRE: 27 15:56:58 -26.981117 0.283688
FIRE: 28 15:56:58 -27.004699 0.232056
FIRE: 29 15:56:58 -27.023875 0.129160
FIRE: 30 15:56:58 -27.023665 0.065418
FIRE: 31 15:56:58 -27.023932 0.063243
FIRE: 32 15:56:58 -27.024426 0.058997
FIRE: 33 15:56:58 -27.025071 0.052890
FIRE: 34 15:56:58 -27.025778 0.045230
FIRE: 35 15:56:58 -27.026455 0.036435
FIRE: 36 15:56:58 -27.027033 0.030872
FIRE: 37 15:56:58 -27.027485 0.033525
FIRE: 38 15:56:58 -27.027863 0.046407
FIRE: 39 15:56:58 -27.028209 0.055288
FIRE: 40 15:56:58 -27.028598 0.057928
FIRE: 41 15:56:58 -27.029081 0.052454
FIRE: 42 15:56:58 -27.029611 0.037632
FIRE: 43 15:56:58 -27.029996 0.013477
FIRE: 44 15:56:58 -27.029974 0.017129
FIRE: 45 15:56:58 -27.029982 0.016447
FIRE: 46 15:56:58 -27.029995 0.015113
FIRE: 47 15:56:58 -27.030013 0.013191
FIRE: 48 15:56:58 -27.030031 0.010771
FIRE: 49 15:56:58 -27.030046 0.007969
FIRE: 50 15:56:58 -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:
[7]:
! 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:
[8]:
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
[8]:
Trajectories¶
We can also save the trajectory during optimisation:
[9]:
! janus geomopt --arch mace_mp --struct ../data/NaCl-deformed.xyz --write-traj --no-tracker
/home/runner/work/janus-core/janus-core/.venv/lib/python3.12/site-packages/codecarbon/core/gpu.py:4: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you.
import pynvml
/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:197: 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)
Using head Default out of ['Default']
Step Time Energy fmax
LBFGS: 0 15:57:04 -24.377063 2.732337
LBFGS: 1 15:57:04 -24.576368 2.152513
LBFGS: 2 15:57:04 -25.093690 0.949845
LBFGS: 3 15:57:04 -25.290212 0.685740
LBFGS: 4 15:57:04 -25.403788 0.951695
LBFGS: 5 15:57:04 -25.465213 1.066671
LBFGS: 6 15:57:04 -25.566446 1.095941
LBFGS: 7 15:57:04 -25.695050 0.931041
LBFGS: 8 15:57:04 -25.856683 0.857034
LBFGS: 9 15:57:04 -26.109079 0.811033
LBFGS: 10 15:57:04 -26.298823 0.545193
LBFGS: 11 15:57:04 -26.353032 0.656387
LBFGS: 12 15:57:04 -26.383904 0.710488
LBFGS: 13 15:57:04 -26.485793 0.790107
LBFGS: 14 15:57:04 -26.643853 0.768081
LBFGS: 15 15:57:05 -26.751857 0.477356
LBFGS: 16 15:57:05 -26.792940 0.329784
LBFGS: 17 15:57:05 -26.815123 0.343673
LBFGS: 18 15:57:05 -26.829194 0.334430
LBFGS: 19 15:57:05 -26.845943 0.317718
LBFGS: 20 15:57:05 -26.912307 0.249075
LBFGS: 21 15:57:05 -26.975682 0.229256
LBFGS: 22 15:57:05 -27.018210 0.150833
LBFGS: 23 15:57:05 -27.023510 0.093603
This creates an additional file, janus_results/NaCl-deformed-traj.extxyz:
[10]:
! 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:
[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]:
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:
[12]:
! janus geomopt --arch mace_mp --struct ../data/NaCl-deformed.xyz --write-traj --opt-cell-lengths --no-tracker
/home/runner/work/janus-core/janus-core/.venv/lib/python3.12/site-packages/codecarbon/core/gpu.py:4: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you.
import pynvml
/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:197: 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)
Using head Default out of ['Default']
Step Time Energy fmax
LBFGS: 0 15:57:10 -24.377063 2.732337
LBFGS: 1 15:57:10 -24.593123 2.143240
LBFGS: 2 15:57:10 -25.178180 0.955170
LBFGS: 3 15:57:11 -25.445502 0.670912
LBFGS: 4 15:57:11 -25.641907 0.919610
LBFGS: 5 15:57:11 -25.735611 1.036999
LBFGS: 6 15:57:11 -25.816918 1.015167
LBFGS: 7 15:57:11 -25.937123 0.774473
LBFGS: 8 15:57:11 -26.075553 0.561448
LBFGS: 9 15:57:11 -26.194652 0.423668
LBFGS: 10 15:57:11 -26.290628 0.391005
LBFGS: 11 15:57:11 -26.323319 0.439570
LBFGS: 12 15:57:11 -26.337274 0.452690
LBFGS: 13 15:57:11 -26.372755 0.457902
LBFGS: 14 15:57:11 -26.436676 0.640085
LBFGS: 15 15:57:11 -26.522182 0.786499
LBFGS: 16 15:57:11 -26.636096 0.839171
LBFGS: 17 15:57:11 -26.785166 0.746960
LBFGS: 18 15:57:11 -26.953112 0.390060
LBFGS: 19 15:57:12 -27.019414 0.155454
LBFGS: 20 15:57:12 -27.024708 0.104189
LBFGS: 21 15:57:12 -27.025735 0.100499
LBFGS: 22 15:57:12 -27.040218 0.057352
As before, we can visualise this trajectory:
[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]:
We can also allow the cell angles to change:
[14]:
! janus geomopt --arch mace_mp --struct ../data/NaCl-deformed.xyz --write-traj --opt-cell-fully --no-tracker
/home/runner/work/janus-core/janus-core/.venv/lib/python3.12/site-packages/codecarbon/core/gpu.py:4: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you.
import pynvml
/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:197: 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)
Using head Default out of ['Default']
Step Time Energy fmax
LBFGS: 0 15:57:17 -24.377063 2.732337
LBFGS: 1 15:57:17 -24.593440 2.145517
LBFGS: 2 15:57:17 -25.187807 0.956555
LBFGS: 3 15:57:17 -25.462335 0.675018
LBFGS: 4 15:57:17 -25.663853 0.932226
LBFGS: 5 15:57:17 -25.757785 1.056063
LBFGS: 6 15:57:17 -25.841303 1.038396
LBFGS: 7 15:57:18 -25.968290 0.806821
LBFGS: 8 15:57:18 -26.115858 0.539264
LBFGS: 9 15:57:18 -26.230485 0.404401
LBFGS: 10 15:57:18 -26.323082 0.351277
LBFGS: 11 15:57:18 -26.354901 0.404205
LBFGS: 12 15:57:18 -26.366593 0.418519
LBFGS: 13 15:57:18 -26.402172 0.433457
LBFGS: 14 15:57:18 -26.468589 0.596828
LBFGS: 15 15:57:18 -26.555830 0.733795
LBFGS: 16 15:57:18 -26.664503 0.782841
LBFGS: 17 15:57:18 -26.800773 0.695400
LBFGS: 18 15:57:18 -26.945555 0.318072
LBFGS: 19 15:57:18 -26.988824 0.103903
LBFGS: 20 15:57:18 -26.991370 0.096582
Visualising this:
[15]:
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
[15]:
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:
[16]:
%%writefile geomopt_config_1.yml
arch: mace_mp
struct: ../data/NaCl-deformed.xyz
opt_cell_fully: True
filter: 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 ExpCellFilter –minimize-kwargs “{‘filter_kwargs’: {‘constant_volume’ : True}” –no-tracker
[17]:
! janus geomopt --config geomopt_config_1.yml
/home/runner/work/janus-core/janus-core/.venv/lib/python3.12/site-packages/codecarbon/core/gpu.py:4: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you.
import pynvml
/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:197: 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)
Using head Default out of ['Default']
Step Time Energy fmax
LBFGS: 0 15:57:24 -24.377063 2.732337
LBFGS: 1 15:57:24 -24.592435 2.306591
LBFGS: 2 15:57:24 -25.267905 1.323804
LBFGS: 3 15:57:24 -25.449977 1.003470
LBFGS: 4 15:57:24 -25.521347 0.823847
LBFGS: 5 15:57:24 -25.573640 0.866332
LBFGS: 6 15:57:24 -25.673315 1.004378
LBFGS: 7 15:57:24 -25.811015 1.004781
LBFGS: 8 15:57:25 -25.975638 0.889642
LBFGS: 9 15:57:25 -26.150934 0.689834
LBFGS: 10 15:57:25 -26.319557 0.690018
LBFGS: 11 15:57:25 -26.467568 0.708275
LBFGS: 12 15:57:25 -26.574768 0.722124
LBFGS: 13 15:57:25 -26.637370 0.637261
LBFGS: 14 15:57:25 -26.708598 0.646450
LBFGS: 15 15:57:25 -26.781501 0.681819
LBFGS: 16 15:57:25 -26.844998 0.708926
LBFGS: 17 15:57:26 -26.874141 0.642122
LBFGS: 18 15:57:26 -26.894946 0.530522
LBFGS: 19 15:57:26 -26.919861 0.583009
LBFGS: 20 15:57:26 -26.954422 0.529294
LBFGS: 21 15:57:26 -26.977868 0.252808
LBFGS: 22 15:57:26 -26.985391 0.153658
LBFGS: 23 15:57:26 -26.987735 0.137245
LBFGS: 24 15:57:26 -26.991006 0.146929
LBFGS: 25 15:57:26 -26.998079 0.227195
LBFGS: 26 15:57:27 -27.010168 0.266027
LBFGS: 27 15:57:27 -27.022858 0.207394
LBFGS: 28 15:57:27 -27.029032 0.073577
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]:
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:
[19]:
! janus geomopt --arch mace_mp --struct ../data/NaCl-deformed.xyz --write-traj --pressure 10 --opt-cell-fully --symmetrize --no-tracker
/home/runner/work/janus-core/janus-core/.venv/lib/python3.12/site-packages/codecarbon/core/gpu.py:4: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you.
import pynvml
/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:197: 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)
Using head Default out of ['Default']
Step Time Energy fmax
LBFGS: 0 15:57:32 -13.238868 2.732337
LBFGS: 1 15:57:32 -13.462923 2.165368
LBFGS: 2 15:57:32 -14.202196 0.816332
LBFGS: 3 15:57:32 -14.482717 0.974135
LBFGS: 4 15:57:32 -14.615201 1.357127
LBFGS: 5 15:57:33 -14.725192 1.523082
LBFGS: 6 15:57:33 -14.965894 1.582837
LBFGS: 7 15:57:33 -15.224801 1.365773
LBFGS: 8 15:57:33 -15.565080 1.478877
LBFGS: 9 15:57:33 -16.012663 1.304065
LBFGS: 10 15:57:33 -16.228606 0.830748
LBFGS: 11 15:57:33 -16.309868 0.857244
LBFGS: 12 15:57:33 -16.420749 0.921891
LBFGS: 13 15:57:33 -16.573410 0.762856
LBFGS: 14 15:57:33 -16.657135 0.352620
LBFGS: 15 15:57:33 -16.675171 0.233915
LBFGS: 16 15:57:33 -16.683179 0.220636
LBFGS: 17 15:57:33 -16.689688 0.228274
LBFGS: 18 15:57:33 -16.704695 0.244224
LBFGS: 19 15:57:34 -16.735181 0.260771
LBFGS: 20 15:57:34 -16.783558 0.280588
LBFGS: 21 15:57:34 -16.825748 0.241423
LBFGS: 22 15:57:34 -16.838071 0.184039
LBFGS: 23 15:57:34 -16.842487 0.164277
LBFGS: 24 15:57:34 -16.845307 0.148875
LBFGS: 25 15:57:34 -16.851921 0.108184
LBFGS: 26 15:57:34 -16.859492 0.071406
Visualising this:
[20]:
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
[20]:
Comparing MACE to Orb¶
Finally, let’s compare to the structure optimised by Orb:
[21]:
%%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
[22]:
! janus geomopt --config geomopt_config_mace.yml
! janus geomopt --config geomopt_config_mace.yml --arch orb --file-prefix janus_results/NaCl-orb
/home/runner/work/janus-core/janus-core/.venv/lib/python3.12/site-packages/codecarbon/core/gpu.py:4: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you.
import pynvml
/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:197: 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)
Using head Default out of ['Default']
Step Time Energy fmax
LBFGS: 0 15:57:39 -13.238868 2.732337
LBFGS: 1 15:57:40 -13.438489 2.154819
LBFGS: 2 15:57:40 -13.964822 0.951229
LBFGS: 3 15:57:40 -14.168259 0.691490
LBFGS: 4 15:57:40 -14.287776 0.967388
LBFGS: 5 15:57:40 -14.353338 1.090650
LBFGS: 6 15:57:40 -14.462727 1.132334
LBFGS: 7 15:57:40 -14.599603 0.979719
LBFGS: 8 15:57:40 -14.777078 0.839389
LBFGS: 9 15:57:40 -15.047300 0.768000
LBFGS: 10 15:57:40 -15.234272 0.506532
LBFGS: 11 15:57:40 -15.289339 0.637933
LBFGS: 12 15:57:40 -15.315515 0.687128
LBFGS: 13 15:57:40 -15.381258 0.722061
LBFGS: 14 15:57:40 -15.501767 0.656301
LBFGS: 15 15:57:40 -15.613928 0.438204
LBFGS: 16 15:57:41 -15.657094 0.267788
LBFGS: 17 15:57:41 -15.672232 0.278590
LBFGS: 18 15:57:41 -15.685886 0.274355
LBFGS: 19 15:57:41 -15.695658 0.276833
LBFGS: 20 15:57:41 -15.737869 0.260571
LBFGS: 21 15:57:41 -15.795216 0.210303
LBFGS: 22 15:57:41 -15.844755 0.192020
LBFGS: 23 15:57:41 -15.855027 0.128712
LBFGS: 24 15:57:41 -15.857020 0.105543
LBFGS: 25 15:57:41 -15.859034 0.082504
/home/runner/work/janus-core/janus-core/.venv/lib/python3.12/site-packages/codecarbon/core/gpu.py:4: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you.
import pynvml
/home/runner/work/janus-core/janus-core/.venv/lib/python3.12/site-packages/orb_models/utils.py:30: UserWarning: Setting global torch default dtype to torch.float32.
warnings.warn(f"Setting global torch default dtype to {torch_dtype}.")
Downloading …ervative-20-omat-20250404.ckpt ━━━━━━━━━━━ 0% 0:00:00 0.0/102.1
Downloading …ervative-20-omat-20250404.ckpt ━━━━━━━━━━━ 0% 0:00:00 0.0/102.1
Downloading …ervative-20-omat-20250404.ckpt ━━━━━━━━━━━ 0% 0:00:00 0.0/102.1
Downloading …ervative-20-omat-20250404.ckpt ━━━━━━━━━━━ 0% 0:00:00 0.0/102.1
Downloading …ervative-20-omat-20250404.ckpt ━━━━━━━━━━━ 0% 0:00:00 0.1/102.1
Downloading …ervative-20-omat-20250404.ckpt ━━━━━━━━━━━ 0% 0:00:00 0.4/102.1
Downloading …ervative-20-omat-20250404.ckpt ━━━━━━━━━━━ 1% 0:00:00 1.0/102.1
Downloading …ervative-20-omat-20250404.ckpt ╸━━━━━━━━━━ 5% 0:00:00 5.3/102.1
Downloading …ervative-20-omat-20250404.ckpt ━╺━━━━━━━━ 11% 0:00:00 11.6/102.1
Downloading …ervative-20-omat-20250404.ckpt ━╸━━━━━━━━ 18% 0:00:00 18.2/102.1
Downloading …ervative-20-omat-20250404.ckpt ━━╺━━━━━━━ 24% 0:00:01 25.0/102.1
Downloading …ervative-20-omat-20250404.ckpt ━━━╺━━━━━━ 31% 0:00:01 31.6/102.1
Downloading …ervative-20-omat-20250404.ckpt ━━━╸━━━━━━ 37% 0:00:01 37.9/102.1
Downloading …ervative-20-omat-20250404.ckpt ━━━━╺━━━━━ 43% 0:00:01 44.1/102.1
Downloading …ervative-20-omat-20250404.ckpt ━━━━╸━━━━━ 49% 0:00:01 50.3/102.1
Downloading …ervative-20-omat-20250404.ckpt ━━━━━╸━━━━ 56% 0:00:01 56.7/102.1
Downloading …ervative-20-omat-20250404.ckpt ━━━━━━╺━━━ 62% 0:00:01 62.9/102.1
Downloading …ervative-20-omat-20250404.ckpt ━━━━━━╸━━━ 68% 0:00:01 69.9/102.1
Downloading …ervative-20-omat-20250404.ckpt ━━━━━━━╺━━ 74% 0:00:01 75.5/102.1
Downloading …ervative-20-omat-20250404.ckpt ━━━━━━━╸━━ 78% 0:00:01 79.4/102.1
Downloading …ervative-20-omat-20250404.ckpt ━━━━━━━━╺━ 83% 0:00:02 85.0/102.1
Downloading …ervative-20-omat-20250404.ckpt ━━━━━━━━╸━ 89% 0:00:02 90.8/102.1
Downloading …ervative-20-omat-20250404.ckpt ━━━━━━━━━╺ 95% 0:00:02 96.9/102.1
Downloading …ervative-20-omat-20250404.ckpt ━━━━━━━━━━━ 100% 0:00:02 102.1/102.1
MB
/home/runner/work/janus-core/janus-core/.venv/lib/python3.12/site-packages/torch/cuda/__init__.py:61: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you.
import pynvml # type: ignore[import]
Step Time Energy fmax
LBFGS: 0 15:58:29 -13.126874 2.678413
LBFGS: 1 15:58:29 -13.316277 2.164414
LBFGS: 2 15:58:29 -13.867301 0.876200
LBFGS: 3 15:58:29 -14.055682 0.921783
LBFGS: 4 15:58:29 -14.168099 1.217233
LBFGS: 5 15:58:29 -14.248955 1.319721
LBFGS: 6 15:58:29 -14.386883 1.258692
LBFGS: 7 15:58:29 -14.545424 0.955374
LBFGS: 8 15:58:29 -14.757332 0.982592
LBFGS: 9 15:58:30 -15.045882 0.785888
LBFGS: 10 15:58:30 -15.197545 0.605819
LBFGS: 11 15:58:30 -15.240164 0.733672
LBFGS: 12 15:58:30 -15.282086 0.798982
LBFGS: 13 15:58:30 -15.395701 0.849327
LBFGS: 14 15:58:30 -15.572037 0.694654
LBFGS: 15 15:58:30 -15.667040 0.304853
LBFGS: 16 15:58:30 -15.696938 0.277427
LBFGS: 17 15:58:30 -15.713296 0.296171
LBFGS: 18 15:58:31 -15.723715 0.292463
LBFGS: 19 15:58:31 -15.741387 0.281043
LBFGS: 20 15:58:31 -15.778401 0.256297
LBFGS: 21 15:58:31 -15.836583 0.220242
LBFGS: 22 15:58:31 -15.895636 0.177559
LBFGS: 23 15:58:31 -15.910929 0.093283
Visualising the final structures:
[23]:
from ase.io import read
from weas_widget import WeasWidget
mace_opt = read("janus_results/NaCl-mace-opt.extxyz")
orb_opt = read("janus_results/NaCl-orb-opt.extxyz")
opt_comparison = [mace_opt, orb_opt]
v=WeasWidget()
v.from_ase(opt_comparison)
v.avr.model_style = 1
v.avr.show_hydrogen_bonds = True
v
[23]: