<commit_before>open BatPervasives
open Utils

(* this is not gradient descent by itself, it depends
   on the cost and step functions *)
let optimize ?max_steps ~epsilon (x0:'a) (cost:'a -> float) (step:int -> 'a -> 'a) =
  let rec loop cur_step prev_cost x =
    if (match max_steps with
	  | Some max_steps when cur_step >= max_steps -> false
	  | _ -> true)
    then
      let new_cost = cost x in
	Printf.printf "cost difference %f\n" (new_cost -. prev_cost);
	if abs_float (new_cost -. prev_cost) < epsilon 
	 then 
	   let _ = Printf.printf "found at step %d\n" cur_step in
	     x
	else loop (cur_step + 1) new_cost (step cur_step x)
    else
      x
  in
    loop 0 (cost x0) (step 0 x0)

let derivate ~epsilon f x = (f (x +. epsilon) -. f (x -. epsilon)) /. (2.0 *. epsilon)

let gradient2 ~epsilon f1 f2 (x1, x2) = 
  let e1, e2 = epsilon in
    (derivate e1 f1 x1, derivate e2 f2 x2)

(* m = number of training examples
   x = training input
   y = output

 *)

let linreg_h2 x theta = 
  theta.(0) +. theta.(1) *. x.(0)

let (!@) x = float (Array.length x)

let linreg_cost xs ys theta =
  assert (Array.length xs = Array.length ys);
  1.0 /. (2.0 *. !@ theta) *. sum (BatArray.map2 (fun x y -> (linreg_h2 x theta -. y) ** 2.0) xs ys)

let linreg_cost2' xs ys theta =
  assert (Array.length theta = 2);
  assert (Array.length xs = Array.length ys);
  [|1.0 /. !@ theta *. sum (BatArray.map2 (fun x y -> linreg_h2 x theta -. y) xs ys);
    1.0 /. !@ theta *. sum (BatArray.map2 (fun x y -> (linreg_h2 x theta -. y) *. x.(0)) xs ys)|]

let linreg ?max_steps ~epsilon theta0 training =
  let xs = Array.map fst training in
  let ys = Array.map snd training in
    assert (Array.length theta0 = 2);	(* other theta sizes not supported yet *)
    assert (Array.length xs.(0) = 1); 	(* other sizes not supported yet *)
    optimize ?max_steps ~epsilon theta0
      (linreg_cost xs ys)
      (fun _nth_step theta ->
	 let step = linreg_cost2' xs ys theta in
	   BatArray.map2 (-.) theta step
      )<commit_msg>Optimize.linreg: generalized to any number of variables<commit_after>open BatPervasives
open Utils

(* this is not gradient descent by itself, it depends
   on the cost and step functions *)
let optimize ?max_steps ~epsilon (x0:'a) (cost:'a -> float) (step:int -> 'a -> 'a) =
  let rec loop cur_step prev_cost x =
    if (match max_steps with
	  | Some max_steps when cur_step >= max_steps -> false
	  | _ -> true)
    then
      let new_cost = cost x in
	Printf.printf "cost difference %f\n" (new_cost -. prev_cost);
	if abs_float (new_cost -. prev_cost) < epsilon 
	 then 
	   let _ = Printf.printf "found at step %d\n" cur_step in
	     x
	else loop (cur_step + 1) new_cost (step cur_step x)
    else
      x
  in
    loop 0 (cost x0) (step 0 x0)

let derivate ~epsilon f x = (f (x +. epsilon) -. f (x -. epsilon)) /. (2.0 *. epsilon)

let gradient2 ~epsilon f1 f2 (x1, x2) = 
  let e1, e2 = epsilon in
    (derivate e1 f1 x1, derivate e2 f2 x2)

(* m = number of training examples
   x = training input
   y = output

 *)

let linreg_h2 x theta = 
  let sum = ref theta.(0) in
    for c = 0 to Array.length x - 1 do
      sum := !sum +. theta.(c + 1) *. x.(c);
    done;
    !sum

let (!@) x = float (Array.length x)

let linreg_cost xs ys theta =
  assert (Array.length xs = Array.length ys);
  1.0 /. (2.0 *. !@ theta) *. sum (BatArray.map2 (fun x y -> (linreg_h2 x theta -. y) ** 2.0) xs ys)

let trace_exn label f a =
  try f a 
  with exn -> 
    Printf.printf "exception at %s\n" label;
    raise exn

let linreg_cost' xs ys theta =
  assert (Array.length xs = Array.length ys);
  let m = Array.length xs in
    Array.init (Array.length theta) **> fun j ->
      1.0 /. !@ theta *. sum (Array.init m **> fun i ->
				let x = xs.(i) in
				let y = ys.(i) in
				  (linreg_h2 x theta -. y) *. (if j = 0 then 1.0 else x.(j - 1))
			     )

let linreg ?max_steps ~epsilon alpha theta0 training =
  let xs = Array.map fst training in
  let ys = Array.map snd training in
    optimize ?max_steps ~epsilon theta0
      (linreg_cost xs ys)
      (fun _nth_step theta ->
	 let step = Array.map (( *. ) alpha) **> linreg_cost' xs ys theta in
	   BatArray.map2 (-.) theta step
      )
{'fraction_non_alphanumeric': 0.12863268222963317, 'fraction_numerical': 0.022868032396379228, 'mean_word_length': 2.5020850708924103, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 27, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1139337', 'n_tokens_mistral': 1699, 'n_tokens_neox': 1551, 'n_words': 568}