Skip to content

Commit

Permalink
Saving work
Browse files Browse the repository at this point in the history
  • Loading branch information
gmatteo committed Sep 6, 2023
1 parent 3f38b90 commit 911e036
Show file tree
Hide file tree
Showing 4 changed files with 91 additions and 89 deletions.
64 changes: 32 additions & 32 deletions abipy/flowtk/psrepos.py
Original file line number Diff line number Diff line change
Expand Up @@ -139,38 +139,38 @@ 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_latest_pseudos(xc_name: str, ps_type: str = "NC", relativity_type: str = "SR", accuracy: str = "standard") -> PseudoTable:
"""
Args:
xc_name:
ps_type:
relativity_type
accuracy:
"""
if ps_type == "NC":
version = "0.4"
ps_generator, project_name = OncvpspRepo.ps_generator, OncvpspRepo.project_name
repo_name = f"{ps_generator}-{xc_name}-{relativity_type}-{project_name}v{version}"
#if relativity_type == "SR":
# repo_name = {
# "PBE": "ONCVPSP-PBE-SR-PDv0.4",
# "PBEsol": "ONCVPSP-PBEsol-SR-PDv0.4",
# "LDA": "ONCVPSP-LDA-SR-PDv0.4",
# }[xc_name]
#else
# repo_name = {
# "PBE": "ONCVPSP-PBE-FR-PDv0.4",
# "PBEsol": "ONCVPSP-PBEsol-FR-PDv0.4",
# "LDA": "ONCVPSP-LDA-FR-PDv0.4",
# }[xc_name]

elif ps_type == "PAW":
raise NotImplementedError(f"Invalid {ps_type=}")

else:
raise ValueError(f"Invalid {ps_type=}")

return get_repo_from_name(repo_name).get_pseudos(accuracy)
#def get_latest_pseudos(xc_name: str, ps_type: str = "NC", relativity_type: str = "SR", accuracy: str = "standard") -> PseudoTable:
# """
# Args:
# xc_name:
# ps_type:
# relativity_type
# accuracy:
# """
# if ps_type == "NC":
# version = "0.4"
# ps_generator, project_name = OncvpspRepo.ps_generator, OncvpspRepo.project_name
# repo_name = f"{ps_generator}-{xc_name}-{relativity_type}-{project_name}v{version}"
# #if relativity_type == "SR":
# # repo_name = {
# # "PBE": "ONCVPSP-PBE-SR-PDv0.4",
# # "PBEsol": "ONCVPSP-PBEsol-SR-PDv0.4",
# # "LDA": "ONCVPSP-LDA-SR-PDv0.4",
# # }[xc_name]
# #else
# # repo_name = {
# # "PBE": "ONCVPSP-PBE-FR-PDv0.4",
# # "PBEsol": "ONCVPSP-PBEsol-FR-PDv0.4",
# # "LDA": "ONCVPSP-LDA-FR-PDv0.4",
# # }[xc_name]
#
# elif ps_type == "PAW":
# raise NotImplementedError(f"Invalid {ps_type=}")
#
# else:
# raise ValueError(f"Invalid {ps_type=}")
#
# return get_repo_from_name(repo_name).get_pseudos(accuracy)


def get_installed_repos_and_root(dirpath: Optional[str] = None) -> tuple[list[PseudosRepo], str]:
Expand Down
98 changes: 51 additions & 47 deletions abipy/ml/aseml.py
Original file line number Diff line number Diff line change
Expand Up @@ -977,8 +977,8 @@ def __init__(self, *args, **kwargs):

from collections import deque
maxlen = 2
self.__correct_forces_algo = CORRALGO.delta
self.__correct_stress_algo = CORRALGO.delta
self.__correct_forces_algo = CORRALGO.none
self.__correct_stress_algo = CORRALGO.none
self.__abi_forces_list = deque(maxlen=maxlen)
self.__abi_stress_list = deque(maxlen=maxlen)
self.__abi_atoms_list = deque(maxlen=maxlen)
Expand Down Expand Up @@ -1022,25 +1022,24 @@ def store_abi_forstr_atoms(self, abi_forces, abi_stress, atoms):
self.__abi_stress_list.append(abi_stress)
self.__abi_atoms_list.append(atoms.copy())

# Compute ML forces ans stresses for the input atoms.
self.reset()
# Compute ML forces and stresses for the input atoms.
old_forces_algo = self.set_correct_forces_algo(CORRALGO.none)
old_stress_algo = self.set_correct_stress_algo(CORRALGO.none)
self.reset()
ml_forces = self.get_forces(atoms=atoms)
ml_stress = self.get_stress(atoms=atoms)
#print(f"{ml_forces=}"); print(f"{ml_stress=}")
self.reset()
self.__ml_forces_list.append(ml_forces)
self.__ml_stress_list.append(ml_stress)
self.set_correct_forces_algo(old_forces_algo)
self.set_correct_stress_algo(old_stress_algo)
self.reset()

def fmt_vec3(vec3) -> str:
return "{:.6e} {:.6e} {:.6e}".format(*vec3)
def fmt_vec6(vec6) -> str:
return "{:.6e} {:.6e} {:.6e} {:.6e} {:.6e} {:.6e}".format(*vec6)

if self.__verbose:
def fmt_vec3(vec3) -> str:
return "{:.6e} {:.6e} {:.6e}".format(*vec3)
def fmt_vec6(vec6) -> str:
return "{:.6e} {:.6e} {:.6e} {:.6e} {:.6e} {:.6e}".format(*vec6)

from ase.stress import full_3x3_to_voigt_6_stress
print("abi_stress6:", fmt_vec6(full_3x3_to_voigt_6_stress(abi_stress)))
print("ml_stress6: ", fmt_vec6(full_3x3_to_voigt_6_stress(ml_stress)))
Expand Down Expand Up @@ -1139,7 +1138,7 @@ class CalcBuilder:
Supports different backends defined by `name` string.
Possible formats are:
1) nn_type e.g. m3net
1) nn_type e.g. m3gnet
2) nn_type:model_name
3) nn_type@filepath
"""
Expand Down Expand Up @@ -1185,6 +1184,9 @@ def reset(self) -> None:
def get_calculator(self, reset=False) -> Calculator:
"""
Return ASE calculator with ML potential.
Args:
reset: True if internal cache should be reset.
"""
if reset: self.reset()

Expand Down Expand Up @@ -1249,7 +1251,7 @@ class MyAlignnCalculator(_MyMlCalculator, AlignnAtomwiseCalculator):
"""Add abi_forces and abi_stress"""

model_name = default_path() if self.model_name is None else self.model_name
return AlignnAtomwiseCalculator(path=model_name)
return MyAlignnAtomwiseCalculator(path=model_name)

#if self.nn_type == "quip":
# try:
Expand Down Expand Up @@ -1562,8 +1564,8 @@ def __init__(self, atoms: Atoms, relax_mode, fmax, pressure, steps, optimizer, n
fmax: tolerance for relaxation convergence. Here fmax is a sum of force and stress forces.
pressure: Target pressure.
steps: max number of steps for relaxation.
optimizer: name of the ASE optimizer to use.
nn_name:
optimizer: name of the ASE optimizer to use for relaxation.
nn_name: String defining the NN potential. See also CalcBuilder.
verbose: Verbosity level.
"""
super().__init__(workdir, prefix)
Expand Down Expand Up @@ -1642,13 +1644,13 @@ def __init__(self, atoms: Atoms, temperature, timestep, steps, loginterval,
ensemble, nn_name, verbose, workdir, prefix=None):
"""
Args:
atoms:
temperature:
atoms: ASE atoms.
temperature: Temperatur in K
timestep:
steps:
loginterval:
ensemble:
nn_name:
nn_name: String defining the NN potential. See also CalcBuilder.
verbose: Verbosity level.
workdir:
prefix:
Expand Down Expand Up @@ -1794,7 +1796,7 @@ def __init__(self, atoms_list: list[Atoms], nn_name, verbose,
"""
Args:
atoms_list: List of ASE atoms
nn_name:
nn_name: String defining the NN potential. See also CalcBuilder.
verbose: Verbosity level.
"""
super().__init__(workdir, prefix)
Expand Down Expand Up @@ -1842,17 +1844,17 @@ def __init__(self, initial_atoms: Atoms, final_atoms: Atoms,
nn_name, verbose, workdir, prefix=None):
"""
Args:
initial_atoms
final_atoms:
nimages:
initial_atoms: initial ASE atoms.
final_atoms: final ASE atoms.
nimages: Number of images.
neb_method:
climb:
optimizer:
relax_mode:
fmax:
pressure:
nn_name:
verbose:
optimizer: name of the ASE optimizer to use for relaxation.
relax_mode: "ions" to relax ions only, "cell" for ions + cell, "no" for no relaxation.
fmax: tolerance for relaxation convergence. Here fmax is a sum of force and stress forces.
pressure: Target pressure
nn_name: String defining the NN potential. See also CalcBuilder.
verbose: Verbosity level.
workdir:
prefix:
"""
Expand Down Expand Up @@ -1992,16 +1994,16 @@ def __init__(self, atoms_list: list[Atoms], nimages, neb_method, climb, optimize
nn_name, verbose, workdir, prefix=None):
"""
Args:
atoms_list:
nimages:
atoms_list: List of ASE atoms.
nimages: Number of NEB images.
neb_method:
climb:
optimizer:
relax_mode:
fmax:
relax_mode: "ions" to relax ions only, "cell" for ions + cell, "no" for no relaxation.
fmax: tolerance for relaxation convergence. Here fmax is a sum of force and stress forces.
pressure:
nn_name:
verbose:
nn_name: String defining the NN potential. See also CalcBuilder.
verbose: Verbosity level.
workdir:
prefix:
"""
Expand Down Expand Up @@ -2132,15 +2134,15 @@ def __init__(self, structure, max_ns, optimizer, relax_mode, fmax, pressure, ste
workdir, prefix=None):
"""
Args:
structure:
structure: Abipy Structure object or any object that can be converted to structure.
max_ns:
optimizer:
relax_mode:
fmax:
pressure:
steps:
nn_name:
verbose:
nn_name: String defining the NN potential. See also CalcBuilder.
verbose: Verbosity level.
workdir:
prefix:
"""
Expand Down Expand Up @@ -2263,8 +2265,10 @@ def run(self) -> None:
self._finalize()


class MlPhonons(_MlBase):
"""Compute phonons with ASE and ML potential."""
class MlAsePhonons(_MlBase):
"""
Compute phonons with ASE and ML potential.
"""

def __init__(self, atoms: Atoms, supercell, kpts, asr, nqpath,
relax_mode, fmax, pressure, steps, optimizer, nn_name,
Expand All @@ -2273,15 +2277,15 @@ def __init__(self, atoms: Atoms, supercell, kpts, asr, nqpath,
Args:
atoms: ASE atoms.
supercell: tuple with supercell dimension.
kpts:
kpts: K-mesh for phonon-DOS
asr: Enforce acoustic sum-rule.
nqpath: Number of q-point along the q-path.
relax_mode:
fmax:
steps:
optimizer:
nn_name,
verbose:
relax_mode: "ions" to relax ions only, "cell" for ions + cell, "no" for no relaxation.
fmax: tolerance for relaxation convergence. Here fmax is a sum of force and stress forces.
verbose: whether to print stdout.
optimizer: name of the ASE optimizer to use for relaxation.
nn_name: String defining the NN potential. See also CalcBuilder.
verbose: Verbosity level.
workdir:
prefix:
"""
Expand Down Expand Up @@ -2326,7 +2330,7 @@ def to_string(self, verbose=0):
return s

def run(self) -> None:
"""Run MlPhonons."""
"""Run AseMlPhonons."""
#self.pickle_dump()
workdir = self.workdir
calculator = self.get_calculator()
Expand Down
6 changes: 2 additions & 4 deletions abipy/ml/ml_relax.py
Original file line number Diff line number Diff line change
Expand Up @@ -246,14 +246,13 @@ def run(self, workdir=None, prefix=None):
print(f"\nBegin ABINIT + {self.nn_name} hybrid relaxation")
if self.xc_name == "PBE":
print(f"Starting from ML-optimized Atoms as {self.xc_name=}")
atoms = ml_opt.atoms.copy()
atoms = abisanitize_atoms(atoms)
atoms = abisanitize_atoms(ml_opt.atoms.copy())
else:
print(f"Starting from initial Atoms as {self.xc_name=}")
atoms = self.initial_atoms.copy()

count, abiml_nsteps, ml_nsteps = 0, 0, 0
count_max = 15
count, abiml_nsteps, ml_nsteps = 0, 0, 0
t_start = time.time()
while count <= count_max:
count += 1
Expand All @@ -268,7 +267,6 @@ def run(self, workdir=None, prefix=None):

# Store ab-initio forces/stresses in the ML calculator and attach it to atoms.
ml_calc.store_abi_forstr_atoms(gs.forces, gs.stress, atoms)
#ml_calc.reset()
atoms.calc = ml_calc

opt_kws = dict(
Expand Down
12 changes: 6 additions & 6 deletions abipy/scripts/abiml.py
Original file line number Diff line number Diff line change
Expand Up @@ -321,7 +321,7 @@ def mneb(ctx, filepaths, nn_name,
@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("--kpts", "-k", nargs=3, type=int, default=(4, 4, 4), show_default=True, help="K-mesh for phonon-DOS")
@click.option('--asr/--no-asr', default=True, show_default=True,
help="Restore the acoustic sum rule on the force constants.")
@click.option('--nqpath', default=100, type=int, show_default=True, help="Number of q-points along the q-path")
Expand All @@ -347,11 +347,11 @@ def aseph(ctx, filepath, nn_name, supercell, kpts, asr, nqpath,
abiml.py.py aseph -nn m3gnet [...]
"""
ml_ph = aseml.MlPhonons(filepath, supercell, kpts, asr, nqpath,
relax_mode, fmax, pressure, steps, optimizer, nn_name,
verbose, workdir, prefix="_aseph_")
print(ml_ph.to_string(verbose=verbose))
ml_ph.run()
ml_aseph = aseml.MlAsePhonons(filepath, supercell, kpts, asr, nqpath,
relax_mode, fmax, pressure, steps, optimizer, nn_name,
verbose, workdir, prefix="_aseph_")
print(ml_aseph.to_string(verbose=verbose))
ml_aseph.run()
return 0


Expand Down

0 comments on commit 911e036

Please sign in to comment.