Skip to content

Instantly share code, notes, and snippets.

@EntilZha
Created October 6, 2015 06:35
Show Gist options
  • Star 1 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save EntilZha/faaa11e4b05afcba68e0 to your computer and use it in GitHub Desktop.
Save EntilZha/faaa11e4b05afcba68e0 to your computer and use it in GitHub Desktop.
Rust code trying to implement the EM algorithm for a simple example
fn main() {
let observations = [5, 9, 8, 4, 7];
let (theta_a, theta_b) = em(0.5, 0.5, observations);
println!("ThetaA: {}, ThetaB: {}", theta_a, theta_b);
}
fn em(theta_a_0: f64, theta_b_0: f64, observations: [i32; 5]) -> (f64, f64) {
let mut theta_a = theta_a_0;
let mut theta_b = theta_b_0;
let a_probs = observations.iter().map(|x| binomial(theta_a, *x, 10));
let b_probs = observations.iter().map(|x| binomial(theta_b, *x, 10));
let norms: Vec<f64> = a_probs.zip(b_probs).map(|x| x.0 + x.1).collect();
let a_norm_probs = a_probs.zip(&norms).map(|x| x.0 / x.1);
let b_norm_probs = b_probs.zip(&norms).map(|x| x.0 / x.1);
(0.0, 0.0)
}
fn binomial(p: f64, heads: i32, n: i32) -> f64 {
p.powi(heads) * (1.0 - p).powi(n - heads)
}
@EntilZha
Copy link
Author

EntilZha commented Oct 6, 2015

Compile error:

$ cargo build
   Compiling EMAlgorithm v0.1.0 (file:///Users/pedro/Documents/Code/EMAlgorithm)
src/main.rs:13:24: 13:31 error: use of moved value: `a_probs` [E0382]
src/main.rs:13     let a_norm_probs = a_probs.zip(&norms).map(|x| x.0 / x.1);
                                      ^~~~~~~
src/main.rs:12:27: 12:34 note: `a_probs` moved here because it has type `core::iter::Map<core::slice::Iter<'_, i32>, [closure@src/main.rs:10:43: 10:72 theta_a:&f64]>`, which is non-copyable
src/main.rs:12     let norms: Vec<f64> = a_probs.zip(b_probs).map(|x| x.0 + x.1).collect();
                                         ^~~~~~~
src/main.rs:14:24: 14:31 error: use of moved value: `b_probs` [E0382]
src/main.rs:14     let b_norm_probs = b_probs.zip(&norms).map(|x| x.0 / x.1);
                                      ^~~~~~~
src/main.rs:12:39: 12:46 note: `b_probs` moved here because it has type `core::iter::Map<core::slice::Iter<'_, i32>, [closure@src/main.rs:11:43: 11:72 theta_b:&f64]>`, which is non-copyable
src/main.rs:12     let norms: Vec<f64> = a_probs.zip(b_probs).map(|x| x.0 + x.1).collect();
                                                     ^~~~~~~
error: aborting due to 2 previous errors
Could not compile `EMAlgorithm`.

Trying to make code look similar to my scala implementation:

bject EMAlgorithm {

  val flips = 10

  def binomial(p: Double, heads: Int, n: Int): Double = {
    math.pow(p, heads) * math.pow(1 - p, n - heads)
  }

  def expected_counts(p: Double, obs: Int): Experiment = {
    val heads = obs
    val tails = flips - obs
    Experiment(p * heads, tails * p)
  }

  case class Experiment(heads: Double, tails: Double)

  def em(theta_a_0: Double, theta_b_0: Double, iter_max: Int, observations: List[Int]): (Double, Double) = {
    var theta_a = theta_a_0
    var theta_b = theta_b_0
    for (i <- 0 until iter_max) {
      println(s"i= $i ThetaA: $theta_a ThetaB: $theta_b")

      // E Step

      // Calculate the probability of observation using Binomial Distribution
      val a_probs = observations.map(h => binomial(theta_a, h, 10))
      val b_probs = observations.map(h => binomial(theta_b, h, 10))
      val norms = a_probs.zip(b_probs).map{case (a, b) => a + b}
      val a_norm = a_probs.zip(norms).map{case (p, total) => p / total}
      val b_norm = b_probs.zip(norms).map{case (p, total) => p / total}

      // Compute sufficient statistics for ThetaA and ThetaB
      val a_expected_counts: Experiment = a_norm
        .zip(observations)
        .map{ case (p, obs) => expected_counts(p, obs)}
        .reduce((e0, e1) => Experiment(e0.heads + e1.heads, e0.tails + e1.tails))
      val b_expected_counts = b_norm
        .zip(observations)
        .map(t => expected_counts(t._1, t._2))
        .reduce((e0, e1) => Experiment(e0.heads + e1.heads, e0.tails + e1.tails))

      // M Step

      // Using expected counts, set parameters again
      theta_a = a_expected_counts.heads / (a_expected_counts.heads + a_expected_counts.tails)
      theta_b = b_expected_counts.heads / (b_expected_counts.heads + b_expected_counts.tails)
    }
    (theta_a, theta_b)
  }

  def main(args: Array[String]) {
    println("Running EM Algorithm")
    val observations = List(5, 9, 8, 4, 7)
    val theta = em(0.6, 0.5, 20, observations)
    println(s"Theta: $theta")
  }
}

@EntilZha
Copy link
Author

EntilZha commented Oct 6, 2015

Why would this have a compile error on inference? If you put types Vec<f64> for those it works, but don't know why I should have to:

$ cargo build
   Compiling EMAlgorithm v0.1.0 (file:///Users/pedro/Documents/Code/EMAlgorithm)
src/main.rs:12:27: 12:41 error: the type of this value must be known in this context
src/main.rs:12     let norms: Vec<f64> = a_probs.iter().zip(b_probs).map(|x| x.0 + x.1).collect();
                                         ^~~~~~~~~~~~~~
src/main.rs:12:63: 12:66 error: the type of this value must be known in this context
src/main.rs:12     let norms: Vec<f64> = a_probs.iter().zip(b_probs).map(|x| x.0 + x.1).collect();
                                                                             ^~~
note: in expansion of closure expansion
src/main.rs:12:59: 12:72 note: expansion site
src/main.rs:13:59: 13:62 error: the type of this value must be known in this context
src/main.rs:13     let a_norm_probs = a_probs.iter().zip(&norms).map(|x| x.0 / x.1);
                                                                         ^~~
note: in expansion of closure expansion
src/main.rs:13:55: 13:68 note: expansion site
src/main.rs:14:59: 14:62 error: the type of this value must be known in this context
src/main.rs:14     let b_norm_probs = b_probs.iter().zip(&norms).map(|x| x.0 / x.1);
                                                                         ^~~
note: in expansion of closure expansion
src/main.rs:14:55: 14:68 note: expansion site
error: aborting due to 4 previous errors
Could not compile `EMAlgorithm`.
fn em(theta_a_0: f64, theta_b_0: f64, observations: [i32; 5]) -> (f64, f64) {
    let mut theta_a = theta_a_0;
    let mut theta_b = theta_b_0;
    let a_probs = observations.iter().map(|x| binomial(theta_a, *x, 10)).collect();
    let b_probs = observations.iter().map(|x| binomial(theta_b, *x, 10)).collect();
    let norms: Vec<f64> = a_probs.iter().zip(b_probs.iter()).map(|x| x.0 + x.1).collect();
    let a_norm_probs = a_probs.iter().zip(&norms).map(|x| x.0 / x.1);
    let b_norm_probs = b_probs.iter().zip(&norms).map(|x| x.0 / x.1);
    (0.0, 0.0)
}

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment