sgd.ml

type sample = {
  x : float;
  y : float
}

type model = {
  theta : float;
  beta  : float
}

let make_sample s =
  { x = fst s; y = snd s }

let score m x =
  m.theta *. x + m.beta

let step_parameters m s alpha =
  let y_hat = score m s.x in
  { theta = m.theta - alpha * (y_hat - s.y) * s.x;
    beta = m.beta - alpha * (y_hat - s.y)
  }

let print_models ms =
  BatEnum.iter (fun m -> Printf.printf "theta: %f beta: %f\n" m.theta m.beta) ms

(* Simple regression example from http://stattrek.com/regression/regression-example.aspx *)
let x = BatArray.enum [|95.; 85.; 80.; 70.; 60.|] in
let y = BatArray.enum [|85.; 95.; 70.; 65.; 70.|] in
let samples = BatEnum.map (BatEnum.combine (x, y)) in
(* Initial parameter values *)
let m_0 = { theta = Random.float 1.0; beta = Random.float 1.0 } in
let shuffled_samples = BatRandom.shuffle samples in
let alpha = 0.1 in
let iter_1 = BatEnum.fold (fun ms s -> (step_parameters (hd ms) s alpha) :: ms) [m_0] shuffled_samples in
print_models iter_1