AtomRef Updates

[lbluque]

A few edits to the AtomRef class:

• added max_z as optional kwarg for cases where offsets have not been fitted yet.
• Removed all use of numpy
• Use lstsq instead of explicitly pseudo-inverse
• Made property_offset a buffer to avoid errors when running on GPU and with lightning modules
• Save precomputed identity matrix for onehot vectors as buffer.

[lbluque]

The errors have to do with "Bad serialized model". I don't really understand how to fix those.

I also keep getting this error when running pre-commit:

error: TOML parse error at line 69, column 3
   |
69 |   "PERF203", # try-except-in-loop
   |   ^^^^^^^^^
Unknown rule selector: `PERF203`

Any ideas on how to get rid of it?

[shyuep]

I think it is a version of ruff that is the problem for the PERF. For bad serialized model, I think your fork may have cached an old model. Maybe mgl clear --yes command to remove all models first?

[lbluque]

Thanks @shyuep !

The "bad serialized" error was actually from my changes. I had to do a few things to fix it:

• use strict=False when using load_state_dict in IOMixin.load.
• check if property_offset is a list or ndarray in AtomRef constructor for backcompat with pretrained models.

I also added a raise from the BaseException leading to the "Bad serialized" error because otherwise it can be unclear what the error actually is.

Pls have a look and lmk your thoughts

[codecov]

Codecov Report

Patch coverage: 94.73% and no project coverage change.

Comparison is base (e825b75) 98.06% compared to head (dbd8fdf) 98.06%.
Report is 1 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #158   +/-   ##
=======================================
  Coverage   98.06%   98.06%           
=======================================
  Files          28       28           
  Lines        1807     1810    +3     
=======================================
+ Hits         1772     1775    +3     
  Misses         35       35           

Files Changed	Coverage Δ
matgl/layers/_atom_ref.py	97.61% <94.11%> (-2.39%)	⬇️
matgl/utils/io.py	100.00% <100.00%> (ø)

... and 1 file with indirect coverage changes

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[shyuep]

Thanks. I think this is fine. I will merge.