FEAT: Add RMoK

nbs/docs/capabilities/01_overview.ipynb

@@ -34,6 +34,7 @@
     "|`NHITS` | `AutoNHITS` | MLP | Univariate | Direct | F/H/S | \n",
     "|`NLinear` | `AutoNLinear` | MLP | Univariate | Direct  | - | \n",
     "|`PatchTST` | `AutoPatchTST` | Transformer | Univariate | Direct  | - | \n",
+    "|`RMoK` | `AutoRMoK` | KAN | Multivariate | Direct | - |\n",
     "|`RNN` | `AutoRNN` | RNN | Univariate | Recursive  | F/H/S | \n",
     "|`SOFTS` | `AutoSOFTS` | MLP | Multivariate | Direct  | - | \n",
     "|`StemGNN` | `AutoStemGNN` | GNN | Multivariate | Direct  | - | \n",

nbs/models.ipynb

@@ -66,6 +66,7 @@
     "from neuralforecast.models.itransformer import iTransformer\n",
     "\n",
     "from neuralforecast.models.kan import KAN\n",
+    "from neuralforecast.models.rmok import RMoK\n",
     "\n",
     "from neuralforecast.models.stemgnn import StemGNN\n",
     "from neuralforecast.models.hint import HINT\n",
@@ -2300,6 +2301,14 @@
     "model.fit(dataset=dataset)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "8efc1692",
+   "metadata": {},
+   "source": [
+    "## C. KAN-Based"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -2444,7 +2453,7 @@
    "id": "fd705a56",
    "metadata": {},
    "source": [
-    "## C. Transformer-Based"
+    "## D. Transformer-Based"
    ]
   },
   {
@@ -3427,7 +3436,7 @@
    "id": "57d6cb1f",
    "metadata": {},
    "source": [
-    "## D. CNN Based"
+    "## E. CNN Based"
    ]
   },
   {
@@ -3573,7 +3582,7 @@
    "id": "e6fd22c7",
    "metadata": {},
    "source": [
-    "## E. Multivariate"
+    "## F. Multivariate"
    ]
   },
   {
@@ -4680,6 +4689,158 @@
     "model.fit(dataset=dataset)"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ab15c4b6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#| export\n",
+    "class AutoRMoK(BaseAuto):\n",
+    "\n",
+    "    default_config = {\n",
+    "        \"input_size_multiplier\": [1, 2, 3, 4, 5],\n",
+    "        \"h\": None,\n",
+    "        \"n_series\": None,\n",
+    "        \"taylor_order\": tune.choice([3, 4, 5]),\n",
+    "        \"jacobi_degree\": tune.choice([4, 5, 6]),\n",
+    "        \"wavelet_function\": tune.choice(['mexican_hat', 'morlet', 'dog', 'meyer', 'shannon']),\n",
+    "        \"learning_rate\": tune.loguniform(1e-4, 1e-1),\n",
+    "        \"scaler_type\": tune.choice([None, 'robust', 'standard', 'identity']),\n",
+    "        \"max_steps\": tune.choice([500, 1000, 2000]),\n",
+    "        \"batch_size\": tune.choice([32, 64, 128, 256]),\n",
+    "        \"loss\": None,\n",
+    "        \"random_seed\": tune.randint(1, 20),\n",
+    "    }\n",
+    "\n",
+    "    def __init__(self,\n",
+    "                 h,\n",
+    "                 n_series,\n",
+    "                 loss=MAE(),\n",
+    "                 valid_loss=None,\n",
+    "                 config=None, \n",
+    "                 search_alg=BasicVariantGenerator(random_state=1),\n",
+    "                 num_samples=10,\n",
+    "                 refit_with_val=False,\n",
+    "                 cpus=cpu_count(),\n",
+    "                 gpus=torch.cuda.device_count(),\n",
+    "                 verbose=False,\n",
+    "                 alias=None,\n",
+    "                 backend='ray',\n",
+    "                 callbacks=None):\n",
+    "        \n",
+    "        # Define search space, input/output sizes\n",
+    "        if config is None:\n",
+    "            config = self.get_default_config(h=h, backend=backend, n_series=n_series)                 \n",
+    "\n",
+    "        # Always use n_series from parameters, raise exception with Optuna because we can't enforce it\n",
+    "        if backend == 'ray':\n",
+    "            config['n_series'] = n_series\n",
+    "        elif backend == 'optuna':\n",
+    "            mock_trial = MockTrial()\n",
+    "            if ('n_series' in config(mock_trial) and config(mock_trial)['n_series'] != n_series) or ('n_series' not in config(mock_trial)):\n",
+    "                raise Exception(f\"config needs 'n_series': {n_series}\")           \n",
+    "\n",
+    "        super(AutoRMoK, self).__init__(\n",
+    "              cls_model=RMoK, \n",
+    "              h=h,\n",
+    "              loss=loss,\n",
+    "              valid_loss=valid_loss,\n",
+    "              config=config,\n",
+    "              search_alg=search_alg,\n",
+    "              num_samples=num_samples, \n",
+    "              refit_with_val=refit_with_val,\n",
+    "              cpus=cpus,\n",
+    "              gpus=gpus,\n",
+    "              verbose=verbose,\n",
+    "              alias=alias,\n",
+    "              backend=backend,\n",
+    "              callbacks=callbacks,            \n",
+    "        )\n",
+    "\n",
+    "    @classmethod\n",
+    "    def get_default_config(cls, h, backend, n_series):\n",
+    "        config = cls.default_config.copy()        \n",
+    "        config['input_size'] = tune.choice([h * x \\\n",
+    "                        for x in config[\"input_size_multiplier\"]])\n",
+    "\n",
+    "        # Rolling windows with step_size=1 or step_size=h\n",
+    "        # See `BaseWindows` and `BaseRNN`'s create_windows\n",
+    "        config['step_size'] = tune.choice([1, h])\n",
+    "        del config[\"input_size_multiplier\"]\n",
+    "        if backend == 'optuna':\n",
+    "            # Always use n_series from parameters\n",
+    "            config['n_series'] = n_series\n",
+    "            config = cls._ray_config_to_optuna(config)           \n",
+    "\n",
+    "        return config         "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "111d8d3b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "show_doc(AutoRMoK, title_level=3)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2073d4aa",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%capture\n",
+    "# Use your own config or AutoRMoK.default_config\n",
+    "config = dict(max_steps=1, val_check_steps=1, input_size=12, learning_rate=1e-2)\n",
+    "model = AutoRMoK(h=12, n_series=1, config=config, num_samples=1, cpus=1)\n",
+    "\n",
+    "# Fit and predict\n",
+    "model.fit(dataset=dataset)\n",
+    "y_hat = model.predict(dataset=dataset)\n",
+    "\n",
+    "# Optuna\n",
+    "model = AutoRMoK(h=12, n_series=1, config=None, backend='optuna')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ebe2c500",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#| hide\n",
+    "# Check Optuna\n",
+    "assert model.config(MockTrial())['h'] == 12\n",
+    "\n",
+    "# Unit test to test that Auto* model contains all required arguments from BaseAuto\n",
+    "test_args(AutoRMoK, exclude_args=['cls_model']) \n",
+    "\n",
+    "# Unit test for situation: Optuna with updated default config\n",
+    "my_config = AutoRMoK.get_default_config(h=12, n_series=1, backend='optuna')\n",
+    "def my_config_new(trial):\n",
+    "    config = {**my_config(trial)}\n",
+    "    config.update({'max_steps': 1, 'val_check_steps': 1, 'input_size': 12, 'learning_rate': 1e-1})\n",
+    "    return config\n",
+    "\n",
+    "model = AutoRMoK(h=12, n_series=1, config=my_config_new, backend='optuna', num_samples=1, cpus=1)\n",
+    "model.fit(dataset=dataset)\n",
+    "\n",
+    "# Unit test for situation: Ray with updated default config\n",
+    "my_config = AutoRMoK.get_default_config(h=12, n_series=1, backend='ray')\n",
+    "my_config['max_steps'] = 1\n",
+    "my_config['val_check_steps'] = 1\n",
+    "my_config['input_size'] = 12\n",
+    "my_config['learning_rate'] = 1e-1\n",
+    "model = AutoRMoK(h=12, n_series=1, config=my_config, backend='ray', num_samples=1, cpus=1)\n",
+    "model.fit(dataset=dataset)"
+   ]
+  },
   {
    "attachments": {},
    "cell_type": "markdown",