Skip to content

Step Function

In this basic tutorial, we fit distributional kernel regression model to data that comes from a step function with heterskedastic noise. The ability to naturally capture heterskedastic noise is one of the biggest benefits of distributional kernel regression when compared to Gaussian processes.

For this problem, we generate the data points according to: $$ x\sim\mathcal{U}(0,10) $$ $$ y\sim\mathcal{N}(h(x),0.1\lvert x-5\rvert) $$ where \(h\) is the step function defined as: $$ h(x) = \begin{cases}1 & \mathrm{if}~x\geq 6 \\ 0 & \mathrm{else}\end{cases} $$ We start by importing all necessary packages for the tutorial:

import torch
from dkregression.kernels import RBF
from dkregression.likelihoods import UnivariateGaussianLikelihood
from dkregression.cross_validation import CrossValidation
from dkregression import DKR
import matplotlib.pyplot as plt
In the next step we generate 50 training data points stored in X and Y. Note that for consistency between uni and multivariate problems, both X and Y are always required to be two-dimensional tensors, even if the problem is one-dimensional. For our problem, both X and Y will be of size (50,1). In addition we also generate the query values Xq which we will use later on for plotting.
Xq = torch.linspace(0,10,200).unsqueeze(-1)
X = torch.distributions.uniform.Uniform(torch.Tensor([0]),torch.Tensor([10])).sample((50,))
Y = torch.where(X>=6,1.,0.) + 0.1*torch.abs(X-5)*torch.randn((50,1))
Next, we set up the DKR model. We DKR object requires a kernel, an observation likelihood, and a cross-validation setup for fitting the kernel hyperparameters. In this tutorial, we use the commonly used radial basis function (RBF) kernel which is also known as squared exponential kernel. The RBF kernel has one hyperparameter: the scale length \(\ell\). To provide an initial guess for the scale length, it is necessary to pass in the X values of the dataset to the RBF object when creating it. The initial guess is used later on by the fit method as an inital value for the optimizer. More details about the heuristic used to estimated \ell can be found in the reference section of the RBF kernel. Neither the observation likelihood, not the cross validation configuration setup require any additional configuration for the purpose of this tutorial.
kernel = RBF(X)
likelihood = UnivariateGaussianLikelihood()
cv = CrossValidation()
model = DKR(kernel,likelihood,cv)
No we can finally fit the model to the data X and Y and run inference on the model with Xq.
model.fit(X,Y)
Y_pred = model.predict(Xq)
Finally, we plot the results.
plt.scatter(X.flatten(),Y.flatten(),s=5)
plt.plot(Xq.flatten(),
         torch.where(Xq>=6,1.,0.).flatten(),
         "--",
         label="Step function")
plt.plot(Xq.flatten(),
         Y_pred["mu"].flatten(),
         c="#485B6E",
         label=r"Model $\mu$")
plt.fill_between(Xq.flatten(),
                 Y_pred["mu"].flatten() - 1.96*Y_pred["sigma"].flatten(), 
                 Y_pred["mu"].flatten() + 1.96*Y_pred["sigma"].flatten(),
                 color="#485B6E",
                 alpha=0.3,
                 label="Model 95% CI")
plt.xlabel(r"$x$")
plt.ylabel(r"$y$")
plt.legend()
plt.tight_layout()
plt.show()
Step Function Step Function