Quantile Regressor interaction

quantile.py

%matplotlib ipympl
import numpy as np
import matplotlib.pyplot as plt
import ipywidgets as widgets
from IPython.display import display

def symmetrize_xy(ax, ref="min"):
    x_min, x_max = ax.get_xlim()
    y_min, y_max = ax.get_ylim()
    x_range = x_max - x_min
    y_range = y_max - y_min
    max_range = max(x_range, y_range)
    if ref=='min':
        ax.set_xlim((x_min, x_min + max_range))
        ax.set_ylim((y_min, y_min + max_range))
    elif ref=="center":
        ax.set_xlim((x_min-max_range/2, x_min + max_range/2))
        #ax.set_ylim((y_min-max_range/2, y_min + max_range/2))


# Set random seed for reproducibility
np.random.seed(42)

class QuantileRegressorVisualiser:
    K = 3
    _N = 100
    def __init__(self, n_samples=100, q=0.5):
        self.beta0_min = -10*self.K
        self.beta0_max = 10*self.K
        self.beta1_min = -10*self.K
        self.beta1_max = 10*self.K
        
        self.n_samples = n_samples
        self.q = q

        self.x = np.random.uniform(low=2, high=10, size=self.n_samples)
        self.x_ = np.linspace(0, 10)
        self.noise = self.x/2 * np.random.normal(loc=0, scale=1, size=self.n_samples)
        self.y = 2 * self.x + 5 + self.noise
        self.fig, self.axes = plt.subplots(1,3)
        self.scatter = self.axes[0].scatter(self.x, self.y, alpha=0.7, label='Data points')

        self.axes[0].plot(self.x_, self.x_*2+5-self.x_/4, label="-sigma/2")
        self.axes[0].set_xlabel('Feature (x)')
        self.axes[0].set_ylabel('Target (y)')
        #self.axes[0].set_title('Toy Dataset for Quantile Regression')
        self.axes[0].legend()
        self.axes[0].grid(True)
        self.axes[0].set_ylim(0)
        self.axes[0].set_aspect('equal')

        self.beta_0 = 5
        self.beta_1 = 1

        self.scatter_resids = self.axes[1].scatter(self.loss_x, self.loss_y, label="PB residuals")
        self.fit_line       = self.axes[0].plot(self.x_, self.beta_1*self.x_ + self.beta_0)[0]
        self.axes[1].set_ylim(0)
        self.axes[1].set_xlim(-self.x__max, self.x__max)
        self.pb_line = self.axes[1].plot(self.x__, self.x__y, label=f"PB(q={self.q:.2f})", color="r")[0]
        self.legend = self.axes[1].legend()
        self.axes[1].set_aspect('equal')
        
        symmetrize_xy(self.axes[1], ref="min")
        self.axes[1].fill_between(self.x__, np.abs(self.x__), self.axes[1].get_ylim()[1], color='gray', alpha=0.3)  # Fill area above y=abs(x) with gray
        self.axes[1].set_title("Pinball loss and residuals")

        #symmetrize_xy(self.axes[0], ref="center")

        self.axes[2].grid()
        self.axes[2].set_xlim(self.beta0_min*1.1, self.beta0_max*1.1)
        self.axes[2].set_ylim(self.beta1_min*1.1, self.beta1_max*1.1)
        self.axes[2].set_aspect('equal')
        self.axes[2].set_title('Loss map')
        self.axes[2].set_xlabel(r"$\beta_0$ (offset)")
        self.axes[2].set_ylabel(r"$\beta_1$ (slope)")

        self.q_slider = widgets.FloatSlider(value=q, min=0.01, max=0.99, step=0.01, description='q:')
        self.beta0_slider = widgets.FloatSlider(value=self.beta_0, min=self.beta0_min, max=self.beta0_max, step=0.02, description='Beta0:')
        self.beta1_slider = widgets.FloatSlider(value=self.beta_1, min=self.beta1_min, max=self.beta1_max, step=0.02, description='Beta1:')
        self.q_slider.observe(self.update_q, 'value')
        self.beta0_slider.observe(self.update_beta0, 'value')
        self.beta1_slider.observe(self.update_beta1, 'value')
        #output = widgets.Output()
        #output.children = [self.fig]
        display(widgets.HBox([widgets.VBox([self.q_slider, self.beta0_slider, self.beta1_slider])]))

        self.explored_beta = []
        self.losses = []
        self.explored_beta.append(self.beta)
        self.losses.append(self.loss_y.mean())

        self.explored_beta_sc = self.axes[2].scatter(*np.array(self.explored_beta).T, c=self.losses, cmap='jet', alpha=0.2)
        self.beta_point = self.axes[2].scatter(*self.beta, color="r")

        Xs, Ys = np.meshgrid(np.linspace(self.beta0_min, self.beta0_max, self._N), np.linspace(self.beta1_min, self.beta1_max, self._N))
        y_preds = Ys.flatten().reshape(-1,1) * self.x + Xs.flatten().reshape(-1,1)
        self._Zs = (self.q * np.maximum(self.y - y_preds, 0) + (1-self.q) * np.maximum(y_preds - self.y, 0)).mean(axis=1).reshape(self._N, self._N)
        self.axes[2].contourf(Xs, Ys, self._Zs, 200, cmap='jet')
        self.fig.tight_layout()

    @property
    def beta(self): return self.beta_0, self.beta_1

    def update_q(self, change):
        self.q = change.new
        self.legend.texts[1].set_text(f"PB(q={self.q:.2f})")
        self.explored_beta = []
        self.losses = []
        self.explored_beta.append(self.beta)
        self.losses.append(self.loss_y.mean())
        self.update()

    def update_beta0(self, change):
        self.beta_0 = change.new
        self.explored_beta.append(self.beta)
        self.losses.append(self.loss_y.mean())
        self.update()

    def update_beta1(self, change):
        self.beta_1 = change.new
        self.explored_beta.append(self.beta)
        self.losses.append(self.loss_y.mean())
        self.update()

    def update(self):
        self.explored_beta_sc.remove()
        betas = self.explored_beta #np.unique(self.explored_beta, axis=0).T
        PB_losses = self.losses
        self.explored_beta_sc = self.axes[2].scatter(*np.array(betas).T, c=PB_losses, cmap='jet', zorder=10, alpha=0.2, vmin=self._Zs.min(), vmax=self._Zs.max())
        self.current = self.axes[2].scatter(*np.array(betas).T, c=PB_losses)
        
        self.scatter_resids.set_offsets(np.array([self.loss_x, self.loss_y]).T)
        self.beta_point.set_offsets(np.array(self.beta).T)
        self.fit_line.set_xdata(self.x_)
        self.fit_line.set_ydata(self.beta_1*self.x_ + self.beta_0)
        self.pb_line.set_xdata(self.x__)
        self.pb_line.set_ydata(self.x__y)
        
    @property
    def x__max(self): return np.max(self.axes[1].get_xlim())
    @property
    def x__(self): return np.linspace(-self.x__max, self.x__max, 200)
    @property
    def x__y(self): return self.q * np.maximum(self.x__, 0) + ( 1-self.q) * np.maximum(-self.x__, 0)
    @property
    def y_pred(self): return (self.beta_1 * self.x + self.beta_0)
    @property
    def loss_x(self): return self.y - self.y_pred
    @property
    def loss_y(self): return self.q * np.maximum(self.y - self.y_pred, 0) + (1-self.q) * np.maximum(self.y_pred - self.y, 0)

q = QuantileRegressorVisualiser()