New structure implemented, test passes. Now I will have to add a new test with different teams to verify that it works.

src/lib.rs

@@ -72,34 +72,29 @@ fn draw_margin(p: f64, beta: f64, total_players: f64) -> f64 {
     math::icdf((p + 1.0) / 2.0) * total_players.sqrt() * beta
 }
 
-pub fn rate<R>(rating_groups: &[(R, u16)], ranks: &[u16]) -> Vec<Rating>
+pub fn rate<R>(ratings: &[(R, u16)], ranks: &[u16]) -> Vec<Rating>
 where
     R: Rateable,
 {
     // TODO Validate rating_groups is orderded in teams.
     // TODO Validate that teams are orderd after rank.
 
-    let flatten_ratings = rating_groups
-        .iter()
-        .map(|group| &group.0)
-        .collect::<Vec<_>>();
-
-    let size = flatten_ratings.len();
-
     let tau_sqr = TAU.powi(2);
     let beta_sqr = BETA.powi(2);
 
     let mut variable_arena = VariableArena::new();
 
-    // ---------------------------------------
-
-    let ratings = rating_groups;
-
     let rating_count = ratings.len();
     let team_count = ranks.len();
 
     let rating_vars = (0..rating_count).map(|_| variable_arena.create()).collect::<Vec<_>>();
-    let perf_vars = (0..rating_count).map(|_| variable_arena.create()).collect::<Vec<_>>();
+    let perf_vars = ratings
+        .iter()
+        .map(|(_, team)| {
+            (variable_arena.create(), *team)
+        })
+        .collect::<Vec<_>>();
+
     let team_perf_vars = (0..team_count).map(|_| variable_arena.create()).collect::<Vec<_>>();
     let team_diff_vars = (0..team_count - 1).map(|_| variable_arena.create()).collect::<Vec<_>>();
 
@@ -107,7 +102,7 @@ where
 
     let rating_layer = rating_vars
         .iter()
-        .zip(flatten_ratings.iter())
+        .zip(ratings.iter().map(|(rating, _)| rating))
         .map(|(rating_var, rating)| {
             let gaussian = Gaussian::from_mu_sigma(rating.mu(), (rating.sigma().powi(2) + tau_sqr).sqrt());
 
@@ -124,7 +119,7 @@ where
 
     let perf_layer = rating_vars
         .iter()
-        .zip(perf_vars.iter())
+        .zip(perf_vars.iter().map(|(variable, _)| variable))
         .map(|(rating_var, perf)| {
             factor_id += 1;
 
@@ -138,165 +133,111 @@ where
         })
         .collect::<Vec<_>>();
 
-    // let team_perf_layer = team_perf_vars.iter();
+    let team_perf_layer = team_perf_vars
-
+        .iter()
-    /*
+        .enumerate()
-        def build_team_perf_layer():
+        .map(|(i, variable)| {
-            for team, team_perf_var in enumerate(team_perf_vars):
-                if team > 0:
-                    start = team_sizes[team - 1]
-                else:
-                    start = 0
-                end = team_sizes[team]
-                child_perf_vars = perf_vars[start:end]
-                coeffs = flatten_weights[start:end]
-                yield SumFactor(team_perf_var, child_perf_vars, coeffs)
-    */
-
-    // ---------------------------------------
-
-    let mut ss = Vec::new();
-    let mut ps = Vec::new();
-    let mut ts = Vec::new();
-    let mut ds = Vec::new();
-
-    for _ in 0..size {
-        ss.push(variable_arena.create());
-        ps.push(variable_arena.create());
-        ts.push(variable_arena.create());
-    }
-
-    for _ in 0..size - 1 {
-        ds.push(variable_arena.create());
-    }
-
-    let mut factor_id = 0;
-
-    let mut skill = Vec::new();
-
-    for (i, rating) in flatten_ratings.iter().enumerate() {
-        let variable = ss[i];
-        let gaussian =
-            Gaussian::from_mu_sigma(rating.mu(), (rating.sigma().powi(2) + tau_sqr).sqrt());
-
-        skill.push(PriorFactor::new(
-            &mut variable_arena,
-            factor_id,
-            variable,
-            gaussian,
-        ));
-
             factor_id += 1;
-    }
 
-    let mut skill_to_perf = Vec::new();;
+            let team = perf_vars
+                .iter()
+                .filter(|(_, team)| *team as usize == i)
+                .map(|(variable, _)| *variable)
+                .collect::<Vec<_>>();
 
-    for i in 0..size {
+            let team_count = team.len();
-        skill_to_perf.push(LikelihoodFactor::new(
+
+            SumFactor::new(
                 &mut variable_arena,
                 factor_id,
-            ss[i],
+                *variable,
-            ps[i],
+                team,
-            beta_sqr,
+                vec![1.0; team_count],
-        ));
+            )
+        })
+        .collect::<Vec<_>>();
 
+    let team_diff_layer = team_diff_vars
+        .iter()
+        .zip(team_perf_vars.windows(2))
+        .map(|(variable, teams)| {
             factor_id += 1;
-    }
 
-    let mut perf_to_team = Vec::new();;
+            SumFactor::new(
-
-    for i in 0..size {
-        perf_to_team.push(SumFactor::new(
                 &mut variable_arena,
                 factor_id,
-            ts[i],
+                *variable,
-            vec![ps[i]],
+                teams.to_vec(),
-            vec![1.0],
-        ));
-
-        factor_id += 1;
-    }
-
-    let mut team_diff = Vec::new();;
-
-    for i in 0..size - 1 {
-        team_diff.push(SumFactor::new(
-            &mut variable_arena,
-            factor_id,
-            ds[i],
-            vec![ts[i], ts[i + 1]],
                 vec![1.0, -1.0],
-        ));
+            )
 
-        factor_id += 1;
+        })
-    }
+        .collect::<Vec<_>>();
-
-    let mut trunc = Vec::new();
 
     let epsilon = draw_margin(DRAW_PROBABILITY, BETA, 2.0);
 
-    for i in 0..size - 1 {
+    let trunc_layer = team_diff_vars
-        trunc.push(TruncateFactor::new(
+        .iter()
+        .enumerate()
+        .map(|(i, variable)| {
+            factor_id += 1;
+
+            TruncateFactor::new(
                 &mut variable_arena,
                 factor_id,
-            ds[i],
+                *variable,
                 epsilon,
                 ranks[i] == ranks[i + 1],
-        ));
+            )
 
-        factor_id += 1;
+        })
-    }
+        .collect::<Vec<_>>();;
 
-    for factor in &skill {
+    for factor in &rating_layer {
         factor.start(&mut variable_arena);
     }
 
-    for factor in &skill_to_perf {
+    for factor in &perf_layer {
         factor.update_value(&mut variable_arena);
     }
 
-    for factor in &perf_to_team {
+    for factor in &team_perf_layer {
         factor.update_sum(&mut variable_arena);
     }
 
     for _ in 0..5 {
-        for factor in &team_diff {
+        for factor in &team_diff_layer {
             factor.update_sum(&mut variable_arena);
         }
 
-        for factor in &trunc {
+        for factor in &trunc_layer {
             factor.update(&mut variable_arena);
         }
 
-        for factor in &team_diff {
+        for factor in &team_diff_layer {
             factor.update_term(&mut variable_arena, 0);
             factor.update_term(&mut variable_arena, 1);
         }
     }
 
-    for factor in &perf_to_team {
+    for factor in &team_perf_layer {
         factor.update_term(&mut variable_arena, 0);
     }
 
-    for factor in &skill_to_perf {
+    for factor in &perf_layer {
         factor.update_mean(&mut variable_arena);
     }
 
-    let mut ratings = Vec::new();
+    rating_vars
-
+        .iter()
-    for i in 0..size {
+        .map(|variable| variable_arena.get(*variable).unwrap().get_value())
-        let value = variable_arena
+        .map(|value| {
-            .get(ss[i])
+            Rating {
-            .map(|variable| variable.get_value())
-            .unwrap();
-
-        ratings.push(Rating {
                 mu: value.mu(),
                 sigma: value.sigma(),
-        });
             }
-
+        })
-    ratings
+        .collect::<Vec<_>>()
 }
 
 pub fn quality<R>(rating_groups: &[&[R]]) -> f64