Numpy-like index/slice accessors for Clojure

core.clj

(ns numpy-wannabe.core
  (:gen-class)
  (:require [clojure.edn :as edn]
            [clojure.string :as str]))

;;;;;;;; Helper Functions
(defn gsub [matcher replacer value]
  (str/replace value matcher replacer))

(defn nths [row, columns]
  (keep-indexed #(if (.contains columns %1) %2) row))

(defn multi-dimensional? [coll]
  (some coll? coll))

;;;;;;; Input normalization
(defn normalize-slice [index coll]
  (->>  index
    ; ensure we have 3 components
    (gsub #"^([^:]*):([^:]*)$" "$1:$2:")
    ; default start=0
    (gsub #"^:.*" #(str "0" %))
    ; default end=collection.count
    (gsub #"^(.*)::(.*)$" (str "$1:" (count coll) ":$2"))
    ; default step=1
    (gsub #"^(.*):$"  "$1:1")))

(defn parse-slice [coll index]
  (->>
    ; let's get input to "start:end:step" format
    (-> index (normalize-slice coll) (str/split #":"))

    ; convert to ints
    (map edn/read-string)

    ; numpy's start:end:step matches clojure's range function, perfectly!
    (apply range)))

(defn normalize-index [coll neg-index]
  (if (neg-int? neg-index)
      (+ neg-index (count coll))
      neg-index))

(defn np-get-one-dimensional-value [coll index]
  (nth coll index))

(defn coll-of-coll? [coll] (and (coll? coll) (coll? (first coll))))

(defn slicing-index? [index] (string? index))

(defn plain-index? [index] (int? index))

(defn nested-index-2d? [coll index] (and (coll? index) (coll-of-coll? coll)))
(defn nested-index-1d? [coll index] (and (coll? index) (not (coll-of-coll? coll))))
(defn minimal-index? [index] (int? index))

(defn transpose [xy]
  (apply map vector xy))

(defn np-get-slice [coll index](nths coll (parse-slice coll index)))

(defn np-get-column-list [coll index]
  (apply map vector (mapv #(nth coll %) index)))


(defn np-get-value [coll index]
  (if-not (neg-int? index)
    (cond
      (slicing-index? index) (np-get-slice coll index)
      (plain-index? index) (np-get-one-dimensional-value coll index)

      (nested-index-2d? coll index) (np-get-column-list coll index)
      (nested-index-1d? coll index) (mapv #(np-get-value coll %) index)
    :else (throw (Throwable. (str "Unknown index type " index " " coll))))

    (recur coll (normalize-index coll index))))

(defn np-get-r [coll index dimension]
  (if (and (coll? index) (empty? index))
    coll
    (do
      (recur
        (np-get-value
          (if (and (coll-of-coll? coll) (= dimension 0) (identity (slicing-index? (first index))))
              (transpose coll)
              coll)
               (first index))
        (rest index)
        (inc dimension)
        ))))

(defn np-get [coll index]
  (np-get-r
    coll
    (if-not (coll? index) [index] index)
    0))

core_test.clj

(ns numpy-wannabe.core-test
  (:require [clojure.test :refer :all]
            [numpy-wannabe.core :refer :all]))

(def one-d [:a :b :c :d :e])

(def two-d [
  [0 1 2 3 4 5]
  [9 8 7 6 5 4]
  [11 22 33 44 55 66]])

(def three-d [
  [ [:a :b]
    [:c :d]]

  [ [:z :y]
    [:x :w]]
  ])

; INDEXING / SLICING
(deftest index-test
  (testing "1d numpy-style indexing/slicing"
    ; can't use 'are' macro w/ dynamic index argument list
    (are [indexes data expected] (= (np-get data indexes) expected)
      ; convenience for simple accessors without a wrapping vector
      3       one-d   :d
      -1      one-d   :e
      ":"     one-d   one-d
      "1:4"   one-d   [:b :c :d]
      "1:4"   one-d   [:b :c :d]
      ":4"    one-d   [:a :b :c :d]
      "1:"    one-d   [:b :c :d :e]
      "0:6:2" one-d   [:a :c :e]
      1       two-d  [9 8 7 6 5 4]

      [3]       one-d   :d
      [-1]      one-d   :e
      [1]       two-d  [9 8 7 6 5 4]
      [1 2]     two-d  7
      [":"]     one-d   one-d


      [[3 2]]   one-d   [:d :c]
      [[3]]     one-d   [:d]
      [1 [1 2]]   two-d   [8 7]
      [1 ":"] two-d   [9 8 7 6 5 4]
      [":" 1] two-d   [1 8 22]


      [1 1 0]   three-d :x

      [":" [1 2]] two-d [[1 2] [8 7] [22 33]]
      )))