Skip to content

Commit

Permalink
Merge pull request #74 from QuantumChemist/main
Browse files Browse the repository at this point in the history 
Fully Integrate all MLIPs
  • Loading branch information
@JaGeo
JaGeo authored Jul 2, 2024
2 parents 1e5313d + 4da9dc3 commit 57b2404
Show file tree
Hide file tree
Showing 12 changed files with 497 additions and 146 deletions.
7 changes: 5 additions & 2 deletions autoplex/auto/phonons/flows.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -150,6 +150,7 @@ class CompleteDFTvsMLBenchmarkWorkflow(Maker):
atomwise_regularization_list: list | None = None
soap_delta_list: list | None = None
n_sparse_list: list | None = None
benchmark_kwargs: dict = field(default_factory=dict)

def make(
self,
Expand Down Expand Up @@ -263,7 +264,7 @@ def make(
mlip_hyper=ml_hyper,
).make(
species_list=isoatoms.output["species"],
isolated_atoms_energy=isoatoms.output["energies"],
isolated_atoms_energies=isoatoms.output["energies"],
fit_input=fit_input,
split_ratio=split_ratio,
f_max=f_max,
Expand Down Expand Up @@ -311,6 +312,7 @@ def make(
symprec=self.symprec,
phonon_displacement_maker=self.phonon_displacement_maker,
dft_references=dft_references,
**self.benchmark_kwargs,
)
flows.append(complete_bm)
bm_outputs.append(complete_bm.output)
Expand Down Expand Up @@ -343,7 +345,7 @@ def make(
mlip_type=ml_model, mlip_hyper=ml_hyper
).make(
species_list=isoatoms.output["species"],
isolated_atoms_energy=isoatoms.output["energies"],
isolated_atoms_energies=isoatoms.output["energies"],
fit_input=fit_input,
split_ratio=split_ratio,
f_max=f_max,
Expand Down Expand Up @@ -379,6 +381,7 @@ def make(
symprec=self.symprec,
phonon_displacement_maker=self.phonon_displacement_maker,
dft_references=dft_references,
**self.benchmark_kwargs,
)
flows.append(complete_bm)
bm_outputs.append(complete_bm.output)
Expand Down
41 changes: 36 additions & 5 deletions autoplex/auto/phonons/jobs.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,9 @@ def complete_benchmark( # this function was put here to prevent circular import
symprec,
phonon_displacement_maker: BaseVaspMaker,
dft_references=None,
relax_maker_kwargs: dict | None = None,
static_maker_kwargs: dict | None = None,
**ml_phonon_maker_kwargs,
):
"""
Construct a complete flow for benchmarking the MLIP fit quality using a DFT based phonon structure.
Expand Down Expand Up @@ -76,6 +79,12 @@ def complete_benchmark( # this function was put here to prevent circular import
Maker used to compute the forces for a supercell.
dft_references:
a list of DFT reference files containing the PhononBSDOCDoc object. Default None.
relax_maker_kwargs: dict
Keyword arguments that can be passed to the RelaxMaker.
static_maker_kwargs: dict
Keyword arguments that can be passed to the StaticMaker.
ml_phonon_maker_kwargs: dict
Keyword arguments that can be passed to the MLPhononMaker.
"""
jobs = []
collect_output = []
Expand All @@ -84,16 +93,34 @@ def complete_benchmark( # this function was put here to prevent circular import
if min_length >= 18:
phonon_displacement_maker = TightDFTStaticMakerBigSupercells()
for suffix in ["", "_wo_sigma", "_phonon", "_rand_struc"]:
if Path(Path(ml_path) / f"gap_file{suffix}.xml").exists():
# TODO: this needs to beextended for the other MLIPs
# _wo_sigma", "_phonon", "_rand_struc" only available for GAP at the moment
if ml_model == "GAP":
ml_potential = Path(ml_path) / f"gap_file{suffix}.xml"
elif ml_model == "J-ACE":
raise UserWarning("No atomate2 ACE.jl PhononMaker implemented.")
elif ml_model in ["M3GNET"]:
ml_potential = Path(ml_path.join(suffix)) / "training"
# M3GNet requires path
# also need to find a different solution for separated fit then
elif ml_model in ["NEQUIP"]:
ml_potential = Path(ml_path) / f"deployed_nequip_model{suffix}.pth"
else: # MACE
ml_potential = Path(ml_path) / f"MACE_model{suffix}.model"

if Path(ml_potential).exists():
add_data_ml_phonon = MLPhononMaker(
min_length=min_length,
relax_maker_kwargs=relax_maker_kwargs,
static_maker_kwargs=static_maker_kwargs,
).make_from_ml_model(
structure=benchmark_structure,
ml_model=ml_path,
suffix=suffix,
ml_model=ml_model,
potential_file=ml_potential,
**ml_phonon_maker_kwargs,
)
jobs.append(add_data_ml_phonon)

# DFT benchmark reference preparations
if dft_references is None and benchmark_mp_ids is not None:
if (
benchmark_mp_ids[ibenchmark_structure] in mp_ids
Expand Down Expand Up @@ -322,18 +349,22 @@ def get_iso_atom(structure_list: list[Structure]):
list of pymatgen Structure objects
"""
jobs = []
iso_atoms_dict = {}
all_species = list(
{specie for s in structure_list for specie in s.types_of_species}
)

isoatoms = IsoAtomMaker().make(all_species=all_species)
jobs.append(isoatoms)

for i, species in enumerate(all_species):
iso_atoms_dict.update({species.number: isoatoms.output["energies"][i]})

flow = Flow(
jobs,
{
"species": all_species,
"energies": isoatoms.output["energies"],
"energies": iso_atoms_dict,
"dirs": isoatoms.output["dirs"],
},
)
Expand Down
155 changes: 114 additions & 41 deletions autoplex/data/phonons/flows.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,12 @@
ForceFieldStaticMaker,
GAPRelaxMaker,
GAPStaticMaker,
M3GNetRelaxMaker,
M3GNetStaticMaker,
MACERelaxMaker,
MACEStaticMaker,
NequipRelaxMaker,
NequipStaticMaker,
)
from atomate2.vasp.flows.core import DoubleRelaxMaker
from atomate2.vasp.flows.phonons import PhononMaker
Expand All @@ -27,6 +33,7 @@
from pymatgen.core import Molecule, Site

from autoplex.data.common.jobs import generate_randomized_structures
from autoplex.data.phonons.utils import ml_phonon_maker_preparation

__all__ = [
"DFTPhononMaker",
Expand Down Expand Up @@ -332,17 +339,17 @@ class MLPhononMaker(FFPhononMaker):
High-throughput electronic band structure calculations:
Challenges and tools. Computational Materials Science,
49(2), 299-312. doi:10.1016/j.commatsci.2010.05.010.
We will however use seekpath and primitive structures
We will, however, use seekpath and primitive structures
as determined by from phonopy to compute the phonon band structure
bulk_relax_maker : .ForceFieldRelaxMaker or None
bulk_relax_maker: .ForceFieldRelaxMaker or None
A maker to perform a tight relaxation on the bulk.
Set to ``None`` to skip the
bulk relaxation
static_energy_maker : .ForceFieldStaticMaker or None
static_energy_maker: .ForceFieldStaticMaker or None
A maker to perform the computation of the DFT energy on the bulk.
Set to ``None`` to skip the
static energy computation
phonon_displacement_maker : .ForceFieldStaticMaker or None
phonon_displacement_maker: .ForceFieldStaticMaker or None
Maker used to compute the forces for a supercell.
generate_frequencies_eigenvectors_kwargs : dict
Keyword arguments passed to :obj:`generate_frequencies_eigenvectors`.
Expand All @@ -364,6 +371,10 @@ class MLPhononMaker(FFPhononMaker):
in the future
store_force_constants: bool
if True, force constants will be stored
relax_maker_kwargs: dict
Keyword arguments that can be passed to the RelaxMaker.
static_maker_kwargs: dict
Keyword arguments that can be passed to the StaticMaker.
"""

name: str = "ml phonon"
Expand All @@ -374,7 +385,7 @@ class MLPhononMaker(FFPhononMaker):
)
)
phonon_displacement_maker: ForceFieldStaticMaker | None = field(
default_factory=lambda: GAPStaticMaker(name="ml phonon static")
default_factory=lambda: GAPStaticMaker(name="gap phonon static")
)
static_energy_maker: ForceFieldStaticMaker | None = field(
default_factory=lambda: GAPStaticMaker()
Expand All @@ -383,11 +394,18 @@ class MLPhononMaker(FFPhononMaker):
generate_frequencies_eigenvectors_kwargs: dict = field(
default_factory=lambda: {"units": "THz", "tol_imaginary_modes": 1e-1}
)
relax_maker_kwargs: dict = field(default_factory=dict)
static_maker_kwargs: dict = field(default_factory=dict)
relax_maker_kwargs: dict | None = field(default_factory=dict)
static_maker_kwargs: dict | None = field(default_factory=dict)

@job
def make_from_ml_model(self, structure, ml_model, suffix: str = "", **make_kwargs):
def make_from_ml_model(
self,
structure,
potential_file,
ml_model: str = "GAP",
calculator_kwargs: dict | None = None,
**make_kwargs,
):
"""
Maker for GAP phonon jobs.
Expand All @@ -397,10 +415,13 @@ def make_from_ml_model(self, structure, ml_model, suffix: str = "", **make_kwarg
A pymatgen structure. Please start with a structure
that is nearly fully optimized as the internal optimizers
have very strict settings!
ml_model : str
Complete path to MLIP file(s).
suffix:
Train, test and MLIP suffix ("", "_wo_sigma", "_phonon", "_rand_struc").
ml_model: str
ML model to be used. Default is GAP.
potential_file :
Complete path to MLIP file(s)
Train, test and MLIP files (+ suffixes "", "_wo_sigma", "_phonon", "_rand_struc").
calculator_kwargs :
Keyword arguments for the ASE Calculator.
make_kwargs :
Keyword arguments for the PhononMaker.
Expand All @@ -409,41 +430,93 @@ def make_from_ml_model(self, structure, ml_model, suffix: str = "", **make_kwarg
PhononMaker jobs.
"""
ml_model = ml_model + f"/gap_file{suffix}.xml"
if self.bulk_relax_maker is not None:
br = self.bulk_relax_maker
self.bulk_relax_maker = br.update_kwargs(
update={
"calculator_kwargs": {
"args_str": "IP GAP",
"param_filename": str(ml_model),
},
**self.relax_maker_kwargs,
if ml_model == "GAP":
if calculator_kwargs is None:
calculator_kwargs = {
"args_str": "IP GAP",
"param_filename": str(potential_file),
}

ml_prep = ml_phonon_maker_preparation(
bulk_relax_maker=self.bulk_relax_maker,
phonon_displacement_maker=self.phonon_displacement_maker,
static_energy_maker=self.static_energy_maker,
calculator_kwargs=calculator_kwargs,
relax_maker_kwargs=self.relax_maker_kwargs,
static_maker_kwargs=self.static_maker_kwargs,
)
if self.phonon_displacement_maker is not None:
ph_disp = self.phonon_displacement_maker
self.phonon_displacement_maker = ph_disp.update_kwargs(
update={
"calculator_kwargs": {
"args_str": "IP GAP",
"param_filename": str(ml_model),
},
**self.static_maker_kwargs,

elif ml_model == "J-ACE":
raise UserWarning("No atomate2 ACE.jl PhononMaker implemented.")

elif ml_model == "NEQUIP":
if calculator_kwargs is None:
calculator_kwargs = {
"model_path": str(potential_file),
"device": "cuda",
}
else:
calculator_kwargs.update({"model_path": str(potential_file)})

ml_prep = ml_phonon_maker_preparation(
bulk_relax_maker=NequipRelaxMaker(
relax_cell=True, relax_kwargs={"interval": 500}
),
phonon_displacement_maker=NequipStaticMaker(
name="nequip phonon static"
),
static_energy_maker=NequipStaticMaker(),
calculator_kwargs=calculator_kwargs,
relax_maker_kwargs=self.relax_maker_kwargs,
static_maker_kwargs=self.static_maker_kwargs,
)
if self.static_energy_maker is not None:
stat_en = self.static_energy_maker
self.static_energy_maker = stat_en.update_kwargs(
update={
"calculator_kwargs": {
"args_str": "IP GAP",
"param_filename": str(ml_model),
},
**self.static_maker_kwargs,
}

elif ml_model == "M3GNET":
if calculator_kwargs is None:
calculator_kwargs = {"path": str(potential_file)}

ml_prep = ml_phonon_maker_preparation(
bulk_relax_maker=M3GNetRelaxMaker(
relax_cell=True, relax_kwargs={"interval": 500}
),
phonon_displacement_maker=M3GNetStaticMaker(
name="m3gnet phonon static"
),
static_energy_maker=M3GNetStaticMaker(),
calculator_kwargs=calculator_kwargs,
relax_maker_kwargs=self.relax_maker_kwargs,
static_maker_kwargs=self.static_maker_kwargs,
)

else: # MACE
if calculator_kwargs is None:
calculator_kwargs = {"model": str(potential_file), "device": "cuda"}
elif "model" in calculator_kwargs:
calculator_kwargs.update(
{"default_dtype": "float64"}
) # Use float64 for geometry optimization.
else:
calculator_kwargs.update(
{"model": str(potential_file), "default_dtype": "float64"}
)

ml_prep = ml_phonon_maker_preparation(
bulk_relax_maker=MACERelaxMaker(
relax_cell=True, relax_kwargs={"interval": 500}
),
phonon_displacement_maker=MACEStaticMaker(name="mace phonon static"),
static_energy_maker=MACEStaticMaker(),
calculator_kwargs=calculator_kwargs,
relax_maker_kwargs=self.relax_maker_kwargs,
static_maker_kwargs=self.static_maker_kwargs,
)

(
self.bulk_relax_maker,
self.phonon_displacement_maker,
self.static_energy_maker,
) = ml_prep

flow = self.make(structure=structure, **make_kwargs)
return Response(replace=flow, output=flow.output)

Expand Down
Loading

0 comments on commit 57b2404

Please sign in to comment.