Skip to content
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

BUG: custom CHGNET model in PhononFlow throws jsanitize error #781

Closed
badw opened this issue Mar 19, 2024 · 27 comments
Closed

BUG: custom CHGNET model in PhononFlow throws jsanitize error #781

badw opened this issue Mar 19, 2024 · 27 comments

Comments

@badw
Copy link

badw commented Mar 19, 2024

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....

@JaGeo
Copy link
Member

JaGeo commented Mar 19, 2024

@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
Copy link
Contributor

naik-aakash commented Mar 19, 2024

@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.

@badw
Copy link
Author

badw commented Mar 19, 2024

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

@JaGeo
Copy link
Member

JaGeo commented Mar 19, 2024

@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?

@JaGeo
Copy link
Member

JaGeo commented Mar 19, 2024

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

Yes, many things changed recently. Good to check out the latest version. If it persists, we will check.

@badw
Copy link
Author

badw commented Mar 19, 2024

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

@JaGeo
Copy link
Member

JaGeo commented Mar 19, 2024

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.

@rogeriog
Copy link

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)

@JaGeo
Copy link
Member

JaGeo commented Mar 19, 2024

Yeah, sorry. We need to investigate. But you are also invited to do so.

@utf
Copy link
Member

utf commented Mar 19, 2024

Can you try downgrading to emmet-core==0.78.0rc4?

E.g.,

pip install emmet-core==0.78.0rc4

@badw
Copy link
Author

badw commented Mar 19, 2024

for me emmet-core==0.78.90rc4 didn't rectify the problem

@rogeriog
Copy link

rogeriog commented Mar 19, 2024

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  

@naik-aakash
Copy link
Contributor

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)

@JaGeo
Copy link
Member

JaGeo commented Mar 19, 2024

@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

@matthewkuner
Copy link
Collaborator

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?

@JaGeo
Copy link
Member

JaGeo commented Mar 20, 2024

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

@JaGeo
Copy link
Member

JaGeo commented Mar 20, 2024

Something similar to:

calculator = NequIPCalculator.from_deployed_model(

But this should probably be reviewed and merged first: #722

@bowen-bd
Copy link

bowen-bd commented Mar 20, 2024

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?

@naik-aakash
Copy link
Contributor

naik-aakash commented Mar 21, 2024

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?

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.

@bowen-bd
Copy link

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?

@JaGeo
Copy link
Member

JaGeo commented Mar 21, 2024

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?

@JaGeo
Copy link
Member

JaGeo commented Mar 21, 2024

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?

Yes, I think this should work. We just need to implement it like this.

@bowen-bd
Copy link

Just made the commit

If this is not urgent, we will release it in the next release

@badw
Copy link
Author

badw commented Mar 21, 2024

Thanks! This isn't particularly urgent for me, but I can play around a little bit and see what I manage!

@esoteric-ephemera
Copy link
Contributor

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)

@esoteric-ephemera
Copy link
Contributor

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

@JaGeo
Copy link
Member

JaGeo commented May 11, 2024

i will close this as the new emmet version has been released. If it shows up again, let us know

@JaGeo JaGeo closed this as completed May 11, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

No branches or pull requests

8 participants