utils.scm 6.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200
  1. (library (array-utils)
  2. (export array-len-in-dim
  3. arrays->hash-table
  4. array-display
  5. array-map
  6. array-next-index
  7. array-cell-ref-vec
  8. array-index-of
  9. array-indices-of)
  10. (import
  11. (except (rnrs base)
  12. let-values
  13. map
  14. error
  15. vector-map)
  16. (only (guile)
  17. lambda* λ
  18. current-output-port
  19. ;; arrays
  20. array-shape
  21. array-ref
  22. array-cell-ref
  23. array-map!
  24. array-rank
  25. ;; display
  26. display
  27. simple-format
  28. current-output-port)
  29. ;; lists
  30. (srfi srfi-1)
  31. ;; vectors
  32. (srfi srfi-43)
  33. ;; hash tables
  34. (srfi srfi-69)
  35. (ice-9 arrays))
  36. (define array-len-in-dim
  37. (λ (arr dim)
  38. "Return the array length of a multi-dimensional array ARR in the
  39. specified dimension DIM."
  40. (let* ([shape (array-shape arr)]
  41. [dim-min-max (list-ref shape dim)])
  42. (+ (- (second dim-min-max)
  43. (first dim-min-max))
  44. 1))))
  45. (define arrays->hash-table
  46. (lambda* (keys-arr vals-arrs #:optional (equal-func equal?))
  47. (let ([rows (array-len-in-dim keys-arr 0)]
  48. [cols (array-len-in-dim keys-arr 1)]
  49. [table (make-hash-table equal-func)])
  50. (let iter-rows ([row-ind 0])
  51. (let iter-cols ([col-ind 0])
  52. (cond
  53. [(< row-ind rows)
  54. (cond
  55. [(< col-ind cols)
  56. (hash-table-set! table
  57. (array-ref keys-arr row-ind col-ind)
  58. (array-ref vals-arrs row-ind col-ind))
  59. (iter-cols (+ col-ind 1))]
  60. [else (iter-rows (+ row-ind 1))])]
  61. [else table]))))))
  62. (define array-display
  63. (lambda* (landscape
  64. #:optional (port (current-output-port))
  65. #:key (formatter (λ (elem) elem)))
  66. (let ([rows (array-len-in-dim landscape 0)]
  67. [cols (array-len-in-dim landscape 1)])
  68. (let iter-rows ([row-ind 0])
  69. (let iter-cols ([col-ind 0])
  70. (cond
  71. [(>= row-ind rows) 'done]
  72. [(>= col-ind cols)
  73. (display "\n" (current-output-port))
  74. (iter-rows (+ row-ind 1))]
  75. [else
  76. (display (formatter (array-cell-ref landscape row-ind col-ind)) port)
  77. (iter-cols (+ col-ind 1))]))))))
  78. (define array-map
  79. (λ (proc src-arr)
  80. (define target-arr (array-copy src-arr))
  81. (array-map! target-arr proc src-arr)
  82. target-arr))
  83. (define array-next-index
  84. (λ (shape-vec indices-vec max-dim)
  85. "Increment one of indices in INDICES-VEC for which the
  86. following conditions are true:
  87. 1. There is no index in INDICES-VEC that sits at a later
  88. position (higher index) than MAX-DIM of INDICES-VEC and is
  89. not yet at its maximum. The maximum is specified by the
  90. SHAPE-VEC. Each index in INDICES-VEC has a corresponding
  91. minimum and maximum in SHAPE-VEC at the same position.
  92. After incrementing the index in INDICES-VEC, all later
  93. indices (at a higher index of INDICES-VEC) are set to their
  94. corresponding minimum to get a correct indices vector."
  95. (cond
  96. [(>= max-dim 0)
  97. (let ([index (vector-ref indices-vec max-dim)]
  98. [max-for-index (second (vector-ref shape-vec max-dim))])
  99. (cond
  100. [(< index max-for-index)
  101. ;; Copy the vector to not mutate argument.
  102. (let ([updated-indices-vec (vector-copy indices-vec)])
  103. ;; Increase index at position.
  104. (vector-set! updated-indices-vec
  105. max-dim
  106. (+ (vector-ref updated-indices-vec max-dim) 1))
  107. ;; Set later indices to their corresponding minimum.
  108. (let ([indices-vec-len (vector-length indices-vec)])
  109. (let iter ([dim° (+ max-dim 1)])
  110. (cond
  111. [(< dim° indices-vec-len)
  112. (let ([minimum-ind-val (first (vector-ref shape-vec dim°))])
  113. (vector-set! updated-indices-vec dim° minimum-ind-val)
  114. (iter (+ dim° 1)))]
  115. [else updated-indices-vec]))))]
  116. [else
  117. ;; Increment next higher dimension.
  118. (array-next-index shape-vec
  119. indices-vec
  120. (- max-dim 1))]))]
  121. [else #f])))
  122. (define array-cell-ref-vec
  123. (λ (arr indices-vec)
  124. "array-cell-ref takes a vector of indices INDICES-VEC, whose
  125. length depends on how many dimensions the array ARR
  126. has (what its rank is)."
  127. (let ([vec-len (vector-length indices-vec)])
  128. (let iter ([index-into-indices-vec° 0] [result arr])
  129. (cond
  130. [(< index-into-indices-vec° vec-len)
  131. (let ([cell-index (vector-ref indices-vec index-into-indices-vec°)])
  132. (iter (+ index-into-indices-vec° 1)
  133. (array-cell-ref result cell-index)))]
  134. [else result])))))
  135. (define array-index-of
  136. (lambda* (arr pred #:optional (start-indices #f))
  137. "Return the index of the first element in ARR which
  138. satisfies the predicate PRED."
  139. (let* ([shape (array-shape arr)]
  140. [shape-vec (list->vector shape)]
  141. [rank (array-rank arr)]
  142. [initial-indices
  143. (if start-indices
  144. start-indices
  145. (vector-map (λ (_i elem) (car elem)) shape-vec))])
  146. (let iter ([indices° initial-indices])
  147. (cond
  148. [indices°
  149. (cond
  150. ;; If the array element satisfies the
  151. ;; predicate, return the indices of the
  152. ;; element.
  153. [(pred (array-cell-ref-vec arr indices°))
  154. indices°]
  155. [else
  156. ;; Potential optimization: Make better use of
  157. ;; the rank argument, so that array-next-index
  158. ;; does not have to search unnecessarily for
  159. ;; the index to increment.
  160. (iter (array-next-index shape-vec indices° (- rank 1)))])]
  161. ;; No index found at which the predicate would be
  162. ;; satisfied.
  163. [else #f])))))
  164. (define array-indices-of
  165. (λ (arr pred)
  166. "Return a list of vectors of indices into the given array ARR which
  167. satisfy the given predicate PRED, for an array of arbitrary rank."
  168. (let* ([shape (array-shape arr)]
  169. [shape-vec (list->vector shape)]
  170. [rank (array-rank arr)]
  171. [initial-indices (vector-map (λ (_i elem) (car elem)) shape-vec)])
  172. (let iter ([indices° initial-indices])
  173. (cond
  174. [indices°
  175. (cond
  176. [(pred (array-cell-ref-vec arr indices°))
  177. (cons indices°
  178. (iter (array-next-index shape-vec indices° (- rank 1))))]
  179. [else (iter (array-next-index shape-vec indices° (- rank 1)))])]
  180. [else '()]))))))