From 6fb84b9b13b06e65e1972334ae9af4e3b17333b6 Mon Sep 17 00:00:00 2001 From: gmatteo Date: Thu, 31 Aug 2023 23:32:52 +0200 Subject: [PATCH] Saving work --- abipy/flowtk/psrepos.py | 15 ++ abipy/ml/aseml.py | 463 ++++++++++++++++++++++++++------------ abipy/ml/ml_relax.py | 3 +- abipy/ml/relax_scanner.py | 143 +++++++----- abipy/scripts/abiml.py | 188 ++++++---------- dev_scripts/mlrun.py | 4 +- 6 files changed, 495 insertions(+), 321 deletions(-) diff --git a/abipy/flowtk/psrepos.py b/abipy/flowtk/psrepos.py index b3b215a8c..05731ab6e 100644 --- a/abipy/flowtk/psrepos.py +++ b/abipy/flowtk/psrepos.py @@ -139,6 +139,21 @@ def get_repo_from_name(repo_name: str) -> PseudosRepo: raise KeyError(f"Couldn't find {repo_name} in the list of registered repos:\n{all_names}") +#def get_pseudos(xc_name, ps_type: str, relativity_type: str = "SR", accuracy="standard") +# if ps_type == "NC": +# repo_name = { +# "PBE": "ONCVPSP-PBE-SR-PDv0.4", +# "PBEsol": "ONCVPSP-PBEsol-SR-PDv0.4", +# "LDA": "ONCVPSP-LDA-SR-PDv0.4", +# }[xc_name] +# elif ps_type == "PAW": +# raise ValueError(f"Invalid {ps_type=}") +# else: +# raise ValueError(f"Invalid {ps_type=}") +# pseudos = get_repo_from_name(repo_name).get_pseudos(accuracy) +# return pseudos + + def get_installed_repos_and_root(dirpath: Optional[str] = None) -> tuple[list[PseudosRepo], str]: """ Return (all_repos, dirpath) diff --git a/abipy/ml/aseml.py b/abipy/ml/aseml.py index 115ddd86a..079f4c51d 100644 --- a/abipy/ml/aseml.py +++ b/abipy/ml/aseml.py @@ -39,7 +39,8 @@ from ase.md.nptberendsen import NPTBerendsen, Inhomogeneous_NPTBerendsen from ase.md.nvtberendsen import NVTBerendsen from abipy.core import Structure -from abipy.tools.iotools import workdir_with_prefix +import abipy.core.abinit_units as abu +from abipy.tools.iotools import workdir_with_prefix, PythonScript from abipy.tools.printing import print_dataframe from abipy.abio.enums import StrEnum, EnumMixin from abipy.tools.plotting import (set_axlims, add_fig_kwargs, get_ax_fig_plt, get_axarray_fig_plt, set_grid_legend, @@ -53,10 +54,10 @@ def nprocs_for_ntasks(nprocs, ntasks, title=None) -> int: """ Return the number of procs to be used in a multiprocessing Pool. - If negative or None, use all procs in the system. + If negative or None, use hlaf the procs in the system. """ if nprocs is None or nprocs <= 0: - nprocs = max(1, os.cpu_count()) + nprocs = max(1, os.cpu_count() // 2) else: nprocs = int(nprocs) @@ -74,11 +75,19 @@ def nprocs_for_ntasks(nprocs, ntasks, title=None) -> int: class RX_MODE(EnumMixin, StrEnum): # StrEnum added in 3.11 - """Relaxation mode string flags.""" + """ + Relaxation mode string flags. + """ no = "no" ions = "ions" cell = "cell" + #@classmethod + #def from_ionmov_optcell(cls, ionmov: int, optcell: int) -> RX_MODE: + # cls.no + # cls.ions + # cls.cell + def to_ase_atoms(structure: PmgStructure, calc=None) -> Atoms: """Convert pymatgen structure to ASE atoms. Optionally, attach a calculator.""" @@ -105,7 +114,7 @@ def abisanitize_atoms(atoms: Atoms, **kwargs) -> Atoms: Call abisanitize, return new Atoms instance. """ structure = Structure.as_structure(atoms) - new_structure = structure.abisanitize(**kwargs) + new_structure = structure.abi_sanitize(**kwargs) return to_ase_atoms(get_atoms(new_structure), calc=atoms.calc) @@ -241,7 +250,7 @@ def pressure(self) -> float: return -self.stress.trace() / 3 def get_voigt_stress(self): - "xx, yy, zz, yz, xz, xy" + """xx, yy, zz, yz, xz, xy""" from ase.stress import full_3x3_to_voigt_6_stress return full_3x3_to_voigt_6_stress(self.stress) @@ -311,7 +320,6 @@ def zip_sort(xs, ys): return xs[isort].copy(), ys[isort].copy() - def diff_stats(xs, ys): abs_diff = np.abs(ys - xs) return AttrDict( @@ -322,18 +330,18 @@ def diff_stats(xs, ys): ) -def linear_fit_ax(ax, xs, ys, with_labels=False, with_ideal_line=False) -> tuple[float]: +def linear_fit_ax(ax, xs, ys, fontsize, with_label=True, with_ideal_line=False) -> tuple[float]: """ """ - amat = np.vstack([xs, np.ones(len(xs))]).T - alpha, beta = np.linalg.lstsq(amat, ys, rcond=None)[0] - label = r"Linear fit $\alpha={:.2f}$".format(alpha) - ax.plot(xs, alpha*xs + beta, 'r', label=label if with_labels else None) + from scipy.stats import linregress + result = linregress(xs, ys) + label = r"Linear fit $\alpha={:.2f}$, $r^2$={:.2f}".format(result.slope, result.rvalue**2) + ax.plot(xs, result.slope*xs + result.intercept, 'r', label=label if with_label else None) if with_ideal_line: # Plot y = x line ax.plot([xs[0], xs[-1]], [ys[0], ys[-1]], color='k', linestyle='-', - linewidth=1, label='Ideal' if with_labels else None) - return alpha, beta + linewidth=1, label='Ideal' if with_label else None) + return result def make_square_axes(ax_mat): @@ -358,6 +366,14 @@ class AseResultsComparator: ALL_VOIGT_COMPS = "xx yy zz yz xz xy".split() + @classmethod + def pickle_load(cls, workdir): + """ + Reconstruct the object from a pickle file located in workdir. + """ + with open(Path(workdir) / f"{cls.__name__}.pickle", "rb") as fh: + return pickle.load(fh) + @classmethod def from_ase_results(cls, keys: list[str], results_list: list[list[AseResults]]): """ @@ -396,9 +412,9 @@ def __init__(self, structure, keys, ene_list, forces_list, stress_list): """ self.structure = structure self.keys = keys - self.ene_list = ene_list # [nkeys, nsteps] - self.forces_list = forces_list # [nkeys, nsteps, natom, 3] - self.stress_list = stress_list # [nkeys, nsteps, 6] + self.ene_list = ene_list # [nkeys, nsteps] + self.forces_list = forces_list # [nkeys, nsteps, natom, 3] + self.stress_list = stress_list # [nkeys, nsteps, 6] # Consistency check nkeys = len(self) @@ -553,11 +569,11 @@ def plot_energies(self, fontsize=8, **kwargs): stats = diff_stats(xs, ys) ax = ax_mat[irow, icol] ax.scatter(xs, ys, marker="o") - linear_fit_ax(ax, xs, ys) ax.grid(True) ax.set_xlabel(f"{key1} energy", fontsize=fontsize) ax.set_ylabel(f"{key2} energy", fontsize=fontsize) - #ax.legend(loc="best", shadow=True, fontsize=fontsize) + linear_fit_ax(ax, xs, ys, fontsize=fontsize, with_label=True) + ax.legend(loc="best", shadow=True, fontsize=fontsize) if irow == 0: ax.set_title(f"{key1}/{key2} MAE: {stats.MAE:.6f}", fontsize=fontsize) @@ -584,8 +600,9 @@ def plot_forces(self, fontsize=8, **kwargs): stats = diff_stats(xs, ys) ax = ax_mat[irow, icol] ax.scatter(xs, ys, marker="o") - linear_fit_ax(ax, xs, ys) ax.grid(True) + linear_fit_ax(ax, xs, ys, fontsize=fontsize, with_label=True) + ax.legend(loc="best", shadow=True, fontsize=fontsize) f_tex = f"$F_{direction}$" if icol == 0: ax.set_ylabel(f"{key2} {f_tex}", fontsize=fontsize) @@ -616,7 +633,8 @@ def plot_stresses(self, fontsize=6, **kwargs): stats = diff_stats(xs, ys) ax = ax_mat[irow, icol] ax.scatter(xs, ys, marker="o") - linear_fit_ax(ax, xs, ys) + linear_fit_ax(ax, xs, ys, fontsize=fontsize, with_label=True) + ax.legend(loc="best", shadow=True, fontsize=fontsize) ax.grid(True) s_tex = "$\sigma_{%s}$" % voigt_comp if icol == 0: @@ -635,7 +653,7 @@ def plot_energies_traj(self, delta_mode=True, fontsize=6, markersize=2, **kwargs Plot energies along the trajectory. Args: - delta_mode: True to plot difference instead of absolute values. + delta_mode: True to plot differences instead of absolute values. """ key_pairs = self.get_key_pairs() nrows, ncols = 1, len(key_pairs) @@ -654,12 +672,12 @@ def plot_energies_traj(self, delta_mode=True, fontsize=6, markersize=2, **kwargs ax.plot(np.abs(e1 - e2), marker="o", markersize=markersize) ax.set_yscale("log") else: - ax.plot(e1, marker="o", color="red", legend=key1, markersize=markersize) - ax.plot(e2, marker="o", color="blue", legend=key2, markersize=markersize) + ax.plot(e1, marker="o", color="red", label=key1, markersize=markersize) + ax.plot(e2, marker="o", color="blue", label=key2, markersize=markersize) set_grid_legend(ax, fontsize, xlabel='trajectory', ylabel="$|\Delta_E|$" if delta_mode else "$E$", - grid=True, legend=not delta_mode, legend_loc="left", + grid=True, legend_loc="upper left", title=f"{key1}/{key2} MAE: {stats.MAE:.6f} eV") head = "$\Delta$-Energy in eV" if delta_mode else "Energy in eV" @@ -673,7 +691,7 @@ def plot_forces_traj(self, delta_mode=True, fontsize=6, markersize=2, **kwargs): Plot forces along the trajectory. Args: - delta_mode: True to plot difference instead of absolute values. + delta_mode: True to plot differences instead of absolute values. """ # Fx,Fy,Fx along rows, pairs along columns. key_pairs = self.get_key_pairs() @@ -705,7 +723,6 @@ def plot_forces_traj(self, delta_mode=True, fontsize=6, markersize=2, **kwargs): style = dict(marker=marker_idir[idir], markersize=markersize, color=atom1_cmap(float(iatom) / self.natom)) abs_delta = np.abs(f1_tad[:,iatom,idir] - f2_tad[:,iatom,idir]) - #print(f"{abs_delta=}") zero_values = zero_values or np.any(abs_delta == 0.0) ax.plot(abs_delta, **style, label=f"$\Delta {fp_tex}$" if iatom == 0 else None) else: @@ -713,10 +730,11 @@ def plot_forces_traj(self, delta_mode=True, fontsize=6, markersize=2, **kwargs): color=atom1_cmap(float(iatom) / self.natom)) f2_style = dict(marker=marker_idir[idir], markersize=markersize, color=atom2_cmap(float(iatom) / self.natom)) + with_label = (iatom, idir) == (0,0) ax.plot(f1_tad[:,iatom,idir], **f1_style, - label=f"${key1}\, {fp_tex}$" if (iatom, idir) == (0,0) else None) + label=f"${key1}\, {fp_tex}$" if with_label else None) ax.plot(f2_tad[:,iatom,idir], **f2_style, - label=f"${key2}\, {fp_tex}$" if (iatom, idir) == (0,0) else None) + label=f"${key2}\, {fp_tex}$" if with_label else None) if delta_mode: ax.set_yscale("log" if not zero_values else "symlog") @@ -725,7 +743,6 @@ def plot_forces_traj(self, delta_mode=True, fontsize=6, markersize=2, **kwargs): grid=True, legend=not delta_mode, legend_loc="upper left", ylabel=f"$|\Delta {fp_tex}|$" if delta_mode else f"${fp_tex}$") - #for ax in ax_mat.flat: ax.set_box_aspect(1) head = "$\Delta$-forces in eV/Ang" if delta_mode else "Forces in eV/Ang" if "title" not in kwargs: fig.suptitle(f"{head} for {self.structure.latex_formula}") @@ -737,7 +754,7 @@ def plot_stress_traj(self, delta_mode=True, markersize=2, fontsize=6, **kwargs): Plot stresses along the trajectory. Args: - delta_mode: True to plot difference instead of absolute values. + delta_mode: True to plot differences instead of absolute values. """ # Sxx,Syy,Szz,... along rows, pairs along columns. key_pairs = self.get_key_pairs() @@ -748,6 +765,7 @@ def plot_stress_traj(self, delta_mode=True, markersize=2, fontsize=6, **kwargs): ) marker_voigt = {"xx": ">", "yy": "<", "zz": "^", "yz": 1, "xz": 2, "xy":3} + for icol, (key1, key2) in enumerate(key_pairs): # Plot (delta of) stresses along the trajectory. for iv, voigt_comp in enumerate(self.ALL_VOIGT_COMPS): @@ -759,6 +777,7 @@ def plot_stress_traj(self, delta_mode=True, markersize=2, fontsize=6, **kwargs): voigt_comp_tex = "{" + voigt_comp + "}" s1, s2 = self.xy_stress_for_keys(key1, key2, voigt_comp, sort=False) s_style = dict(marker=marker_voigt[voigt_comp], markersize=markersize) + if delta_mode: abs_delta_stress = np.abs(s1 - s2) ax.plot(abs_delta_stress, **s_style, label=f"$|\Delta \sigma_{voigt_comp_tex}|$") @@ -927,22 +946,31 @@ def silence_tensorflow() -> None: class _MyCalculatorMixin: """ - Add _delta_forces and _delta_stress attributes to an ASE calculator. + Add _delta_forces_list and _delta_stress_list internal attributes to an ASE calculator. Extend `calculate` method so that forces and stresses are corrected accordingly. """ + def set_delta_forces(self, delta_forces): """F_Abinitio - F_ML""" - self._delta_forces = delta_forces + if getattr(self, "_delta_forces_list", None) is None: + self._delta_forces_list = [] + self._delta_forces_list.append(delta_forces) def get_delta_forces(self): - return getattr(self, "_delta_forces", None) + delta_forces_list = getattr(self, "_delta_forces_list", None) + if delta_forces_list is not None: return delta_forces_list[-1] + return None def set_delta_stress(self, delta_stress): """S_Abinitio - S_ML""" - self._delta_stress = delta_stress + if getattr(self, "_delta_stress_list", None) is None: + self._delta_stress_list = [] + self._delta_stress_list.append(delta_stress) def get_delta_stress(self): - return getattr(self, "_delta_stress", None) + delta_stress_list = getattr(self, "_delta_stress_list", None) + if delta_stress_list is not None: return delta_stress_list[-1] + return None def calculate( self, @@ -966,8 +994,8 @@ def calculate( forces = self.results["forces"] delta_forces = self.get_delta_forces() if delta_forces is not None: - forces += delta_forces #print("Updating forces with delta_forces:\n", forces) + forces += delta_forces self.results.update( forces=forces, ) @@ -976,8 +1004,8 @@ def calculate( stress = self.results["stress"] delta_stress = self.get_delta_stress() if delta_stress is not None: - stress += delta_stress #print("Updating stress with delta_stress:\n", stress) + stress += delta_stress self.results.update( stress=stress, ) @@ -999,24 +1027,31 @@ class CalcBuilder: """ ALL_NN_TYPES = [ - "m3gnet_old", "m3gnet", + "matgl", "chgnet", "alignn", + #"quip", ] def __init__(self, name: str, **kwargs): self.name = name - self._model = None - def _get_nntype_model_name(self): - if ":" not in self.name: - return self.name, None - nn_type, model_name = self.name.split(":") - return nn_type, model_name + # Extract nn_type and model_name from name + self.nn_type, self.model_name = name, None + if ":" in name: + self.nn_type, self.model_name = name.split(":") + + if self.nn_type not in self.ALL_NN_TYPES: + raise ValueError(f"Invalid {name=}, it should be in {self.ALL_NN_TYPES=}") + + self._model = None def __str__(self): - return f"{self.__class__.__name__}: name: {self.name}" + if self.model_name is not None: + return f"{self.__class__.__name__} nn_type: {self.nn_type}, model_name: {self.model_name}" + else: + return f"{self.__class__.__name__} nn_type: {self.nn_type}" # pickle support. def __getstate__(self): @@ -1030,9 +1065,8 @@ def get_calculator(self) -> Calculator: """ Return ASE calculator with ML potential. """ - nn_type, model_name = self._get_nntype_model_name() - if nn_type == "m3gnet_old": + if self.nn_type == "m3gnet": # m3gnet legacy version. if self._model is None: silence_tensorflow() @@ -1042,16 +1076,16 @@ def get_calculator(self) -> Calculator: raise ImportError("m3gnet not installed. Try `pip install m3gnet`.") from exc if self._model is None: - assert model_name is None + assert self.model_name is None self._model = Potential(M3GNet.load()) class MyM3GNetCalculator(M3GNetCalculator, _MyCalculatorMixin): - """Add delta_forces""" + """Add delta_forces and delta_stress""" return MyM3GNetCalculator(potential=self._model) - if nn_type == "m3gnet": - # m3gnet new version. See https://github.com/materialsvirtuallab/matgl + if self.nn_type == "matgl": + # See https://github.com/materialsvirtuallab/matgl try: import matgl from matgl.ext.ase import M3GNetCalculator @@ -1059,15 +1093,15 @@ class MyM3GNetCalculator(M3GNetCalculator, _MyCalculatorMixin): raise ImportError("matgl not installed. Try `pip install matgl`.") from exc if self._model is None: - model_name = "M3GNet-MP-2021.2.8-PES" if model_name is None else model_name + model_name = "M3GNet-MP-2021.2.8-PES" if self.model_name is None else self.model_name self._model = matgl.load_model(model_name) class MyM3GNetCalculator(M3GNetCalculator, _MyCalculatorMixin): - """Add delta_forces""" + """Add delta_forces and delta_stress""" return MyM3GNetCalculator(potential=self._model) - if nn_type == "chgnet": + if self.nn_type == "chgnet": try: from chgnet.model.dynamics import CHGNetCalculator from chgnet.model.model import CHGNet @@ -1075,38 +1109,39 @@ class MyM3GNetCalculator(M3GNetCalculator, _MyCalculatorMixin): raise ImportError("chgnet not installed. Try `pip install chgnet`.") from exc if self._model is None: - assert model_name is None + assert self.model_name is None self._model = CHGNet.load() class MyCHGNetCalculator(CHGNetCalculator, _MyCalculatorMixin): - """Add delta_forces""" + """Add delta_forces and delta_stress""" return MyCHGNetCalculator(model=self._model) - if nn_type == "alignn": + if self.nn_type == "alignn": try: from alignn.ff.ff import AlignnAtomwiseCalculator, default_path except ImportError as exc: raise ImportError("alignn not installed. See https://github.com/usnistgov/alignn") from exc class MyAlignnCalculator(AlignnAtomwiseCalculator, _MyCalculatorMixin): - """Add delta_forces""" + """Add delta_forces and delta_stress""" - model_name = default_path() if model_name is None else model_name + model_name = default_path() if self.model_name is None else self.model_name return AlignnAtomwiseCalculator(path=model_name) - #if nn_type == "quip": + #if self.nn_type == "quip": # try: # from quippy.potential import Potential # except ImportError as exc: # raise ImportError("quippy not installed. Try `pip install quippy-ase`.\n" + # "See https://github.com/libAtoms/QUIP") from exc - # class MyPotential(Potential, _MyCalculatorMixin): - # """Add delta_forces""" + # class MyQuipPotential(Potential, _MyCalculatorMixin): + # """Add delta_forces and delta_stress""" + # assert self.model_name is None # args_str = "" - # return Potential(args_str="") + # return MyQuipPotential(args_str="") raise ValueError(f"Invalid {self.name=}") @@ -1118,7 +1153,7 @@ class _MlBase: and object persistence via pickle. """ @classmethod - def pickle_load(cls, workdir) -> RelaxScanner: + def pickle_load(cls, workdir): """ Reconstruct the object from a pickle file located in workdir. """ @@ -1140,28 +1175,6 @@ def pickle_dump(self): with open(self.workdir / f"{self.__class__.__name__}.pickle", "wb") as fh: pickle.dump(self, fh) - def set_delta_forces_stress_from_abiml_nc(self, filepath: str) -> None: - """ - Read ab-initio forces from a netcdf file produced by ABINIT and use - these values to set the delta corrections in the calculator. - """ - if os.path.basename(filepath) != "ABIML_RELAX_IN.nc": return - - # Read structure, forces and stresses from the nc file produced by ABINIT. - from abipy.iotools import ETSF_Reader - with ETSF_Reader(filepath) as reader: - abi_forces = reader.read_value("cart_forces") - abi_stress = reader.read_value("cart_stress") - structure = reader.read_structure() - - # Compute ML forces for this structure. - atoms = to_ase_atoms(structure, calc=self.get_calculator()) - ml_forces = atoms.get_forces() - ml_stress = atoms.get_stress() - - # Set delta forces/stresses so that the next invokation to get_calculator include the deltas - self.set_delta_forces_stress(abi_forces - ml_forces, abi_stress - ml_stress) - def set_delta_forces_stress(self, delta_forces, delta_stress) -> None: """Set the value of the delta corrections.""" self.delta_forces = delta_forces @@ -1171,6 +1184,10 @@ def __str__(self): # Delegated to the subclass. return self.to_string() + @lazy_property + def calc_builder(self): + return CalcBuilder(self.nn_name) + def get_calculator_name(self) -> tuple[Calculator, str]: return self.get_calculator(), self.calc_builder.name @@ -1181,6 +1198,7 @@ def get_calculator(self) -> Calculator: if self.delta_forces is not None: #print("Setting delta_forces:\n", self.delta_forces) calc.set_delta_forces(self.delta_forces) + if self.delta_stress is not None: #print("Setting delta_stress:\n", self.delta_stress) calc.set_delta_stress(self.delta_stress) @@ -1302,7 +1320,144 @@ class MlRelaxer(_MlBase): Relax structure with ASE and ML-potential. """ - def __init__(self, atoms: Atoms, relax_mode, fmax, pressure, steps, optimizer, calc_builder, verbose, + @classmethod + def from_abinit_yaml_file(cls, filepath: str, workdir=None, prefix=None) -> MlRelaxer: + """ + Build object from a YAML file produced by ABINIT in hybrid relaxation mode. + """ + # Read yaml file produced by Abinit: + # + # iteration_state: {dtset: 1, itime: 1, icycle: 1, } + # comment : Summary of ground state results + # lattice_vectors: + # - [ -5.1690735, -5.1690735, 0.0000000, ] + # - [ -5.1690735, 0.0000000, -5.1690735, ] + # - [ 0.0000000, -5.1690735, -5.1690735, ] + # lattice_lengths: [ 7.31017, 7.31017, 7.31017, ] + # lattice_angles: [ 60.000, 60.000, 60.000, ] # degrees, (23, 13, 12) + # lattice_volume: 2.7622826E+02 + # convergence: {deltae: -1.926E-11, res2: 3.761E-10, residm: 1.588E-05, diffor: null, } + # etotal : -8.46248947E+00 + # entropy : 0.00000000E+00 + # fermie : 1.42500714E-01 + # cartesian_stress_tensor: # hartree/bohr^3 + # - [ 3.06384355E-06, 0.00000000E+00, 0.00000000E+00, ] + # - [ 0.00000000E+00, 3.06384355E-06, 0.00000000E+00, ] + # - [ 0.00000000E+00, 0.00000000E+00, 3.06384355E-06, ] + # pressure_GPa: -9.0141E-02 + # xred : + # - [ 5.0000E-01, 5.0000E-01, 5.0000E-01, Si] + # - [ 2.5000E-01, 2.5000E-01, 2.5000E-01, Si] + # cartesian_forces: # hartree/bohr + # - [ 5.08705549E-32, 5.08705549E-32, -1.52611665E-31, ] + # - [ -5.08705549E-32, -5.08705549E-32, 1.52611665E-31, ] + # force_length_stats: {min: 1.68718543E-31, max: 1.68718543E-31, mean: 1.68718543E-31, } + # format_version: 1 + # natom: 2 + # ionmov: 1 + # optcell: 2 + # nn_name: matgl + # prtvol: 1 + + from ruamel.yaml import YAML + with open(os.path.expanduser(filepath), "rt") as fh: + doc = YAML().load(fh) #; print(doc) + + format_version = doc.pop("format_version") + natom = doc.pop("natom") + ntypat = doc.pop("ntypat") + typat = np.array(doc.pop("typat"), dtype=int) + znucl = np.array(doc.pop("znucl"), dtype=float) + rprim = np.array(doc.pop("lattice_vectors")) + xred = np.array(doc.pop("xred")) + xred = np.array(xred[:,:3], dtype=float) + abi_cart_forces = np.array(doc.pop("cartesian_forces")) * abu.Ha_eV / abu.Bohr_Ang + abi_cart_stresses = np.array(doc.pop("cartesian_stress_tensor")) * abu.Ha_eV / (abu.Bohr_Ang**3) + + ionmov = doc.pop("ionmov") + optcell = doc.pop("optcell") + iatfix = doc.pop("iatfix") # [3,natom] array or None if unconstrained. + strtarget = np.array(doc.pop("strtarget"), dtype=float) + nn_name = doc.pop("nn_name", default="chgnet") + verbose = doc.pop("prtvol") + + structure = Structure.from_abivars( + acell=3*[1.0], + rprim=rprim, + typat=typat, + xred=xred, + ntypat=ntypat, + znucl=znucl, + ) #; print(structure) + + atoms = structure.to_ase_atoms() + if iatfix is not None: + raise NotImplementedError() + #aseml.fix_atoms(atoms, fix_inds=fix_inds, fix_symbols=fix_symbols) + + ###################################################################### + # Consistency check as not all the Abinit options are supported by ASE + ###################################################################### + relax_mode = RX_MODE.cell if optcell != 0 else RX_MODE.ions + + allowed_optcells = (0, 2) + if optcell not in allowed_optcells: + raise ValueError(f"{optcell=} not in {allowed_optcells=}") + + # Target pressure is taken from strtarget. + # The components of the stress tensor are stored in a.u. according to: + # (1,1) → 1; (2,2) → 2; (3,3) → 3; (2,3) → 4; (3,1) → 5; (1,2) → 6. + pressure = -strtarget[0] * abu.HaBohr3_GPa + if np.any(strtarget[:3] != strtarget[0]): + raise ValueError(f"Only hydrostatic stress in strtarget is supported. {strtarget=}") + if np.any(strtarget[3:] != 0.0): + raise ValueError(f"Off diagonal components in strtarget are not supported. {strtarget=}") + + # Set internal parameters according to yaml file and build object. + fmax, steps, optimizer = 0.01, 500, "BFGS" + + new = cls(atoms, relax_mode, fmax, pressure, steps, optimizer, nn_name, verbose, + workdir=workdir, prefix=prefix) + + # Set delta forces and delta stress if script is called by ABINIT. + + #new.set_delta_forces_stress_from_abiml_nc(filepath) + ## Compute ML forces for this structure. + #atoms = to_ase_atoms(structure, calc=new.get_calculator()) + #ml_cart_forces = atoms.get_forces() + #ml_cart_stress = atoms.get_stress() + + ## Set delta forces/stresses so that the next invokation to get_calculator include the deltas + #new.set_delta_forces_stress(abi_cart_forces - ml_cart_forces, abi_stress - ml_cart_stress) + + return new + + def write_output_file_for_abinit(self) -> Path: + """ + Write output file with results in a format that can be parsed by ABINIT. + Return path to output file. + """ + filepath = self.get_path("ABI_MLRELAXER.out", "Output file for hybrid relaxation with ABINIT.") + format_version = 1 + + def fmt_vec3(vec) -> str: + return "{:.12e} {:.12e} {:.12e}".format(*vec) + + with open(filepath, "wt") as fh: + fh.write("%i # format_version\n" % format_version) + fh.write("%i # natom\n" % len(self.atoms)) + # Write lattice vectors. + rprimd = self.atoms.cell.array * abu.Ang_Bohr + for i in range(3): + fh.write("%s # lattice vector %i\n" % (fmt_vec3(rprimd[i]), i+1)) + # Write relaxed fractional coordinates. + fh.write("xred\n") + for atom in self.atoms: + fh.write(fmt_vec3(atom.scaled_position) + "\n") + + return filepath + + def __init__(self, atoms: Atoms, relax_mode, fmax, pressure, steps, optimizer, nn_name, verbose, workdir, prefix=None): """ Args: @@ -1312,7 +1467,7 @@ def __init__(self, atoms: Atoms, relax_mode, fmax, pressure, steps, optimizer, c pressure: Target pressure. steps: max number of steps for relaxation. optimizer: name of the ASE optimizer to use. - calc_builder: + nn_name: verbose: Verbosity level. """ super().__init__(workdir, prefix) @@ -1323,11 +1478,9 @@ def __init__(self, atoms: Atoms, relax_mode, fmax, pressure, steps, optimizer, c self.steps = steps self.optimizer = optimizer self.pressure = pressure - self.calc_builder = calc_builder + self.nn_name = nn_name self.verbose = verbose - self.atoms.calc = self.get_calculator() - def to_string(self, verbose=0) -> str: """String representation with verbosity level `verbose`.""" return f"""\ @@ -1339,7 +1492,7 @@ def to_string(self, verbose=0) -> str: steps = {self.steps} optimizer = {self.optimizer} pressure = {self.pressure} - calculator = {self.calc_builder} + nn_name = {self.nn_name} workdir = {self.workdir} verbose = {self.verbose} @@ -1349,10 +1502,11 @@ def to_string(self, verbose=0) -> str: """ - def run(self) -> None: + def run(self): """Run structural relaxation.""" #self.pickle_dump() workdir = self.workdir + self.atoms.calc = self.get_calculator() print(f"Relaxing structure with relax mode: {self.relax_mode} ...") relax_kws = dict(calculator=self.atoms.calc, @@ -1377,17 +1531,18 @@ def run(self) -> None: label = "xdatcar with structural relaxation" write_vasp_xdatcar(self.get_path("XDATCAR", label), relax.traj, label=label) - #from abipy.core.structure import StructDiff - #diff = StructDiff(["INPUT", "ML RELAXED"], [initial_atoms, self.atoms]) + # Write output file for Abinit + self.write_output_file_for_abinit() self._finalize() + return relax class MlMd(_MlBase): """Perform MD calculations with ASE and ML potential.""" def __init__(self, atoms: Atoms, temperature, timestep, steps, loginterval, - ensemble, calc_builder, verbose, workdir, prefix=None): + ensemble, nn_name, verbose, workdir, prefix=None): """ Args: atoms: @@ -1396,7 +1551,7 @@ def __init__(self, atoms: Atoms, temperature, timestep, steps, loginterval, steps: loginterval: ensemble: - calc_builder: + nn_name: verbose: Verbosity level. workdir: prefix: @@ -1408,11 +1563,9 @@ def __init__(self, atoms: Atoms, temperature, timestep, steps, loginterval, self.steps = steps self.loginterval = loginterval self.ensemble = ensemble - self.calc_builder = calc_builder + self.nn_name = nn_name self.verbose = verbose - self.atoms.calc = self.get_calculator() - def to_string(self, verbose=0) -> str: """String representation with verbosity level `verbose`.""" return f"""\ @@ -1438,6 +1591,7 @@ def run(self) -> None: """Run MD""" #self.pickle_dump() workdir = self.workdir + self.atoms.calc = self.get_calculator() traj_file = self.get_path("md.traj", "ASE MD trajectory") logfile = self.get_path("md.log", "ASE MD log file") @@ -1538,17 +1692,17 @@ class MlGsList(_MlNebBase): Inherits from _MlNebBase so that we can reuse postprocess_images and read_neb_data. """ - def __init__(self, atoms_list: list[Atoms], calc_builder, verbose, + def __init__(self, atoms_list: list[Atoms], nn_name, verbose, workdir, prefix=None): """ Args: atoms_list: List of ASE atoms - calc_builder: + nn_name: verbose: Verbosity level. """ super().__init__(workdir, prefix) self.atoms_list = atoms_list - self.calc_builder = calc_builder + self.nn_name = nn_name self.verbose = verbose def to_string(self, verbose=0) -> str: @@ -1557,9 +1711,9 @@ def to_string(self, verbose=0) -> str: {self.__class__.__name__} parameters: - calculator = {self.calc_builder} - workdir = {self.workdir} - verbose = {self.verbose} + nn_name = {self.nn_name} + workdir = {self.workdir} + verbose = {self.verbose} """ @@ -1588,7 +1742,7 @@ class MlNeb(_MlNebBase): def __init__(self, initial_atoms: Atoms, final_atoms: Atoms, nimages, neb_method, climb, optimizer, relax_mode, fmax, pressure, - calc_builder, verbose, workdir, prefix=None): + nn_name, verbose, workdir, prefix=None): """ Args: initial_atoms @@ -1600,7 +1754,7 @@ def __init__(self, initial_atoms: Atoms, final_atoms: Atoms, relax_mode: fmax: pressure: - calc_builder: + nn_name: verbose: workdir: prefix: @@ -1618,7 +1772,7 @@ def __init__(self, initial_atoms: Atoms, final_atoms: Atoms, RX_MODE.validate(self.relax_mode) self.fmax = fmax self.pressure = pressure - self.calc_builder = calc_builder + self.nn_name = nn_name self.verbose = verbose def to_string(self, verbose=0) -> str: @@ -1634,7 +1788,7 @@ def to_string(self, verbose=0) -> str: pressure = {self.pressure} relax_mode = {self.relax_mode} fmax = {self.fmax} - calculator = {self.calc_builder} + nn_name = {self.nn_name} workdir = {self.workdir} verbose = {self.verbose} @@ -1698,8 +1852,7 @@ def run(self) -> None: nebtraj_file = str(workdir / "neb.traj") logfile = self.get_path("neb.log", "Log file of NEB calculation.") optimizer = opt_class(neb, trajectory=nebtraj_file, logfile=logfile) - - #print("Starting NEB algorithm with optimizer:", optimizer, "...") + print("Starting NEB algorithm with optimizer:", opt_class, "...") optimizer.run(fmax=self.fmax) # To read the last nimages atoms e.g. 5: read('neb.traj@-5:') @@ -1739,7 +1892,7 @@ class MultiMlNeb(_MlNebBase): """ def __init__(self, atoms_list: list[Atoms], nimages, neb_method, climb, optimizer, relax_mode, fmax, pressure, - calc_builder, verbose, workdir, prefix=None): + nn_name, verbose, workdir, prefix=None): """ Args: atoms_list: @@ -1750,7 +1903,7 @@ def __init__(self, atoms_list: list[Atoms], nimages, neb_method, climb, optimize relax_mode: fmax: pressure: - calc_builder: + nn_name: verbose: workdir: prefix: @@ -1765,7 +1918,7 @@ def __init__(self, atoms_list: list[Atoms], nimages, neb_method, climb, optimize RX_MODE.validate(self.relax_mode) self.fmax = fmax self.pressure = pressure - self.calc_builder = calc_builder + self.nn_name = nn_name self.verbose = verbose def to_string(self, verbose=0) -> str: @@ -1781,7 +1934,7 @@ def to_string(self, verbose=0) -> str: relax_mode = {self.relax_mode} fmax = {self.fmax} pressure = {self.pressure} - calculator = {self.calc_builder} + nn_name = {self.nn_name} workdir = {self.workdir} verbose = {self.verbose} """ @@ -1800,7 +1953,7 @@ def run(self) -> None: for i in range(len(atoms_list) - 1): ml_neb = MlNeb(atoms_list[i], atoms_list[i+1], self.nimages, self.neb_method, self.climb, self.optimizer, - self.relax_mode, self.fmax, self.pressure, self.calc_builder, self.verbose, camp_dirs[i]) + self.relax_mode, self.fmax, self.pressure, self.nn_name, self.verbose, camp_dirs[i]) ml_neb.run() # Read energies from json files and remove first/last point depending on CAMP index.. @@ -1878,7 +2031,7 @@ class MlOrderer(_MlBase): """ Order a disordered structure using pymatgen and ML potential. """ - def __init__(self, structure, max_ns, optimizer, relax_mode, fmax, pressure, steps, calc_builder, verbose, + def __init__(self, structure, max_ns, optimizer, relax_mode, fmax, pressure, steps, nn_name, verbose, workdir, prefix=None): """ Args: @@ -1889,7 +2042,7 @@ def __init__(self, structure, max_ns, optimizer, relax_mode, fmax, pressure, ste fmax: pressure: steps: - calc_builder: + nn_name: verbose: workdir: prefix: @@ -1903,7 +2056,7 @@ def __init__(self, structure, max_ns, optimizer, relax_mode, fmax, pressure, ste self.fmax = fmax self.pressure = pressure self.steps = steps - self.calc_builder = calc_builder + self.nn_name = nn_name self.verbose = verbose def to_string(self, verbose=0) -> str: @@ -1918,7 +2071,7 @@ def to_string(self, verbose=0) -> str: fmax = {self.fmax} pressure = {self.pressure} steps = {self.steps} - calculator = {self.calc_builder} + nn_name = {self.nn_name} workdir = {self.workdir} verbose = {self.verbose} @@ -2017,7 +2170,7 @@ class MlPhonons(_MlBase): """Compute phonons with ASE and ML potential.""" def __init__(self, atoms: Atoms, supercell, kpts, asr, nqpath, - relax_mode, fmax, pressure, steps, optimizer, calc_builder, + relax_mode, fmax, pressure, steps, optimizer, nn_name, verbose, workdir, prefix=None): """ Args: @@ -2030,7 +2183,7 @@ def __init__(self, atoms: Atoms, supercell, kpts, asr, nqpath, fmax: steps: optimizer: - calc_builder, + nn_name, verbose: workdir: prefix: @@ -2047,7 +2200,7 @@ def __init__(self, atoms: Atoms, supercell, kpts, asr, nqpath, self.pressure = pressure self.steps = steps self.optimizer = optimizer - self.calc_builder = calc_builder + self.nn_name = nn_name self.verbose = verbose def to_string(self, verbose=0): @@ -2065,7 +2218,7 @@ def to_string(self, verbose=0): steps = {self.steps} optimizer = {self.optimizer} pressure = {self.pressure} - calculator = {self.calc_builder} + nn_name = {self.nn_name} workdir = {self.workdir} verbose = {self.verbose} @@ -2145,17 +2298,17 @@ class MlCompareWithAbinitio(_MlNebBase): Compare ab-initio energies, forces and stresses with ML results. """ - def __init__(self, filepath, nn_names, traj_range, verbose, workdir, prefix=None): + def __init__(self, filepaths, nn_names, traj_range, verbose, workdir, prefix=None): """ Args: - filepath: File produced by the ab-initio code with energies, forces and stresses. + filepaths: List of file produced by the ab-initio code with energies, forces and stresses. nn_names: String or list of strings defining the NN potential. traj_range: Trajectory range. None to include all steps. verbose: Verbosity level. workdir: Working directory. """ super().__init__(workdir, prefix) - self.filepath = filepath + self.filepaths = list_strings(filepaths) self.traj_range = traj_range self.nn_names = list_strings(nn_names) self.verbose = verbose @@ -2166,7 +2319,7 @@ def to_string(self, verbose=0) -> str: {self.__class__.__name__} parameters: - filepath = {self.filepath} + filepaths = {self.filepaths} traj_range = {self.traj_range} nn_names = {self.nn_names} workdir = {self.workdir} @@ -2174,18 +2327,24 @@ def to_string(self, verbose=0) -> str: """ - def _get_abinitio_results(self) -> list[AseResults]: + def get_abinitio_results(self) -> list[AseResults]: + results = [] + for filepath in self.filepaths: + results.extend(self._get_results_filepath(filepath)) + return results + + def _get_results_filepath(self, filepath) -> list[AseResults]: """ Extract ab-initio results from self.filepath according to the file extension. """ - basename = os.path.basename(self.filepath) + basename = os.path.basename(filepath) abi_results = [] from fnmatch import fnmatch if basename.endswith("_HIST.nc"): # Abinit HIST file produced by a structural relaxation. from abipy.dynamics.hist import HistFile - with HistFile(self.filepath) as hist: + with HistFile(filepath) as hist: # etotals in eV units. etotals = hist.etotals if self.traj_range is None: self.traj_range = range(0, len(hist.etotals), 1) @@ -2213,7 +2372,10 @@ def get_energy_step(step: dict) -> float: return total_energy from pymatgen.io.vasp.outputs import Vasprun - vasprun = Vasprun(self.filepath) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + vasprun = Vasprun(filepath) + num_steps = len(vasprun.ionic_steps) if self.traj_range is None: self.traj_range = range(0, num_steps, 1) for istep, step in enumerate(vasprun.ionic_steps): @@ -2225,7 +2387,7 @@ def get_energy_step(step: dict) -> float: abi_results.append(r) return abi_results - raise ValueError(f"Don't know how to extract data from {self.filepath}") + raise ValueError(f"Don't know how to extract data from: {filepath=}") def run(self, nprocs, print_dataframes=True) -> AseResultsComparator: """ @@ -2235,7 +2397,7 @@ def run(self, nprocs, print_dataframes=True) -> AseResultsComparator: workdir = self.workdir labels = ["abinitio"] - abi_results = self._get_abinitio_results() + abi_results = self.get_abinitio_results() results_list = [abi_results] ntasks = len(abi_results) @@ -2258,11 +2420,36 @@ def run(self, nprocs, print_dataframes=True) -> AseResultsComparator: comp = AseResultsComparator.from_ase_results(labels, results_list) + # Write pickle file for object persistence. + with open(self.workdir / f"{comp.__class__.__name__}.pickle", "wb") as fh: + pickle.dump(comp, fh) + + py_path = self.workdir / "analyze.py" + print("Writing python script to analyze the results in:", py_path.name) + with PythonScript(py_path) as script: + script.add_text(""" +def main(): + from abipy.ml.aseml import AseResultsComparator + c = AseResultsComparator.pickle_load(".") + with_stress = True + from abipy.tools.plotting import Exposer + exposer = "mpl" # or "panel" + with Exposer.as_exposer(exposer) as e: + e(c.plot_energies(show=False)) + e(c.plot_forces(delta_mode=True, show=False)) + e(c.plot_energies_traj(delta_mode=True, show=False)) + e(c.plot_energies_traj(delta_mode=False, show=False)) + if with_stress: + e(c.plot_stresses(delta_mode=True, show=False)) + e(c.plot_forces_traj(delta_mode=True, show=False)) + e(c.plot_stress_traj(delta_mode=True, show=False)) +""").add_main() + if print_dataframes: # Write dataframes to disk in CSV format. - forces_df = comp.get_forces_dataframe() #; print(forces_df) + forces_df = comp.get_forces_dataframe() self.write_df(forces_df, "cart_forces.csv", info="CSV file with cartesian forces.") - stress_df = comp.get_stress_dataframe() #; print(stress_df) + stress_df = comp.get_stress_dataframe() self.write_df(stress_df, "voigt_stress.csv", info="CSV file with cartesian stresses in Voigt notation") self._finalize() diff --git a/abipy/ml/ml_relax.py b/abipy/ml/ml_relax.py index dfd4ea5cc..c9dc16f00 100644 --- a/abipy/ml/ml_relax.py +++ b/abipy/ml/ml_relax.py @@ -246,7 +246,7 @@ def run(self, workdir=None, prefix=None): ml_calc = CalcBuilder(self.nn_name).get_calculator() print(f"\nBegin ABINIT + {self.nn_name} hybrid relaxation") - if self.xc.is_gga_family == "GGA": + if self.xc == "PBE": print(f"Starting from ML-optimized Atoms as {self.xc=}") atoms = ml_opt.atoms.copy() atoms = abisanitize_atoms(atoms) @@ -346,6 +346,7 @@ def run(self, workdir=None, prefix=None): }[xc] print(f"Using {repo_name=}") pseudos = get_repo_from_name(repo_name).get_pseudos("standard") + #pseudos = get_pseudos(xc) from ase.build import bulk atoms = bulk('Si') diff --git a/abipy/ml/relax_scanner.py b/abipy/ml/relax_scanner.py index 6be6e69f8..e217b64ba 100644 --- a/abipy/ml/relax_scanner.py +++ b/abipy/ml/relax_scanner.py @@ -1,5 +1,5 @@ """ -Objects to perform ASE calculations with machine-learned potentials. +Objects to perform ASE calculations with machine-learning potentials. """ from __future__ import annotations @@ -10,19 +10,19 @@ import itertools import warnings import dataclasses +import shutil import numpy as np import pandas as pd -#import abipy.tools.duck as duck from pathlib import Path from multiprocessing import Pool -#from typing import Type, Any, Optional, Union +#from typing import Any from monty.json import MontyEncoder from monty.functools import lazy_property from monty.collections import dict2namedtuple from pymatgen.core.lattice import Lattice from pymatgen.util.coord import pbc_shortest_vectors - +from ase.io.vasp import write_vasp # write_vasp_xdatcar, from abipy.core import Structure from abipy.tools.plotting import add_fig_kwargs, get_ax_fig_plt #, get_axarray_fig_plt, from abipy.tools.iotools import workdir_with_prefix, PythonScript, ShellScript @@ -219,7 +219,7 @@ def _make_directory(self, start, stop) -> Path: newdir.mkdir() return newdir - def get_atoms_with_frac_coords(self, frac_coords, with_fixex_atoms=True) -> Atoms: + def get_atoms_with_frac_coords(self, frac_coords, with_fixed_atoms=True) -> Atoms: """ Return Atoms instance with frac_coords at site index `isite`. By default, Fixed Contraints are applied to all atoms except isite. @@ -227,10 +227,20 @@ def get_atoms_with_frac_coords(self, frac_coords, with_fixex_atoms=True) -> Atom structure = self.initial_structure.copy() structure._sites[self.isite].frac_coords = np.array(frac_coords) atoms = get_atoms(structure) - if with_fixex_atoms: + if with_fixed_atoms: fix_atoms(atoms, fix_inds=[i for i in range(len(atoms)) if i != self.isite]) return atoms + def get_structure_with_two_frac_coords(self, frac_coords1, frac_coords2) -> Structure: + """ + Return Structure instance with frac_coords at site index `isite`. + """ + new_structure = self.initial_structure.copy() + species = new_structure._sites[self.isite].species + new_structure._sites[self.isite].frac_coords = np.array(frac_coords1) + new_structure.insert(self.isite, species, np.array(frac_coords2)) + return new_structure + def run_start_count(self, start: int, count: int) -> Path: """ Invoke ASE to relax `count` structures starting from index `start`. @@ -316,13 +326,15 @@ def main(): #rsa.histplot() # Tolerances for pairs detection. - ediff_tol, dist_tol = 1e-3, 3.5 + #ediff_tol, dist_tol = 1e-3, 3.5 # For max values. + ediff_tol, dist_tol = (0, 1e-3), (0.5, 3.5) # For ranges. # Find pairs and save them to file rsa.pairs_enediff_dist(ediff_tol, dist_tol, neb_method=None) # Find pairs and compute static transition path. - # NO NEB HERE, just linear interpolation between final and end points and energy calculations. + # NO NEB HERE, just linear interpolation between final and end points + # followed by total energy calculations. #rsa.pairs_enediff_dist(ediff_tol, dist_tol, neb_method="static") # Find pairs and compute transition path with NEB. @@ -410,6 +422,19 @@ def __init__(self, entries: list[Entry], scanner: RelaxScanner, verbose: int = 0 self.scanner = scanner self.verbose = verbose + # Look for VASP templates. + self.jobsh_path = self.workdir / "job.sh" + self.incar_path = self.workdir / "INCAR" + self.potcar_path = self.workdir / "POTCAR" + self.kpoints_path = self.workdir / "KPOINTS" + print("has_vasp_inputs", self.has_vasp_inputs()) + + def has_vasp_inputs(self) -> bool: + """Return True if all the input files required to run VASP exist.""" + return True + return self.jobsh_path.exists() and self.incar_path.exists() and \ + self.potcar_path.exists() and self.kpoints_path.exists() + @property def workdir(self): return self.scanner.workdir @@ -421,7 +446,7 @@ def workdir(self): @lazy_property def df(self) -> pd.DataFrame: """ - Dataframe with the total energy in eV and the relaxed coords + Dataframe with the total energy in eV and the relaxed coordinates of the atomic site that has been relaxed. """ dict_list = [] @@ -439,7 +464,7 @@ def df(self) -> pd.DataFrame: df = pd.DataFrame(dict_list).sort_values(by="energy", ignore_index=True) - # Add medata + # Add metadata df.attrs['lattice_matrix'] = entries[0].structure.lattice.matrix df.attrs['structure'] = entries[0].structure return df @@ -455,14 +480,6 @@ def to_string(self, verbose=0) -> str: def lattice(self): return Lattice(self.df.attrs["lattice_matrix"]) - #def get_emin_emax(self, fact=0.001) -> tuple[float, float]: - # """ - # Return min and max energy. The range is extended by .1% on each side. - # """ - # emin, emax = self.df['energy'].min(), self.df['energy'].max() - # emin -= fact * abs(emin); emax += fact * abs(emax) - # return emin, emax - def pairs_enediff_dist(self, ediff_tol=1e-3, dist_tol=3.5, neb_method=None, nprocs=-1) -> list[Pair]: """ Find pairs (i.e. relaxed configurations) that differ in energy less than ediff_tol @@ -470,8 +487,11 @@ def pairs_enediff_dist(self, ediff_tol=1e-3, dist_tol=3.5, neb_method=None, npro (minimum-image convention is applied). Args: - ediff_tol: Energy difference in eV. - dist_tol: Tolerance on site distance in Ang. + ediff_tol: Energy difference in eV. Tuple for range, scalar for max value. + dist_tol: Tolerance on site distance in Ang. Tuple for range, scalar for max value. + neb_method: None to print pairs only, "static" to compute energy profile along static path + or ASE neb method to perform NEB calculation. + nprocs: Number of processes for Multiprocessing parallelism. Return: list of Pair objects. """ @@ -484,19 +504,33 @@ def adiff_matrix(vec): aediff_mat = adiff_matrix(energy) xreds = self.df[["xred0", "xred1", "xred2"]].to_numpy(dtype=float) + # Define ranges + if isinstance(ediff_tol, (tuple, list)): + ediff_min, ediff_max = ediff_tol + else: + ediff_min, ediff_max = 0.0, float(ediff_tol) + + if isinstance(dist_tol, (tuple, list)): + dist_min, dist_max = dist_tol + else: + dist_min, dist_max = 0.0, float(dist_tol) + # Find pairs. pairs = [] inds = np.triu_indices_from(aediff_mat) for i, j in zip(inds[0], inds[1]): - if aediff_mat[i,j] > ediff_tol or i == j: continue + if i == j or not (ediff_max >= aediff_mat[i,j] >= ediff_min): continue _, d2 = pbc_shortest_vectors(self.lattice, xreds[i], xreds[j], return_d2=True) dist = np.sqrt(float(d2)) - if dist > dist_tol: continue + if not (dist_max >= dist >= dist_min): continue pair = Pair(index1=i, index2=j, ediff=energy[j] - energy[i], dist=dist, frac_coords1=xreds[i], frac_coords2=xreds[j]) pairs.append(pair) - print(f"Found {len(pairs)} pair(s) with {ediff_tol=} eV and {dist_tol=} Ang.") + print(f"Found {len(pairs)} pair(s) with {ediff_min=}, {ediff_max=} eV and {dist_min=}, {dist_max=} Ang.") + if not pairs: + print("Returning immediately from pairs_enediff_dist") + return pairs if neb_method is None: # Build dataframe with pairs info. @@ -507,6 +541,7 @@ def adiff_matrix(vec): nprocs = nprocs_for_ntasks(nprocs, len(pairs), title=f"Computing transition energies for each pair with {neb_method=}") + # Run'em all if nprocs == 1: d_list = [self.run_pair(pair, neb_method=neb_method) for pair in pairs] else: @@ -517,7 +552,8 @@ def adiff_matrix(vec): df = pd.DataFrame(d_list) print(df) - df["ediff_tol"], df["dist_tol"] = ediff_tol, dist_tol + df["ediff_min"], df["ediff_max"] = ediff_min, ediff_max + df["distl_min"], df["dist_max"] = dist_min, dist_max df = df.sort_values(by=["ediff", "dist"], ignore_index=True) print_dataframe(df) path = self.workdir / f"{neb_method=}_data.csv".replace("'", "") @@ -529,7 +565,7 @@ def adiff_matrix(vec): def run_pair(self, pair: Pair, neb_method="static", nimages=14, climb=False) -> dict: """ Perform NEB calculation for the given pair. Return dictionary with results. - NB: Contraints are enforces during the NEB + NB: Contraints are enforced during the NEB. Args: pair: Info on Pair. @@ -538,8 +574,9 @@ def run_pair(self, pair: Pair, neb_method="static", nimages=14, climb=False) -> nimages: Number of images climb: True to use climbing images. """ - # Get atoms with Fixed constraints. + # Get atoms with fixed constraints. scanner = self.scanner + delta_frac = pair.frac_coords2 - pair.frac_coords1 initial_atoms = scanner.get_atoms_with_frac_coords(pair.frac_coords1) final_atoms = scanner.get_atoms_with_frac_coords(pair.frac_coords2) @@ -560,22 +597,41 @@ def run_pair(self, pair: Pair, neb_method="static", nimages=14, climb=False) -> #interpolate(images, mic=False, apply_constraint=False) # NB: It's not a NEB ML object but it provides a similar API. - ml_neb = MlGsList(neb.images, calc_builder, self.verbose, - workdir=self.workdir / f"GSLIST/pair_{pair_str}") + my_workdir = self.workdir / f"GSLIST/pair_{pair_str}" + ml_neb = MlGsList(neb.images, scanner.nn_name, self.verbose, + workdir=my_workdir) + + # Write POSCAR file with the two sites so that we can visualize it with e.g. Vesta. + structure_with_two_sites = scanner.get_structure_with_two_frac_coords(pair.frac_coords1, pair.frac_coords2) + structure_with_two_sites.to(filename=str(my_workdir/ "TWO_SITES_POSCAR.vasp")) else: # Real NEB stuff + my_workdir = self.workdir / f"NEB/pair_{pair_str}" ml_neb = MlNeb(initial_atoms, final_atoms, nimages, neb_method, climb, scanner.optimizer_name, relax_mode, scanner.fmax, scanner.pressure, - calc_builder, self.verbose, - workdir=self.workdir / f"NEB/pair_{pair_str}") + scanner.nn_name, self.verbose, + workdir=my_workdir) if self.verbose: print(ml_neb.to_string(verbose=self.verbose)) ml_neb.run() neb_data = ml_neb.read_neb_data() + if self.has_vasp_inputs(): + # Create directories to run VASP. Note symlinks except for jobsh_path. + for ip in range(2): + directory = ml_neb.workdir / f"VASP_JOB_{ip}" + directory.mkdir() + atoms = initial_atoms if ip == 0 else final_atoms + write_vasp(directory / "POSCAR", atoms, label=None) + shutil.copyfile(self.jobsh_path, directory / "job.sh") + (directory / "INCAR").symlink_to(self.incar_path) + (directory / "KPOINTS").symlink_to(self.kpoints_path) + (directory / "POTCAR").symlink_to(self.potcar_path) + (directory / "__INITIALIZED__").touch() + # Build out_data dict. out_data = pair.get_dict4pandas() @@ -597,30 +653,3 @@ def histplot(self, ax=None, **kwargs): import seaborn as sns sns.histplot(self.df, x="energy", ax=ax) return fig - - -#def linear_path_with_extra_points(isite, initial, final, nimages, nextra, step) -> list[Atoms]: -# -# # Get cart coords -# p0 = initial[isite].position -# p1 = final[isite].position -# -# dvers_10 = (p1 - p0) / np.linalg.norm(p1 - p0) -# -# from ase.neb import interpolate -# first_segment = [] -# for i in range(nextra - 1): -# new = initial.copy() -# new[isite].position = p0 - i * step * dvers_10 -# first_segment.append(new) -# -# images += [initial.copy() for i in range(nimages - 2)] -# -# for i in range(nextra): -# new = final.copy() -# new[isite].position = p1 + i * step * dvers_10 -# images.append(new) -# -# interpolate(images, mic=False, interpolate_cell=False, use_scaled_coord=False, apply_constraint=None) -# -# return first_segment + images + last_segment diff --git a/abipy/scripts/abiml.py b/abipy/scripts/abiml.py index 73e1d5851..d2a71d3f7 100755 --- a/abipy/scripts/abiml.py +++ b/abipy/scripts/abiml.py @@ -84,7 +84,7 @@ def add_neb_opts(f): """Add CLI options for NEB calculations.""" f = click.option("--nimages", "-n", default=14, type=click.IntRange(3, None), show_default=True, help='Number of NEB images including initial/final points. Must be >= 3')(f) - f = click.option("--relax", "-r", default="ions", show_default=True, type=click.Choice(["no", "ions", "cell"]), + f = click.option("--relax-mode", "-r", default="ions", show_default=True, type=click.Choice(["no", "ions", "cell"]), help="Relax initial and final structure. Use `cell` to relax ions and cell, " + "`ions` to relax atomic positions only, `no` to disable relaxation")(f) f = click.option("--fmax", default=0.01, type=float, show_default=True, help='Stopping criterion.')(f) @@ -112,14 +112,19 @@ def add_workdir_verbose_opts(f): return f +def add_nn_name_opt(f): + """Add CLI options to select NN potential.""" + f = click.option("--nn-name", "-nn", default="chgnet", show_default=True, + type=click.Choice(aseml.CalcBuilder.ALL_NN_TYPES), + help='ML potential to be used')(f) + return f + + @click.group() @click.pass_context -@click.option("--nn-name", "-nn", default="m3gnet", show_default=True, - type=click.Choice(aseml.CalcBuilder.ALL_NN_TYPES), - help='ML potential to be used') @click.option("--seaborn", "-sns", default=None, show_default=True, help='Use seaborn settings. Accept value defining context in ("paper", "notebook", "talk", "poster").') -def main(ctx, nn_name, seaborn): +def main(ctx, seaborn): """Script to perform calculations with ML potentials.""" ctx.ensure_object(dict) @@ -130,19 +135,16 @@ def main(ctx, nn_name, seaborn): sns.set(context=seaborn, style='darkgrid', palette='deep', font='sans-serif', font_scale=1, color_codes=False, rc=None) - global CALC_BUILDER - CALC_BUILDER = aseml.CalcBuilder(nn_name) - ctx.obj["nn_name"] = nn_name - @main.command() @herald @click.pass_context @click.argument("filepath", type=str) +@add_nn_name_opt @add_relax_opts @add_constraint_opts @add_workdir_verbose_opts -def relax(ctx, filepath, +def relax(ctx, filepath, nn_name, relax_mode, fmax, pressure, steps, optimizer, fix_inds, fix_symbols, workdir, verbose): @@ -159,17 +161,30 @@ def relax(ctx, filepath, To change the ML potential, use e.g.: - abiml.py.py -nn chgnet relax [...] + abiml.py.py relax -nn m3gnet [...] """ atoms = aseml.get_atoms(filepath) aseml.fix_atoms(atoms, fix_inds=fix_inds, fix_symbols=fix_symbols) ml_relaxer = aseml.MlRelaxer(atoms, relax_mode, fmax, pressure, steps, optimizer, - CALC_BUILDER, verbose, workdir, prefix="_relax_") + nn_name, verbose, workdir, prefix="_relax_") - # Set Delta forces and delta stress if script is called by ABINIT - ml_relaxer.set_delta_forces_stress_from_abiml_nc(filepath) + print(ml_relaxer.to_string(verbose=verbose)) + ml_relaxer.run() + return 0 + +@main.command() +@herald +@click.pass_context +@click.argument("filepath", type=str) +@add_workdir_verbose_opts +def abinit_relax(ctx, filepath, + workdir, verbose): + """ + Interact with ABINIT in hybrid relaxation mode. + """ + ml_relaxer = aseml.MlRelaxer.from_abinit_yaml_file(filepath) print(ml_relaxer.to_string(verbose=verbose)) ml_relaxer.run() return 0 @@ -179,6 +194,7 @@ def relax(ctx, filepath, @herald @click.pass_context @click.argument("filepath", type=str) +@add_nn_name_opt @click.option('--temperature', "-t", default=600, type=float, show_default=True, help='Temperature in Kelvin') @click.option('--timestep', "-ts", default=1, type=float, show_default=True, help='Timestep in fs.') @click.option('--steps', "-s", default=1000, type=int, show_default=True, help='Number of timesteps.') @@ -187,7 +203,7 @@ def relax(ctx, filepath, type=click.Choice(["nvt", "npt", "npt_berendsen"]), help='Ensemble e.g. nvt, npt.') @add_constraint_opts @add_workdir_verbose_opts -def md(ctx, filepath, +def md(ctx, filepath, nn_name, temperature, timestep, steps, loginterval, ensemble, fix_inds, fix_symbols, workdir, verbose): @@ -204,13 +220,13 @@ def md(ctx, filepath, To change the ML potential, use e.g.: - abiml.py.py -nn chgnet md [...] + abiml.py.py md -nn m3gnet [...] """ # See https://github.com/materialsvirtuallab/m3gnet#molecular-dynamics atoms = aseml.get_atoms(filepath) aseml.fix_atoms(atoms, fix_inds=fix_inds, fix_symbols=fix_symbols) - ml_md = aseml.MlMd(atoms, temperature, timestep, steps, loginterval, ensemble, CALC_BUILDER, verbose, + ml_md = aseml.MlMd(atoms, temperature, timestep, steps, loginterval, ensemble, nn_name, verbose, workdir, prefix="_md_") print(ml_md.to_string(verbose=verbose)) ml_md.run() @@ -221,10 +237,11 @@ def md(ctx, filepath, @herald @click.pass_context @click.argument("filepaths", nargs=2, type=str) +@add_nn_name_opt @add_neb_opts @add_constraint_opts @add_workdir_verbose_opts -def neb(ctx, filepaths, +def neb(ctx, filepaths, nn_name, nimages, relax_mode, fmax, pressure, optimizer, neb_method, climb, fix_inds, fix_symbols, workdir, verbose @@ -243,7 +260,7 @@ def neb(ctx, filepaths, To change the ML potential, use e.g.: - abiml.py.py -nn chgnet neb [...] + abiml.py.py neb -nn m3gnet [...] """ initial_atoms = aseml.get_atoms(filepaths[0]) aseml.fix_atoms(initial_atoms, fix_inds=fix_inds, fix_symbols=fix_symbols) @@ -252,7 +269,7 @@ def neb(ctx, filepaths, ml_neb = aseml.MlNeb(initial_atoms, final_atoms, nimages, neb_method, climb, optimizer, relax_mode, fmax, pressure, - CALC_BUILDER, verbose, workdir, prefix="_neb_") + nn_name, verbose, workdir, prefix="_neb_") print(ml_neb.to_string(verbose=verbose)) ml_neb.run() return 0 @@ -262,10 +279,11 @@ def neb(ctx, filepaths, @herald @click.pass_context @click.argument("filepaths", nargs=-1, type=str) # , help="Files with structures") +@add_nn_name_opt @add_neb_opts @add_constraint_opts @add_workdir_verbose_opts -def mneb(ctx, filepaths, +def mneb(ctx, filepaths, nn_name, nimages, relax_mode, fmax, pressure, optimizer, neb_method, climb, fix_inds, fix_symbols, workdir, verbose): @@ -283,7 +301,7 @@ def mneb(ctx, filepaths, To change the ML potential, use e.g.: - abiml.py.py -nn chgnet mneb [...] + abiml.py.py mneb -nn m3gnet [...] """ # Fix atoms atoms_list = [aseml.get_atoms(p) for p in filepaths] @@ -291,100 +309,17 @@ def mneb(ctx, filepaths, aseml.fix_atoms(atoms, fix_inds=fix_inds, fix_symbols=fix_symbols) mneb = aseml.MultiMlNeb(atoms_list, nimages, neb_method, climb, optimizer, relax_mode, fmax, pressure, - CALC_BUILDER, verbose, workdir, prefix="_mneb_") + nn_name, verbose, workdir, prefix="_mneb_") print(mneb.to_string(verbose=verbose)) mneb.run() return 0 -@main.command() -@herald -@click.pass_context -@click.argument("filepaths", nargs=2, type=str) -@click.option("--nimages", "-n", default=7, type=click.IntRange(3, None), show_default=True, - help='Number of images. Must be >= 3') -#@click.option("--method", default="aseneb", show_default=True, type=str, help="NEB method") -#@click.option("--climb", is_flag=True, help="Use a climbing image (default is no climbing image).") -@add_workdir_verbose_opts -def tsaseneb(ctx, filepaths, nimages, workdir, verbose): -#def run_tsaseneb(ctx, initial, final, nimages, fmax, pressure, relax_mode, method, climb): - """ - NEB calculations with TSASE and ML potential. - """ - initial, final = filepaths - initial = os.path.abspath(initial) - final = os.path.abspath(final) - os.chdir(workdir) - - # Read initial and final structures - initial = Structure.from_file(initial, primitive=False) - final = Structure.from_file(final, primitive=False) - - """ - GP: Example take from the tsase examples folder - -------------------- - ssneb optimization example - output: 1. fe.out - contais force, energy and distance between images - the final results can be plotted by nebspline.pl in vtstscripts - 2. i.CON: i is the number of image - geomety of each image, in vasp POSCAR format - - If parallel=True is set, the ssneb calculation is parallelized over images. - The command to execute this script needs to be: - - mpirun -np 5 python ssneb.py - - where 5 is the number of intermidiate images (nim-2) - """ - # use my ssneb to pass directly the initial structures - try: - from tsase import neb - from tsase.neb import ssneb - except ImportError as exc: - raise - - import glob - ''' - #-------------- uncomment if parallelize over image ------------------ - # parallelize over images with mpi4py - from mpi4py import MPI - ''' - p1 = aseml.to_ase_atoms(initial, calculator=CALC_BUILDER.get_calculator()) - p2 = aseml.to_ase_atoms(final, calculator=CALC_BUILDER.get_calculator()) - - # external stress applied in the unit of GPa - stress = np.zeros((3, 3)) - #stress[2,2] = -1.0 #negative is tension - - # no climbing image first - # band = ssneb(p1, p2, numImages = nim, express = stress) - band = ssneb(p1, p2, numImages=nimages, express=stress) - #band = neb.ssneb(p1, p2, numImages=nim, method='ci', express=stress) - - ''' - #-------------- substitute for the above line if parallelize over image ------------------ - band = neb.ssneb(p1, p2, numImages = nim, method = 'ci', express = stress, parallel = True) - ''' - - # to restart, uncomment the following lines which read the previous optimized images into the band - #restart = False - #if restart: - # for i in range(1, nim-1): - # b = images[i] - # band.path[i].set_positions(b.get_positions()) - # band.path[i].set_cell(b.get_cell()) - - opt = neb.fire_ssneb(band, maxmove=1.0, dtmax=0.05, dt=0.05) - #opt = neb.qm_ssneb(band, maxmove=0.10, dt=0.05) - opt.minimize(forceConverged=0.3, maxIterations=5000) - return 0 - - @main.command() @herald @click.pass_context @click.argument("filepath", type=str) +@add_nn_name_opt @click.option("--supercell", "-s", nargs=3, type=int, default=(4, 4, 4), show_default=True, help="Supercell") @click.option("--kpts", "-k", nargs=3, type=int, default=(4, 4, 4), show_default=True, help="K-mesh for ph-DOS") @click.option('--asr/--no-asr', default=True, show_default=True, @@ -392,7 +327,7 @@ def tsaseneb(ctx, filepaths, nimages, workdir, verbose): @click.option('--nqpath', default=100, type=int, show_default=True, help="Number of q-points along the q-path") @add_relax_opts @add_workdir_verbose_opts -def aseph(ctx, filepath, supercell, kpts, asr, nqpath, +def aseph(ctx, filepath, nn_name, supercell, kpts, asr, nqpath, relax_mode, fmax, pressure, steps, optimizer, workdir, verbose): """ Compute phonon band structure and DOS with ML potential. @@ -410,10 +345,10 @@ def aseph(ctx, filepath, supercell, kpts, asr, nqpath, To change the ML potential, use e.g.: - abiml.py.py -nn chgnet aseph [...] + abiml.py.py aseph -nn m3gnet [...] """ ml_ph = aseml.MlPhonons(filepath, supercell, kpts, asr, nqpath, - relax_mode, fmax, pressure, steps, optimizer, CALC_BUILDER, + relax_mode, fmax, pressure, steps, optimizer, nn_name, verbose, workdir, prefix="_aseph_") print(ml_ph.to_string(verbose=verbose)) ml_ph.run() @@ -424,10 +359,12 @@ def aseph(ctx, filepath, supercell, kpts, asr, nqpath, @herald @click.pass_context @click.argument("filepath", type=str) +@add_nn_name_opt @click.option("--max-ns", "-m", default=100, type=int, show_default=True, help='Max number of structures') @add_relax_opts @add_workdir_verbose_opts -def order(ctx, filepath, max_ns, relax_mode, fmax, pressure, steps, optimizer, workdir, verbose): +def order(ctx, filepath, nn_name, + max_ns, relax_mode, fmax, pressure, steps, optimizer, workdir, verbose): """ Generate ordered structures from CIF with partial occupancies. @@ -441,7 +378,7 @@ def order(ctx, filepath, max_ns, relax_mode, fmax, pressure, steps, optimizer, w Based on: https://matgenb.materialsvirtuallab.org/2013/01/01/Ordering-Disordered-Structures.html """ ml_orderer = aseml.MlOrderer(filepath, max_ns, optimizer, relax_mode, fmax, pressure, - steps, CALC_BUILDER, verbose, workdir, prefix="_order_") + steps, nn_name, verbose, workdir, prefix="_order_") print(ml_orderer.to_string(verbose=verbose)) ml_orderer.run() return 0 @@ -451,13 +388,14 @@ def order(ctx, filepath, max_ns, relax_mode, fmax, pressure, steps, optimizer, w @herald @click.pass_context @click.argument("filepath", type=str) +@add_nn_name_opt @click.option("-isite", "--isite", required=True, help='Index of atom to displace or string with chemical element to be added to input structure.') @click.option("--mesh", type=int, default=4, show_default=True, help='Mesh size along the smallest cell size.') @add_relax_opts @add_nprocs_opt @add_workdir_verbose_opts -def scan_relax(ctx, filepath, +def scan_relax(ctx, filepath, nn_name, isite, mesh, relax_mode, fmax, pressure, steps, optimizer, nprocs, @@ -470,12 +408,15 @@ def scan_relax(ctx, filepath, Usage example: \b - abiml.py.py scan_relax FILE -isite 0 --mesh 4 # Move first atom in the structure - abiml.py.py --nn-name chgnet scan_relax FILE -isite H # Add H to the structure read from FILE. + abiml.py.py scan-relax FILE -isite 0 --mesh 4 # Move first atom in the structure + abiml.py.py scan-relax FILE -isite H # Add H to the structure read from FILE. where `FILE` is any file supported by abipy/pymatgen e.g. netcdf files, Abinit input, POSCAR, xsf, etc. + + To change the ML potential, use e.g.: + + abiml.py.py scan-relax -nn m3gnet [...] """ - nn_name = ctx.obj["nn_name"] structure = Structure.from_file(filepath) from abipy.ml.relax_scanner import RelaxScanner @@ -492,7 +433,7 @@ def scan_relax(ctx, filepath, @main.command() @herald @click.pass_context -@click.argument("filepath", type=str) +@click.argument('filepaths', type=str, nargs=-1) @click.option("-nns", '--nn-names', type=str, multiple=True, show_default=True, help='ML potentials to be used', #default=["m3gnet", "chgnet"], default=["chgnet"]) @@ -503,7 +444,7 @@ def scan_relax(ctx, filepath, help='Plotting backend: mpl for matplotlib, panel for web-based') @add_nprocs_opt @add_workdir_verbose_opts -def compare(ctx, filepath, +def compare(ctx, filepaths, nn_names, traj_range, exposer, @@ -516,23 +457,24 @@ def compare(ctx, filepath, Usage example: \b - abiml.py.py compare FILE --nn-names m3gnet --nn-names chgnet + abiml.py.py compare FILE --nn-names matgl --nn-names chgnet where `FILE` can be either a _HIST.nc or a VASPRUN file. """ traj_range = cli.range_from_str(traj_range) - ml_comp = aseml.MlCompareWithAbinitio(filepath, nn_names, traj_range, verbose, workdir, prefix="_compare_") + ml_comp = aseml.MlCompareWithAbinitio(filepaths, nn_names, traj_range, verbose, workdir, prefix="_compare_") print(ml_comp) c = ml_comp.run(nprocs=nprocs) with_stress = True from abipy.tools.plotting import Exposer - with Exposer.as_exposer(exposer, title=os.path.basename(filepath)) as e: + with Exposer.as_exposer(exposer, title=" ".join(os.path.basename(p) for p in filepaths)) as e: e(c.plot_energies(show=False)) - e(c.plot_forces(show=False)) + e(c.plot_forces(delta_mode=True, show=False)) e(c.plot_energies_traj(delta_mode=True, show=False)) + e(c.plot_energies_traj(delta_mode=False, show=False)) if with_stress: - e(c.plot_stresses(show=False)) + e(c.plot_stresses(delta_mode=True, show=False)) e(c.plot_forces_traj(delta_mode=True, show=False)) e(c.plot_stress_traj(delta_mode=True, show=False)) diff --git a/dev_scripts/mlrun.py b/dev_scripts/mlrun.py index 898b96fc7..b5dc304d5 100755 --- a/dev_scripts/mlrun.py +++ b/dev_scripts/mlrun.py @@ -46,8 +46,8 @@ def add_workdir_verbose_opts(f): @click.command() @click.argument("filepath", type=str) -@click.option("--nn-name", "-nn", default="m3gnet", show_default=True, - type=click.Choice(["m3gnet_old", "m3gnet", "chgnet"]), help='ML potential') +@click.option("--nn-name", "-nn", default="chgnet", show_default=True, + type=click.Choice(["m3gnet", "matgl", "chgnet"]), help='ML potential') @add_relax_opts @click.option("-k", "--kppa", default=1000, type=float, show_default=True, help="Defines the sampling of BZ mesh (k-points per atom)")