From 06d6b0b0aad9c6d689ab0bdb2edf17e6e4f90cf9 Mon Sep 17 00:00:00 2001 From: hjung Date: Tue, 30 Jul 2024 18:57:02 +0000 Subject: [PATCH 1/4] MACE descriptor added --- wfl/descriptors/mace_desc.py | 65 ++++++++++++++++++++++++++++++++++++ 1 file changed, 65 insertions(+) create mode 100644 wfl/descriptors/mace_desc.py diff --git a/wfl/descriptors/mace_desc.py b/wfl/descriptors/mace_desc.py new file mode 100644 index 00000000..277be757 --- /dev/null +++ b/wfl/descriptors/mace_desc.py @@ -0,0 +1,65 @@ +import numpy as np + +from wfl.autoparallelize import autoparallelize, autoparallelize_docstring +from wfl.utils.parallel import construct_calculator_picklesafe +from ase.atoms import Atoms + + +def _calc_autopara_wrappable(atoms, calculator, key, invariants_only=True, per_atom=False, normalize=True, force=False): + """ + MACE descriptor + + Parameters + ---------- + atoms : ase.atoms.Atoms / list(Atoms) + input configurations + calculator : calculator + In this case MACE calculator + key: str + key in Atoms.info (``not per_atom``) or Atoms.arrays (``per_atom``) to store information + invariants_only : bool, default True + get only invariant part of the descriptor + per_atom: bool + calculate a local (per-atom) descriptor, as opposed to global (per-config) + normalize: bool, default True + normalize final vector (e.g. if contributions from multiple descriptors were concatenated) + force: bool, default False + overwrite key if already exists + + Returns + ------- + atoms : Atoms or list(Atoms) + Input configurations with descriptors in info/arrays + """ + + calculator = construct_calculator_picklesafe(calculator) + + if isinstance(atoms, Atoms): + at_list = [atoms] + else: + at_list = atoms + + descriptor_encoded_atoms = [] + for i, at in enumerate(at_list): + if key in at.info: + if force: + del at.info[key] + else: + raise RuntimeError(f'Got info key {key} which already exists, pass force=True to overwrite') + + if per_atom: + at.info[key] = calculator.get_descriptors(at, invariants_only=invariants_only) + else: + descriptor = np.average(calculator.get_descriptors(at, invariants_only=invariants_only), axis=0) + if normalize: + at.info[key] = descriptor / np.linalg.norm(descriptor) + else: + at.info[key] = descriptor + + descriptor_encoded_atoms.append(at) + + return descriptor_encoded_atoms + +def calculate(*args, **kwargs): + return autoparallelize(_calc_autopara_wrappable, *args, **kwargs) +autoparallelize_docstring(calculate, _calc_autopara_wrappable, "Atoms") From a2a158f730cbfe3e6de8e7c1e1f31f7cd9e8e22c Mon Sep 17 00:00:00 2001 From: hjung Date: Tue, 30 Jul 2024 18:59:26 +0000 Subject: [PATCH 2/4] small change --- wfl/descriptors/mace_desc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/wfl/descriptors/mace_desc.py b/wfl/descriptors/mace_desc.py index 277be757..9140a4ec 100644 --- a/wfl/descriptors/mace_desc.py +++ b/wfl/descriptors/mace_desc.py @@ -19,7 +19,7 @@ def _calc_autopara_wrappable(atoms, calculator, key, invariants_only=True, per_a key in Atoms.info (``not per_atom``) or Atoms.arrays (``per_atom``) to store information invariants_only : bool, default True get only invariant part of the descriptor - per_atom: bool + per_atom: bool, default False calculate a local (per-atom) descriptor, as opposed to global (per-config) normalize: bool, default True normalize final vector (e.g. if contributions from multiple descriptors were concatenated) From 5425862c9ece7b326464ffee6a8be98afb12b7ba Mon Sep 17 00:00:00 2001 From: hjung Date: Tue, 8 Oct 2024 13:12:06 +0000 Subject: [PATCH 3/4] add attach function and abort for neb.py --- wfl/generate/neb.py | 24 ++++++++++++++++++++---- 1 file changed, 20 insertions(+), 4 deletions(-) diff --git a/wfl/generate/neb.py b/wfl/generate/neb.py index 1f2be063..1d936ee4 100644 --- a/wfl/generate/neb.py +++ b/wfl/generate/neb.py @@ -13,8 +13,8 @@ def _run_autopara_wrappable(list_of_images, calculator, fmax=5e-2, steps=1000, - traj_step_interval=1, traj_subselect=None, skip_failures=True, - results_prefix='last_op__neb_', verbose=False, logfile=None, update_config_type="append", + traj_step_interval=1, traj_subselect=None, skip_failures=True, attach_kwargs=None, attach_interval=None, + results_prefix='last_op__neb_', verbose=False, logfile=None, update_config_type="append", abort_check=None, **neb_kwargs): """runs a structure optimization. By default calculator properties will be stored in keys prefixed with "last_op__neb_", which may be overwritten by next operation. @@ -37,6 +37,10 @@ def _run_autopara_wrappable(list_of_images, calculator, fmax=5e-2, steps=1000, Currently implemented: "last_converged", which takes the last config, if converged. skip_failures: bool, default True just skip optimizations that raise an exception + attach_kwargs : function, default None + function to be attached to optimizer + attach_interval : int, default None + interval on which function is executed during optimization results_prefix: str, default "last_op__neb_" prefix to info/arrays keys where calculator properties will be stored Will overwrite any other properties that start with same "__", so that by @@ -46,6 +50,9 @@ def _run_autopara_wrappable(list_of_images, calculator, fmax=5e-2, steps=1000, optimisation logs are not printed unless this is True update_config_type: ["append" | "overwrite" | False], default "append" whether/how to add at.info['neb_config_type'] to at.info['config_type'] + abort_check : default None, + wfl.generate.md.abort_base.AbortBase - derived class that + checks the NEB snapshots and aborts the simulation on some condition. neb_kwargs keyword arguments for DyNEB and FIRE @@ -54,7 +61,9 @@ def _run_autopara_wrappable(list_of_images, calculator, fmax=5e-2, steps=1000, list(Atoms) trajectories """ logfile = neb_kwargs.get("logfile", None) - + if attach_kwargs is not None: + attach_constructor = attach_kwargs.pop("attach_function", None) + if logfile is None and verbose: logfile = "-" @@ -84,9 +93,16 @@ def process_step(): cur_images.append(new_config) traj.append(cur_images) - + + if abort_check is not None: + if abort_check.stop(opt): + raise RuntimeError(f"NEB was stopped by the NEB checker function {abort_check.__class__.__name__}") opt.attach(process_step, interval=traj_step_interval) + if attach_kwargs is not None: + attach_kwargs["neb"] = neb + opt.attach(attach_constructor, interval=attach_interval, **attach_kwargs) + # preliminary value final_status = 'unconverged' From 2b32fc59e5b1045f551af8766cb805f1a5d3d049 Mon Sep 17 00:00:00 2001 From: hjung Date: Tue, 8 Oct 2024 13:21:34 +0000 Subject: [PATCH 4/4] mace_desc.py deleted --- wfl/descriptors/mace_desc.py | 65 ------------------------------------ 1 file changed, 65 deletions(-) delete mode 100644 wfl/descriptors/mace_desc.py diff --git a/wfl/descriptors/mace_desc.py b/wfl/descriptors/mace_desc.py deleted file mode 100644 index 9140a4ec..00000000 --- a/wfl/descriptors/mace_desc.py +++ /dev/null @@ -1,65 +0,0 @@ -import numpy as np - -from wfl.autoparallelize import autoparallelize, autoparallelize_docstring -from wfl.utils.parallel import construct_calculator_picklesafe -from ase.atoms import Atoms - - -def _calc_autopara_wrappable(atoms, calculator, key, invariants_only=True, per_atom=False, normalize=True, force=False): - """ - MACE descriptor - - Parameters - ---------- - atoms : ase.atoms.Atoms / list(Atoms) - input configurations - calculator : calculator - In this case MACE calculator - key: str - key in Atoms.info (``not per_atom``) or Atoms.arrays (``per_atom``) to store information - invariants_only : bool, default True - get only invariant part of the descriptor - per_atom: bool, default False - calculate a local (per-atom) descriptor, as opposed to global (per-config) - normalize: bool, default True - normalize final vector (e.g. if contributions from multiple descriptors were concatenated) - force: bool, default False - overwrite key if already exists - - Returns - ------- - atoms : Atoms or list(Atoms) - Input configurations with descriptors in info/arrays - """ - - calculator = construct_calculator_picklesafe(calculator) - - if isinstance(atoms, Atoms): - at_list = [atoms] - else: - at_list = atoms - - descriptor_encoded_atoms = [] - for i, at in enumerate(at_list): - if key in at.info: - if force: - del at.info[key] - else: - raise RuntimeError(f'Got info key {key} which already exists, pass force=True to overwrite') - - if per_atom: - at.info[key] = calculator.get_descriptors(at, invariants_only=invariants_only) - else: - descriptor = np.average(calculator.get_descriptors(at, invariants_only=invariants_only), axis=0) - if normalize: - at.info[key] = descriptor / np.linalg.norm(descriptor) - else: - at.info[key] = descriptor - - descriptor_encoded_atoms.append(at) - - return descriptor_encoded_atoms - -def calculate(*args, **kwargs): - return autoparallelize(_calc_autopara_wrappable, *args, **kwargs) -autoparallelize_docstring(calculate, _calc_autopara_wrappable, "Atoms")