hylangでnumpyっぽいインデキシング

はじめに

numpyがlispから使えるようにならないかなあみたいな夢もみたことあるくらいですが、最近hylangのおかげであっさり実現しました。
しかし、色々使ってみると、matlabから輸入された構文だけあってS式との相性は今のところ良いようには感じません。
その中でも特に気に入らないところがあったので今回もマクロの力で簡単に直しました。

問題

まずはもとのpythonコードから確認してください。

import numpy as np
ar = np.random.random((10, 10, 10, 10))
print(ar[:, 1:3, :5, 6:].shape)

こんな感じです。期待される出力は(10, 2, 5, 4)です。

これをhylangで書くとこうなります。

(import (numpy :as np))
(setv ar (np.random.random '(10 10 10 10)))
(print (. (get ar [(slice None) (slice 1 3) (slice None 5) (slice 6 None)]) shape))

どうですか。ハゲそうですよね。
直します。

理想的な文法を妄想（デザイン）する

マクロ名はngetとかで良いですかね。文法はもとのgetをベースに下のように書ければ満足です。

(import (numpy :as np))
(setv ar (np.random.random '(10 10 10 10)))
(print (. (nget ar : 1:3 :5 6:) shape))

要するに、

=> (parse-indexing ':)
'(None None)
=> (parse-indexing '1:3)
'(1 3)
=> (parse-indexing ':5)
'(None 3)
=> (parse-indexing '6:)
'(6 None)

みたいな関数があればこれをマクロ展開時に各引数に適用して貼り付ければ終わりです。
作るのは簡単そうですね。

実装

実装は以下です。CL化したhylangでプロトタイプを作って、無理してhylang標準の関数とマクロだけで書きなおしました。
標準だとあるべき関数とマクロが無さすぎてキレイに書く方法が全然わかりません。

(defn parse-indexing [sym]
  (if (not (in ":" sym)) 
      sym
      (do
        (setv splited (.split sym ":"))
        (list (map (fn [el]
                     (if (or (not el) (= el "\ufdd0"))
                         None
                         (do
                           (setv iel (try
                                       (int el)
                                       (except [e Exception] 
                                         None)))
                           (if iel
                               iel
                               (HySymbol el)))))
                   splited)))))

(defmacro nget [ar &rest indices]
  `(get ~ar ~(list-comp
               (if (or (symbol? i) (keyword? i)) 
                   `(slice ~@(parse-indexing i))
                   i)
               [i indices])))

使ってみる

=> (. (nget ar :  1:3 :5 6:) shape)
ar[[slice(None, None), slice(1, 3), slice(None, 5), slice(6, None)]].shape
(10, 2, 5, 4)

無事動きました。

まだちょっと問題が

今回はシンボルに文法を埋め込んで自前でパースすることにしましたが、これだとatom以外のものは入れられません。
そのため以下みたいな式を含むインデキシングは省略できません。

import numpy as np
ar = np.random.random((10, 10, 10, 10))
print(ar[(a+1):(b+2), 1:3, :5, 6:].shape)

あ、read上書きしなきゃ無理かもと思いましたが、意外と簡単に解決できます。
他の言語と同様にlispでも文字列リテラルの中だけは聖域です。parserが止まっているので書けないものはありません。文字列の中は楽園やでぇ。。。
以下が完成品です。

(import hy.lex.tokenize)

(defn parse-indexing [sym]
  (if (not (in ":" sym))
      sym
      (do
        (setv splited (.split sym ":"))
        (list (map (fn [el]
                     (if (or (not el) (= el "\ufdd0"))
                         None
                         (do
                           (setv iel (try
                                       (int el)
                                       (except [e Exception] 
                                         None)))
                           (if iel
                               iel
                               (HySymbol el)))))
                   splited)))))

(defn parse-str-indexing [str-i]
  (setv splited (.split str-i ":"))
  (list (map (fn [i]
               (setv i (.strip i))
               (if i
                   (get (hy.lex.tokenize i) 0)
                   None)) splited)))

(defmacro nget [ar &rest indices]
  `(get ~ar ~(list-comp
               (cond
                 [(or (symbol? i) (keyword? i)) `(slice ~@(parse-indexing i))]
                 [(string? i) `(slice ~@(parse-str-indexing i))]
                 [True i])               
               [i indices])))

使ってみます。

=> (setv ar (np.random.random '(10 10 10 10 10))
         a 1
         b 2)
from hy import HyExpression, HyInteger
ar = np.random.random(HyExpression(((((([] + [HyInteger(10)]) + [HyInteger(10)]) + [HyInteger(10)]) + [HyInteger(10)]) + [HyInteger(10)])))
a = 1
b = 2
None 
=> (. (nget ar : 1:3 :5 6: "(+ a 1):(* b 2)") shape)
ar[[slice(None, None), slice(1, 3), slice(None, 5), slice(6, None), slice((a + 1), (b * 2))]].shape
(10, 2, 5, 4, 2)

完成！