[DRAFT] Arbitrary Attribute Support #354
Conversation
add a few tests and rename to extra_x_attributes get value strict static state can now handle arbitrary things update model interface update validate model outputs proper runner code handle deprecated model attributes fix state init revert test_state wrote tests for state more fixes more fixes better set attribute validation more tests more progress towrads better attributes assert we are only concatenating states with the same attribute names fix concat state
|
|
||
|
|
||
| # TODO: we should put this logic inside __init_subclass__ of Modelinterface to | ||
| # automatically validate the model outputs when the model is subclassed. |
There was a problem hiding this comment.
I think we should do this in another PR because this Pr is already kinda long, and automatically enforcing validate_model_outputs (by putting this logic inside init_subclass) will require a refactor of all models (which will make this PR longer). Since we are backwards compatible with older models that just returns energy/forces/stress it's okay to not do this enforcement now.
|
|
||
| # Create a new instance of the same class | ||
| return state_class(**concatenated) | ||
| new_state = state_class(**concatenated) |
There was a problem hiding this comment.
this line of code is the main reason why I put up this feature as a draft. If we just pass in these extra attributes (e.g. energy) into the constructor of the state_class, we will error out since the init function doesn't expect it. So to compensate, I am just setting these extra attributes after we've initialized the state_class.
This works - but is obviously a bit longer (since we need to distinctly specify get_attrs_for_scope for "only_extra" and "only_default". It'd be nice if we didn't need to do this and the constructor just understood which attributes should go into extra or not. maybe the extra dictionary isn't even the right way to do this.
Thoughts are very welcome here!
| yield attr_name, state.get(attr_name) | ||
|
|
||
|
|
||
| def _filter_attrs_by_mask( |
There was a problem hiding this comment.
This function is what's causing the tests to fail. I have yet to migrate get_attrs_for_scope to set the attributes differently for default vs extra attributes.
|
Can you say a little more about the used case here? Are you sure we need arbitrary properties? Reporters support arbitrary properties, so we can get them out as long as they are a function of the model and state. If integrators/optimizers need them, it will be easier to build on them if they are defined in code. If we do need things like charge and spin, it might be sufficient to just add charge and spin to the SimState as optional attributes. Until the diversity of properties supported by models becomes overwhelming, that should work fine. Perhaps I'll regret saying this, but there are only so many atomic properties that models can support. If we feel it becomes overwhelming, we can figure out a more abstractable solution down the road. With that in mind, I'd suggest either opening a new PR that adds charge/spin to SimState or just adding it to the model interface for MACE and Fairchem. |
|
@orionarcher Yeah I wanted to offer a solution that was more general but I'm fine with just using #373 for now. I kinda don't like the solution in this PR because making a general solution feels a bit clunky. |
2 participants
Summary
This is the first draft for supporting arbitrary attributes inside of SimState. The tests don't pass because this is half-complete. We may make some changes which will affect this code.
The main idea is to introduce 3 new dictionaries inside SimState to track all of the extra attributes each state may want to store.
When we define a new model, the output of this model needs to specify the kind of each attribute it outputs (if it's an atom-wise, system-wise, or global attribute) - this allows us to add models that work on any set of properties (not just energy, forces, and stress)
Note: I want the constraints PR #294 to go in before this one because these 2 PRs change fundamental parts of torchsim and I'd rather rebase this one ontop of the constraints PR.
Checklist
Before a pull request can be merged, the following items must be checked: