Adding spin state embedding

[chrisiacovella]

Pull Request Summary

This PR adds in embedding for spin state, following the same approach used for total charge. That is, the per_system property is replicated to match the length of the per-atom property tensor and these are sent through a linear layer. '

Note, as with total_charge, forces and dipole_moment, S is optional. If not specified, it will simply be populated with tensor of zero (with appropriate shape) when we set up the properties. Presumably, one can then use spin state embedding with any dataset (regardless if that property is available), without the code failing...although since every value is zero, it shouldn't really impact anything.

Key changes

Notable points that this PR has either accomplished or will accomplish.

• Spin state is now in the PropertyNames class as "S"
• To be able to to embed this, per_system_spin_state is now in NNPInput.

Associated Issue(s)

• spin state embedding #353

Pull Request Checklist

• Issue(s) raised/addressed and linked
• Includes appropriate unit test(s)
• Appropriate docstring(s) added/updated
• Appropriate .rst doc file(s) added/updated
• PR is ready for review

[codecov-commenter]

Codecov Report

❌ Patch coverage is 81.81818% with 4 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.61%. Comparing base (9d51d94) to head (25e92f1).
⚠️ Report is 346 commits behind head on main.

❌ Your project status has failed because the head coverage (81.61%) is below the target coverage (90.00%). You can increase the head coverage or adjust the target coverage.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[Copilot]

Pull Request Overview

This PR integrates spin state embedding into modelforge following the established pattern for total charge, forces, and dipole moment. Key changes include:

• Adding a new property "S" and its corresponding unit mapping in PropertyNames.
• Introducing and propagating the per_system_spin_state tensor in NNPInput and related modules.
• Updating tests, dataset defaults, documentation, and API calls to include spin state handling.

Reviewed Changes

Copilot reviewed 13 out of 13 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
modelforge/utils/prop.py	Added spin state ("S") property and updated NNPInput with spin state.
modelforge/tests/test_dataset.py	Updated tests to reflect new spin state property in properties list.
modelforge/tests/data/dataset_defaults/fe_ii.toml	Included "spin_multiplicities" in the properties_of_interest list.
modelforge/potential/potential.py	Updated input signature for spin state in forward_for_jit_inference.
modelforge/potential/featurization.py	Added logic to append per_system_spin_state to the embedding tensor.
modelforge/jax.py	Propagated the spin state tensor into the JAX conversion workflow.
modelforge/dataset/parameters.py	Included spin state ("S") in the properties definitions.
modelforge/dataset/fe_II.py	Added spin state mapping in dataset and properties association.
modelforge/dataset/dataset.py	Updated loading and collation of spin state property in dataset.
modelforge-openmm/modelforge/openmm/utils.py	Added spin state to NNPInput; note default value differs from core code.
modelforge-openmm/modelforge/openmm/potential.py	Updated potential wrapper to forward the spin state tensor.
docs/potentials.rst	Updated documentation to mention usage of per_system_spin_state.