guile-ploy/ploy/slices.scm

; (c) Daniel Llorens 2012-2015
; Array operations, after J & others.

; This library is free software; you can redistribute it and/or modify it under
; the terms of the GNU General Public License as published by the Free
; Software Foundation; either version 3 of the License, or (at your option) any
; later version.

; @todo reductions: look at J/repa/numpy/PDL/SAC?; tensordot/tensorsolve.

(define-module (ploy slices))
(import (ice-9 optargs) (srfi srfi-26) (srfi srfi-1) (srfi srfi-11)
        (srfi srfi-9) (srfi srfi-8) (ploy basic) (ploy assert) (ploy ploy)
        (ploy as-array) (ploy reduce) (ploy cat) (ice-9 control))

; from (ploy basic)
(re-export $. $ tally array-dim pk-shape pk=>)

; from (ploy ploy)
(re-export rank verb w/rank array-atom array-type*
           ply/t ply ply! ply-n/o from J range amend! out/t out ravel
           reshape reshape-strict rollaxis out/t out
           i. iota. linspace. linspace-m.)

; from (ploy reduce)
(re-export over over/t folda folda/t foldb foldb/t dot cdot)

; from (ploy cat)
(re-export cat icat)

; from (ploy as-array)
(re-export arraylike-dimensions as-array)

; ---------------------------------------------------------------------
; joining arrays. @todo Put cat in here, or put these in cat.scm.
; ---------------------------------------------------------------------

; name is after J monad ; but w/o padding or replication.
(define (raze a)
; @todo use over + tally, but over must support rank > 1.
  (let* ((m (tally a))
         (n (folda (verb (lambda (n a) (+ n (tally a))) '() 0 '_) 0 a))
         (b (apply make-typed-array (or (zero? m) (array-type (array-ref a 0)))
                   *unspecified* n (cdr ($ (array-cell-ref a 0))))))
    (let loop ((i 0) (st 0))
; @TODO need amend! or...
      (if (= i m)
        b
        (let ((ai (array-cell-ref a i)))
          (amend! b ai (J (tally ai) st))
          (loop (+ 1 i) (+ st (tally ai))))))))

(export raze)

; @todo See how one would do this in J.
(define (tile a . sb)
  "tile a . sb

   Replicate array a over dimensions sb.

   sb indicate axes matching those of a. That is, the first element of sb indicates
   how many times to repeat a's first dimension, and so on.

   If a has shape [a0 a1 ... a(n-1)] and sb is [b0 b1 ... b(m-1)] with n>=m,
   the result has shape [a0*b0 a1*b1 ... a(m-1)*b(m-1) a(m) ... a(n-1)]."

  (let ((rb (length sb))
        (ra (rank a)))
    (assert (<= rb ra) "bad ranks")
    (apply reshape
      (apply transpose-array
        (out (verb (lambda (b a) a) '() 0 0) (apply reshape #f sb) a)
        (append (iota rb 0 2) (iota rb 1 2) (iota (- ra rb) (+ rb rb))))
      (let ((sa ($ a)))
        (append (map * sb sa) (drop sa rb))))))

(export tile)

; ---------------------------------------------------------------------
; explode arrays, for when we need to loop using map, etc.
; ---------------------------------------------------------------------

(define (explode cell-rank A)
  (if (or (zero? cell-rank) (= (rank A) cell-rank))
    A
    (ply/t #t (verb values '() cell-rank) A)))

(export explode)

; ---------------------------------------------------------------------
; axis operations.
; ---------------------------------------------------------------------

(define (squeeze. A)
  (apply reshape A (filter (cut not= 1 <>) ($ A))))

(export squeeze.)

; @todo have this in Guile.
(define* (reverse. A #:optional (i 0))
  "reverse. A [i]

   Reverse array over axis i, by default the first (i=0)"
  (apply make-shared-array
         A
         (lambda j (append (take j i) (list (- ($. A i) (list-ref j i) 1)) (drop j (+ i 1))))
         ($ A)))

(export reverse.)

; -----------------------------
; cant
; -----------------------------

; #(a b c d ...) -> #2((a b c) (b c d) ...)
; the name is from one of Iverson's papers (A dict. of APL?)
; c cannot be 0, but one can use reshape in that case.
; @todo apply to first axis?
; @TODO Check c:is_c_order? (ra-large.H:is_c_order) against this.
(define cant
  (case-lambda
    ((a c n)
     (make-shared-array a
                        (lambda (i j) (list (+ (* i n) (* j 1))))
                        (+ (floor/ (- (tally a) c) n) 1) c))
    ((a c)
     (cant a c 1))))

(export cant)

; ---------------------------------------------------------------------
; roll
; ---------------------------------------------------------------------

; @bug dubious use of from for set!.
; @todo A case for loop fusion, also w/rank use.
(define (roll n a)
  (if (zero? n)
    a
    (let* ((b (apply make-typed-array (array-type a) *unspecified* ($ a)))
           (s (tally a))
           (n (modulo n s)))
      (array-copy! (from a (J (- s n) 0)) (from b (J (- s n) n)))
      (array-copy! (from a (J n (- s n))) (from b (J n 0)))
      b)))

(export roll)

; ---------------------------------------------------------------------
; family is sort, partial-sort, median, min-by, max-by.
; ---------------------------------------------------------------------

(define (max-by l <)
  (assert (pair? l) "bad args to max-by")
  (fold (lambda (s b) (if (< b s) s b)) (car l) (cdr l)))

(define (min-by l <)
  (assert (pair? l) "bad args to min-by")
  (fold (lambda (s b) (if (< s b) s b)) (car l) (cdr l)))

(export max-by min-by)

; sort i according to the order of (from b i).
; this is a shortcut for (sort-by i (from b i) <).
; @todo There's a (sort-indices-by-norm) in (array symmetrizer). See to replace it.
(define sort-indices-by.
  (case-lambda
    ((i b) (sort-indices-by. i b <))
    ((i b <) (stable-sort* i (lambda (i j) (< (array-cell-ref b i) (array-cell-ref b j)))))))

; sort over items.
(define sort.
  (case-lambda
    ((a) (sort. a <))
    ((a <) (ply/t (array-type* a) (verb (cut array-cell-ref a <>) (cdr ($ a)) 0)
                     (sort-indices-by. (i. (tally a)) a <)))))

; sort a according to the order of the corresponding elements of b.
(define sort-by.
  (case-lambda
    ((a b) (sort-by. a b <))
    ((a b <)
     (assert (>= (tally b) (tally a)) "b too small to grade a")
     (ply/t (array-type* a) (verb (cut array-cell-ref a <>) (cdr ($ a)) 0)
               (sort-indices-by. (i. (tally a)) b <)))))

; @todo unfortunately stable-sort doesn't work on uniform arrays.
(define stable-sort*
  (case-lambda
    ((a) (stable-sort* a <))
    ((a <)
     (if (or (eq? #t (array-type a)) (list? a))
       (stable-sort a <)
       (let ((A (array-copy #t a)))
         (stable-sort! A <)
         (array-copy (array-type a) A))))))

(export sort. sort-by. sort-indices-by.)

; ---------------------------------------------------------------------
; variant selectors
; ---------------------------------------------------------------------

; @todo Use (over) when it supports rank>1.
(define (count. pred? a)
;  (over + (lambda (a) (if (pred? a) 1 0)) a)
  (folda (verb (lambda (c a) (+ c (if (pred? a) 1 0))) '() 0 '_) 0 a))

(define copy./t
  (case-lambda
    ((type w b)
     (let* ((o (apply make-typed-array type *unspecified*
                      (cons (count. values w) (cdr ($ b)))))
            (j 0))
       (array-slice-for-each 1
                             (lambda (w b)
                               (when (array-cell-ref w)
                                 (array-cell-set! o (array-cell-ref b) j)
                                 (set! j (+ 1 j))))
                             w b)
       o))
    ((type b)
     (array-copy type b))))

(define copy.
  (case-lambda
    ((w b) (copy./t (array-type* b) w b))
    ((b) (copy./t (array-type* b) b))))

; this is the same as (ply (verb array-cell-ref #f '_ 0) a i), unless there are repeated indices.
(define (copy-i. i a)
  (copy. (let ((l (make-array #f (tally a))))
           (array-for-each (lambda (i) (array-set! l #t i)) i)
           l)
         a))

(define (filter. pred? a)
  (let ((pred? (if (verb? pred?) pred? (verb pred? '() -1))))
    (copy. (ply/t 'b pred? a) a)))

(define (filter-map. f a)
  (let ((fa (ply (if (verb? f) f (verb f '() -1)) a)))
    (copy. fa fa)))

(export filter. filter-map.)

(define (remove-i. i a)
  (if (zero? (tally i))
    a
    (let ((i (sort. i <))
          (b (apply make-typed-array (array-type* a) *unspecified*
                    (cons (- (tally a) (tally i)) (cdr ($ a))))))
      (let loop ((k 0) (a0 0) (b0 0))
        (if (= k (tally i))
          (let ((ik (tally a)))
            (array-copy! (from a (J (- ik a0) a0)) (from b (J (- ik a0) b0)))
            b)
          (let* ((ik (array-cell-ref i k))
                 (d (- ik a0)))
            (if (zero? d)
              (loop (+ 1 k) (+ 1 ik) (+ b0 d))
              (begin
                (array-copy! (from a (J d a0)) (from b (J d b0)))
                (loop (+ 1 k) (+ 1 ik) (+ b0 d))))))))))

(define drop.
  (case-lambda
    ((a) (drop. a 1))
    ((a n) (from a (range n (tally a))))))

(export count. copy. copy-i. remove-i. drop. copy./t)

; @todo neither over nor fold is a good fit for every. any. index  .
; @todo over should support identity...
(define (every. pred? a)
  (reset (over (case-lambda ((c a) a) ((a) a) (() #t))
               (lambda (a) (let ((x (pred? a))) (or x (shift k x))))
               a)))

; @todo see every.
(define (any. pred? a)
  (reset (over (case-lambda ((c a) a) ((a) a) (() #f))
               (lambda (a) (let ((x (pred? a))) (and x (shift k x))))
               a)))

; @todo maybe rename to index-first.
; @todo foldb carries the index, redundant.
; @todo pred? what if pred? is a verb.
(define (index pred? a)
  (reset (folda (verb (lambda (i a) (or (and (pred? a) (shift k i)) (+ 1 i)))
                      '() 0 '_)
                0 a)
; J dyad i. would leave the result of the fold.
         #f))

; @TODO see (index). ply/t carries the index, (i. (tally a)) is redundant.
; @TODO pred? should be a verb...
; @TODO a variable rank version, returning multi-indices.
(define (indices pred? . a)
  (assert (procedure? pred?) "indices can't take verb, sorry")
  (copy. (apply ply/t 'b (apply verb pred? '() (make-list (length a) -1)) a)
         (i. (tally (car a)))))

(export every. any. index indices)

; should be (set! (from y x) (iota (size x))) without having to create i.
; Maybe i ~ TensorIndex<0>.
; cf gradeup for inverting permutations.
(define* (invert-index x #:optional (m (+ 1 (over max x))))
  "invert-index x #:optional (m 1+(max x))
   return y such that y[x[i]] = i. Range of y is 0..(m-1)."
  (let ((y (make-vector m #f))
        (n (tally x)))
    (let loop ((i 0))
      (cond ((= i n) y)
            (else (array-set! y i (array-ref x i))
                  (loop (+ 1 i)))))))

(export invert-index)

(define (first. x) (from x 0))
(define (last. x) (from x (- (tally x) 1)))

(export first. last.)

; ---------------------------------------------------------------------
; conversion or casting.
; ---------------------------------------------------------------------

; @todo f64/c64 uniform vectors are also bytevectors, so try to cast instead.
; @todo ply with vector constructor is wasteful.
; @todo maybe keep the last axis, but must fix users.
(define (complex->f64 c)
  (apply reshape
   (ply/t 'f64 (verb (lambda (x) (vector (real-part x) (imag-part x))) '(2) 0) c)
   (append (drop-right ($ c) 1) (list (* 2 (last ($ c)))))))

(define (real->c64 r)
  (assert (even? (last ($ r))) "last dimension must be even")
  (let ((r2 (apply reshape r (append (drop-right ($ r) 1) (list (/ (last ($ r)) 2) 2)))))
   (ply/t 'c64 (verb (lambda (x) (make-rectangular (array-ref x 0) (array-ref x 1))) '() 1) r2)))

(export complex->f64 real->c64)

; ---------------------------------------------------------------------
; common verbs, but @todo maybe not the right place.
; ---------------------------------------------------------------------

(define vnorm. (verb v-norm '() 1))
(define sqrm. (verb sqrm-reduce '() 1))

(export vnorm. sqrm.)
(re-export v-norm sqrm-reduce)