开发者

Summing out over custom axes

开发者 https://www.devze.com 2023-01-25 04:46 出处:网络
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]

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]]
  ]
0

精彩评论

暂无评论...
验证码 换一张
取 消

关注公众号