lib/groupie.rb in groupie-0.1.1 vs lib/groupie.rb in groupie-0.2.2

@@ -10,42 +10,63 @@
 
   def [](group)
     @groups[group] ||= Group.new(group)
   end
 
-  def classify(entry)
+  def classify(entry, strategy=:sum)
     results = {}
     total_count = @groups.inject(0) do |sum, name_group|
       group = name_group.last
-      sum + group.count(entry)
+      count = group.count(entry)
+      if strategy==:sum
+        sum += count
+      elsif strategy==:sqrt
+        sum += Math::sqrt(count)
+      elsif strategy==:log
+        sum += Math::log10(count) if count > 0
+      else
+        raise "Invalid strategy: #{strategy}"
+      end
+      next sum
     end
     return results if 0 == total_count
 
     @groups.each do |name, group|
       count = group.count(entry)
+      if strategy==:sum
+        # keep count
+      elsif strategy==:sqrt
+        count = Math::sqrt(count)
+      elsif strategy==:log
+        count = Math::log10(count) if count > 0
+      else
+        raise "Invalid strategy: #{strategy}"
+      end
       results[name] = count > 0 ? count.to_f / total_count : 0.0
     end
     return results
   end
 
   # Classify a text by taking the average of all word classifications.
-  def classify_text(words)
+  def classify_text(words, strategy=:sum)
+    hits = 0
     group_score_sums = words.inject({}) do |results, word|
-      word_results = classify(word)
+      word_results = classify(word, strategy)
+      next results if word_results.empty?
+      hits += 1
       results.merge(word_results) do |key, old, new|
         old + new
       end
     end
 
-    words_count = words.size.to_f
     averages={}
     group_score_sums.each do |group, sum|
-      averages[group] = sum / words_count
+      averages[group] = hits > 0 ? sum / hits : 0
     end
 
     averages
   end
-  
+
   def self.version
     File.read(File.join(File.dirname(File.expand_path(__FILE__)), "..", "VERSION")).strip
   end
 end
\ No newline at end of file