Visualize Stochastic Gradient Descent using Contour plot in Python

Question

I tried to implement the stochastic gradient descent method and apply it to my build dataset. The data set follows a linear regression ( wx + b = y). The process has also somehow converged towards the appropriate values. What causes me difficulties is plotting the associated contour plot. The code of the plot is in the 2nd box.

#%% Packages
import numpy as np
import matplotlib.pyplot as plt
import random

#%% Prepare Dataset

# Generate data set along f + random noises 
f = lambda x: x*7+13 + np.random.randn(len(x))*np.random.normal(10, 33, len(x))
x = np.random.rand(500000)*100 # 500 random choosen inputs 
y = f(x) # thier according outputs
xy = np.vstack((x,y))

# obsere distribution
plt.scatter(x,y,alpha = 0.3, linewidths = 0.3) 
plt.show()

# simple regression and evaluation metric
h = lambda b,w,x: b + w*x
loss = lambda b,w, x_i, y_i: 0.5*(h(b,w,x_i)-y_i)**2
dloss = lambda b,w, x_i, y_i: np.array([h(b,w,x_i)-y_i, x_i*(h(b,w,x_i)-y_i)])

#%% Implementation of sgd

# Initailistion values
wb_old = np.array([0.,0.])
wb_new = np.array([1.,1.])
grad = np.array([0,0])

alpha = 0.00005 # step size
loss_history = [] # memorizes the loss of each iteration
avrgloss_history = []
weight_history = [] # memorizes the weights of each iteration
inpout_history = []
itmax = 5000000 # max amount of iteration
sumloss = 0
#Implementation of sgd
for it in range(itmax):
    wb_old = wb_new
    random_index = random.randrange(0,len(y)) # generate random index 
    (inp,out) = x[random_index], y[random_index] # randomly picking one data point from the data 
    grad = dloss(wb_old[0],wb_old[1],inp,out) # calculate gradient at that point

    wb_new = wb_old - alpha * grad # iteration instruction
    
    l = loss(wb_old[0],wb_old[1],inp,out)
    loss_history.append(l) # safe loss in history
    weight_history.append(wb_new)
    inpout_history.append((inp,out))   
    sumloss += l

    if (it+1) % 100000 == 0: # aprox. "Stopwatch"
            #print(it)
            #print(l)
            avrgloss_history.append(sumloss/100000)
            sumloss = 0   
            #print(l)


print("Local minimum occurs where:")
print("b =", wb_new[0])
print("w =", wb_new[1])

The stochastic gradient descent now yields b= 12.951 and w = 6.787, which is not too far from the values already set.

#%%

loss_history = np.asarray(loss_history)
inpout = np.asarray(inpout_history)

levels =  np.sort(loss_history)[::50000]
levels = np.sort(np.hstack((levels, np.array([0, 100, 300, 600, 1000, 5000, 15000, 50000, 150000, 1000000]))))

x = np.linspace(0, 10, 60)
y = np.linspace(0, 26, 60)
lo = np.zeros(shape=(x.size, y.size))

for i, value1 in enumerate(x):
    #print(i)
    for j, value2 in enumerate(y):       
        lo[i, j] = loss(value1,value2, inpout[i,0], inpout[i,1])

plt.figure(figsize=(20,20))
cp = plt.contour(x, y, lo, levels, colors='black', linestyles='dashed', linewidths=0.001)
plt.clabel(cp, inline=1, fontsize=8)
cp = plt.contourf(x, y, lo, levels, alpha=.7)
plt.title("Contour Plot of SGD")
plt.xlabel("b")
plt.ylabel("w")
plt.show()

Now I would like to visualise the behaviour of the method in a contourplot. Unfortunately, this did not work so well for me. I am sure that the problem lies with the "inpout" variable and that I have to choose the values differently - unfortunately, I don't know how. Can someone please help me with this? Thanks in advance

The resulting plot