[Atomate2] How to decode a TaskDocument from queried dictionaries?

I’m trying to load an atomate2 TaskDocument from the dictionary I queried from JobStore, is there any method to help me do so? For example, in pymatgen, every object is basically a child class of monty.MSONable, and they are supposed to have from_dict and as_dict methods. However, in atomate2, output documents seem to be pydantic.BaseModel, which means if one of the attributes for the BaseModel is supposed to be another BaseModel, you can not naively use BaseModel(**dict) to initialize it.

I’ve tried MontyDecoder as follows, and it did not work for me at all:

After I reviewed the codes in Monty.json, I discovered that line 429 of monty/json.py did not check whether self.process_decoded has been applied recusively to subdictionary “data”, while the “data” dict of a pydantic model might still have encoded pydantic models or MSONables as its values. This happens exactly in atomate2 TaskDocument, where taskdocument.calcs_reversed is a list of Calculation BaseModels. This iterative check should be added to ensure that atomate2 TaskDocument’s are re-loadable.

Just tried pydantic.parse_obj_as(), and this is what I got:

When I tried to print out the force constants:


It seems that phonon information is saved into blobs, but the blob infos are not cleared before trying to load data with TaskDocument. It must be replaced with the correct format by specifying load=True when calling jobstore.query().

In summary, to correctly load a TaskDocument from store, one needs to use:

data = store.query_one(query_dict, load=True)
doc = parse_obj_as(TaskDocument, data)

Hope this will be helpful for new users.

Case closed.