lib/eps/naive_bayes.rb in eps-0.2.0 vs lib/eps/naive_bayes.rb in eps-0.2.1

@@ -1,12 +1,12 @@
 module Eps
   class NaiveBayes < BaseEstimator
     attr_reader :probabilities
 
     def initialize(probabilities: nil, target: nil)
-      @probabilities = probabilities if probabilities
-      @target = target if target
+      @probabilities = probabilities
+      @target = target
     end
 
     def train(*args)
       super
 
@@ -20,23 +20,23 @@
         x = @x.dup
         x.each_with_index do |xi, i|
           xi[@target] = @y[i]
         end
         keys.each do |k|
-          conditional[k] = {}
+          conditional[k.to_s] = {}
           x.group_by { |xi| xi[@target] }.each do |group, xs|
             v = xs.map { |xi| xi[k] }
 
             if categorical?(v[0])
               # TODO apply smoothing
               # apply smoothing only to
               # 1. categorical features
               # 2. conditional probabilities
               # TODO more efficient count
-              conditional[k][group] = group_count(v)
+              conditional[k.to_s][group] = group_count(v)
             else
-              conditional[k][group] = {mean: mean(v), stdev: stdev(v)}
+              conditional[k.to_s][group] = {mean: mean(v), stdev: stdev(v)}
             end
           end
         end
       end
 
@@ -46,19 +46,17 @@
       }
     end
 
     # TODO better summary
     def summary(extended: false)
-      @summary_str ||= begin
-        str = String.new("")
-        probabilities[:prior].each do |k, v|
-          str += "#{k}: #{v}\n"
-        end
-        str += "\n"
-        str += "accuracy: %d%%\n" % [(100 * accuracy).round]
-        str
+      str = String.new("")
+      probabilities[:prior].each do |k, v|
+        str += "#{k}: #{v}\n"
       end
+      str += "\n"
+      str += "accuracy: %d%%\n" % [(100 * accuracy).round]
+      str
     end
 
     def accuracy
       self.class.metrics(predict(@x), @y)[:accuracy]
     end
@@ -184,11 +182,11 @@
 
     private
 
     def _predict(x)
       x.map do |xi|
-        probs = calculate_class_probabilities(xi)
+        probs = calculate_class_probabilities(stringify_keys(xi))
         # deterministic for equal probabilities
         probs.sort_by { |k, v| [-v, k.to_s] }[0][0]
       end
     end
 
@@ -233,8 +231,16 @@
 
     def stdev(arr)
       m = mean(arr)
       sum = arr.inject(0) { |accum, i| accum + (i - m)**2 }
       Math.sqrt(sum / (arr.length - 1).to_f)
+    end
+
+    def stringify_keys(h)
+      o = {}
+      h.each do |k, v|
+        o[k.to_s] = v
+      end
+      o
     end
   end
 end