guile-ploy/ploy/ploy.scm

; (c) Daniel Llorens - 2012-2015
; Array operations.

; 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.

; @bug array-copy! is lax wrt copies (dest can be larger) which lets some errors
; through. Should use something stricter, or fix array-copy!.

(define-module (ploy ploy))
(import (ice-9 optargs) (srfi srfi-26) (srfi srfi-1) (srfi srfi-11) (ice-9 match)
        (srfi srfi-9) (srfi srfi-8) (ploy basic) (ploy assert))

(re-export $. $ tally array-dim rank array-type* i. iota. linspace. linspace-m.)
(re-export ravel reshape reshape-strict)

; Not for users, b/c array-first may return #0(.).

(define (array-first a)
  (apply make-shared-array a
         (lambda x (cons 0 x))
         (cdr ($ a))))
(define (array-rest a)
  (apply make-shared-array a
         (lambda x (cons (+ (car x) 1) (cdr x)))
         (cons (- (tally a) 1) (cdr ($ a)))))

; -------------------------------------------------------------------
; Guile array iteration ops
; -------------------------------------------------------------------
;
;                         write-arg  ordered/any-order  indices/element
; array-for-each              R             O                  E
; array-map!                  W             A                  E
; array-index-map!            W             A                  I
; array-map-in-order!         W             O                  E
;
; probably need array-index-map-in-order!, array-for-each-any-order!
; (see ply-n/o!), something else?

; -------------------------------------------------------------------
; declare verb ranks
; -------------------------------------------------------------------

; A necessary first step in establishing the use of ranks is to define the ranks
; of all existing functions. Because of the treatment of degenerate cases, it is
; possible to do this in a manner compatible with existing definitions by simply
; assigning infinite ranks to all. However, it is generally more useful to
; assign the lowest rank which will ensure compatibility with existing
; definitions, since this leads to simpler definitions of the functions, and
; greater utilization of the general scheme of extension. In particular, the
; scalar functions can, and will, be assigned rank 0.
;   ---Iverson 1983, 'Rationalized APL'

; The emphasis on examples using first axis variants of functions and operators
; is intentional. Intermediate and last axis definitions fall out logically and
; simply from use of the Rank Operator in conjunction with first axis variants,
; while the converse is not true. For example [phi] can be written as [theta]"1,
; whereas [theta] cannot be expressed in terms of [phi].
;   ---Bernecky1988, 'An introduction to function rank'

(define-record-type <verb>
  (make-verb op oshape ri) verb?
  (op verb-op)
  (oshape verb-oshape)
  (ri verb-ri))

(define (valid-ranks? ri)
  (every (lambda (ri) (or (integer? ri) (eq? '_ ri))) ri))

; declare verb ranks.
; @TODO A way to reshape to match, either with cycle (easier ---that's what
; reshape does) or filler (??).
(define verb
  (case-lambda
    "verb {op | op oshape | op oshape . ri}

   Construct verb from procedure op taking args of ranks ri.
   oshape can be - a function (arg0-shape ...) -> output shape.
                 - a fixed output shape.
                 - #f meaning 'arbitrary'."
    ((op)
     (unless (procedure? op) (throw 'bad-op-0 op))
     (make-verb op #f #f))
    ((op oshape)
     (unless (procedure? op) (throw 'bad-op-1 op))
     (make-verb op oshape #f))
    ((op oshape . ri)
     (unless (procedure? op) (throw 'bad-op-2 op))
     (unless (valid-ranks? ri) (throw 'bad-input-ranks ri))
     (make-verb op oshape ri))))

; modify a verb to take other cell ranks; rank conjunction.
; @TODO Make it easier to change just some ranks of op.
; @TODO Deduce output rank.
(define (w/rank op . ri)
  (make-verb (cond ((procedure? op) (verb op))
                   ((verb? op) op)
                   (else (throw 'bad-op-3 op)))
             #f
             (if (valid-ranks? ri)
               ri
               (throw 'bad-input-ranks ri))))

(define (verb-actual-ri op . ra)
  (match (verb-ri op)
    (#f (make-list (length ra) 0))
    (ri (unless (= (length ra) (length ri)) (throw 'bad-number-of-args op ri (length ra)))
        (map (lambda (ri ra)
               (let ((r (cond ((eq? '_ ri) ra)
                              ((negative? ri) (+ ri ra))
                              (else ri))))
                 (unless (<= 0 r ra) (throw 'bad-rank ri ra))
                 r))
             ri ra))))

(define (verb-final? op)
  (procedure? (verb-op op)))

(export verb? verb w/rank verb-op verb-ri verb-actual-ri verb-final? verb-oshape)

; -------------------------------------------------------------------
; array-map over already matching frames
; -------------------------------------------------------------------

; the fundamental array ops.
(define array-map/frame!
  (case-lambda
; optimizations. @TODO wish it was automatic; would apply elsewhere.
   ((o f op a0)
    (if (= (length f) (rank o) (rank a0))
      (array-map! o op a0)
      (array-slice-for-each (length f)
                            (lambda (o a0)
                              (array-cell-set! o (op (array-cell-ref a0))))
                            o a0)))
   ((o f op a0 a1)
    (let ((oo o))
      (if (= (length f) (rank o) (rank a0) (rank a1))
        (array-map! o op a0 a1)
        (array-slice-for-each (length f)
                              (lambda (o a0 a1)
                                (array-cell-set! o (op (array-cell-ref a0) (array-cell-ref a1))))
                              o a0 a1))))
   ((o f op . a)
    (if (apply = (length f) (rank o) (map rank a))
      (apply array-map! o op a)
      (apply array-slice-for-each (length f)
             (lambda (o . a)
               (array-cell-set! o (apply op (map array-cell-ref a))))
             o a)))))

(define (array-atom A)
  "array-atom A

  Return an atom of A. If A is a zero-size slice of a root array, this function
  may return a result that is not in A, but is in the root array."
  (if (array? A)
    (array-ref (shared-array-root A) 0)
    A))

; convert nested array to straight array, first nesting only. We may need this
; on the output of array-map/frame! if op output shape can't be computed.
(define (collapse-array type a)
  (let ((b (and (positive? (apply * ($ a)))
                (let ((b (array-atom a)))
                  (and (array? b) b)))))
    (cond (b (let ((o (apply make-typed-array type *unspecified* (append ($ a) ($ b)))))
               (array-slice-for-each (rank a)
                                     (lambda (o a) (array-cell-set! o (array-ref a)))
                                     o a)
               o))
          ((eq? (array-type a) type) a)
          (else (array-copy type a)))))

(define (array-map/frame type oshape f op . a)
  (cond ((null? f)
         (let ((o (apply op a)))
           (if (and (array? o) (not (eq? (array-type o) type)))
             (array-copy type o)
             o)))
        ((or (null? oshape) (pair? oshape))
         (let ((o (apply make-typed-array type *unspecified* (append f oshape))))
           (apply array-map/frame! o f op a)
           o))
        (else
; output cell may not be scalar, so can't use uniform type before collapse.
         (collapse-array
          type
          (let ((o (apply make-array *unspecified* f)))
            (apply array-map/frame! o f op a)
            o)))))

; without output. @TODO don't need in order traversal, despite for-each.
(define (array-map/frame-n/o f op . a)
  (cond
   ((null? f)
    (apply op a))
   ((apply = (length f) (map rank a))
    (apply array-for-each op a))
   (else
    (apply array-slice-for-each (length f) (lambda a (apply op (map array-cell-ref a))) a))))

(export collapse-array array-atom array-map/frame! array-map/frame)

; -------------------------------------------------------------------
; match frames
; -------------------------------------------------------------------

; Rank-extend A with cell rank r to frame f.
; [(a skipped k prefix) ( (-r)-frame of a/k... ) ### (a/k r-cell ...)]
; [(a skipped k prefix) ( f ...................... ) (a/k r-cell ...)]
; where ### are the extended axes.
; @todo the scalar cases should be handled specially so that this isn't needed.
(define (match-frame a f r k)
  "Rank-extend a with cell rank r to frame f."
  (cond ((length=? (- (rank a) k r) f)
         a)
        ((not (array? a))
         (match-frame (make-array a) f r k))
        (else
         (apply make-shared-array a
                (lambda i (append (take i k) (take (drop i k) (- (rank a) k r)) (take-right i r)))
                (append (take ($ a) k) f (take-right ($ a) r))))))

(define (prefix-frame a r k)
  "Frame common to all arrays a with cell ranks r, ignoring first k axes."
  (fold (lambda (a r f)
          (let ((fa (drop-right! (drop ($ a) k) r)))
            (let loop ((s f) (sa fa))
              (cond ((null? sa) f)
                    ((null? s) fa)
                    ((= (car s) (car sa)) (loop (cdr s) (cdr sa)))
                    (else (error "shape clash" ($ a) r f k))))))
        '() a r))

(define (nested-op-frames op k . a)
  "Match argument frames for all nested w/rank ops."
  (let loop ((op op) (ff '()) (a a) (k k))
    (let* ((r (apply verb-actual-ri op (map (lambda (a) (- (rank a) k)) a)))
           (f (prefix-frame a r k))
           (a (map! (cut match-frame <> f <> k) a r)))
      (if (verb-final? op)
        (values (let ((vo (verb-oshape op)))
                  (if (procedure? vo)
                    (apply vo (map (lambda (a r) (take-right ($ a) r)) a r))
                    vo))
                (concatenate! (reverse! (cons f ff)))
                (verb-op op)
; ply doesn't use r, but fold does.
                r a)
        (loop (verb-op op) (cons f ff) a (+ k (length f)))))))

; -------------------------------------------------------------------
; interface
; -------------------------------------------------------------------

; @todo reduce scalar arguments with a closure.
(define (ply/t type op . a)
  (let ((op (if (verb? op) op (verb op))))
    (receive (oshape f op ri a) (apply nested-op-frames op 0 a)
      (apply array-map/frame type oshape f op a))))

(define (ply op . a)
  (apply ply/t (array-type* (car a)) op a))

(export prefix-frame match-frame nested-op-frames ply/t ply)

; -------------------------------------------------------------------
; ply with reference argument.
; -------------------------------------------------------------------
; See (ammend!) http://www.jsoftware.com/help/jforc/modifying_an_array_m.htm#_Toc191734471
; y[m] = x ... x m} y ... where $x must be a suffix of $(m{ y).
; See amend! below and how it relates to this.

;; 0 (<1 1)} 4 4 $ 5
;; 0 1 2 3 (0)} 4 4 $ 5
;; 0 (0)} 4 4 $ 5
;; 0 1 (0)} 4 4 $ 5 -> error
;; 1 2 (1 1;2 2)} 4 4 $ 5
;; 1 2 (1 1;1 1)} 4 4 $ 5 -> undefined

; also look at expand-X in /box.

(define (ply! o op . a)
  (let ((op (if (verb? op) op (verb op))))
    (receive (oshape f op ri a) (apply nested-op-frames op 0 a)
      (if oshape
        (let ((suffix (take-right ($ o) (length oshape))))
          (if (equal? oshape suffix)
            (let ((oo (apply reshape o #f suffix)))
              (apply array-map/frame! (array-first oo) f op a)
              (array-copy! (apply reshape (array-first oo) (- (tally oo) 1) suffix)
                           (array-rest oo)))
            (error "mismatched arguments")))
        (let* ((ss (apply array-map/frame (array-type o) oshape f op a))
               (oshape ($ ss)))
          (let ((suffix (take-right ($ o) (length oshape))))
            (if (equal? oshape suffix)
              (let ((oo (apply reshape o #f suffix)))
                (array-copy! (apply reshape ss (tally oo) suffix)
                             oo))
              (error "mismatched arguments")))))
      o)))

(define (ply!! o op . a)
  (let ((op (if (verb? op) op (verb op))))
; @TODO not using oshape, should maybe compute it separately.
    (receive (oshape f op ri o-a) (apply nested-op-frames op 0 o a)
      (apply array-map/frame! (car o-a) f op (cdr o-a))
      (car o-a))))

; (ply! (make-array 0 10) (verb (const 2) '()))
; (ply! (make-array 0 10) (const 2))

(define (ply-n/o op . a)
  (let ((op (if (verb? op) op (verb op))))
; @TODO not using oshape, should maybe compute it separately.
    (receive (oshape f op ri a) (apply nested-op-frames op 0 a)
      (apply array-map/frame-n/o f op a))))

(export ply! ply!! ply-n/o)

; -------------------------------------------------------------------
; selection
; -------------------------------------------------------------------

(define-record-type <J>
  (make-J n from step) J?
  (n     J-n)
  (from  J-from)
  (step  J-step))

(define J
  (case-lambda (() #f)
               ((n) (make-J n 0 1))
               ((n from) (make-J n from 1))
               ((n from step) (make-J n from step))))

(define (J-index J i)
  (+ (J-from J) (* (J-step J) i)))

(define range
  (case-lambda ((from til)
                (range from til 1))
               ((from til step)
                (assert (eq? (>= til from) (>= step 0)))
                (J (floor/ (- til from) step) from step))))

; return the numbers in [0 .. k) not in l, assumed sorted. @todo vectors here.
(define (complement-list k l)
  (let loop ((i 0) (l l) (r '()))
    (cond ((= i k) r)
          ((null? l) (append! r (iota (- k i) i)))
          ((< i (car l)) (loop (+ (car l) 1) (cdr l) (append! r (iota (- (car l) i) i))))
          ((= i (car l)) (loop (+ i 1) (cdr l) r))
          (else (error "bad arguments")))))

(define-record-type <parsed-array-arg>
  (make-paa axis rank arg) paa?
  (axis paa-axis)
  (rank paa-rank)
  (arg  paa-arg))

(define (parse-from-args a . x)
  (define (checked-cdr d) (if (pair? d) (cdr d) (throw 'bad-rank)))
  (let*-values
; c is ((axis rank arg)  ...) for array args.
; b is array after applying shortcuts.
; s is dimensions of b.
      (((c s) (let loop ((c '()) (s '()) (i 0) (x x) (d ($ a)))
                (match x
                  (()
                   (values (reverse! c) (append! (reverse! s) d)))
                  (((? boolean?) . xtail)
                   (loop c (cons (car d) s) (1+ i) xtail (checked-cdr d)))
                  (((? integer?) . xtail)
                   (loop c s i xtail (checked-cdr d)))
                  (((? J? x) . xtail)
                   (loop c (cons (J-n x) s) (1+ i) xtail (checked-cdr d)))
                  (((? array? x) . xtail)
                   (loop (cons (make-paa i (rank x) x) c) (cons (car d) s)
                         (1+ i) xtail (checked-cdr d))))))
       ((b) (apply make-shared-array
              a
              (lambda u
                (let loop ((u u) (x x) (r '()))
                  (match x
                    (() (append! (reverse! r) u))
                    (((or (? boolean?) (? array?)) . xtail)
                     (loop (cdr u) xtail (cons (car u) r)))
                    (((? integer? x) . xtail)
                     (loop u xtail (cons x r)))
                    (((? J? x) . xtail)
                     (loop (cdr u) xtail (cons (J-index x (car u)) r))))))
              s))
       ((lc) (length c))
       ((lb) (array-rank b)))
    (values c s b lc lb)))

(define (forward-backward b c lb lc)
  (let* ((cx (map paa-axis c))
         (forward (gradeup (append cx (complement-list (array-rank b) cx))))
         (backward
          (let loop ((cc '()) (ww '()) (c c) (i 0) (j 0))
            (cond ((null? c)
                   (append! cc (reverse! ww) (iota (- lb i) j)))
                  ((= (paa-axis (car c)) i)
                   (loop (append! cc (iota (paa-rank (car c)) j))
                         ww (cdr c) (+ i 1) (+ j (paa-rank (car c)))))
                  (else
                   (loop cc (cons j ww) c (+ i 1) (+ j 1)))))))
    (values forward backward)))

(define (from a . x)
  "from A . index-list

   Rectangular selection operator. Each index applies to one axis of the array
   A. The indices can be either integers or integer arrays of any rank. The
   result will have rank equal to the sum of the ranks of the indices. The
   following shortcuts for common rank-1 indices are accepted:

   #t                      - the whole axis, i.e. #(lbnd lbnd+1 .. ubnd-1 ubnd).
   (J n [from=0 [step=1]]) - #(from from+step .. from+(n-1)*step)
   (range a b [step=1])    - #(a a+step ... [a+floor((b-a)/step)*step])

   If every index is either a scalar or one of the shortcuts above and the
   result has rank > 0, the result will be a shared array over the same storage
   as A.

   If index-list does not cover the rank of A, the missing indices are taken to
   be #t. It is an error if index-list is longer than the rank of A.

   If the result x has rank 0, (from) will return x, never #0(x).

   Examples:

   (from A) => A, if the rank of A > 0, but x if A is #0(x)
   (from A #t) => A, if the rank of A > 0, otherwise error
   (apply from A [list of (rank A) scalars]) => (apply array-ref A [list ...])
   (from (i. 2 3) 0) => #(0 1 2)
   (from (i. 2 3) 0 #t) => #(0 1 2), i.e. row 0
   (from (i. 2 3) #t 0) => #(0 3), i.e. column 0
   (from #(#(0 1) #(2)) 1) => #(2))
   (from #(#(0 1) #(2)) 1 0) => error
   (from (i. 10 2) (J 2 2) #t) => #2((4 5) (6 7))
   (from #(1 2 3) #2((0 1) (1 2) (2 0))) => #2((1 2) (2 3) (3 1))
   (from (i. 10 10) #(3 4) #(7 9 2)) => #2((37 39 32) (47 49 42))"

  (receive (c s b lc lb) (apply parse-from-args a x)
; selection without need of transpose.
    (define (prefix-from b)
      (let ((oshape (drop ($ b) lc)))
        (apply ply/t (array-type* b)
               (let loop ((i 0) (cr (map paa-rank c)))
                 (if (null? cr)
                   (apply verb (lambda x (apply array-cell-ref b x)) oshape (make-list lc 0))
                   (apply w/rank
                     (loop (+ 1 i) (cdr cr))
                     (append! (make-list (+ 1 i) 0) (cdr cr)))))
               (map paa-arg c))))
    (cond ((null? s) (array-ref b))
; optimizations.
          ((null? c) b)
          ((= lc lb) (prefix-from b))
          (else
; move explicit indexing to the front of the shape of b, so that it can be done w/rank.
           (receive (forward backward) (forward-backward b c lb lc)
             (apply transpose-array
               (prefix-from (apply transpose-array b forward))
               backward))))))

(export <J> J? J J-index J-n paa-arg paa-rank paa-axis parse-from-args from range)

(define (extend-left z s)
  "extend-left z s

   Broadcast array z to shape s. The shape of z must be a suffix of s."
  (assert (equal? (take-right s (rank z)) ($ z)) "cannot extend to shape")
  (cond ((null? s) (if (array? z) (array-ref z) z))
        ((array? z) (apply make-shared-array z (lambda i (take-right i (array-rank z))) s))
        (else (apply make-shared-array (make-array z) (const '()) s))))

(define (extend-right z s)
  "extend-left z s

   Broadcast array z to shape s. The shape of z must be a prefix of s."
  (assert (equal? (take s (rank z)) ($ z)) "cannot extend to shape")
  (cond ((null? s) (if (array? z) (array-ref z) z))
        ((array? z) (apply make-shared-array z (lambda i (take i (array-rank z))) s))
        (else (apply make-shared-array (make-array z) (const '()) s))))

; @TODO Factor out more pieces shared with (from).
; @TODO Consider in the light of (ply!).
(define (amend! a z . x)
  "amend! A z indices ... -> modified A

   Rectangular assignment operator.

   Think of this as (array-copy! z (from A indices ...)), except that it will
   work for any indices that work with (from), even for those that
   (from A indices) doesn't return an assignable array for.

   ($ z) must be a suffix of ($ (from A indices ...)); if (from A indices ...)
   has higher rank than z, each of the (rank z)-cells of (from A indices) will
   be assigned z. This behavior is after J amend.

   Examples:

   (define A (array-copy #(10 20 30 40)))
   (amend! A #(a b c) #(3 0 2))
   => A: #(b 20 c a)

   @TODO others..."

  (receive (c s b lc lb) (apply parse-from-args a x)
; selection without need of transpose.
    (define (prefix-amend! b z)
      (apply ply-n/o
             (let loop ((i 0) (cr (map paa-rank c)))
               (if (null? cr)
                 (apply verb (lambda (z . x) (apply array-cell-set! b z x)) #f '_ (make-list lc 0))
                 (apply w/rank
                   (loop (+ 1 i) (cdr cr))
                   (- (car cr))
                   (append! (make-list (+ 1 i) 0) (cdr cr)))))
             z (map paa-arg c)))
    (cond
; optimization.
     ((and (null? c) (not (array? z))) (array-fill! b z))
     ((= lc lb)
      (prefix-amend! b (extend-left z (concatenate (map (lambda (c) ($ (paa-arg c))) c)))))
     (else
; move explicit indexing to the front of the shape of b, so that it can be done w/rank.
      (receive (forward backward) (forward-backward b c lb lc)
        (prefix-amend!
         (apply transpose-array b forward)
         (apply transpose-array
; expected-z-shape: ---($b[i0])--($b[i1])--- will become ---($i0)--($i1)---.
; @TODO compute along [backward]. @TODO compute only if required.
           (extend-left z
                        (let ((sz (map list ($ b))))
                          (for-each (lambda (c) (list-set! sz (paa-axis c) ($ (paa-arg c)))) c)
                          (concatenate sz)))
           (gradeup backward))))))
    a))

(export extend-left extend-right amend!)

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

(define* (rollaxis a s #:optional (d -1))
  "rollaxis a s [d: -1]

   Transpose axis s of array a to d by shifting the axes in between. Negative
   axes count from the end.

   ($ (rollaxis (i. 2 3 4 5)  0 -1)) => (3 4 5 2)
   ($ (rollaxis (i. 2 3 4 5) -1  0)) => (5 2 3 4)
   ($ (rollaxis (i. 2 3 4 5)  0  2)) => (3 4 2 5)
   ($ (rollaxis (i. 2 3 4 5)  3  1)) => (2 5 3 4)
   "
  (let ((s (if (negative? s) (+ s (rank a)) s))
        (d (if (negative? d) (+ d (rank a)) d)))
; transpose-array args are the destination of each source axis.
    (cond ((< s d)
           (apply transpose-array a (append! (iota s)
                                             (list d)
                                             (iota (- d s) s)
                                             (iota (- (rank a) d 1) (+ d 1)))))
          ((> s d)
           (apply transpose-array a (append! (iota d)
                                             (iota (- s d) (+ d 1))
                                             (list d)
                                             (iota (- (rank a) s 1) (+ s 1)))))
          (else a))))

(export rollaxis)

; ---------------------------------------------------------------------
; outer product
; @todo Define a verb instead of an operation, then do (ply (out ..) a ..) (or ply!).
; But need input ranks.
; ---------------------------------------------------------------------

(define (out/t type op . a)
  (let* ((op (if (verb? op) op (verb op)))
         (ra (map rank a))
         (ri (apply verb-actual-ri op ra)))
    (let ((n (length a)))
      (apply
       ply/t type
       (let loop ((i 1))
         (if (>= i n)
           op
           (apply w/rank (loop (+ i 1)) (append (take ri i) (drop ra i)))))
       a))))

(define (out op . a) (apply out/t (array-type* (car a)) op a))

(export out/t out)