lib/svmkit/model_selection/cross_validation.rb in svmkit-0.2.3 vs lib/svmkit/model_selection/cross_validation.rb in svmkit-0.2.4

@@ -19,22 +19,28 @@
 
       # Return the splitter that divides dataset.
       # @return [Splitter]
       attr_reader :splitter
 
+      # Return the evaluator that calculates score.
+      # @return [Evaluator]
+      attr_reader :evaluator
+
       # Return the flag indicating whether to caculate the score of training dataset.
       # @return [Boolean]
       attr_reader :return_train_score
 
       # Create a new evaluator with cross-validation method.
       #
       # @param estimator [Classifier] The classifier of which performance is evaluated.
       # @param splitter [Splitter] The splitter that divides dataset to training and testing dataset.
+      # @param evaluator [Evaluator] The evaluator that calculates score of estimator results.
       # @param return_train_score [Boolean] The flag indicating whether to calculate the score of training dataset.
-      def initialize(estimator: nil, splitter: nil, return_train_score: false)
+      def initialize(estimator: nil, splitter: nil, evaluator: nil, return_train_score: false)
         @estimator = estimator
         @splitter = splitter
+        @evaluator = evaluator
         @return_train_score = return_train_score
       end
 
       # Perform the evalution of given classifier with cross-validation method.
       #
@@ -47,11 +53,11 @@
       #   * :test_score (Array<Float>) The scores of testing dataset for each split.
       #   * :train_score (Array<Float>) The scores of training dataset for each split. This option is nil if
       #     the return_train_score is false.
       def perform(x, y)
         # Initialize the report of cross validation.
-        report = {test_score: [], train_score: nil, fit_time: []}
+        report = { test_score: [], train_score: nil, fit_time: [] }
         report[:train_score] = [] if @return_train_score
         # Evaluate the estimator on each split.
         @splitter.split(x, y).each do |train_ids, test_ids|
           # Split dataset into training and testing dataset.
           feature_ids = !kernel_machine? || train_ids
@@ -62,11 +68,16 @@
           # Fit the estimator.
           start_time = Time.now.to_i
           @estimator.fit(train_x, train_y)
           # Calculate scores and prepare the report.
           report[:fit_time].push(Time.now.to_i - start_time)
-          report[:test_score].push(@estimator.score(test_x, test_y))
-          report[:train_score].push(@estimator.score(train_x, train_y)) if @return_train_score
+          if @evaluator.nil?
+            report[:test_score].push(@estimator.score(test_x, test_y))
+            report[:train_score].push(@estimator.score(train_x, train_y)) if @return_train_score
+          else
+            report[:test_score].push(@evaluator.score(test_y, @estimator.predict(test_x)))
+            report[:train_score].push(@estimator.score(train_x, @estimator.predict(train_x))) if @return_train_score
+          end
         end
         report
       end
 
       private