Equivalent of numpy's newaxis

Question

Numpy has a newaxis object that allows you to insert a new dimension of length 1 into an array. So after y = x[:, numpy.newaxis, :] we have y[i, 0, j] == x[i, j].
Is there something similar in Mathematica? I'd like something easier to use than ArrayReshape that requires knowing the size of the array.

swish · Accepted Answer

Like this maybe:
NewAxis /: HoldPattern[part : Part[array_, ___, NewAxis, ___]] := 
  With[{i = Rest[List @@ Unevaluated[part]]},
    Part[
      ArrayReshape[array, Fold[Insert[#1, 1, #2] &, Dimensions[array], Position[i, NewAxis]]], 
      Sequence @@ (i /. NewAxis -> All)
    ]
]

x = Array[a, {3, 3}];
y = x[[All, NewAxis]];
y[[All, 1, All]] == x
(* True *)

Michael E2 · Answer

Here's what seems a simple way with ArrayReshape, with typical Mathematica syntax. I don't really understand the OP's objection to ArrayReshape, and I think that perhaps the flexibility of it was not completely appreciated.
ClearAll[insertNewAxis];
insertNewAxis[array_?ArrayQ, level_] := 
 ArrayReshape[array, Insert[Dimensions@array, 1, level]];

Examples:
array = Table[10 i + j, {i, 3}, {j, 4}]
(*  {{11, 12, 13, 14}, {21, 22, 23, 24}, {31, 32, 33, 34}}  *)

insertNewAxis[array, 2]
% // Dimensions
(*
  {{{11, 12, 13, 14}}, {{21, 22, 23, 24}}, {{31, 32, 33, 34}}}
  {3, 1, 4}
*)

insertNewAxis[array, 1]
% // Dimensions
(*
  {{{11, 12, 13, 14}, {21, 22, 23, 24}, {31, 32, 33, 34}}}
  {1, 3, 4}
*)

insertNewAxis[array, 3]
% // Dimensions
(*
  {{{11}, {12}, {13}, {14}}, {{21}, {22}, {23}, {24}}, {{31}, {32}, {33}, {34}}}
  {3, 4, 1}
*)

kglr · Answer

ClearAll[insertNewDim]
insertNewDim[a_?ArrayQ, l_] := Map[List, a, {l - 1}]

Using Michael E2's example array:
array = {{11, 12, 13, 14}, {21, 22, 23, 24}, {31, 32, 33, 34}};

{#, ## &[#, Dimensions@#] &@insertNewDim[array, #]} & /@ {1, 2, 3} // Grid

Equivalent of numpy's newaxis

3 Answers

Add your own answers!

Ask a Question