Add Nadaraya-Watson kernel regression with Gaussian kernel

fortuna/kernel_regression/__init__.py

@@ -0,0 +1 @@
+from fortuna.kernel_regression.nadaraya_watson import NadarayaWatsonKernelRegressor

fortuna/kernel_regression/kernels/gaussian.py

@@ -0,0 +1,7 @@
+import jax.numpy as jnp
+
+from fortuna.typing import Array
+
+
+def gaussian_kernel(x: Array) -> Array:
+    return jnp.exp(-0.5 * x**2) / jnp.sqrt(2 * jnp.pi)

fortuna/kernel_regression/nadaraya_watson.py

@@ -0,0 +1,92 @@
+from typing import Callable
+
+from jax import vmap
+import jax.numpy as jnp
+
+from fortuna.kernel_regression.kernels.gaussian import gaussian_kernel
+from fortuna.typing import Array
+
+
+class NadarayaWatsonKernelRegressor:
+    def __init__(
+        self,
+        train_inputs: Array,
+        train_targets: Array,
+        kernel: Callable[[Array], Array] = gaussian_kernel,
+    ):
+        """
+        A Nadaraya-Watson kernel regressor.
+
+        Parameters
+        ----------
+        kernel: Callable[[Array], Array]
+            A kernel.
+        train_inputs: Array
+            Training inputs.
+        train_targets: Array
+            Training targets.
+        """
+        self.kernel = kernel
+
+        if (train_inputs.ndim > 1) or (train_targets.ndim > 1):
+            raise ValueError(
+                "Both `train_inputs` and `train_targets` must be one-dimensional arrays."
+            )
+        if len(train_inputs) != len(train_targets):
+            raise ValueError(
+                "`train_inputs` and `train_targets` must have the same length."
+            )
+        self.train_inputs = jnp.copy(train_inputs)
+        self._mean_train_targets = jnp.mean(train_targets)
+        self._std_train_targets = jnp.std(train_targets)
+        self.standardized_train_targets = (
+            jnp.copy(train_targets) - self._mean_train_targets
+        ) / self._std_train_targets
+
+    def predict(self, inputs: Array, bandwidth: float = 0.1):
+        """
+        Predict the target
+
+        Parameters
+        ----------
+        inputs: Array
+            Inputs.
+        bandwidth: float
+            Kernel bandwidth.
+
+        Returns
+        -------
+        Predictions for the given inputs.
+        """
+        """"""
+        if inputs.ndim > 1:
+            raise ValueError("`inputs` must be a one-dimensional array.")
+        inputs = jnp.copy(inputs)
+        kernels = vmap(
+            lambda x: self.evaluate_scaled_kernel(
+                x - self.train_inputs, bandwidth=bandwidth
+            )
+        )(inputs)
+        m = jnp.sum(kernels * self.standardized_train_targets[None], axis=1) / jnp.sum(
+            kernels, axis=1
+        )
+        m *= self._std_train_targets
+        m += self._mean_train_targets
+        return m
+
+    def evaluate_scaled_kernel(self, inputs: Array, bandwidth: float) -> Array:
+        """
+        Given a kernel :math:`K(x)`, Evaluate the scaled kernel :math:`K_h(x) := \frac{1}{h}K\left(\frac{x}{h}\right)`.
+
+        Parameters
+        ----------
+        inputs: Array
+            Inputs.
+        bandwidth: Array
+            Bandwidth.
+
+        Returns
+        -------
+        Evaluation of the scaled kernel.
+        """
+        return self.kernel(inputs / bandwidth) / bandwidth

tests/fortuna/test_kernel_regression.py

@@ -0,0 +1,37 @@
+import unittest
+
+from jax import random
+import jax.numpy as jnp
+import numpy as np
+
+from fortuna.kernel_regression.nadaraya_watson import NadarayaWatsonKernelRegressor
+
+
+class TestKernelRegression(unittest.TestCase):
+    def test_nadaraya_watson(self):
+        train_x = random.normal(random.PRNGKey(0), shape=(3,))
+        train_y = random.normal(random.PRNGKey(1), shape=(3,))
+        eval_x = random.normal(random.PRNGKey(2), shape=(4,))
+
+        kr = NadarayaWatsonKernelRegressor(train_inputs=train_x, train_targets=train_y)
+        preds = kr.predict(inputs=eval_x)
+        assert preds.shape == (4,)
+
+        kr = NadarayaWatsonKernelRegressor(
+            train_inputs=np.array(train_x), train_targets=np.array(train_y)
+        )
+        preds = kr.predict(inputs=np.array(eval_x))
+        assert preds.shape == (4,)
+
+        with self.assertRaises(ValueError):
+            NadarayaWatsonKernelRegressor(
+                train_inputs=train_x, train_targets=train_y[None]
+            )
+        with self.assertRaises(ValueError):
+            NadarayaWatsonKernelRegressor(
+                train_inputs=train_x[None], train_targets=train_y
+            )
+        with self.assertRaises(ValueError):
+            NadarayaWatsonKernelRegressor(
+                train_inputs=jnp.concatenate((train_x, train_y)), train_targets=train_y
+            )