Median quantile finder

finding_median.py

import numpy as np
import torch as th
import torch.nn as nn
import matplotlib.pyplot as plt

from torch.optim import Adam


arr = np.arange(100)


class QR(nn.Module):
    def __init__(self):
        super().__init__()
        self.x = nn.Parameter(data=th.from_numpy(np.array([1], dtype=np.float32)))

qr = QR()
opt = Adam(params=qr.parameters(), lr=0.1)

estimated_median, real_median, losses = [], [], []
for i in range(2000):
    loss = th.mean(th.abs(th.from_numpy(arr)-qr.x))
    opt.zero_grad()
    loss.backward()
    opt.step()
    losses.append(loss.item())
    estimated_median.append(qr.x.clone().detach())
    real_median.append(np.median(arr))
    print(f"step: {i}, loss:{loss.item()}, estimated median: {qr.x.clone().detach()}, real meadian: {np.median(arr)}")
    
plt.plot(range(len(estimated_median)), estimated_median, label='estimated_median')
plt.plot(range(len(estimated_median)), real_median, label='median')
plt.plot(range(len(estimated_median)), losses, label='loss')
plt.legend()
plt.show()