Summing out over custom axes
Starting with an Array
arr
, and a List
of axes lst
of length k
, what's a good way to sum out values over axes specified in lst
? When lst={1,2,...,m}
, this would be the same as Nest[Total,arr,m]
Example:
arr = Array[a, {2, 3, 4}];
The开发者_如何学编程n f[arr,{1}]
would have dimensions {3,4}
, f[arr,{2}]
will have dimensions {2,4}
, f[arr,{2,3}]
will have dimensions {2}
, f[arr,{1,2,3}]
will have head Plus
and dimensions {}
Does Fold[Total[#, {#2}]&, arr, lst]
do what you want?
UPDATE
How about this?
f[arr_, lst_] :=
Fold[Total[#, {#2}] &, arr, Sort[lst, Greater]]
(and a tip o' the hat to @belisarius =) )
Does this do what you want?
f[arr_, coords_] :=
With[{perm =
Ordering[
Join[coords, Complement[Range[TensorRank[arr]], coords]]]},
Total[Transpose[arr, perm], Length[coords]]
]
精彩评论