BUG: serialization of ForceFieldMixin subclasses with dict `force_field_name`
In order to use a calculator for a Force field that is not included among the standard ones (e.g. the original chgnet) the suggested solution is to pass
```
{
"@module": "chgnet.model.dynamics",
"@callable": "CHGNetCalculator"
}
```
to the `force_field_name` in the subclasses of `ForceFieldMixin`.
https://github.com/materialsproject/atomate2/blob/44fd9b1e739f779dbb8530127f07613c8d51bdc8/src/atomate2/forcefields/utils.py#L247
The problem is that, while this works with the generated Job object, it fails when the object is serialized. The following example:
```python
from atomate2.forcefields.jobs import ForceFieldRelaxMaker
from jobflow import run_locally
from pymatgen.core import Structure
from monty.json import MontyDecoder, MontyEncoder
import json
mgo_structure = Structure(
lattice=[[0, 2.13, 2.13], [2.13, 0, 2.13], [2.13, 2.13, 0]],
species=["Mg", "O"],
coords=[[0, 0, 0], [0.5, 0.5, 0.5]],
)
flow = ForceFieldRelaxMaker(force_field_name = {"@module": "chgnet.model.dynamics", "@callable": "CHGNetCalculator"}).make(mgo_structure)
flow = json.loads(json.dumps(flow, cls=MontyEncoder), cls=MontyDecoder)
run_locally(flow)
```
fails with
> ValueError: Could not create ASE calculator for MLFF.Forcefield.
This leads to a failure in jobflow-remote, that needs to serialize the job in the DB, but I suppose the same would happen with Fireworks.
The primary reason is that the `calculator_meta` is set at creation time, but not serialized due to the presence of `field(init=False)`
https://github.com/materialsproject/atomate2/blob/44fd9b1e739f779dbb8530127f07613c8d51bdc8/src/atomate2/forcefields/utils.py#L178
Simply removing the `init=False` does not solve the issue though, because upon deserialization of the Job the `{"@module": "chgnet.model.dynamics", "@callable": "CHGNetCalculator"}` is automatically deserialized by the `MontyDecoder` and the code only expects either an `MLFF` or a `dict` for `calculator_meta`.
I managed to make it work by allowing `calculator_meta` to also be a `Calculator` subclass. I can open a PR if this seems a reasonable solution, but maybe there are better options?
1 条评论