probabilities branch merged

1 files changed, 18 insertions, 20 deletions

diff --git a/lib/classification.rb b/lib/classification.rb
index 0de8726..93b4f0f 100644
--- a/lib/classification.rb
+++ b/lib/classification.rb
@@ -5,29 +5,27 @@ module OpenTox
 
       def self.weighted_majority_vote compound, params
         neighbors = params[:neighbors]
-        weighted_sum = {}
-        sim_sum = 0.0
-        confidence = 0.0
-        # see ~/src/pubchem-read-across/application.rb:353
-        neighbors.each do |row|
-          sim = row["tanimoto"]
-          row["toxicities"][params[:prediction_feature_id].to_s].each do |act|
-            weighted_sum[act] ||= 0
-            weighted_sum[act] += sim
+        feature_id = params[:prediction_feature_id].to_s
+        sims = {}
+        neighbors.each do |n|
+          sim = n["tanimoto"]
+          n["toxicities"][feature_id].each do |act|
+            sims[act] ||= []
+            sims[act] << sim
           end
         end
-        case weighted_sum.size
-        when 1
-          return {:value => weighted_sum.keys.first, :confidence => weighted_sum.values.first/neighbors.size.abs}
-        when 2
-          sim_sum = weighted_sum[weighted_sum.keys[0]]
-          sim_sum -= weighted_sum[weighted_sum.keys[1]]
-          sim_sum > 0 ? prediction = weighted_sum.keys[0] : prediction = weighted_sum.keys[1] 
-          confidence = (sim_sum/neighbors.size).abs 
-          return {:value => prediction,:confidence => confidence}
-        else
-          bad_request_error "Cannot predict more than 2 classes, multinomial classifications is not yet implemented. Received classes were: '#{weighted.sum.keys}'"
+        sim_all = sims.collect{|a,s| s}.flatten
+        sim_sum = sim_all.sum
+        sim_max = sim_all.max
+        probabilities = {}
+        sims.each do |a,s|
+          probabilities[a] = s.sum/sim_sum
         end
+        probabilities = probabilities.collect{|a,p| [a,sim_max*p]}.to_h
+        p_max = probabilities.collect{|a,p| p}.max
+        prediction = probabilities.key(p_max)
+        {:value => prediction,:probabilities => probabilities}
+
       end
     end
   end