Skip to content

Commit

Permalink
feat(distributed): support ids in predict (#454)
Browse files Browse the repository at this point in the history
  • Loading branch information
@jmoralez
jmoralez authored Nov 22, 2024
1 parent 2342dea commit c833952
Show file tree
Hide file tree
Showing 3 changed files with 111 additions and 60 deletions.
39 changes: 28 additions & 11 deletions mlforecast/distributed/forecast.py
View file
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
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
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

0 comments on commit c833952

Please sign in to comment.