-
Notifications
You must be signed in to change notification settings - Fork 89
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Add forcefield schemas/makers to atomate2 #322
Conversation
I tested this as 1) a standalone job, and 2) as a two-part flow with the CHGNet prerelax job feeding into a VASP job. Both worked. |
People may want to use CHGNet makers as an input in the VASP flows that currently require inputted variables to be of type BaseVaspMaker. Would it be reasonable to change this to allow CHGNetRelaxMakers as well? (or perhaps some base FFRelaxMaker class that I could make) |
Codecov Report
Additional details and impacted files@@ Coverage Diff @@
## main #322 +/- ##
==========================================
+ Coverage 65.57% 65.97% +0.40%
==========================================
Files 72 74 +2
Lines 7018 7113 +95
Branches 896 904 +8
==========================================
+ Hits 4602 4693 +91
- Misses 2150 2152 +2
- Partials 266 268 +2
|
I did not add the new dependency for chgnet/pytorch--unsure where in the |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Hi @matthewkuner, thanks very much for this! It is in a great state for an initial PR.
I've added a bunch of comments. They are mainly style changes, but I had some more thoughts about the task document (I've been thinking about this more since we spoke on slack).
I think it would be nice to add:
- A
CHGNetStaticMaker
that will just calculate forces but not do a relaxation. After you've made the changes I suggested this should be very quick to implement by subclassingCHGNetRelaxMaker
. Just something like:
@dataclass
class CHGNetStaticMaker(CHGNetRelaxMaker):
"""
Maker to calculate forces and stresses using the CHGNet force field.
Parameters
----------
name: str
The job name.
relax_kwargs : dict
Keyword arguments that will get passed to :obj:`StructOptimizer.relax`.
optimizer_kwargs : dict
Keyword arguments that will get passed to :obj:`StructOptimizer()`.
task_document_kwargs : dict
Keyword arguments that will get passed to :obj:`ForceFieldTaskDocument()`.
"""
name: str = "CHGNet static"
relax_kwargs: dict = field(default_factory=dict)
optimizer_kwargs: dict = field(default_factory={"relax_cell": False, "steps": 1})
task_document_kwargs: dict = field(default_factory=dict)
- Can you add tests. CHGNet should be fast enough to run in the testing environment. Just use silicon and limit steps to 10
As for where to add the dependency. You should add it to the strict
set of dependencies with a specified version number (the latest one). You can also add a new set of dependencies called chgnet
which isn't pinned to a specific version. E.g., see how this is done for the lobster
set:
Line 44 in 4321620
lobster = ["lobsterpy"] |
Thanks again, and please get in touch if you have any questions.
src/atomate2/forcefields/schemas.py
Outdated
trajectory: dict = Field( | ||
None, description="Step-by-step trajectory of the structural relaxation." | ||
) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Can we try and emulate the ionic_steps
document in the VASP Task document. See https://github.com/materialsproject/emmet/blob/2289affb31d9b18aa82e49484ca772a2ce80e934/emmet-core/emmet/core/vasp/calculation.py#LL305C5-L305C114
and
So basically, it is a list of dicts rather than multiple dicts of lists. Rather than the funky e_fr_energy
etc in the IonicStep linked, you can just have energy
. You can also remove electronic_steps
and add magmoms
.
This will involve reorganising the output of CHGNet somewhat but will bring the outputs into line with VASP, CP2K, etc.
Note that we should also store the structure at each step. You can convert an ase atoms object using:
from pymatgen.io.ase import AseAtomsAdaptor
structure = AseAtomsAdaptor.get_structure(atoms)
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
so you'd want one of the output fields to be ionic_steps
, which is a list. That list will contain as many dicts as are specifies, which would include things like energy
, structure
, etc.? Sure I can do that.
Separately--considering that the linked VASP Task Document doesn't even let users specify which parts of the trajectory to keep, I don't think we should either. I'll just replace the tuple/list keep_info
with a bool keep_trajectory
. Let me know if you have any objections.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Actually, it appears that users don't even have the option to not store the trajectory for MD jobs. Seems pointless to exclude it for FF relaxations in that case. I'll just remove the tag entirely.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
RE: converting the Atoms
objects to structures. It appears that CHGNet only outputs a single Atoms
object, rather than one for each step.
update: I've been trying to convert the 'atom_positions'
and 'cells'
outputs from results['trajectory']
, but am not having success.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think having the keep_info tag is actually useful and just because its not there in VASP doesn't mean we can't add it here. So I would keep that in and forget about keep_trajectory
for now. However, one we add in MD through force fields, we will want to add trajectory support (a latter PR).
You can generate a new structure from the positions and cells along the lines of:
species = atoms.get_chemical_symbols()
for pos, cell in zip(atom_positions, cells):
structure = Structure(cell, species, positions, coords_are_cartesian=False)
Double check whether the atom_positions
are in cartesian coordinates and update coords_are_cartesian
accordingly.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Ah, I didn't know about the coords_are_cartesian
input, and missed it when looking through the documentation. That fixed it.
Co-authored-by: Alex Ganose <utf@users.noreply.github.com>
Adding a brief note that pymatgen has started offering an API that may make the code a bit simpler: https://pymatgen.org/pymatgen.core.structure.html#pymatgen.core.structure.Structure.relax I believe there’s an open issue to add chgnet as a relaxer too, I think @janosh is involved. |
Yes, I think CHGnet would make a great addition to |
@mkhorton would implementing CHGNet into pymatgen's structure objects make this Flow redundant? Or are they distinct use-cases? |
I don't think that function makes anything redundant, I think you we could just use it inside your job rather than constructing the relaxer manually (once chgnet has been implemented) |
What do you mean by
|
@matthewkuner Both are useful! |
@matthewkuner missed a "use" from my message. But effectively your job would simplify from @job(output_schema=FFStructureRelaxDocument)
def make(self, structure: Structure):
"""
Perform a relaxation of a structure using CHGNet.
Parameters
----------
structure : .Structure
A pymatgen structure.
"""
from chgnet.model import StructOptimizer
relaxer = StructOptimizer(**self.optimizer_kwargs)
result = relaxer.relax(
structure, relax_cell=self.relax_cell, **self.relax_kwargs
)
ff_structure_relax_doc = FFStructureRelaxDocument.from_chgnet_result(
structure,
self.relax_cell,
self.relax_kwargs,
self.optimizer_kwargs,
result,
self.keep_info,
)
return ff_structure_relax_doc To @job(output_schema=FFStructureRelaxDocument)
def make(self, structure: Structure):
"""
Perform a relaxation of a structure using CHGNet.
Parameters
----------
structure : .Structure
A pymatgen structure.
"""
result = structure.relax(
structure, relax_cell=self.relax_cell, **self.relax_kwargs
)
ff_structure_relax_doc = FFStructureRelaxDocument.from_chgnet_result(
structure,
self.relax_cell,
self.relax_kwargs,
self.optimizer_kwargs,
result,
self.keep_info,
)
return ff_structure_relax_doc |
@matthewkuner are you using the latest jobflow version? I should have fixed this in v0.1.11 |
@utf I was using v.0.1.9. I'll check it again with the latest version of jobflow soon. |
@utf I think this is getting very close to complete. Let me know if you have any other things you'd like me to add/change. |
This is very cool! Looking forward to using this in my own work! @utf, could we get a new version release after this is merged if you don't mind? |
Thanks very much for this! |
I just read through some of the changes here as part of #362 and want to say this is a really excellent contribution! Thanks for the work you put into this @matthewkuner! 🙏 |
Looking great to me so far! Are you planning to replicate the unit tests we have in |
@rkingsbury I wrote my own few tests (as I didn't know pymatgen had a m3gnet method /tests until I had written my own tests), so I wasn't planning to. Is it important? |
Ooops, sorry @matthewkuner , I posted that on the wrong PR! Please disregard. (Although I am excited to see this, too) |
Summary
Additional dependencies introduced
TODO
Comments
Based on what Shyue Ping told me here, I think we should wait until M3GNet is fully moved to the matgl repository, because that will the location for future development of M3GNet. This will also make it so that users will only need pytorch for both CHGNet and M3GNet (as opposed to needing TensorFlow too, as TensorFlow powers the older M3GNet implementation).
EDIT: m3gnet support was later added to Atomate2 in #380