datatype_test.clj

(ns tech.v2.datatype-test
  (:require [clojure.test :refer :all]
            [tech.v2.datatype :as dtype]
            [tech.v2.datatype.base :as base]
            [tech.v2.datatype.casting :as casting]
            [tech.v2.datatype.primitive]
            [tech.v2.datatype.list]
            [tech.parallel.for :as parallel-for]
            [tech.v2.datatype.functional :as dfn]
            [tech.v2.datatype.boolean-op :as boolean-op])
  (:import [java.nio FloatBuffer]))


(deftest raw-copy-with-mutable-lazy-sequence
  ;;It is important that raw copy can work with a lazy sequence of double buffers where
  ;;the same buffer is being filled for ever member of the lazy sequence.  This is an
  ;;easy optimization to make that cuts down the memory usage when reading from datasets
  ;;by a fairly large amount.
  (let [input-seq (partition 10 (range 100))
        input-ary (double-array 10)
        ;;Note the input ary is being reused.  Now imagine the input ary is a large
        ;;and we want a batch size of 50.  This is precisely the use case
        ;;where we don't want to allocate dynamically an entire batch worth of data
        ;;every time but copy one image at a time into the packed buffer for upload
        ;;to the gpu.
        double-array-seq (map (fn [data]
                                (dtype/copy-raw->item! data input-ary 0)
                                input-ary)
                              input-seq)
        output-doubles (double-array 100)]
    (dtype/copy-raw->item! double-array-seq output-doubles 0)
    (is (= (vec output-doubles) (mapv double (flatten input-seq))))))


(defn basic-copy
  [src-fn dest-fn src-dtype dst-dtype]
  (try
    (let [ary (dtype/make-container src-fn src-dtype (range 10))
          buf (dtype/make-container dest-fn dst-dtype 10)
          retval (dtype/make-jvm-container :float64 10)]
      ;;copy starting at position 2 of ary into position 4 of buf 4 elements
      (dtype/copy! ary 2 buf 4 4)
      (dtype/copy! buf 0 retval 0 10)
      (is (= [0 0 0 0 2 3 4 5 0 0]
             (mapv int (dtype/->vector retval)))
          (pr-str {:src-fn src-fn :dest-fn dest-fn
                   :src-dtype src-dtype :dst-dtype dst-dtype}))
      (is (= 10 (dtype/ecount buf))))
    (catch Throwable e
      (println {:src-fn src-fn :dest-fn dest-fn
                :src-dtype src-dtype :dst-dtype dst-dtype})
      (throw e))))


(def create-functions [:typed-buffer :native-buffer :list :sparse])


(deftest generalized-copy-test
  (->> (for [src-container create-functions
             dst-container create-functions
             src-dtype casting/numeric-types
             dst-dtype casting/numeric-types]
         (basic-copy src-container dst-container src-dtype dst-dtype))
       dorun))


(deftest array-of-array-support
  (let [^"[[D" src-data (make-array (Class/forName "[D") 5)
        _ (doseq [idx (range 5)]
            (aset src-data idx (double-array (repeat 10 idx))))
        dst-data (float-array (* 5 10))]
    ;;This should not hit any slow paths.
    (dtype/copy-raw->item! src-data dst-data 0)
    (is (= (vec (float-array (flatten (map #(repeat 10 %) (range 5)))))
           (vec dst-data)))))


(deftest out-of-range-data-causes-exception
  (is (thrown? Throwable (dtype/copy! (int-array [1000 2000 3000 4000])
                                      (byte-array 4)))))


(deftest out-of-range-data-does-not-cause-exception-if-unchecked
  (let [byte-data
        (dtype/copy! (int-array [1000 2000 3000 4000]) 0
                     (byte-array 4) 0
                     4 {:unchecked? true})]
    (is (= [-24 -48 -72 -96]
           (vec byte-data)))))


(deftest offset-buffers-should-copy-correctly
  (let [fbuf (-> (dtype/make-buffer-of-type :float32 (range 10))
                 (dtype/sub-buffer 3))
        result-buf (dtype/make-jvm-container :float32 (dtype/ecount fbuf))]
    (dtype/copy! fbuf result-buf)
    (is (= (dtype/get-value fbuf 0)
           (float 3)))
    (is (= (vec (drop 3 (range 10)))
           (mapv long (dtype/->vector result-buf))))))


(deftest primitive-types-are-typed
  (doseq [[cls dtype] [[(byte 1) :int8]
                       [(short 1) :int16]
                       [(int 1) :int32]
                       [(long 1) :int64]
                       [(float 1) :float32]
                       [(double 1) :float64]
                       [(boolean false) :boolean]]]
    (is (= dtype (dtype/get-datatype cls)))))


(set! *warn-on-reflection* true)
(set! *unchecked-math* :warn-on-boxed)

(deftest copy-time-test
  (testing "Run perf regression of times spent to copy data"
    (let [num-items (long 100000)
          src-data (float-array (range num-items))
          dst-data (float-array num-items)
          array-copy (fn []
                       (parallel-for/parallel-for
                        idx num-items
                        (aset dst-data idx (aget src-data idx))))
          src-buf (FloatBuffer/wrap src-data)
          dst-buf (FloatBuffer/wrap dst-data)
          buffer-copy (fn []
                        (let [
                              ;; Curiously, uncommenting this gets a far faster result.
                              ;; But it isn't at all practical.
                              ;; src-buf (FloatBuffer/wrap src-data)
                              ;; dst-buf (FloatBuffer/wrap dst-data)
                              ]
                          (parallel-for/parallel-for
                           idx num-items
                           (.put dst-buf idx (.get src-buf idx)))))

          dtype-copy (fn []
                       (dtype/copy! src-buf 0 dst-buf 0 num-items))

          unchecked-dtype-copy (fn []
                                 (dtype/copy! src-buf 0 dst-buf 0 num-items
                                              {:unchecked? true}))

          make-array (fn []
                       (dtype/make-container :java-array :float32 dst-buf))
          marshal-buf (int-array num-items)
          ;;If you have to do a marshalling copy then exploiting parallelism will be
          ;;your best bet.  It costs a lot to marshal across datatypes, esp. int->float.
          marshal-copy (fn []
                         (dtype/copy! src-data 0
                                      marshal-buf 0
                                      num-items
                                      {:unchecked? true}))
          fns {:array-copy array-copy
               :buffer-copy buffer-copy
               :dtype-copy dtype-copy
               :unchecked-dtype-copy unchecked-dtype-copy
               :make-array make-array
               :marshal-copy marshal-copy}
          run-timed-fns (fn []
                          (->> fns
                               (map (fn [[fn-name time-fn]]
                                      [fn-name (with-out-str
                                                 (time
                                                  (dotimes [iter 400]
                                                    (time-fn))))]))
                               (into {})))
          warmup (run-timed-fns)
          times (run-timed-fns)]
      (println times))))


(deftest nil-has-nil-shape
  (is (= nil (dtype/shape nil)))
  (is (= 0 (dtype/ecount nil)))
  (is (= 0 (dtype/shape->ecount (dtype/shape nil)))))


(deftest boolean-support
  (is (= [0 1 0]
         (-> (dtype/make-container :java-array  :int8 [false true false])
             (dtype/->vector))))


  (is (= [0 1 0]
         (-> (dtype/make-container :java-array :int8
                                   (boolean-array [false true false]))
             (dtype/->vector))))


  (is (= [0 1 0]
         (-> (dtype/copy! (boolean-array [false true false]) 0
                          (dtype/make-jvm-container :int8 3) 0
                          3
                          {:unchecked? true})
             (dtype/->vector))))


  (is (= [false true false]
         (-> (dtype/make-jvm-container :boolean [0 1 0])
             (dtype/->vector))))


  (is (= [false true false]
         (-> (dtype/make-jvm-container :boolean [0.0 0.01 0.0])
             (dtype/->vector))))

  (is (= true (dtype/cast 10 :boolean)))

  (is (= false (dtype/cast 0 :boolean)))

  (is (= 1 (dtype/cast true :int16)))

  (is (= 0.0 (dtype/cast false :float64))))


(deftest list-regression
  (is (= (Class/forName "[D")
         (type (dtype/->array-copy
                (dtype/make-container :list :float64 10)))))
  (is (= (Class/forName "[D")
         (type (dtype/->array-copy
                (-> (dtype/make-container :list :float64 10)
                    (dtype/sub-buffer 2 2)))))))


(deftest nested-array-things-have-appropriate-shape
  (is (= [5 5]
         (->> (repeatedly 5 #(double-array 5))
              (into-array)
              dtype/shape)))
  (is (= 25
         (->> (repeatedly 5 #(double-array 5))
              (into-array)
              dtype/ecount)))
  (is (= [5 5]
         (->> (repeatedly 5 #(range 5))
              (into-array)
              dtype/shape)))
  (is (= 25
         (->> (repeatedly 5 #(range 5))
              (into-array)
              dtype/ecount))))


(deftest base-math-sanity
  (is (= 0.0 (-> (dfn/- (range 10) (range 10))
                 (dtype/->reader :float64)
                 (dfn/pow 2)
                 (dfn/reduce-+))))

  (is (= 0.0 (-> (dfn/- (into-array (range 10)) (range 10))
                 (dtype/->reader :float64)
                 (dfn/pow 2)
                 (dfn/reduce-+)))))


(deftest nan-insanity
  (let [test-floats (float-array [0 ##NaN 2 4 ##NaN])
        nan-floats (float-array (repeat 5 ##NaN))]
    (is (= [1 4]
           (-> (dfn/argfilter #(= 0 (Float/compare % Float/NaN)) test-floats)
               vec)))
    (is (= [1 4]
           (->> test-floats
                (dfn/argfilter (boolean-op/make-boolean-unary-op
                                :float32
                                (= 0 (Float/compare x Float/NaN))))
               vec)))
    (is (= [false true false false true]
           (dfn/nan? test-floats)))

    (is (= [false true false false true]
           (dfn/eq test-floats Float/NaN)))

    (is (= [1 4]
           (->> (dfn/eq test-floats nan-floats)
                (dfn/argfilter identity)
                vec)))))


(deftest round-and-friends
  (is (= [2.0 3.0 4.0 5.0]
         (vec (dfn/round [2.2 2.8 3.5 4.6]))))
  (is (= [2.0 2.0 3.0 4.0]
         (vec (dfn/floor [2.2 2.8 3.5 4.6]))))
  (is (= [3.0 3.0 4.0 5.0]
         (vec (dfn/ceil [2.2 2.8 3.5 4.6]))))
  (is (= [2.2 2.8 3.5 4.6]
         (vec (dfn/abs [2.2 2.8 3.5 4.6]))))
  (is (= [1 4 9 16]
         (vec (dfn/sq (byte-array [1 2 3 4])))))
  (is (dfn/equals [0.381 0.540 0.675 0.785]
                  (dfn/atan2 [2 3 4 5] 5)
                  0.001))
  (is (thrown? Throwable (dfn/rem (double-array [1 2 3 4] 2))))
  (is (= [1 0 1 0]
         (dfn/rem [1 2 3 4] 2))))


(deftest generic-lists
  (is (= (vec (range 200 255))
         (-> (dtype/make-container :list :uint8 (range 200 255))
             dtype/->vector))))


(deftest desc-stats-unsigned-types
  (is (= {:min 0.0
          :max 255.0}
         (-> (dtype/make-container :typed-buffer :uint8 (range 256))
             (dfn/descriptive-stats)
             (select-keys [:min :max])))))


(deftest reader-sorting
  (doseq [datatype (->> casting/base-host-datatypes
                        (remove #(= :boolean %)))]
    (let [breader (dtype/->reader (dtype/make-container :java-array datatype
                                                        (reverse (range 10))))]
      (is (= (vec (range 10))
             (vec (map int (sort breader))))))))


(deftest object->string-reader-cast
  (is (= :string
         (-> (dtype/->reader ["a" "b"] :string)
             (dtype/get-datatype)))))


(deftest object->uint16
  (is (= :uint16
         (-> (dtype/->reader [1 2] :uint16)
             (dtype/get-datatype)))))


(defn copy-raw->item-time-test-1
  []
  (let [n-elems 10000
        dst-ary (float-array n-elems)
        src-data (map float (range n-elems))]
    (time (dotimes [iter 1000]
            (dtype/copy-raw->item! src-data dst-ary)))))


(defn copy-raw->item-time-test-2
  []
  (let [n-elems 10000
        n-seq-items 100
        dst-ary (float-array (* n-elems n-seq-items))
        src-data (map float-array (repeat n-seq-items (range n-elems)))]
    (time (dotimes [iter 100]
            (dtype/copy-raw->item! src-data dst-ary)))))


(deftest arggroup-by-test
  ;;The operation returns values in object space
  ;;if it operates as :object datatype.  This is the default.
  (let [{truevals true
         falsevals false}
        (dfn/arggroup-by even? (range 20))]
    (is (= truevals
           (vec (filter even? (range 20)))))
    (is (= falsevals
           (vec (remove even? (range 20))))))


  ;;Since the arguments is a long reader, the operation
  ;;operates and returns values in long space.
  (let [{truevals 1
         falsevals 0}
        (dfn/arggroup-by even? (range 20) {:datatype :int64})]
    (is (= truevals
           (vec (filter even? (range 20)))))
    (is (= falsevals
           (vec (remove even? (range 20)))))))


(deftest argpartition-by-test
  ;;The operation returns values in object space
  ;;if it operates as :object datatype.  This is the default.
  (is (= [[0 (range 5)]
          [1 (range 5 10)]
          [2 (range 10 15)]
          [3 (range 15 20)]]
         (vec (dfn/argpartition-by #(quot (long %) 5) (range 20))))))



(deftest typed-buffer-destructure
  (let [[a b c] (dtype/make-container :typed-buffer :int64 [1 2 3])]
    (is (= [a b c]
           [1 2 3]))))


(deftest simple-mse
  (let [predictions (list 1 2 3 4 5)
        labels (list 3 4 5 6 7)]
    (is (= 20.0
           (-> (dfn/- predictions labels)
               (dfn/pow 2)
               (dfn/reduce-+))))))