Conversation
There was a problem hiding this comment.
Pull request overview
Introduces optional Numba-accelerated compute kernels and refactors the simulation runner to reduce per-step Python overhead during ODE integration.
Changes:
- Added
utils/numba_kernels.pywithmaybe_njitwrappers and kernels for slip, torque demand, and pulley force/torque calculations. - Refactored
simulation_runner.pyto cache model references and use precomputed lookup tables for CVT ratio/derivative and engine torque. - Updated pulley and slip model implementations to call the new kernels; added an optional
numbaextra and documented installation.
Reviewed changes
Copilot reviewed 8 out of 8 changed files in this pull request and generated 5 comments.
Show a summary per file
| File | Description |
|---|---|
cvtModel/src/cvt_simulator/utils/numba_kernels.py |
New optional JIT-friendly kernels (with fallback when numba is absent). |
cvtModel/src/cvt_simulator/simulation_runner.py |
Major runtime hot-path refactor: caching + LUTs + kernel calls + event functions moved to instance methods. |
cvtModel/src/cvt_simulator/models/slip_model.py |
Replaces Python/Numpy-side slip + torque-demand computations with kernel calls. |
cvtModel/src/cvt_simulator/models/pulley/secondary_pulley_torque_reactive.py |
Uses kernels for helix force + max torque computation. |
cvtModel/src/cvt_simulator/models/pulley/pulley_interface.py |
Uses kernel for radial force computation. |
cvtModel/src/cvt_simulator/models/pulley/primary_pulley_flyweight.py |
Uses kernels for flyweight force + max torque computation. |
cvtModel/setup.py |
Adds extras_require["numba"]. |
cvtModel/README.md |
Documents optional .[dev,numba] install for acceleration. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| import numpy as np | ||
|
|
There was a problem hiding this comment.
numpy is imported but not used in this module, which will trigger flake8 F401 in CI. Remove the unused import numpy as np (or use it if intentionally needed).
| import numpy as np |
| def maybe_njit(*args, **kwargs): | ||
| return njit(*args, **kwargs) | ||
|
|
||
| except ImportError: # pragma: no cover - exercised when numba is installed |
There was a problem hiding this comment.
The # pragma: no cover note here is misleading: this except ImportError path runs when numba is not installed. Either remove the pragma or update the comment so coverage expectations/documentation match actual behavior.
| except ImportError: # pragma: no cover - exercised when numba is installed | |
| except ImportError: # pragma: no cover - exercised only when numba is not installed |
| self._pulley_calc_shift_max = min(MAX_SHIFT, secondary_max_shift) | ||
|
|
||
|
|
||
|
|
There was a problem hiding this comment.
This blank line contains trailing whitespace, which will trigger flake8 W293. Remove the extra spaces so the line is truly empty.
| ) | ||
|
|
||
| net_radial = primary_radial - secondary_radial | ||
| friction = self._cvt_shift_model._frictional_force(net_radial, shift_velocity) |
There was a problem hiding this comment.
simulation_runner is calling the private method CvtShiftModel._frictional_force(...) directly. This tightly couples the runner to an internal implementation detail; consider exposing a public method (or reusing get_breakdown) so changes inside CvtShiftModel don't silently break the runner.
| friction = self._cvt_shift_model._frictional_force(net_radial, shift_velocity) | |
| def _compute_shift_friction(model, normal_force, axial_velocity): | |
| """ | |
| Compute the shift friction using the CVT shift model. | |
| Prefer a public `frictional_force` method if available, and fall back | |
| to the private `_frictional_force` for backward compatibility. | |
| """ | |
| friction_method = getattr(model, "frictional_force", None) | |
| if friction_method is None: | |
| friction_method = getattr(model, "_frictional_force", None) | |
| if friction_method is None: | |
| raise AttributeError( | |
| "CvtShiftModel instance does not provide a friction computation method." | |
| ) | |
| return friction_method(normal_force, axial_velocity) | |
| friction = _compute_shift_friction( | |
| self._cvt_shift_model, | |
| net_radial, | |
| shift_velocity, | |
| ) |
| @maybe_njit(cache=True, fastmath=True) | ||
| def slip_relative_speed_kernel( | ||
| primary_angular_velocity: float, | ||
| secondary_angular_velocity: float, | ||
| cvt_ratio: float, | ||
| ) -> float: | ||
| return primary_angular_velocity - (secondary_angular_velocity * cvt_ratio) | ||
|
|
||
|
|
||
| @maybe_njit(cache=True, fastmath=True) | ||
| def slip_coupling_torque_kernel( | ||
| relative_speed: float, | ||
| torque_demand: float, | ||
| t_max_capacity: float, | ||
| slip_speed_smoothing: float, | ||
| ) -> tuple[float, bool]: | ||
| coulomb_torque = t_max_capacity * math.tanh(relative_speed / slip_speed_smoothing) | ||
| alpha = min(max(abs(relative_speed) / slip_speed_smoothing, 0.0), 1.0) | ||
| torque_demand_clamped = min(max(torque_demand, -t_max_capacity), t_max_capacity) | ||
| coupling_torque = (1.0 - alpha) * torque_demand_clamped + alpha * coulomb_torque | ||
| return coupling_torque, alpha > 0.0 | ||
|
|
There was a problem hiding this comment.
This PR introduces several new kernel functions that are now part of the core dynamics (torque demand, slip coupling, etc.), but there are no unit tests validating their numerical behavior against known values / the previous Python implementations. Adding focused tests for these kernels would help catch regressions (e.g., smoothing/threshold edge cases, sign conventions, clamping).
There was a problem hiding this comment.
As far as I understand, numba's main purpose is to optimize all the operations that aren't numpy operations, with the JIT compliling into C-like machine code. Solid and easy approach, alternative is to make it all with numpy numbers throughout instead of the current mix.
Happy to use this approach if the numba stuff is fixed to just be added at the top of existing functions instead of moving them all to one file
There was a problem hiding this comment.
I like the idea of the @maybe_njit(cache=True, fastmath=True) decorator, but redefining all functions in this one spot is a real bad idea. It would be much better to just use that annotation throughout the codebase.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 7 out of 7 changed files in this pull request and generated 3 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| exp_term = math.exp(mu_effective * wrap_angle) | ||
| capstan_term = (exp_term - 1.0) / (exp_term + 1.0) | ||
| radial_force_term = total_radial * radius / math.sin(wrap_angle / 2.0) |
There was a problem hiding this comment.
Inconsistent use of math.exp and math.sin in this numba kernel. In the _primary_flyweight_force_kernel above (lines 32-33), you use math.atan and math.tan, but this file also imports numpy as np. For consistency and to ensure numba can properly optimize all kernels in this file, consider using either all math functions or all np functions within numba kernels. The primary_flyweight_force_kernel uses math, so this kernel should also use math for consistency.
| def maybe_njit(*args, **kwargs): | ||
| return njit(*args, **kwargs) | ||
|
|
||
| except ImportError: # pragma: no cover - exercised when numba is installed |
There was a problem hiding this comment.
The pragma comment is misleading. It states "exercised when numba is installed" but this is the except ImportError block, which executes when numba is NOT installed. The comment should either be "exercised when numba is not installed" or be moved to line 1-7 which is exercised when numba is installed.
| except ImportError: # pragma: no cover - exercised when numba is installed | |
| except ImportError: # pragma: no cover - exercised when numba is not installed |
| ) -> tuple[float, bool]: | ||
| coulomb_torque = t_max_capacity * math.tanh(relative_speed / slip_speed_smoothing) | ||
| alpha = min(max(abs(relative_speed) / slip_speed_smoothing, 0.0), 1.0) | ||
| torque_demand_clamped = min(max(torque_demand, -t_max_capacity), t_max_capacity) | ||
| coupling_torque = (1.0 - alpha) * torque_demand_clamped + alpha * coulomb_torque | ||
| return coupling_torque, alpha > 0.0 |
There was a problem hiding this comment.
The second return value (boolean) from _coupling_torque_kernel is unused. If this value is not needed for any future functionality, consider simplifying the kernel to only return the coupling_torque. This would make the API cleaner and avoid confusion about the unused return value.
| ) -> tuple[float, bool]: | |
| coulomb_torque = t_max_capacity * math.tanh(relative_speed / slip_speed_smoothing) | |
| alpha = min(max(abs(relative_speed) / slip_speed_smoothing, 0.0), 1.0) | |
| torque_demand_clamped = min(max(torque_demand, -t_max_capacity), t_max_capacity) | |
| coupling_torque = (1.0 - alpha) * torque_demand_clamped + alpha * coulomb_torque | |
| return coupling_torque, alpha > 0.0 | |
| ) -> float: | |
| coulomb_torque = t_max_capacity * math.tanh(relative_speed / slip_speed_smoothing) | |
| alpha = min(max(abs(relative_speed) / slip_speed_smoothing, 0.0), 1.0) | |
| torque_demand_clamped = min(max(torque_demand, -t_max_capacity), t_max_capacity) | |
| coupling_torque = (1.0 - alpha) * torque_demand_clamped + alpha * coulomb_torque | |
| return coupling_torque |
gr812b
left a comment
gr812b
left a comment
There was a problem hiding this comment.
Is there a reason you can't add those @ decorators to the functions that already exist? Ideally it would be best if we can just paste those @ above the existing functions instead of separating anything out if you catch my drift
cvtModel/setup.py
Outdated
| "numba": [ | ||
| "numba" | ||
| ] |
There was a problem hiding this comment.
If numba is to be used in production it should likely just be added to the top default packages
Description
Optimizations using numba kernels
Replaced Python-side calculations in
primary_pulley_flyweight.py,secondary_pulley_torque_reactive.py,pulley_interface.py, andslip_model.pywith calls to new Numba-accelerated kernels for force, torque, and slip computations.Updated
setup.pyto add anumbaoptional dependency group, and documented the optional installation for runtime acceleration inREADME.md.Refactored
simulation_runner.pyto cache model references and precompute lookup tables for CVT ratio, ratio derivatives, and engine torque.Changed event function references in the simulation runner to use instance methods instead of global functions.