Make system_idx non-optional in SimState [1/2]#231
Conversation
WalkthroughThe Changes
Estimated code review effort🎯 2 (Simple) | ⏱️ ~8 minutes Poem
Note ⚡️ Unit Test Generation is now available in beta!Learn more here, or try it out under "Finishing Touches" below. 📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (1)
✅ Files skipped from review due to trivial changes (1)
✨ Finishing Touches
🧪 Generate unit tests
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
torch_sim/state.py (1)
85-105: Constructor implementation looks good, but consider the TODO.The constructor signature correctly maintains backward compatibility while enforcing non-optional
system_idxat the class level. The docstring is comprehensive and accurate.However, the TODO comment on line 90 raises a valid point about potential confusion with positional arguments.
Consider making the constructor keyword-only to prevent argument order mistakes:
def __init__( self, + *, positions: torch.Tensor, masses: torch.Tensor, cell: torch.Tensor, pbc: bool, # noqa: FBT001 # TODO(curtis): maybe make the constructor be keyword-only (it can be easy to confuse positions vs masses, etc.) atomic_numbers: torch.Tensor, system_idx: torch.Tensor | None = None, ) -> None:This would require all arguments to be passed by name, eliminating the confusion mentioned in the TODO.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
torch_sim/state.py(5 hunks)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (29)
- GitHub Check: test-examples (examples/scripts/7_Others/7.6_Compare_ASE_to_VV_FIRE.py)
- GitHub Check: test-examples (examples/scripts/6_Phonons/6.2_QuasiHarmonic_MACE.py)
- GitHub Check: test-examples (examples/scripts/7_Others/7.4_Velocity_AutoCorrelation.py)
- GitHub Check: test-examples (examples/scripts/2_Structural_optimization/2.5_MACE_UnitCellFilter_Gradient_Descen...
- GitHub Check: test-examples (examples/scripts/2_Structural_optimization/2.4_MACE_FIRE.py)
- GitHub Check: test-examples (examples/scripts/4_High_level_api/4.1_high_level_api.py)
- GitHub Check: test-examples (examples/scripts/3_Dynamics/3.12_MACE_NPT_Langevin.py)
- GitHub Check: test-examples (examples/scripts/3_Dynamics/3.8_MACE_NPT_Nose_Hoover.py)
- GitHub Check: test-examples (examples/scripts/5_Workflow/5.3_Elastic.py)
- GitHub Check: test-examples (examples/scripts/3_Dynamics/3.5_MACE_NVT_Nose_Hoover.py)
- GitHub Check: test-examples (examples/tutorials/autobatching_tutorial.py)
- GitHub Check: test-examples (examples/tutorials/high_level_tutorial.py)
- GitHub Check: test-examples (examples/tutorials/low_level_tutorial.py)
- GitHub Check: test-model (macos-14, 3.12, lowest-direct, mace, tests/test_optimizers_vs_ase.py)
- GitHub Check: test-model (macos-14, 3.11, highest, mattersim, tests/models/test_mattersim.py)
- GitHub Check: test-model (macos-14, 3.11, highest, metatomic, tests/models/test_metatomic.py)
- GitHub Check: test-model (macos-14, 3.11, highest, orb, tests/models/test_orb.py)
- GitHub Check: test-model (macos-14, 3.11, highest, mace, tests/test_elastic.py)
- GitHub Check: test-model (ubuntu-latest, 3.12, lowest-direct, orb, tests/models/test_orb.py)
- GitHub Check: test-model (ubuntu-latest, 3.12, lowest-direct, mace, tests/test_optimizers_vs_ase.py)
- GitHub Check: test-model (ubuntu-latest, 3.12, lowest-direct, sevenn, tests/models/test_sevennet.py)
- GitHub Check: test-model (ubuntu-latest, 3.12, lowest-direct, fairchem, tests/models/test_fairchem.py)
- GitHub Check: test-model (ubuntu-latest, 3.12, lowest-direct, mace, tests/test_elastic.py)
- GitHub Check: test-model (ubuntu-latest, 3.11, highest, mace, tests/test_elastic.py)
- GitHub Check: test-model (ubuntu-latest, 3.11, highest, orb, tests/models/test_orb.py)
- GitHub Check: test-model (ubuntu-latest, 3.11, highest, sevenn, tests/models/test_sevennet.py)
- GitHub Check: test-model (ubuntu-latest, 3.11, highest, fairchem, tests/models/test_fairchem.py)
- GitHub Check: test-core (ubuntu-latest, 3.12, lowest-direct)
- GitHub Check: build-docs
🔇 Additional comments (6)
torch_sim/state.py (6)
10-10: LGTM!The removal of the
fieldimport is appropriate since it's no longer needed after switching to a manual constructor.
25-25: LGTM!Using
init=Falseis the correct approach for implementing a custom constructor while maintaining dataclass benefits.
50-51: LGTM!The type annotation change correctly reflects that
system_idxis now non-optional in the class definition, improving type safety as intended by the PR.Also applies to: 83-83
106-133: LGTM!The initialization and validation logic is well-implemented with proper error handling and clear error messages.
134-146: LGTM!The system_idx initialization correctly provides backward compatibility by defaulting to zeros when not provided, while the validation ensures data integrity.
150-157: LGTM!The cell shape validation is comprehensive and provides clear error messages.
curtischong
changed the title
SimStateSimState [1/3]
curtischong
changed the title
SimState [1/3]SimState [1/2]
| cell=state.cell, | ||
| pbc=state.pbc, | ||
| atomic_numbers=atomic_numbers, | ||
| system_idx=state.system_idx, |
There was a problem hiding this comment.
types caught this bug!
2 participants
Summary
This is part of my effort to type the codebase. system_idx is only optional during initialization so in the dataclass definition for
SimState, I changed it fromsystem_idx: torch.Tensor | Nonetosystem_idx: torch.TensorMore importantly: This PR is the precursor to the one where I fix the SimState concatenation issue #232 (because that PR has a check that disallows
torch.Tensor | Noneattributes. By making system_idx non-optional, I can remove the| Nonein its definitionTo also allow users to pass in a
Nonesystem_idxin the constructor, I initialize the dataclass withinit=Falseand wrote my own constructor forSimStateChecklist
Before a pull request can be merged, the following items must be checked:
Run ruff on your code.
We highly recommended installing the pre-commit hooks running in CI locally to speedup the development process. Simply run
pip install pre-commit && pre-commit installto install the hooks which will check your code before each commit.Summary by CodeRabbit