Chapter 3, Linear Regression

# load the packages
using Plots
using Flux
using Flux: throttle

# generate 100 data points scattered around a straight line and plot
regX = rand(100)
regY = 50 .+ 100 * regX + 2 * randn(100);
scatter(regX, regY, fmt = :png, legend=:bottomright, label="data")

# linear regression using Julia's linear algebra
X = hcat(ones(length(regX)),regX)
Y = regY
intercept,slope = inv(X'*X)*(X'*Y)

# add the regression line to the data plot
plot!((x) -> intercept + slope * x, 0, 1, label="fit_exact")

data = zip(regX, regY)  # create tuples
data = gpu.(data)
model = Dense(1, 1, identity) |> gpu # define the model
loss(x, y) = Flux.mse(model([x]), y) # use mean square error for loss

# define a call back ftn that prints the loss over the whole data set
evalcb = () -> @show(sum([loss(i[1],i[2]) for i in data]))

# use Stocastic Gradient Descent to optimise the model weights
# opt = SGD(Flux.params(modelReg), 0.1)
opt = Descent(0.1)

# train the model on 10 epochs
for i = 1:10 
  Flux.train!(loss, params(model), data, opt, cb = throttle(evalcb, 10)) 
end

(θ, bias) = Tracker.data.(Flux.params(model))

plot!((x) ->  bias[1] + θ[1] * x, 0, 1, label="fit_flux")

# generate data near a plane in two dimensions
using Plots
regX = 2 .* (rand(500) .- 0.5)
regY = 2 .* (rand(500) .- 0.5)

# Z = f(X,Y)
# regZ = 10 .+ ( ( regX.^2 + regY.^2) .+ 0.1 * (randn(500) .- 0.5) )
regZ = 10 .+ 5 .* regX + 12 .* regY .+ 0.2 * (randn(500) .- 0.5) 
scatter(regX, regY, regZ,  fmt = :png, legend=:bottomright, label="data")

data2D = [([regX[i],regY[i]],regZ[i]) for i in 1:length(regX)]
data2D = gpu.(data2D)
model = Chain(Dense(2, 1, identity)) |> gpu
loss(x, y) = Flux.mse(model(x), y) 
evalcb = () -> @show(sum([loss(i[1],i[2]) for i in data2D]))
opt = ADAM(0.01)
for i=1:10 
  Flux.train!(loss, params(model), data2D, opt, cb = throttle(evalcb, 1000)) end

(θ, bias) = Tracker.data.(Flux.params(model))

x=range(-1,stop=1,length=10)
y=range(-1,stop=1,length=10)
f(x,y) = bias[1] + θ[1] * x + θ[2] * y
plot(x, y, f,st=:wireframe,alpha=0.1,camera=(-30,30),fmt=:png)
scatter!(regX, regY, regZ, colour="blue", fmt = :png, legend=:bottomright, label="data")