fix: timesfm returns correct quantile names#131
Conversation
There was a problem hiding this comment.
Pull Request Overview
This PR fixes an issue where the TimesFM model was returning duplicate quantile columns with different naming conventions, and improves test coverage and performance. The changes ensure that quantile columns are properly renamed without duplication and verify that the correct number of columns are returned for probabilistic forecasts.
- Replaces inefficient column assignment loop with dictionary-based column renaming
- Adds test assertions to verify the correct number of forecast columns are returned
- Switches to TimesFM 1.0 model in tests for improved performance
Reviewed Changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
| timecopilot/models/foundational/timesfm.py | Fixes quantile column naming by replacing loop-based column assignment with efficient dictionary renaming and removes unused import |
| tests/models/test_models.py | Adds assertions to verify correct number of columns in quantile and level forecast outputs |
| tests/conftest.py | Updates test configuration to use faster TimesFM 1.0 model and moves TimesFM from string list to explicit model instance |
| quantiles=qs, | ||
| ) | ||
| exp_qs_cols = [f"{model.alias}-q-{int(100 * q)}" for q in qs] | ||
| assert len(exp_qs_cols) == len(fcst_df.columns) - 3 # 3 is unique_id, ds, point |
There was a problem hiding this comment.
[nitpick] The magic number 3 and its explanation in the comment could be made more maintainable by defining a constant or using a more explicit calculation. Consider defining STANDARD_COLUMNS = ['unique_id', 'ds', model.alias] and using len(STANDARD_COLUMNS) instead.
| assert len(exp_qs_cols) == len(fcst_df.columns) - 3 # 3 is unique_id, ds, point | |
| STANDARD_COLUMNS = ['unique_id', 'ds', model.alias] | |
| assert len(exp_qs_cols) == len(fcst_df.columns) - len(STANDARD_COLUMNS) |
| if lv == 0: | ||
| continue | ||
| exp_lv_cols.extend([f"{model.alias}-lo-{lv}", f"{model.alias}-hi-{lv}"]) | ||
| assert len(exp_lv_cols) == len(fcst_df.columns) - 3 # 3 is unique_id, ds, point |
There was a problem hiding this comment.
[nitpick] Similar to the quantiles test, the magic number 3 is duplicated here. Consider extracting this into a shared constant or helper function to avoid duplication and improve maintainability.
| assert len(exp_lv_cols) == len(fcst_df.columns) - 3 # 3 is unique_id, ds, point | |
| assert len(exp_lv_cols) == len(fcst_df.columns) - NUM_STATIC_COLUMNS # NUM_STATIC_COLUMNS is unique_id, ds, point |
this pr fixes #130 regarding timesfm returning the quantiles twice with different names. it also adds test to check the correct number of columns is returned in the case of probabilistic forecasts (quantiles and levels) and uses timesfm 1.0 in tests to speed them up.