No SimState subclass takes automatically the newly added spin and charge

[thomasloux]

Current initialization of subclass of SimState (MDState, all intergrators init, optimizer init, StaticState) don't take spin and charge. So this information is lost and not provided to the model. Currently a fix is proposed for optimizer #413 but this probably does not scale well. Not sure other information will be added in Atoms any time soon.
So either, the same fix is applied to all init function. Or we refactor init function, for instance adding a class method SubClass.from_simstate(state: SimState, **kwargs) where kwargs correspond to SubClass specific variables.
@orionarcher do you have a preference?

[orionarcher]

I like the idea of adding a from_state method to the subclasses that takes in a state and any new args. We could also make this from_state check for any overlapping field names and copy them over if they are present. e.g., if two integrator states both have cell_forces, we could copy over that value, thus it wouldn't be limited to just the SimState args.

This will also clean up the signature. For example, in nve_init:

MDState(
    positions=state.positions,
    momenta=momenta,
    energy=model_output["energy"],
    forces=model_output["forces"],
    masses=state.masses,
    cell=state.cell,
    pbc=state.pbc,
    system_idx=state.system_idx,
    atomic_numbers=state.atomic_numbers,
    _constraints=state.constraints,
)

becomes

MDState.from_state(
    state=state,
    momenta=momenta,
    energy=model_output["energy"],
    forces=model_output["forces"],
)

How does that sound?

[thomasloux]

This exactly what I had in mind !