Machine Learning

Four classic ML algorithms built from Nx primitives: SVD, broadcasting, reductions, and scalar loops.

Running

dune exec nx/examples/10-machine-learning/pca.exe
dune exec nx/examples/10-machine-learning/kmeans.exe
dune exec nx/examples/10-machine-learning/dbscan.exe
dune exec nx/examples/10-machine-learning/tsne.exe

dbscan.ml

(** DBSCAN density-based clustering.

    Generate two dense clusters with scattered noise, find clusters using
    neighbourhood density, and report cluster sizes and noise count. *)

open Nx

let () =
  let eps = 1.5 in
  let min_samples = 5 in

  (* Two tight blobs plus uniform noise *)
  let c1 = add_s (mul_s (randn Float64 [| 80; 2 |]) 0.6) 3.0 in
  let c2 = sub_s (mul_s (randn Float64 [| 80; 2 |]) 0.6) 3.0 in
  let noise = sub_s (mul_s (rand Float64 [| 20; 2 |]) 14.0) 7.0 in
  let data = concatenate ~axis:0 [ c1; c2; noise ] in
  let n = (shape data).(0) in
  Printf.printf "Data: %d points (eps=%.1f, min_samples=%d)\n\n" n eps
    min_samples;

  (* Pairwise Euclidean distance matrix [n, n] *)
  let diff = sub (expand_dims [ 1 ] data) (expand_dims [ 0 ] data) in
  let dist = sqrt (sum ~axes:[ 2 ] (square diff)) in

  (* Neighbour adjacency and core-point mask *)
  let neighbours = less_equal_s dist eps in
  let counts = sum ~axes:[ 1 ] (cast Float64 neighbours) in
  let core = greater_equal_s counts (Float.of_int min_samples) in

  (* BFS cluster expansion *)
  let labels = Array.make n (-1) in
  let cluster_id = ref 0 in
  for i = 0 to n - 1 do
    if labels.(i) = -1 && item [ i ] core then begin
      let c = !cluster_id in
      incr cluster_id;
      labels.(i) <- c;
      let q = Queue.create () in
      Queue.push i q;
      while not (Queue.is_empty q) do
        let p = Queue.pop q in
        for j = 0 to n - 1 do
          if labels.(j) = -1 && item [ p; j ] neighbours then begin
            labels.(j) <- c;
            if item [ j ] core then Queue.push j q
          end
        done
      done
    end
  done;

  let n_clusters = !cluster_id in
  let n_noise =
    Array.fold_left (fun acc l -> if l = -1 then acc + 1 else acc) 0 labels
  in
  Printf.printf "Clusters found: %d\n" n_clusters;
  Printf.printf "Noise points:   %d\n\n" n_noise;
  for c = 0 to n_clusters - 1 do
    let count =
      Array.fold_left (fun acc l -> if l = c then acc + 1 else acc) 0 labels
    in
    Printf.printf "  Cluster %d: %d points\n" c count
  done

kmeans.ml

(** K-means clustering with kmeans++ initialisation.

    Generate synthetic blobs, cluster them with Lloyd's algorithm, and report
    centroid positions and inertia. *)

open Nx

(* Pairwise squared L2 distances: [n, d] x [k, d] -> [n, k] *)
let sq_distances a b =
  sum ~axes:[ 2 ] (square (sub (expand_dims [ 1 ] a) (expand_dims [ 0 ] b)))

(* Isotropic Gaussian blobs around given centres. *)
let make_blobs ~samples_per_cluster centers =
  let d = (shape centers).(1) in
  let blobs =
    List.init
      (shape centers).(0)
      (fun c ->
        add (randn Float64 [| samples_per_cluster; d |]) (get [ c ] centers))
  in
  shuffle (concatenate ~axis:0 blobs)

(* Kmeans++ initialisation: pick k centres from data. *)
let kmeanspp data k =
  let n = (shape data).(0) in
  let d = (shape data).(1) in
  let centroids = zeros Float64 [| k; d |] in
  let idx = Int32.to_int (item [] (randint Int32 ~high:n [||] 0)) in
  set [ 0 ] centroids (get [ idx ] data);
  for c = 1 to k - 1 do
    let current = slice [ R (0, c); A ] centroids in
    let min_d = min ~axes:[ 1 ] (sq_distances data current) in
    let chosen =
      Int32.to_int (item [] (categorical (log (clamp ~min:1e-30 min_d))))
    in
    set [ c ] centroids (get [ chosen ] data)
  done;
  centroids

let () =
  let true_centers =
    create Float64 [| 3; 2 |] [| 0.0; 0.0; 7.0; 7.0; -5.0; 10.0 |]
  in
  let data = make_blobs ~samples_per_cluster:100 true_centers in
  let n = (shape data).(0) in
  let d = (shape data).(1) in
  let k = 3 in
  Printf.printf "Data: %d points, %d features, %d clusters\n\n" n d k;

  let centroids = kmeanspp data k in
  let labels = ref (zeros Int32 [| n |]) in
  let max_iter = 100 in
  let tol = 1e-6 in
  let converged = ref false in
  let iter = ref 0 in
  while !iter < max_iter && not !converged do
    labels := argmin ~axis:1 (sq_distances data centroids);

    let old = copy centroids in
    for c = 0 to k - 1 do
      let mask = cast Float64 (equal !labels (scalar Int32 (Int32.of_int c))) in
      let count = item [] (sum mask) in
      if count > 0.0 then begin
        let total = sum ~axes:[ 0 ] (mul data (expand_dims [ 1 ] mask)) in
        set [ c ] centroids (div_s total count)
      end
    done;

    let shift = item [] (max (abs (sub centroids old))) in
    converged := shift < tol;
    incr iter
  done;

  Printf.printf "Converged after %d iterations\n\n" !iter;
  Printf.printf "Centroids:\n%s\n" (data_to_string centroids);

  for c = 0 to k - 1 do
    let count =
      item []
        (sum (cast Float64 (equal !labels (scalar Int32 (Int32.of_int c)))))
    in
    Printf.printf "  Cluster %d: %.0f points\n" c count
  done;

  let inertia = item [] (sum (min ~axes:[ 1 ] (sq_distances data centroids))) in
  Printf.printf "\nInertia: %.2f\n" inertia

pca.ml

(** Principal component analysis via SVD.

    Generate synthetic data with known structure, project to lower dimensions,
    and verify the explained variance captures the signal. *)

open Nx
open Nx.Infix

let () =
  (* 200 points in 5D: most variance along axes 0 and 1 *)
  let n = 200 in
  let scale = create Float64 [| 1; 5 |] [| 10.0; 5.0; 1.0; 1.0; 1.0 |] in
  let data = randn Float64 [| n; 5 |] * scale in
  Printf.printf "Data shape: [%d; %d]\n\n" (shape data).(0) (shape data).(1);

  (* Center *)
  let mu = mean ~axes:[ 0 ] ~keepdims:true data in
  let centered = data - mu in

  (* Economy SVD: centered = U diag(S) Vt *)
  let _u, s, vt = svd ~full_matrices:false centered in

  (* Explained variance ratio: s_i^2 / sum(s^2) *)
  let s2 = square s in
  let ratios = s2 /$ item [] (sum s2) in
  Printf.printf "Singular values:          %s\n" (data_to_string s);
  Printf.printf "Explained variance ratio: %s\n" (data_to_string ratios);
  Printf.printf "Cumulative:               %s\n\n"
    (data_to_string (cumsum ratios));

  (* Project to 2 components *)
  let n_components = 2 in
  let components = slice [ R (0, n_components); A ] vt in
  let projected = matmul centered (matrix_transpose components) in
  Printf.printf "Projected shape: [%d; %d]\n"
    (shape projected).(0)
    (shape projected).(1);

  (* Reconstruct and measure error *)
  let reconstructed = matmul projected components + mu in
  let rmse = sqrt (mean (square (data - reconstructed))) in
  Printf.printf "Reconstruction RMSE (2 of 5 components): %.4f\n" (item [] rmse)

tsne.ml

(** t-SNE dimensionality reduction.

    Embed three 10-dimensional clusters into 2D using the exact t-SNE algorithm.
    Reports KL divergence and per-cluster spread. *)

open Nx

(* Pairwise squared distances: [n, d] -> [n, n] *)
let pairwise_sq data =
  let diff = sub (expand_dims [ 1 ] data) (expand_dims [ 0 ] data) in
  sum ~axes:[ 2 ] (square diff)

(* Off-diagonal mask: 1 everywhere except the diagonal. *)
let off_diag n =
  sub (full Float64 [| n; n |] 1.0) (cast Float64 (eye Float64 n))

(* Compute symmetric P matrix via binary search for each row's bandwidth. *)
let compute_p dist_sq ~perplexity =
  let n = (shape dist_sq).(0) in
  let target = Float.log perplexity in
  let p = zeros Float64 [| n; n |] in
  for i = 0 to n - 1 do
    let di = get [ i ] dist_sq in
    let lo = ref 1e-10 in
    let hi = ref 1e4 in
    let row = ref (zeros Float64 [| n |]) in
    for _ = 0 to 50 do
      let sigma = (!lo +. !hi) /. 2.0 in
      let beta = 1.0 /. (2.0 *. sigma *. sigma) in
      let pi = exp (mul_s di (-.beta)) in
      set_item [ i ] 0.0 pi;
      let s = item [] (sum pi) in
      let pi = div_s pi (Float.max s 1e-30) in
      let h = -.item [] (sum (mul pi (log (clamp ~min:1e-30 pi)))) in
      row := pi;
      if h > target then hi := sigma else lo := sigma
    done;
    set [ i ] p !row
  done;
  (* Symmetrise: P = (P + P^T) / (2n) *)
  let p = div_s (add p (matrix_transpose p)) (2.0 *. Float.of_int n) in
  clamp ~min:1e-12 p

let () =
  let n_per = 50 in
  let dim = 10 in
  let perplexity = 20.0 in
  let max_iter = 500 in
  let lr = 100.0 in

  (* Three well-separated clusters in 10D *)
  let c0 = randn Float64 [| n_per; dim |] in
  let c1 = add_s (randn Float64 [| n_per; dim |]) 8.0 in
  let c2 = sub_s (randn Float64 [| n_per; dim |]) 8.0 in
  let data = concatenate ~axis:0 [ c0; c1; c2 ] in
  let n = (shape data).(0) in
  Printf.printf "Data: %d points in %dD, perplexity=%.0f\n\n" n dim perplexity;

  let dist_sq = pairwise_sq data in
  let p = compute_p dist_sq ~perplexity in

  let y = ref (mul_s (randn Float64 [| n; 2 |]) 1e-4) in
  let vel = ref (zeros Float64 [| n; 2 |]) in
  let mask = off_diag n in

  for iter = 1 to max_iter do
    let y_diff = sub (expand_dims [ 1 ] !y) (expand_dims [ 0 ] !y) in
    let y_dsq = sum ~axes:[ 2 ] (square y_diff) in
    let inv_d = mul (div (scalar Float64 1.0) (add_s y_dsq 1.0)) mask in
    let q_sum = Float.max (item [] (sum inv_d)) 1e-30 in
    let q = clamp ~min:1e-12 (div_s inv_d q_sum) in

    let p_eff = if iter <= 100 then mul_s p 4.0 else p in

    (* Gradient: 4 sum_j (p_ij - q_ij)(y_i - y_j)(1+||y_i-y_j||^2)^{-1} *)
    let mult = mul (sub p_eff q) inv_d in
    let grad =
      mul_s (sum ~axes:[ 1 ] (mul (expand_dims [ 2 ] mult) y_diff)) 4.0
    in

    let momentum = if iter <= 100 then 0.5 else 0.8 in
    vel := sub (mul_s !vel momentum) (mul_s grad lr);
    y := add !y !vel;

    if iter = 1 || iter mod 100 = 0 then begin
      let kl = item [] (sum (mul p (log (div p q)))) in
      Printf.printf "  iter %4d  KL = %.4f\n" iter kl
    end
  done;

  Printf.printf "\nPer-cluster spread (mean std of embedded coordinates):\n";
  for c = 0 to 2 do
    let lo = c * n_per in
    let cluster = slice [ R (lo, lo + n_per); A ] !y in
    let sx = item [] (mean (std ~axes:[ 0 ] cluster)) in
    Printf.printf "  Cluster %d: %.4f\n" c sx
  done