Add introspection capabilities on simulation state#232
Conversation
- Also, bring cache invalidation outside of the step call for more control
|
Important Review skippedAuto reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the You can disable this status message by setting the WalkthroughThis PR refactors the event-based simulation pipeline with a state machine approach, introduces snapshot-based introspection capabilities, makes internal physics functions explicitly private, adds JSON-specific serialization hooks, and adjusts deserialization logic for None state handling. It removes a public energy query method and adds corresponding test coverage. Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~50 minutes The PR exhibits heterogeneous changes with multiple interrelated subsystems: a stateful simulation refactor replacing procedural logic, new introspection infrastructure integrating snapshots with caches, visibility boundary changes across physics modules, format-specific serialization paths, and API removals affecting callers. While individual files are reasonably scoped, the interconnected nature of the changes (simulate → snapshots → caches → serialization) and the mix of new logic alongside refactored existing patterns require careful review of integration points and data flow consistency. Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
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. Comment |
|
@coderabbitai review |
✅ Actions performedReview triggered.
|
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #232 +/- ##
==========================================
+ Coverage 44.22% 46.85% +2.62%
==========================================
Files 127 140 +13
Lines 8554 9400 +846
==========================================
+ Hits 3783 4404 +621
- Misses 4771 4996 +225
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
![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 (3)
pooltool/physics/resolve/serialize.py (3)
48-54: Consider reducing duplication between disambiguation functions.The two disambiguation functions differ only in the
SerializeFormatused. You could refactor to a single parameterized function or use a factory pattern to reduce duplication.Example refactor:
-def _disambiguate_model_structuring_yaml(v: dict[str, Any], t: type) -> Any: - return conversion[SerializeFormat.YAML].structure(v, _model_map[t][v["model"]]) - - -def _disambiguate_model_structuring_json(v: dict[str, Any], t: type) -> Any: - return conversion[SerializeFormat.JSON].structure(v, _model_map[t][v["model"]]) +def _make_disambiguate_model_structuring(fmt: SerializeFormat): + def _disambiguate(v: dict[str, Any], t: type) -> Any: + return conversion[fmt].structure(v, _model_map[t][v["model"]]) + return _disambiguateThen in
register_serialize_hooks, use:func=_make_disambiguate_model_structuring(SerializeFormat.YAML),
57-67: Consider using a loop to register structure hooks for both formats.The two structure hook registrations are nearly identical. If you adopt the factory pattern suggested above, you could loop over formats to reduce duplication.
Example:
- conversion.register_structure_hook_func( - check_func=lambda t: t in _model_map, - func=_disambiguate_model_structuring_yaml, - which=(SerializeFormat.YAML,), - ) - - conversion.register_structure_hook_func( - check_func=lambda t: t in _model_map, - func=_disambiguate_model_structuring_json, - which=(SerializeFormat.JSON,), - ) + for fmt in (SerializeFormat.YAML, SerializeFormat.JSON): + conversion.register_structure_hook_func( + check_func=lambda t: t in _model_map, + func=_make_disambiguate_model_structuring(fmt), + which=(fmt,), + )
71-89: Consider reducing duplication in unstructure hook registration.The YAML and JSON unstructure hook registrations are nearly identical. You could nest a format loop inside the model loop to eliminate this duplication.
Example refactor:
for models in _model_map.values(): for model_cls in models.values(): - conversion.register_unstructure_hook( - model_cls, - make_dict_unstructure_fn( - model_cls, - conversion[SerializeFormat.YAML], - _cattrs_include_init_false=True, - ), - which=(SerializeFormat.YAML,), - ) - - conversion.register_unstructure_hook( - model_cls, - make_dict_unstructure_fn( - model_cls, - conversion[SerializeFormat.JSON], - _cattrs_include_init_false=True, - ), - which=(SerializeFormat.JSON,), - ) + for fmt in (SerializeFormat.YAML, SerializeFormat.JSON): + conversion.register_unstructure_hook( + model_cls, + make_dict_unstructure_fn( + model_cls, + conversion[fmt], + _cattrs_include_init_false=True, + ), + which=(fmt,), + )
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (10)
pooltool/events/datatypes.py(1 hunks)pooltool/evolution/event_based/cache.py(4 hunks)pooltool/evolution/event_based/introspection.py(1 hunks)pooltool/evolution/event_based/simulate.py(5 hunks)pooltool/physics/__init__.py(0 hunks)pooltool/physics/evolve/__init__.py(6 hunks)pooltool/physics/resolve/serialize.py(2 hunks)pooltool/system/datatypes.py(0 hunks)tests/evolution/event_based/test_introspection.py(1 hunks)tests/evolution/event_based/test_simulate.py(2 hunks)
💤 Files with no reviewable changes (2)
- pooltool/system/datatypes.py
- pooltool/physics/init.py
⏰ 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). (4)
- GitHub Check: test (3.10, 1.8.4, windows-latest)
- GitHub Check: test (3.13, 1.8.4, windows-latest)
- GitHub Check: test (3.13, 1.8.4, ubuntu-latest)
- GitHub Check: test (3.10, 1.8.4, ubuntu-latest)
🔇 Additional comments (18)
pooltool/events/datatypes.py (1)
257-268: LGTM! More explicit None handling improves robustness.The updated deserialization logic now explicitly handles
Nonestates for both initial and final agent states. This is consistent with the introspection use case where events may be cached before resolution.pooltool/evolution/event_based/cache.py (3)
54-57: LGTM! Copy method correctly duplicates cache state.The implementation creates a new
TransitionCachewith deep-copied events, preventing unintended state sharing during snapshotting.
142-143: LGTM! Correct shallow copy for collision times.Creates a new instance with copied dictionary structure. Since the leaf values are floats (collision times), the shallow copy of inner dicts is appropriate and efficient.
173-203: Edge case is theoretical but unmitigated—consider explicit safeguards.After verifying the codebase, I found no evidence of pipe characters in actual object IDs (ball: "cue", "1", etc.; cushion: "6", "8t"; pocket: "rt"). However, ID fields are unvalidated strings with no documented format constraints. While the "|"-based serialization works correctly in practice, it remains theoretically fragile: if an ID containing "|" were ever created, deserialization would silently corrupt the data.
The comment's suggestion stands: either document that object IDs must exclude "|", or implement robust encoding (e.g., JSON array format as mentioned). Without explicit validation or constraints, the current implementation depends on developer discipline.
pooltool/physics/evolve/__init__.py (2)
33-40: Excellent fix for history immutability!Adding the copy and documenting the no-aliasing contract prevents the bug where resolver operations could corrupt historical ball states through shared array references. The test
test_ball_history_immutabilityvalidates this fix.
79-181: LGTM! Function renaming clarifies internal API.Making these physics helpers private with underscore prefixes is good API hygiene. All internal call sites are correctly updated.
tests/evolution/event_based/test_introspection.py (1)
1-103: Excellent test coverage for introspection capabilities!The test suite comprehensively validates:
- Equivalence with existing
simulatefunction- JSON serialization round-tripping
- Event selection and prospective events
- State progression (pre_evolve → post_evolve → post_resolve)
- Object identity vs. equality semantics
The tests provide strong validation for the new snapshot-based introspection API.
tests/evolution/event_based/test_simulate.py (3)
489-519: LGTM! Test validates critical immutability fix.This test ensures the fix in
evolve_ball_motion(copying the input array) prevents historical ball states from being corrupted by resolver operations. The test case is well-documented and directly validates the regression fix.
522-538: Test logic is correct.The test properly validates
_system_has_energyacross different system states. Note that there's a type annotation issue in the function being tested (see comment on simulate.py).
541-577: LGTM! Validates stick-ball event detection refactor.This test confirms that the refactored event detection properly identifies stick-ball collisions at t=0 as first-class events, rather than treating them as a special case. The test is well-documented and validates the expected event sequence.
pooltool/evolution/event_based/simulate.py (4)
49-131: LGTM! Clean state machine encapsulation.The
_SimulationStateclass nicely encapsulates simulation state and control flow. The separation ofstep(),update_caches(), andevolve()provides clear progression through the simulation loop and supports the introspection use case.
228-239: LGTM! Cleaner simulation loop with state machine.The refactored
simulatefunction is more maintainable with the state machine approach. The loop clearly separates event processing and cache updates.
242-300: LGTM! Enhanced event detection with stick-ball support.The updated signature accepting caches enables introspection while maintaining backward compatibility. The t=0 stick-ball detection treats the initial strike as a first-class event, which is cleaner than the previous special-case handling. The comment clearly explains the optimization strategy.
303-328: LGTM! Stick-ball collision detection with caching.The function correctly detects initial strike conditions (t=0, no ball energy, cue.V0 > 0) and caches the result. The caching prevents redundant checks after the first timestep.
pooltool/evolution/event_based/introspection.py (4)
49-108: LGTM! Correct event reconstruction from cache.The helper properly reconstructs event objects from cached collision times across all event types. The systematic iteration through cache.times ensures all prospective events are captured.
111-165: LGTM! Well-designed snapshot introspection API.The
SimulationSnapshotclass provides clean methods for examining system state at different stages:
get_prospective_events(): all possible events at this steppre_evolve_system(): state before time evolutionpost_evolve_system(): state after time evolution but before resolutionpost_resolve_system(): state after event resolutionThis API enables detailed simulation debugging and analysis.
167-188: LGTM! Clean sequence container with serialization.The
SimulationSnapshotSequenceprovides a simple container with list-like access and built-in JSON serialization support. The design is clean and focused.
191-245: LGTM! Comprehensive snapshot capture orchestration.The
simulate_with_snapshotsfunction correctly captures system state, caches, and events at each simulation step. The sequence of operations (capture pre-evolve state → step → copy caches → update caches) ensures accurate snapshot data for introspection.
| def _system_has_energy(system: System) -> float: | ||
| """Calculate the energy of the system in Joules.""" | ||
| return any( | ||
| bool( | ||
| ptmath.get_ball_energy( | ||
| ball.state.rvw, | ||
| ball.params.R, | ||
| ball.params.m, | ||
| ) | ||
| ) | ||
| for ball in system.balls.values() | ||
| ) |
There was a problem hiding this comment.
Fix type annotation and docstring mismatch.
The function has inconsistent type information:
- Return type annotation:
float - Docstring: "Calculate the energy of the system in Joules" (implies float)
- Actual return:
bool(fromany()) - Function name:
_system_has_energy(suggests boolean) - Usage in tests and code: treated as boolean
Based on the implementation and usage, the function correctly returns a boolean. The annotation and docstring need correction.
Apply this diff:
-def _system_has_energy(system: System) -> float:
- """Calculate the energy of the system in Joules."""
+def _system_has_energy(system: System) -> bool:
+ """Check whether the system has any non-zero energy."""
return any(📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def _system_has_energy(system: System) -> float: | |
| """Calculate the energy of the system in Joules.""" | |
| return any( | |
| bool( | |
| ptmath.get_ball_energy( | |
| ball.state.rvw, | |
| ball.params.R, | |
| ball.params.m, | |
| ) | |
| ) | |
| for ball in system.balls.values() | |
| ) | |
| def _system_has_energy(system: System) -> bool: | |
| """Check whether the system has any non-zero energy.""" | |
| return any( | |
| bool( | |
| ptmath.get_ball_energy( | |
| ball.state.rvw, | |
| ball.params.R, | |
| ball.params.m, | |
| ) | |
| ) | |
| for ball in system.balls.values() | |
| ) |
🤖 Prompt for AI Agents
In pooltool/evolution/event_based/simulate.py around lines 35 to 46, the
function _system_has_energy has mismatched type and docstring: it currently
annotates a float and docstring claims it calculates energy, but it actually
returns a boolean via any() and is used as a boolean. Change the return type
annotation to bool and update the docstring to state it checks whether any ball
in the system has nonzero energy (returns True/False); no logic changes
required.
1 participant
Why?
Introspecting the system state throughout a simulation has been very difficult historically, involving manually inserting breakpoints. This came to a boiling point as I worked towards debugging trajectories in the 3D branch.
The simulation state is now controlled with a stateful class that has a
stepmethod. Thesimulate(...)API remains identical, and now instantiates this_SimulationStepclass and iteratively callsstep.A new module
pooltool.evolution.event_based.introspectionwill track the system state at each step in the simulation process, including ability to introspect the system state, pre-evolve, post-evolve, and post-resolve. For examples, seetests/evolution/event_based/test_introspection.py. Entire simulation trajectories (states at each step) can be serialized/deserialized into JSON format.Summary by CodeRabbit
New Features
Bug Fixes
Refactor
get_system_energy()method from System.Tests