add ensure_sorted (#93)

Author: jmoralez

nbs/processing.ipynb

@@ -33,7 +33,7 @@
     "import re\n",
     "import reprlib\n",
     "import warnings\n",
-    "from typing import Any, Dict, Generator, List, Optional, Tuple, Union\n",
+    "from typing import Any, Dict, Generator, List, NamedTuple, Optional, Tuple, Union\n",
     "\n",
     "import numpy as np\n",
     "import pandas as pd\n",
@@ -242,7 +242,7 @@
    "source": [
     "for engine in engines:\n",
     "    series = generate_series(2, engine=engine)\n",
-    "    x = np.random.rand(series.shape[0])    \n",
+    "    x = np.random.rand(series.shape[0])\n",
     "    series = assign_columns(series, 'x', x)\n",
     "    series = assign_columns(series, ['y', 'z'], np.vstack([x, x]).T)\n",
     "    series = assign_columns(series, 'ones', 1)\n",
@@ -1607,7 +1607,38 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "d3e059e1-8c6f-41b1-b3c0-4ca40adf09e4",
+   "id": "033b0a8a-98b1-486b-9414-1f5ef698f80f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#| export\n",
+    "def ensure_sorted(df: DataFrame, id_col: str, time_col: str) -> DataFrame:\n",
+    "    sort_idxs = maybe_compute_sort_indices(df=df, id_col=id_col, time_col=time_col)\n",
+    "    if sort_idxs is not None:\n",
+    "        df = take_rows(df=df, idxs=sort_idxs)\n",
+    "    return df"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4de12264-0bd1-4eed-935b-7b7fb1cbebc0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#| exporti\n",
+    "class _ProcessedDF(NamedTuple):\n",
+    "    uids: Series\n",
+    "    times: np.ndarray\n",
+    "    data: np.ndarray\n",
+    "    indptr: np.ndarray\n",
+    "    sort_idxs: Optional[np.ndarray]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "62293bd2-b921-40b2-b1af-25f0b8e55006",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1617,7 +1648,7 @@
     "    id_col: str,\n",
     "    time_col: str,\n",
     "    target_col: Optional[str],\n",
-    ") -> Tuple[Series, np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray]]:\n",
+    ") -> _ProcessedDF:\n",
     "    \"\"\"Extract components from dataframe\n",
     "    \n",
     "    Parameters\n",
@@ -1660,7 +1691,7 @@
     "        data = data[sort_idxs]\n",
     "        last_idxs = sort_idxs[last_idxs]\n",
     "    times = df[time_col].to_numpy()[last_idxs]\n",
-    "    return uids, times, data, indptr, sort_idxs    "
+    "    return _ProcessedDF(uids, times, data, indptr, sort_idxs)"
    ]
   },
   {

utilsforecast/_modidx.py

@@ -84,6 +84,8 @@
                                                                                                     'utilsforecast/processing.py'),
                                           'utilsforecast.processing.DataFrameProcessor.process': ( 'processing.html#dataframeprocessor.process',
                                                                                                    'utilsforecast/processing.py'),
+                                          'utilsforecast.processing._ProcessedDF': ( 'processing.html#_processeddf',
+                                                                                     'utilsforecast/processing.py'),
                                           'utilsforecast.processing._ensure_month_ends': ( 'processing.html#_ensure_month_ends',
                                                                                            'utilsforecast/processing.py'),
                                           'utilsforecast.processing._multiply_pl_freq': ( 'processing.html#_multiply_pl_freq',
@@ -111,6 +113,8 @@
                                                                                      'utilsforecast/processing.py'),
                                           'utilsforecast.processing.drop_index_if_pandas': ( 'processing.html#drop_index_if_pandas',
                                                                                              'utilsforecast/processing.py'),
+                                          'utilsforecast.processing.ensure_sorted': ( 'processing.html#ensure_sorted',
+                                                                                      'utilsforecast/processing.py'),
                                           'utilsforecast.processing.fill_null': ( 'processing.html#fill_null',
                                                                                   'utilsforecast/processing.py'),
                                           'utilsforecast.processing.filter_with_mask': ( 'processing.html#filter_with_mask',

utilsforecast/processing.py

@@ -5,14 +5,14 @@
            'filter_with_mask', 'is_nan', 'is_none', 'is_nan_or_none', 'match_if_categorical', 'vertical_concat',
            'horizontal_concat', 'copy_if_pandas', 'join', 'drop_index_if_pandas', 'rename', 'sort', 'offset_times',
            'offset_dates', 'time_ranges', 'repeat', 'cv_times', 'group_by', 'group_by_agg', 'is_in', 'between',
-           'fill_null', 'cast', 'value_cols_to_numpy', 'make_future_dataframe', 'anti_join', 'process_df',
-           'DataFrameProcessor', 'backtest_splits', 'add_insample_levels']
+           'fill_null', 'cast', 'value_cols_to_numpy', 'make_future_dataframe', 'anti_join', 'ensure_sorted',
+           'process_df', 'DataFrameProcessor', 'backtest_splits', 'add_insample_levels']
 
 # %% ../nbs/processing.ipynb 2
 import re
 import reprlib
 import warnings
-from typing import Any, Dict, Generator, List, Optional, Tuple, Union
+from typing import Any, Dict, Generator, List, NamedTuple, Optional, Tuple, Union
 
 import numpy as np
 import pandas as pd
@@ -626,12 +626,27 @@ def anti_join(df1: DataFrame, df2: DataFrame, on: Union[str, List[str]]) -> Data
     return out
 
 # %% ../nbs/processing.ipynb 74
+def ensure_sorted(df: DataFrame, id_col: str, time_col: str) -> DataFrame:
+    sort_idxs = maybe_compute_sort_indices(df=df, id_col=id_col, time_col=time_col)
+    if sort_idxs is not None:
+        df = take_rows(df=df, idxs=sort_idxs)
+    return df
+
+# %% ../nbs/processing.ipynb 75
+class _ProcessedDF(NamedTuple):
+    uids: Series
+    times: np.ndarray
+    data: np.ndarray
+    indptr: np.ndarray
+    sort_idxs: Optional[np.ndarray]
+
+# %% ../nbs/processing.ipynb 76
 def process_df(
     df: DataFrame,
     id_col: str,
     time_col: str,
     target_col: Optional[str],
-) -> Tuple[Series, np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray]]:
+) -> _ProcessedDF:
     """Extract components from dataframe
 
     Parameters
@@ -674,9 +689,9 @@ def process_df(
         data = data[sort_idxs]
         last_idxs = sort_idxs[last_idxs]
     times = df[time_col].to_numpy()[last_idxs]
-    return uids, times, data, indptr, sort_idxs
+    return _ProcessedDF(uids, times, data, indptr, sort_idxs)
 
-# %% ../nbs/processing.ipynb 76
+# %% ../nbs/processing.ipynb 78
 class DataFrameProcessor:
     def __init__(
         self,
@@ -693,7 +708,7 @@ def process(
     ) -> Tuple[Series, np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray]]:
         return process_df(df, self.id_col, self.time_col, self.target_col)
 
-# %% ../nbs/processing.ipynb 81
+# %% ../nbs/processing.ipynb 83
 def _single_split(
     df: DataFrame,
     i_window: int,
@@ -758,7 +773,7 @@ def _single_split(
         )
     return cutoffs, train_mask, valid_mask
 
-# %% ../nbs/processing.ipynb 82
+# %% ../nbs/processing.ipynb 84
 def backtest_splits(
     df: DataFrame,
     n_windows: int,
@@ -790,7 +805,7 @@ def backtest_splits(
         valid = filter_with_mask(df, valid_mask)
         yield cutoffs, train, valid
 
-# %% ../nbs/processing.ipynb 86
+# %% ../nbs/processing.ipynb 88
 def add_insample_levels(
     df: DataFrame,
     models: List[str],