@@ -44,6 +44,31 @@ def _nan_reduce_all(a, c_func, nan_val):
44
44
imag = imag .value
45
45
return real if imag == 0 else real + imag * 1j
46
46
47
+ def _FNSD (dim , dims ):
48
+ if dim >= 0 :
49
+ return int (dim )
50
+
51
+ fnsd = 0
52
+ for i , d in enumerate (dims ):
53
+ if d > 1 :
54
+ fnsd = i
55
+ break
56
+ return int (fnsd )
57
+
58
+ def _rbk_dim (keys , vals , dim , c_func ):
59
+ keys_out = Array ()
60
+ vals_out = Array ()
61
+ rdim = _FNSD (dim , vals .dims ())
62
+ safe_call (c_func (c_pointer (keys_out .arr ), c_pointer (vals_out .arr ), keys .arr , vals .arr , c_int_t (rdim )))
63
+ return keys_out , vals_out
64
+
65
+ def _nan_rbk_dim (a , dim , c_func , nan_val ):
66
+ keys_out = Array ()
67
+ vals_out = Array ()
68
+ rdim = _FNSD (dim , vals .dims ())
69
+ safe_call (c_func (c_pointer (keys_out .arr ), c_pointer (vals_out .arr ), keys .arr , vals .arr , c_int_t (rdim ), c_double_t (nan_val )))
70
+ return keys_out , vals_out
71
+
47
72
def sum (a , dim = None , nan_val = None ):
48
73
"""
49
74
Calculate the sum of all the elements along a specified dimension.
@@ -74,6 +99,34 @@ def sum(a, dim=None, nan_val=None):
74
99
else :
75
100
return _reduce_all (a , backend .get ().af_sum_all )
76
101
102
+
103
+ def sumByKey (keys , vals , dim = - 1 , nan_val = None ):
104
+ """
105
+ Calculate the sum of elements along a specified dimension according to a key.
106
+
107
+ Parameters
108
+ ----------
109
+ keys : af.Array
110
+ One dimensional arrayfire array with reduction keys.
111
+ vals : af.Array
112
+ Multi dimensional arrayfire array that will be reduced.
113
+ dim: optional: int. default: -1
114
+ Dimension along which the sum will occur.
115
+ nan_val: optional: scalar. default: None
116
+ The value that replaces NaN in the array
117
+
118
+ Returns
119
+ -------
120
+ keys: af.Array or scalar number
121
+ The reduced keys of all elements in `vals` along dimension `dim`.
122
+ values: af.Array or scalar number
123
+ The sum of all elements in `vals` along dimension `dim` according to keys
124
+ """
125
+ if (nan_val is not None ):
126
+ return _nan_rbk_dim (keys , vals , dim , backend .get ().af_sum_by_key_nan , nan_val )
127
+ else :
128
+ return _rbk_dim (keys , vals , dim , backend .get ().af_sum_by_key )
129
+
77
130
def product (a , dim = None , nan_val = None ):
78
131
"""
79
132
Calculate the product of all the elements along a specified dimension.
@@ -104,6 +157,33 @@ def product(a, dim=None, nan_val=None):
104
157
else :
105
158
return _reduce_all (a , backend .get ().af_product_all )
106
159
160
+ def productByKey (keys , vals , dim = - 1 , nan_val = None ):
161
+ """
162
+ Calculate the product of elements along a specified dimension according to a key.
163
+
164
+ Parameters
165
+ ----------
166
+ keys : af.Array
167
+ One dimensional arrayfire array with reduction keys.
168
+ vals : af.Array
169
+ Multi dimensional arrayfire array that will be reduced.
170
+ dim: optional: int. default: -1
171
+ Dimension along which the product will occur.
172
+ nan_val: optional: scalar. default: None
173
+ The value that replaces NaN in the array
174
+
175
+ Returns
176
+ -------
177
+ keys: af.Array or scalar number
178
+ The reduced keys of all elements in `vals` along dimension `dim`.
179
+ values: af.Array or scalar number
180
+ The product of all elements in `vals` along dimension `dim` according to keys
181
+ """
182
+ if (nan_val is not None ):
183
+ return _nan_rbk_dim (keys , vals , dim , backend .get ().af_product_by_key_nan , nan_val )
184
+ else :
185
+ return _rbk_dim (keys , vals , dim , backend .get ().af_product_by_key )
186
+
107
187
def min (a , dim = None ):
108
188
"""
109
189
Find the minimum value of all the elements along a specified dimension.
@@ -126,6 +206,28 @@ def min(a, dim=None):
126
206
else :
127
207
return _reduce_all (a , backend .get ().af_min_all )
128
208
209
+ def minByKey (keys , vals , dim = - 1 ):
210
+ """
211
+ Calculate the min of elements along a specified dimension according to a key.
212
+
213
+ Parameters
214
+ ----------
215
+ keys : af.Array
216
+ One dimensional arrayfire array with reduction keys.
217
+ vals : af.Array
218
+ Multi dimensional arrayfire array that will be reduced.
219
+ dim: optional: int. default: -1
220
+ Dimension along which the min will occur.
221
+
222
+ Returns
223
+ -------
224
+ keys: af.Array or scalar number
225
+ The reduced keys of all elements in `vals` along dimension `dim`.
226
+ values: af.Array or scalar number
227
+ The min of all elements in `vals` along dimension `dim` according to keys
228
+ """
229
+ return _rbk_dim (keys , vals , dim , backend .get ().af_min_by_key )
230
+
129
231
def max (a , dim = None ):
130
232
"""
131
233
Find the maximum value of all the elements along a specified dimension.
@@ -148,6 +250,28 @@ def max(a, dim=None):
148
250
else :
149
251
return _reduce_all (a , backend .get ().af_max_all )
150
252
253
+ def maxByKey (keys , vals , dim = - 1 ):
254
+ """
255
+ Calculate the max of elements along a specified dimension according to a key.
256
+
257
+ Parameters
258
+ ----------
259
+ keys : af.Array
260
+ One dimensional arrayfire array with reduction keys.
261
+ vals : af.Array
262
+ Multi dimensional arrayfire array that will be reduced.
263
+ dim: optional: int. default: -1
264
+ Dimension along which the max will occur.
265
+
266
+ Returns
267
+ -------
268
+ keys: af.Array or scalar number
269
+ The reduced keys of all elements in `vals` along dimension `dim`.
270
+ values: af.Array or scalar number
271
+ The max of all elements in `vals` along dimension `dim` according to keys.
272
+ """
273
+ return _rbk_dim (keys , vals , dim , backend .get ().af_max_by_key )
274
+
151
275
def all_true (a , dim = None ):
152
276
"""
153
277
Check if all the elements along a specified dimension are true.
@@ -170,6 +294,28 @@ def all_true(a, dim=None):
170
294
else :
171
295
return _reduce_all (a , backend .get ().af_all_true_all )
172
296
297
+ def allTrueByKey (keys , vals , dim = - 1 ):
298
+ """
299
+ Calculate if all elements are true along a specified dimension according to a key.
300
+
301
+ Parameters
302
+ ----------
303
+ keys : af.Array
304
+ One dimensional arrayfire array with reduction keys.
305
+ vals : af.Array
306
+ Multi dimensional arrayfire array that will be reduced.
307
+ dim: optional: int. default: -1
308
+ Dimension along which the all true check will occur.
309
+
310
+ Returns
311
+ -------
312
+ keys: af.Array or scalar number
313
+ The reduced keys of all true check in `vals` along dimension `dim`.
314
+ values: af.Array or scalar number
315
+ Booleans denoting if all elements are true in `vals` along dimension `dim` according to keys
316
+ """
317
+ return _rbk_dim (keys , vals , dim , backend .get ().af_all_true_by_key )
318
+
173
319
def any_true (a , dim = None ):
174
320
"""
175
321
Check if any the elements along a specified dimension are true.
@@ -192,6 +338,28 @@ def any_true(a, dim=None):
192
338
else :
193
339
return _reduce_all (a , backend .get ().af_any_true_all )
194
340
341
+ def anyTrueByKey (keys , vals , dim = - 1 ):
342
+ """
343
+ Calculate if any elements are true along a specified dimension according to a key.
344
+
345
+ Parameters
346
+ ----------
347
+ keys : af.Array
348
+ One dimensional arrayfire array with reduction keys.
349
+ vals : af.Array
350
+ Multi dimensional arrayfire array that will be reduced.
351
+ dim: optional: int. default: -1
352
+ Dimension along which the any true check will occur.
353
+
354
+ Returns
355
+ -------
356
+ keys: af.Array or scalar number
357
+ The reduced keys of any true check in `vals` along dimension `dim`.
358
+ values: af.Array or scalar number
359
+ Booleans denoting if any elements are true in `vals` along dimension `dim` according to keys.
360
+ """
361
+ return _rbk_dim (keys , vals , dim , backend .get ().af_any_true_by_key )
362
+
195
363
def count (a , dim = None ):
196
364
"""
197
365
Count the number of non zero elements in an array along a specified dimension.
@@ -214,6 +382,28 @@ def count(a, dim=None):
214
382
else :
215
383
return _reduce_all (a , backend .get ().af_count_all )
216
384
385
+ def countByKey (keys , vals , dim = - 1 ):
386
+ """
387
+ Counts non-zero elements along a specified dimension according to a key.
388
+
389
+ Parameters
390
+ ----------
391
+ keys : af.Array
392
+ One dimensional arrayfire array with reduction keys.
393
+ vals : af.Array
394
+ Multi dimensional arrayfire array that will be reduced.
395
+ dim: optional: int. default: -1
396
+ Dimension along which to count elements.
397
+
398
+ Returns
399
+ -------
400
+ keys: af.Array or scalar number
401
+ The reduced keys of count in `vals` along dimension `dim`.
402
+ values: af.Array or scalar number
403
+ Count of non-zero elements in `vals` along dimension `dim` according to keys.
404
+ """
405
+ return _rbk_dim (keys , vals , dim , backend .get ().af_count_by_key )
406
+
217
407
def imin (a , dim = None ):
218
408
"""
219
409
Find the value and location of the minimum value along a specified dimension
0 commit comments