atomate2
atomate2 copied to clipboard
materialsproject
BUG: custom CHGNET model in PhononFlow throws `jsanitize` error
Describe the bug
Including a custom CHGNet model when doing PhononFlow (locally) with CHGNETRelaxMaker and CHGNetStaticMaker results in an error
(from monty.json)
689 return jsanitize(
--> 690 obj.as_dict(),
691 strict=strict,
692 allow_bson=allow_bson,
693 enum_values=enum_values,
694 recursive_msonable=recursive_msonable,
695 )
AttributeError: 'Tensor' object has no attribute 'as_dict'
see full error output below
To Reproduce
This was following the jupyter notebook: https://jageo.github.io/TutorialAtomate2Forcefields/phonon.html
from atomate2.forcefields.jobs import CHGNetRelaxMaker,CHGNetStaticMaker
from pymatgen.core import Structure
from chgnet.model import CHGNet
custom_model = CHGNet.from_file('custom_chgnet_model.pth.tar')
struct = Structure.from_file('teststructure.vasp')
phonon_flow = PhononMaker(
min_length=15.0,
store_force_constants=False,
bulk_relax_maker=CHGNetRelaxMaker(optimizer_kwargs={'model':custom_model},
relax_kwargs={'fmax':0.00001}),
static_energy_maker=CHGNetStaticMaker(optimizer_kwargs={'model':custom_model)
).make(structure=struct)
and then:
from jobflow import run_locally
run_locally(phonon_flow, create_folders=False)
from pymatgen.phonon.bandstructure import PhononBandStructureSymmLine
from pymatgen.phonon.dos import PhononDos
from pymatgen.phonon.plotter import PhononBSPlotter, PhononDosPlotter
from jobflow import SETTINGS
store = SETTINGS.JOB_STORE
store.connect()
Expected behavior
---------------------------------------------------------------------------
AttributeError Traceback (most recent call last)
Cell In[10], line 3
1 from jobflow import run_locally
----> 3 run_locally(phonon_flow, create_folders=False)
5 from pymatgen.phonon.bandstructure import PhononBandStructureSymmLine
6 from pymatgen.phonon.dos import PhononDos
File ~/Library/src/jobflow/src/jobflow/managers/local.py:82, in run_locally(flow, log, store, create_folders, root_dir, ensure_success, allow_external_references, raise_immediately)
79 if log:
80 initialize_logger()
---> 82 flow = get_flow(flow, allow_external_references=allow_external_references)
84 stopped_parents: set[str] = set()
85 errored: set[str] = set()
File ~/Library/src/jobflow/src/jobflow/core/flow.py:865, in get_flow(flow, allow_external_references)
861 flow = Flow(jobs=flow)
863 if not allow_external_references:
864 # ensure that we have all the jobs needed to resolve the reference connections
--> 865 job_references = find_and_get_references(flow.jobs)
866 job_reference_uuids = {ref.uuid for ref in job_references}
867 missing_jobs = job_reference_uuids.difference(set(flow.job_uuids))
File ~/Library/src/jobflow/src/jobflow/core/reference.py:397, in find_and_get_references(arg)
393 if isinstance(arg, (float, int, str, bool)):
394 # argument is a primitive, we won't find a reference here
395 return ()
--> 397 arg = jsanitize(arg, strict=True, enum_values=True, allow_bson=True)
399 # recursively find any reference classes
400 locations = find_key_value(arg, "@class", "OutputReference")
File ~/miniconda3/envs/py311/lib/python3.11/site-packages/monty/json.py:634, in jsanitize(obj, strict, allow_bson, enum_values, recursive_msonable)
632 return obj
633 if isinstance(obj, (list, tuple)):
--> 634 return [
635 jsanitize(i, strict=strict, allow_bson=allow_bson, enum_values=enum_values)
636 for i in obj
637 ]
638 if np is not None and isinstance(obj, np.ndarray):
639 return [
640 jsanitize(i, strict=strict, allow_bson=allow_bson, enum_values=enum_values)
641 for i in obj.tolist()
642 ]
File ~/miniconda3/envs/py311/lib/python3.11/site-packages/monty/json.py:635, in <listcomp>(.0)
632 return obj
633 if isinstance(obj, (list, tuple)):
634 return [
--> 635 jsanitize(i, strict=strict, allow_bson=allow_bson, enum_values=enum_values)
636 for i in obj
637 ]
638 if np is not None and isinstance(obj, np.ndarray):
639 return [
640 jsanitize(i, strict=strict, allow_bson=allow_bson, enum_values=enum_values)
641 for i in obj.tolist()
642 ]
File ~/miniconda3/envs/py311/lib/python3.11/site-packages/monty/json.py:689, in jsanitize(obj, strict, allow_bson, enum_values, recursive_msonable)
680 if pydantic is not None and isinstance(obj, pydantic.BaseModel): # pylint: disable=E1101
681 return jsanitize(
682 MontyEncoder().default(obj),
683 strict=strict,
(...)
686 recursive_msonable=recursive_msonable,
687 )
--> 689 return jsanitize(
690 obj.as_dict(),
691 strict=strict,
692 allow_bson=allow_bson,
693 enum_values=enum_values,
694 recursive_msonable=recursive_msonable,
695 )
File ~/miniconda3/envs/py311/lib/python3.11/site-packages/monty/json.py:648, in jsanitize(obj, strict, allow_bson, enum_values, recursive_msonable)
646 return obj.to_dict()
647 if isinstance(obj, dict):
--> 648 return {
649 str(k): jsanitize(
650 v,
651 strict=strict,
652 allow_bson=allow_bson,
653 enum_values=enum_values,
654 recursive_msonable=recursive_msonable,
655 )
656 for k, v in obj.items()
657 }
658 if isinstance(obj, (int, float)):
659 return obj
File ~/miniconda3/envs/py311/lib/python3.11/site-packages/monty/json.py:649, in <dictcomp>(.0)
646 return obj.to_dict()
647 if isinstance(obj, dict):
648 return {
--> 649 str(k): jsanitize(
650 v,
651 strict=strict,
652 allow_bson=allow_bson,
653 enum_values=enum_values,
654 recursive_msonable=recursive_msonable,
655 )
656 for k, v in obj.items()
657 }
658 if isinstance(obj, (int, float)):
659 return obj
File ~/miniconda3/envs/py311/lib/python3.11/site-packages/monty/json.py:648, in jsanitize(obj, strict, allow_bson, enum_values, recursive_msonable)
646 return obj.to_dict()
647 if isinstance(obj, dict):
--> 648 return {
649 str(k): jsanitize(
650 v,
651 strict=strict,
652 allow_bson=allow_bson,
653 enum_values=enum_values,
654 recursive_msonable=recursive_msonable,
655 )
656 for k, v in obj.items()
657 }
658 if isinstance(obj, (int, float)):
659 return obj
File ~/miniconda3/envs/py311/lib/python3.11/site-packages/monty/json.py:649, in <dictcomp>(.0)
646 return obj.to_dict()
647 if isinstance(obj, dict):
648 return {
--> 649 str(k): jsanitize(
650 v,
651 strict=strict,
652 allow_bson=allow_bson,
653 enum_values=enum_values,
654 recursive_msonable=recursive_msonable,
655 )
656 for k, v in obj.items()
657 }
658 if isinstance(obj, (int, float)):
659 return obj
[... skipping similar frames: <dictcomp> at line 649 (3 times), jsanitize at line 648 (3 times)]
File ~/miniconda3/envs/py311/lib/python3.11/site-packages/monty/json.py:648, in jsanitize(obj, strict, allow_bson, enum_values, recursive_msonable)
646 return obj.to_dict()
647 if isinstance(obj, dict):
--> 648 return {
649 str(k): jsanitize(
650 v,
651 strict=strict,
652 allow_bson=allow_bson,
653 enum_values=enum_values,
654 recursive_msonable=recursive_msonable,
655 )
656 for k, v in obj.items()
657 }
658 if isinstance(obj, (int, float)):
659 return obj
File ~/miniconda3/envs/py311/lib/python3.11/site-packages/monty/json.py:649, in <dictcomp>(.0)
646 return obj.to_dict()
647 if isinstance(obj, dict):
648 return {
--> 649 str(k): jsanitize(
650 v,
651 strict=strict,
652 allow_bson=allow_bson,
653 enum_values=enum_values,
654 recursive_msonable=recursive_msonable,
655 )
656 for k, v in obj.items()
657 }
658 if isinstance(obj, (int, float)):
659 return obj
File ~/miniconda3/envs/py311/lib/python3.11/site-packages/monty/json.py:690, in jsanitize(obj, strict, allow_bson, enum_values, recursive_msonable)
680 if pydantic is not None and isinstance(obj, pydantic.BaseModel): # pylint: disable=E1101
681 return jsanitize(
682 MontyEncoder().default(obj),
683 strict=strict,
(...)
686 recursive_msonable=recursive_msonable,
687 )
689 return jsanitize(
--> 690 obj.as_dict(),
691 strict=strict,
692 allow_bson=allow_bson,
693 enum_values=enum_values,
694 recursive_msonable=recursive_msonable,
695 )
AttributeError: 'Tensor' object has no attribute 'as_dict'
it is also possible i have done this the wrong way....
@naik-aakash did you get something similar recently?
@badw could you list all version numbers of your different codes? I guess it is some unrelated change in another code again. We, unfortunately, have this very often.
@naik-aakash did you get something similar recently?
@badw could you list all version numbers of your different codes? I guess it is some unrelated change in another code again. We, unfortunately, have this very often.
Hi @JaGeo, I did not encounter this yet.
ah right,
atomate2 = v0.0.13
monty = 2023.0.25
jobflow = v0.1.17
I thought I was up to date with atomate2, so will download the latest version as well and try that
@naik-aakash did you get something similar recently? @badw could you list all version numbers of your different codes? I guess it is some unrelated change in another code again. We, unfortunately, have this very often.
Hi @JaGeo, I did encounter this yet.
Did not, I guess, right?
ah right,
atomate2 = v0.0.13 monty = 2023.0.25 jobflow = v0.1.17I thought I was up to date with atomate2, so will download the latest version as well and try that
Yes, many things changed recently. Good to check out the latest version. If it persists, we will check.
I get the same error with 0.0.14
I have attached a tarball with a jupyter notebook that hopefully make it easier to reproduce!
thanks atomate2_error.tar.gz
Likely a recent change in one of the other packages but not atomate2. I am not sure we can get to this fast.
@naik-aakash can you list the version numbers here? Maybe, that already helps.
I am getting jsanitize error too with the example workflow:
from atomate2.vasp.flows.core import RelaxBandStructureMaker
from jobflow import run_locally
from pymatgen.core import Structure
# construct a rock salt MgO structure
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]],
)
# make a band structure flow to optimise the structure and obtain the band structure
bandstructure_flow = RelaxBandStructureMaker().make(mgo_structure)
# run the job
run_locally(bandstructure_flow, create_folders=True)
Produces:
Traceback (most recent call last):
File "/scratch/users/r/g/rgouvea/vasp_runs/atomate_test/workflowatomate.py", line 13, in <module>
bandstructure_flow = RelaxBandStructureMaker().make(mgo_structure)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/ucl/modl/rgouvea/miniconda3/lib/python3.11/site-packages/atomate2/vasp/flows/core.py", line 381, in make
bs_flow = self.band_structure_maker.make(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/ucl/modl/rgouvea/miniconda3/lib/python3.11/site-packages/atomate2/vasp/flows/core.py", line 172, in make
return Flow(jobs, outputs, name=self.name)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/ucl/modl/rgouvea/miniconda3/lib/python3.11/site-packages/jobflow/core/flow.py", line 147, in __init__
self.output = output
^^^^^^^^^^^
File "/home/ucl/modl/rgouvea/miniconda3/lib/python3.11/site-packages/jobflow/core/flow.py", line 272, in output
if contains_flow_or_job(output):
^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/ucl/modl/rgouvea/miniconda3/lib/python3.11/site-packages/jobflow/utils/find.py", line 208, in contains_flow_or_job
obj = jsanitize(obj, strict=True, allow_bson=True)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/ucl/modl/rgouvea/miniconda3/lib/python3.11/site-packages/monty/json.py", line 693, in jsanitize
return {
^
File "/home/ucl/modl/rgouvea/miniconda3/lib/python3.11/site-packages/monty/json.py", line 694, in <dictcomp>
str(k): jsanitize(
^^^^^^^^^^
File "/home/ucl/modl/rgouvea/miniconda3/lib/python3.11/site-packages/monty/json.py", line 734, in jsanitize
return jsanitize(
^^^^^^^^^^
File "/home/ucl/modl/rgouvea/miniconda3/lib/python3.11/site-packages/monty/json.py", line 693, in jsanitize
return {
^
File "/home/ucl/modl/rgouvea/miniconda3/lib/python3.11/site-packages/monty/json.py", line 694, in <dictcomp>
str(k): jsanitize(
^^^^^^^^^^
File "/home/ucl/modl/rgouvea/miniconda3/lib/python3.11/site-packages/monty/json.py", line 672, in jsanitize
return [
^
File "/home/ucl/modl/rgouvea/miniconda3/lib/python3.11/site-packages/monty/json.py", line 673, in <listcomp>
jsanitize(i, strict=strict, allow_bson=allow_bson, enum_values=enum_values)
File "/home/ucl/modl/rgouvea/miniconda3/lib/python3.11/site-packages/monty/json.py", line 672, in jsanitize
return [
^
File "/home/ucl/modl/rgouvea/miniconda3/lib/python3.11/site-packages/monty/json.py", line 673, in <listcomp>
jsanitize(i, strict=strict, allow_bson=allow_bson, enum_values=enum_values)
File "/home/ucl/modl/rgouvea/miniconda3/lib/python3.11/site-packages/monty/json.py", line 735, in jsanitize
obj.as_dict(),
^^^^^^^^^^^
AttributeError: 'VaspObject' object has no attribute 'as_dict'
However the installation example is running fine:
from atomate2.vasp.jobs.core import RelaxMaker
from jobflow import run_locally
from pymatgen.core import Structure
# construct an FCC silicon structure
si_structure = Structure(
lattice=[[0, 2.73, 2.73], [2.73, 0, 2.73], [2.73, 2.73, 0]],
species=["Si", "Si"],
coords=[[0, 0, 0], [0.25, 0.25, 0.25]],
)
# make a relax job to optimise the structure
relax_job = RelaxMaker().make(si_structure)
# run the job
run_locally(relax_job, create_folders=True)
Can you try downgrading to emmet-core==0.78.0rc4?
E.g.,
pip install emmet-core==0.78.0rc4
Can you try downgrading to emmet-core==0.78.0rc4?
E.g.,
pip install emmet-core==0.78.0rc4
Yes, it worked for me now, here is the current setup:
atomate2 0.0.14 pypi_0 pypi
emmet-core 0.78.0rc4 pypi_0 pypi
jobflow 0.1.17 pypi_0 pypi
monty 2024.2.26 pypi_0 pypi
python 3.11.4 h955ad1f_0
Hi @badw , after a bit more checking in detail into this issue, I think the issue is not specific to PhononFlow, but how the forcefield based jobs are generated. When using custom models, this models are also stored in the python objects and as jobflow enforces strict=True in jsanitize, the tensor objects are not being able to converted to dict as they don't have this as_dict method.
I was able to reproduce the same error using the custom model or trying to load locally saved pre-trained chgnet model, even when just when trying to run a relax job.
from jobflow import run_locally
from chgnet.model import CHGNet
from atomate2.forcefields.jobs import CHGNetRelaxMaker
bulk_relax_maker=CHGNetRelaxMaker(relax_kwargs={'fmax':0.1}, steps=3,optimizer_kwargs={'model':custom_model}).make(structure=struct)
run_locally(bulk_relax_maker, create_folders=False)
@janosh @esoteric-ephemera @matthewkuner any ideas?
I am in Munich at the moment and then on Easter holiday. I will not be able to work on this at the moment. But maybe relevant for CHGNet people that the support for fine-tuned models is there as well
I haven't tried using a custom CHGNet model before, nor am I familiar enough with CHGNet to debug this sort of issue. Maybe Janosh (already tagged) or @bowend-ucb would be able to chime in?
I guess adding a link to the model and then loading it in the init of CHGNet instead of directly passing it, should work but would require some refactoring
Something similar to: https://github.com/materialsproject/atomate2/blob/33d747adfda16089d6794019828e3988c6a27a6e/src/atomate2/forcefields/jobs.py#L301
But this should probably be reviewed and merged first: https://github.com/materialsproject/atomate2/pull/722
If I understand correctly, custom_model should be a CHGNet obejct loaded from user's weights.
This doesn't seem to be a issue with CHGNet code?
Is there anything that is suggested to be added to CHGNet functions to resolve this?
If I understand correctly,
custom_modelshould be a CHGNet obejct loaded from user's weights. This doesn't seem to be a issue with CHGNet code? Is there anything that is suggested to be added to CHGNet functions to resolve this?
Hi @BowenD-UCB , I think adding an arg to StructureOptimizer so user can provide path to load locally saved model instead of only CHGNet object should help resolve the issue.
I feel this will help bypass the current issue when jobflow tries to serialize the CHGNet model object, leading to jsanitize errors.
Hi @naik-aakash
The calculator class is more foundamental while relaxer should be a downstream function class.
Will calculator = CHGNetCalculator.from_file(), and relaxer = StructureOptimizer(calculator) work?
Overall, I think this is more of a feature request than a bug as no one has ever used CHGNet in atomate2 with an individual model.
@badw do you maybe want to implement such a solution yourself and ask us or the CHGNet developers for help in case you cannot sort it with the current implementations?
Hi @naik-aakash
The calculator class is more foundamental while relaxer should be a downstream function class. Will
calculator = CHGNetCalculator.from_file(), andrelaxer = StructureOptimizer(calculator) work?
Yes, I think this should work. We just need to implement it like this.
Thanks! This isn't particularly urgent for me, but I can play around a little bit and see what I manage!
This may not be an issue after the refactored workflows are merged, since the interface to a custom calculator would just take JSONable args (like the class/module/callable to load through monty, and the path to a custom model file as a kwarg)
Found the issue: emmet-core PR #944 removed the as_dict method from VaspObject. Every version of emmet-core <= 0.77.1 should work OK, need to get a PR in soon
i will close this as the new emmet version has been released. If it shows up again, let us know