-
Notifications
You must be signed in to change notification settings - Fork 0
/
ndapi.tcl
361 lines (335 loc) · 9.13 KB
/
ndapi.tcl
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
# ndapi.tcl
################################################################################
# Core API for manipulating ND-lists
# Copyright (C) 2024 Alex Baker, ambaker1@mtu.edu
# All rights reserved.
# See the file "LICENSE" in the top level directory for information on usage,
# redistribution, and for a DISCLAIMER OF ALL WARRANTIES.
################################################################################
# GetNDims --
#
# Get dimensionality from ND string (uses regex pattern).
# Either a single digit or with a "D" after.
# e.g. "0" or "0D", or "3" or "3d"
# Returns error if invalid syntax
#
# Syntax:
# GetNDims $nd
#
# Arguments:
# nd Number of dimensions (e.g. 1D, 2D, etc.)
proc ::ndlist::GetNDims {nd} {
if {![IsNDType $nd]} {
return -code error "invalid ND syntax"
}
string trimright $nd {dD}
}
# IsNDType --
#
# Returns whether an input is an ND string
#
# Syntax:
# IsNDType $arg
#
# Arguments:
# arg: Argument to check
proc ::ndlist::IsNDType {arg} {
regexp {^(0|[1-9]\d*)[dD]?$} $arg
}
# ValidateAxis --
#
# Validates axis input
#
# Syntax:
# ValidateAxis $ndims $axis
#
# Arguments:
# ndims Number of dimensions (Inf for arbitrary dimensions)
# axis Axis integer (must be 0-(N-1))
proc ::ndlist::ValidateAxis {ndims axis} {
if {![string is integer -strict $axis]} {
return -code error "expected integer, but got \"$axis\""
}
if {$axis < 0 || $axis >= $ndims} {
return -code error "axis out of range"
}
}
# GetShape --
#
# Private procedure to get list of dimensions of an ND-list along first index
# Returns error if there is a null dimension along a non-zero axis.
#
# Syntax:
# GetShape $ndims $ndlist
#
# Arguments:
# ndims Number of dimensions
# ndlist ND-list to get dimensions from
proc ::ndlist::GetShape {ndims ndlist} {
# Null case
if {[llength $ndlist] == 0} {
return [lrepeat $ndims 0]
}
# Get list of dimensions (along first index)
set dims ""
foreach axis [range $ndims] {
if {[llength $ndlist] == 0} {
return -code error "null dimension along non-zero axis"
}
lappend dims [llength $ndlist]
set ndlist [lindex $ndlist 0]
}
return $dims
}
# IsShape --
#
# Verify that the ND-list is of the specified shape
#
# Syntax:
# IsShape $ndlist $n ...
#
# Arguments:
# ndlist ND-list to check
# n ... Shape of ND-list
proc ::ndlist::IsShape {ndlist args} {
# Scalar base case
if {[llength $args] == 0} {
return 1
}
# Interpret input
set args [lassign $args n]
# Vector base case
if {[llength $ndlist] != $n} {
return 0
}
# Recursion
foreach ndrow $ndlist {
if {![IsShape $ndrow {*}$args]} {
return 0
}
}
return 1
}
# GetMaxShape --
#
# Get maximum dimensions of multiple ND-lists (for expanding)
#
# Syntax:
# GetMaxShape $ndims $arg ...
#
# Arguments:
# ndims Number of dimensions (e.g. 1D, 2D, etc.)
# arg ... ND-lists to get max shape from
proc ::ndlist::GetMaxShape {ndims args} {
set shapes [lmap ndlist $args {GetShape $ndims $ndlist}]
lmap dims [transpose $shapes] {max $dims}
}
# ParseIndices --
#
# Loop through index inputs - returning required information for getting/setting
# Returns index arguments - paired list of index type and index list.
#
# Syntax:
# ParseIndices $dims $index ...
#
# Arguments:
# dims Shape to index into
# index ... Index inputs (e.g. :, {0 3}, 0:10, end*)
proc ::ndlist::ParseIndices {dims args} {
set iArgs ""; # paired list of index type and index list (meaning varies)
foreach dim $dims index $args {
lappend iArgs {*}[ParseIndex $dim $index]
}
return $iArgs
}
# ParseIndex --
#
# Used for parsing index input (i.e. list of indices, range 0:10, etc)
# Returns index type and corresponding values.
#
# Syntax:
# lassign [ParseIndex $n $index] iType iList
#
# Arguments:
# n Size of list
# index Index input (e.g. :, {0 3}, 0:10, end*)
#
# Returns:
# iType Type of index (A, R, L, or S)
# iList List of indices corresponding with type
# A: Empty
# R: Range start and stop
# L: List of indices
# S: Single index (flattens list)
proc ::ndlist::ParseIndex {n index} {
# Check length of input
if {[llength $index] != 1} {
# List of indices (user entered)
return [list L [lmap index $index {Index2Integer $n $index}]]
}
# All index notation
if {$index eq {:}} {
return [list A ""]
}
# Single index notation
if {[string index $index end] eq {*}} {
# Single index notation (flatten along this dimension)
return [list S [Index2Integer $n [string range $index 0 end-1]]]
}
# Single index, not range notation
if {![string match *:* $index]} {
return [list L [Index2Integer $n $index]]
}
# Range index notation
set parts [split $index :]
# Simple range case ($start:$stop)
if {[llength $parts] == 2} {
lassign $parts start stop
set start [Index2Integer $n $start]
set stop [Index2Integer $n $stop]
if {$start == 0 && $stop == ($n - 1)} {
# 0:end case
return [list A ""]
}
# Normal range
return [list R [list $start $stop]]
}
# Skipped range case ($start:$step:$stop)
if {[llength $parts] == 3} {
lassign $parts start step stop
set start [Index2Integer $n $start]
set stop [Index2Integer $n $stop]
if {![string is integer -strict $step]} {
return -code error "expected integer but got \"$step\""
}
# Special case for forward range with step of 1
if {$step == 1 && $start <= $stop} {
if {$start == 0 && $stop == ($n - 1)} {
# 0:1:end case
return [list A ""]
}
# Normal range
return [list R [list $start $stop]]
}
# Special case for reverse range with step of -1
if {$step == -1 && $start >= $stop} {
return [list R [list $start $stop]]
}
# Normal case
return [list L [range $start $stop $step]]
}
return -code error "invalid range index notation: should be \
\"start:stop\" or \"start:step:stop\""
}
# Index2Integer --
#
# Private function, converts end+-integer index format into integer
# Negative indices get converted, such that -1 is end, -2 is end-1, etc.
#
# Syntax:
# Index2Integer $n $index
#
# Arguments:
# n: Length of list to index
# index: Index notation (integer?[+-]integer? or end?[+-]integer?)
proc ::ndlist::Index2Integer {n index} {
# Default case (skip regexp, much faster)
if {[string is integer -strict $index]} {
set i $index
} else {
# Check if index is valid format
set match [regexp -inline {^(end|[+-]?[0-9]+)([+-][0-9]+)?$} $index]
if {[llength $match] == 0} {
return -code error "bad index \"$index\": must be\
integer?\[+-\]integer? or end?\[+-\]integer?"
}
# Convert end to n-1 if needed
set base [lindex $match 1]
if {$base eq {end}} {
set base [expr {$n - 1}]
}
# Handle offset
set offset [lindex $match 2]
if {$offset eq {}} {
set i $base
} else {
set i [expr {$base + $offset}]
}
}
# Handle negative index (from end)
if {$i < 0} {
set i [expr {$i % $n}]
}
# Check if in range
if {$i >= $n} {
return -code error "index out of range"
}
return $i
}
# GetIndexShape --
#
# Get shape of indexed range.
#
# Syntax:
# GetIndexShape $dims $iType $iList ...
#
# Arguments:
# dims Shape to index into
# iType ... Index types
# iList ... Index lists
proc ::ndlist::GetIndexShape {dims args} {
concat {*}[lmap dim $dims {iType iList} $args {
GetIndexDim $dim $iType $iList
}]
}
# GetIndexDim --
#
# Get the size of the indexed range.
#
# Syntax:
# GetIndexDim $n $iType $iList
#
# Arguments:
# n Length of list.
# iType Type of index (returned from ParseIndex).
# iList List corresponding with index type.
proc ::ndlist::GetIndexDim {n iType iList} {
switch $iType {
A { # All indices
return $n
}
R { # Range of indices
lassign $iList start stop
if {$start <= $stop} {
return [expr {$stop - $start + 1}]
} else {
return [expr {$start - $stop + 1}]
}
}
L { # List of indices
return [llength $iList]
}
S { # Single index
return
}
}
}
# UnravelIndex --
#
# Unravel a flat index to its coordinates
#
# Syntax:
# UnravelIndex $i $n ...
#
# Arguments:
# i Flat index into ND list
# n ... Shape of ND list
proc ::ndlist::UnravelIndex {i n args} {
# Base case
if {[llength $args] == 0} {
return [expr {$i % $n}]
}
# Recursion
set N [product $args]
concat [expr {$i / $N}] [UnravelIndex [expr {$i % $N}] {*}$args]
}