Skip to content
Browse files

allow for dataset-set-class to take the name of attribute

  • Loading branch information...
1 parent 93846bc commit 78fc93c2a49b442014996476833a13cd0932cc09 @bmabey bmabey committed Dec 1, 2010
Showing with 13 additions and 8 deletions.
  1. +8 −6 src/clj_ml/data.clj
  2. +5 −2 test/clj_ml/data_test.clj
View
14 src/clj_ml/data.clj
@@ -40,8 +40,10 @@
"Returns the index of an attribute in the attributes definition of an
instance or dataset"
[dataset attr-name]
- (let [attr-name (name attr-name)]
- (find-first #(= attr-name (.name (.attribute dataset %))) (range (.numAttributes dataset)))))
+ (if (number? attr-name)
+ attr-name
+ (let [attr-name (name attr-name)]
+ (find-first #(= attr-name (.name (.attribute dataset %))) (range (.numAttributes dataset))))))
(defn attributes
"Returns the attributes (weka.core.Attribute) of the dataset or instance"
@@ -228,7 +230,7 @@
(defn dataset-get-class
"Returns the index of the class attribute for this dataset"
- [dataset]
+ [^Instances dataset]
(.classIndex dataset))
;; manipulation of instances
@@ -298,12 +300,12 @@ becuase it avoids redundant string interning of the attribute names."
(defn dataset-set-class
"Sets the index of the attribute of the dataset that is the class of the dataset"
- [dataset pos]
- (doto dataset (.setClassIndex pos)))
+ [^Instances dataset index-or-name]
+ (doto dataset (.setClassIndex ^int (index-attr dataset index-or-name))))
(defn dataset-remove-class
"Removes the class attribute from the dataset"
- [dataset]
+ [^Instances dataset]
(doto dataset (.setClassIndex -1)))
(defn dataset-count
View
7 test/clj_ml/data_test.clj
@@ -38,9 +38,12 @@
(let [dataset (make-dataset :test
[:a :b]
2)
- foo(clj-ml.data/dataset-set-class dataset 1)]
+ _ (clj-ml.data/dataset-set-class dataset 1)]
(is (= 1 (.classIndex dataset)))
- (is (= 0 (.classIndex (dataset-set-class dataset 0))))))
+ (is (= 0 (.classIndex (dataset-set-class dataset 0))))
+ (testing "when a string or symbol is passed in"
+ (is (= 1 (.classIndex (dataset-set-class dataset "b"))))
+ (is (= 0 (.classIndex (dataset-set-class dataset :a)))))))
(deftest dataset-count-1
(let [dataset (make-dataset :test

0 comments on commit 78fc93c

Please sign in to comment.
Something went wrong with that request. Please try again.