Adding "global", "per_atom", "per_system" features to SimState

[curtischong]

In an effort to properly type the entire library, I'm proposing to add these attributes to the SimState. This way, we don't need to run infer_property_scope (https://github.com/Radical-AI/torch-sim/blob/main/torch_sim/state.py#L496) and we automatically know if a property belongs globally to the entire simstate, or per-atom/system.

This idea is partially inspired by how orb handles features. They use:

"node_features"
"edge_features"
"system_features"

https://github.com/orbital-materials/orb-models/blob/main/orb_models/forcefield/base.py#L65-L67

Maybe if a state is classified as an "md_state" we'll guarantee exactly which features it'll have? Hopefully by solving this problem, we'll also come up with a good solution to represent md_state.

More discussion in #227, #228

[curtischong]

Note: global is actually a reserved keyword. so we'll probably rename it to "shared". #189

[curtischong]

I think I've got an idea of how it'll look like. We'll make these attributes dictionaries in SimState

"node_features"
"edge_features"
"system_features"
"shared_features" (all systems in the simstate have these features)

Then we can add derived classes like MDState. For these derived classes, they store their data inside these attribute dictionaries
HOWEVER, they will have utility getter/setter functions that reads from these attributes.
e.g. if we want to access velocities, we'll use sim_state.velocities(). This way, we can also handle nice edge cases where the velocities may be None at the start of the simulations inside of these utility functions

perhaps important config data should be stored under shared_features.

[orionarcher]

Modified from a comment in #227, adding for context:

HI @curtischong, this is something I thought about a fair amount and I am happy to revisit. I am not convinced I made the right decision to make everything implicit. I have a few thoughts here:

1. If we did want to make the distinction explicit, I'd advocate for just using a tuple of strings. Then we could just rename infer_property_scope -> return_property_scope and have it return the tuples directly instead of inferring them. Something like this:

class SimState:
    positions: torch.Tensor
    masses: torch.Tensor
    cell: torch.Tensor
    pbc: bool  # TODO: do all calculators support mixed pbc?
    atomic_numbers: torch.Tensor
    system_idx: torch.Tensor | None = field(default=None, kw_only=True)
    _atom_features: tuple[str] = ("positions", "masses", "atomic_numbers")
    _system_features: tuple[str] = ("cell", "system_idx")
    _global_features: tuple[str] = ("pbc")
    
    
    @property
    def atom_features(self) -> torch.Tensor:
        return self._atom_features

    @property
    def atom_features(self) -> torch.Tensor:
        return self._atom_features

    @property
    def atom_features(self) -> torch.Tensor:
        return self._atom_features

    def return_property_scope(self):
        return {"global": self.global_features, "per_atom": self.atom_features", "per_system": self.system_features}

3. If we think this is the right option, let's use atom instead of node.