Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

feat(distributed): support ids in predict #454

Merged
merged 2 commits into from
Nov 22, 2024
Merged

feat(distributed): support ids in predict #454

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
fb42b70
feat(distributed): support ids in predict
jmoralez Nov 22, 2024
aa016aa
update cv schema
jmoralez Nov 22, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 28 additions & 11 deletions mlforecast/distributed/forecast.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@
except ModuleNotFoundError:
RAY_INSTALLED = False
from sklearn.base import clone
from triad import Schema

from mlforecast.core import (
DateFeature,
Expand Down Expand Up @@ -455,31 +456,43 @@ def _predict(
before_predict_callback=None,
after_predict_callback=None,
X_df=None,
ids=None,
schema=None,
) -> Iterable[pd.DataFrame]:
for serialized_ts, _, serialized_valid in items:
valid = cloudpickle.loads(serialized_valid)
if valid is not None:
X_df = valid
ts = cloudpickle.loads(serialized_ts)
if ids is not None:
ids = ts.uids.intersection(ids).tolist()
if not ids:
yield pd.DataFrame(
{
field.name: pd.Series(dtype=field.type.to_pandas_dtype())
for field in schema.values()
}
)
return
res = ts.predict(
models=models,
horizon=horizon,
before_predict_callback=before_predict_callback,
after_predict_callback=after_predict_callback,
X_df=X_df,
ids=ids,
)
if valid is not None:
res = res.merge(valid, how="left")
yield res

def _get_predict_schema(self) -> str:
model_names = self.models.keys()
models_schema = ",".join(f"{model_name}:double" for model_name in model_names)
schema = (
f"{self._base_ts.id_col}:string,{self._base_ts.time_col}:datetime,"
+ models_schema
)
return schema
def _get_predict_schema(self) -> Schema:
ids_schema = [
(self._base_ts.id_col, "string"),
(self._base_ts.time_col, "datetime"),
]
models_schema = [(model, "double") for model in self.models.keys()]
return Schema(ids_schema + models_schema)

def predict(
self,
Expand All @@ -488,6 +501,7 @@ def predict(
after_predict_callback: Optional[Callable] = None,
X_df: Optional[pd.DataFrame] = None,
new_df: Optional[fugue.AnyDataFrame] = None,
ids: Optional[List[str]] = None,
) -> fugue.AnyDataFrame:
"""Compute the predictions for the next `horizon` steps.

Expand All @@ -509,6 +523,8 @@ def predict(
Series data of new observations for which forecasts are to be generated.
This dataframe should have the same structure as the one used to fit the model, including any features and time series data.
If `new_df` is not None, the method will generate forecasts for the new observations.
ids : list of str, optional (default=None)
List with subset of ids seen during training for which the forecasts should be computed.

Returns
-------
Expand Down Expand Up @@ -540,6 +556,8 @@ def predict(
"before_predict_callback": before_predict_callback,
"after_predict_callback": after_predict_callback,
"X_df": X_df,
"ids": ids,
"schema": schema,
},
schema=schema,
engine=self.engine,
Expand Down Expand Up @@ -636,9 +654,8 @@ def cross_validation(
keep_last_n=keep_last_n,
window_info=window_info,
)
schema = (
self._get_predict_schema()
+ f",cutoff:datetime,{self._base_ts.target_col}:double"
schema = self._get_predict_schema() + Schema(
("cutoff", "datetime"), (self._base_ts.target_col, "double")
)
preds = fa.transform(
partition_results,
Expand Down
37 changes: 29 additions & 8 deletions nbs/distributed.forecast.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -98,6 +98,7 @@
"except ModuleNotFoundError:\n",
" RAY_INSTALLED = False\n",
"from sklearn.base import clone\n",
"from triad import Schema\n",
"\n",
"from mlforecast.core import (\n",
" DateFeature,\n",
Expand Down Expand Up @@ -506,29 +507,41 @@
" horizon,\n",
" before_predict_callback=None,\n",
" after_predict_callback=None,\n",
" X_df=None, \n",
" X_df=None,\n",
" ids=None,\n",
" schema=None,\n",
" ) -> Iterable[pd.DataFrame]:\n",
" for serialized_ts, _, serialized_valid in items:\n",
" valid = cloudpickle.loads(serialized_valid)\n",
" if valid is not None:\n",
" X_df = valid\n",
" ts = cloudpickle.loads(serialized_ts)\n",
" if ids is not None:\n",
" ids = ts.uids.intersection(ids).tolist()\n",
" if not ids:\n",
" yield pd.DataFrame(\n",
" {\n",
" field.name: pd.Series(dtype=field.type.to_pandas_dtype())\n",
" for field in schema.values()\n",
" }\n",
" )\n",
" return\n",
" res = ts.predict(\n",
" models=models,\n",
" horizon=horizon,\n",
" before_predict_callback=before_predict_callback,\n",
" after_predict_callback=after_predict_callback,\n",
" X_df=X_df,\n",
" ids=ids,\n",
" )\n",
" if valid is not None:\n",
" res = res.merge(valid, how='left')\n",
" yield res\n",
" \n",
" def _get_predict_schema(self) -> str:\n",
" model_names = self.models.keys()\n",
" models_schema = ','.join(f'{model_name}:double' for model_name in model_names)\n",
" schema = f'{self._base_ts.id_col}:string,{self._base_ts.time_col}:datetime,' + models_schema\n",
" return schema\n",
" def _get_predict_schema(self) -> Schema:\n",
" ids_schema = [(self._base_ts.id_col, 'string'), (self._base_ts.time_col, 'datetime')]\n",
" models_schema = [(model, 'double') for model in self.models.keys()]\n",
" return Schema(ids_schema + models_schema)\n",
"\n",
" def predict(\n",
" self,\n",
Expand All @@ -537,6 +550,7 @@
" after_predict_callback: Optional[Callable] = None,\n",
" X_df: Optional[pd.DataFrame] = None,\n",
" new_df: Optional[fugue.AnyDataFrame] = None,\n",
" ids: Optional[List[str]] = None,\n",
" ) -> fugue.AnyDataFrame:\n",
" \"\"\"Compute the predictions for the next `horizon` steps.\n",
"\n",
Expand All @@ -557,7 +571,9 @@
" new_df : dask or spark DataFrame, optional (default=None)\n",
" Series data of new observations for which forecasts are to be generated.\n",
" This dataframe should have the same structure as the one used to fit the model, including any features and time series data.\n",
" If `new_df` is not None, the method will generate forecasts for the new observations. \n",
" If `new_df` is not None, the method will generate forecasts for the new observations.\n",
" ids : list of str, optional (default=None)\n",
" List with subset of ids seen during training for which the forecasts should be computed. \n",
"\n",
" Returns\n",
" -------\n",
Expand Down Expand Up @@ -589,6 +605,8 @@
" 'before_predict_callback': before_predict_callback,\n",
" 'after_predict_callback': after_predict_callback,\n",
" 'X_df': X_df,\n",
" 'ids': ids,\n",
" 'schema': schema,\n",
" },\n",
" schema=schema,\n",
" engine=self.engine,\n",
Expand Down Expand Up @@ -685,7 +703,10 @@
" keep_last_n=keep_last_n,\n",
" window_info=window_info,\n",
" )\n",
" schema = self._get_predict_schema() + f',cutoff:datetime,{self._base_ts.target_col}:double'\n",
" schema = (\n",
" self._get_predict_schema() + Schema(\n",
" ('cutoff', 'datetime'), (self._base_ts.target_col, 'double'))\n",
" )\n",
" preds = fa.transform(\n",
" partition_results,\n",
" DistributedMLForecast._predict,\n",
Expand Down
95 changes: 54 additions & 41 deletions nbs/docs/getting-started/quick_start_distributed.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -366,32 +366,31 @@
"source": [
"#| hide\n",
"# test num_partitions works properly\n",
"if sys.version_info >= (3, 9):\n",
" num_partitions_test = 4\n",
" test_dd = dd.from_pandas(series, npartitions=num_partitions_test) # In this case we dont have to specify the column\n",
" test_dd['unique_id'] = test_dd['unique_id'].astype(str)\n",
" fcst_np = DistributedMLForecast(\n",
" models=models,\n",
" freq='D',\n",
" target_transforms=[Differences([7])], \n",
" lags=[7],\n",
" lag_transforms={\n",
" 1: [ExpandingMean()],\n",
" 7: [RollingMean(window_size=14)]\n",
" },\n",
" date_features=['dayofweek', 'month'],\n",
" num_threads=1,\n",
" engine=client,\n",
" num_partitions=num_partitions_test\n",
" )\n",
" fcst_np.fit(test_dd)\n",
" test_partition_results_size(fcst_np, num_partitions_test)\n",
" preds_np = fcst_np.predict(7).compute().sort_values(['unique_id', 'ds']).reset_index(drop=True)\n",
" preds = fcst.predict(7, X_df=future).compute().sort_values(['unique_id', 'ds']).reset_index(drop=True)\n",
" pd.testing.assert_frame_equal(\n",
" preds[['unique_id', 'ds']], \n",
" preds_np[['unique_id', 'ds']], \n",
" )"
"num_partitions_test = 4\n",
"test_dd = dd.from_pandas(series, npartitions=num_partitions_test) # In this case we dont have to specify the column\n",
"test_dd['unique_id'] = test_dd['unique_id'].astype(str)\n",
"fcst_np = DistributedMLForecast(\n",
" models=models,\n",
" freq='D',\n",
" target_transforms=[Differences([7])], \n",
" lags=[7],\n",
" lag_transforms={\n",
" 1: [ExpandingMean()],\n",
" 7: [RollingMean(window_size=14)]\n",
" },\n",
" date_features=['dayofweek', 'month'],\n",
" num_threads=1,\n",
" engine=client,\n",
" num_partitions=num_partitions_test\n",
")\n",
"fcst_np.fit(test_dd)\n",
"test_partition_results_size(fcst_np, num_partitions_test)\n",
"preds_np = fcst_np.predict(7).compute().sort_values(['unique_id', 'ds']).reset_index(drop=True)\n",
"preds = fcst.predict(7, X_df=future).compute().sort_values(['unique_id', 'ds']).reset_index(drop=True)\n",
"pd.testing.assert_frame_equal(\n",
" preds[['unique_id', 'ds']], \n",
" preds_np[['unique_id', 'ds']], \n",
")"
]
},
{
Expand Down Expand Up @@ -448,48 +447,48 @@
" <th>0</th>\n",
" <td>id_00</td>\n",
" <td>2002-09-27 00:00:00</td>\n",
" <td>22.267619</td>\n",
" <td>21.835798</td>\n",
" <td>21.722841</td>\n",
" <td>21.725511</td>\n",
" </tr>\n",
" <tr>\n",
" <th>1</th>\n",
" <td>id_00</td>\n",
" <td>2002-09-28 00:00:00</td>\n",
" <td>85.230055</td>\n",
" <td>83.996424</td>\n",
" <td>84.918194</td>\n",
" <td>84.606362</td>\n",
" </tr>\n",
" <tr>\n",
" <th>2</th>\n",
" <td>id_00</td>\n",
" <td>2002-09-29 00:00:00</td>\n",
" <td>168.256154</td>\n",
" <td>163.076652</td>\n",
" <td>162.067624</td>\n",
" <td>163.36802</td>\n",
" </tr>\n",
" <tr>\n",
" <th>3</th>\n",
" <td>id_00</td>\n",
" <td>2002-09-30 00:00:00</td>\n",
" <td>246.712244</td>\n",
" <td>245.827467</td>\n",
" <td>249.001477</td>\n",
" <td>246.422894</td>\n",
" </tr>\n",
" <tr>\n",
" <th>4</th>\n",
" <td>id_00</td>\n",
" <td>2002-10-01 00:00:00</td>\n",
" <td>314.184225</td>\n",
" <td>315.257849</td>\n",
" <td>317.149512</td>\n",
" <td>315.538403</td>\n",
" </tr>\n",
" </tbody>\n",
"</table>\n",
"</div>"
],
"text/plain": [
" unique_id ds DaskXGBForecast DaskLGBMForecast\n",
"0 id_00 2002-09-27 00:00:00 22.267619 21.835798\n",
"1 id_00 2002-09-28 00:00:00 85.230055 83.996424\n",
"2 id_00 2002-09-29 00:00:00 168.256154 163.076652\n",
"3 id_00 2002-09-30 00:00:00 246.712244 245.827467\n",
"4 id_00 2002-10-01 00:00:00 314.184225 315.257849"
"0 id_00 2002-09-27 00:00:00 21.722841 21.725511\n",
"1 id_00 2002-09-28 00:00:00 84.918194 84.606362\n",
"2 id_00 2002-09-29 00:00:00 162.067624 163.36802\n",
"3 id_00 2002-09-30 00:00:00 249.001477 246.422894\n",
"4 id_00 2002-10-01 00:00:00 317.149512 315.538403"
]
},
"execution_count": null,
Expand All @@ -502,6 +501,20 @@
"preds.head()"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "0150de6d-88b5-4513-bd82-c835ba945e79",
"metadata": {},
"outputs": [],
"source": [
"#| hide\n",
"# predict with ids\n",
"ids = np.random.choice(series['unique_id'].unique(), size=10, replace=False)\n",
"preds_ids = fcst.predict(7, X_df=future[future['unique_id'].isin(ids)], ids=ids).compute()\n",
"assert set(preds_ids['unique_id']) == set(ids)"
]
},
{
"cell_type": "code",
"execution_count": null,
Expand Down