Is this really likely to be faster than using ndindex here?

To be honest, most of this function could probably be replaced with a carefully constructed nditer that does the transposes for us, but that's probably not going to increase performance much anyway

Is this assumption (__array_prepare__(x).shape == x.shape) supposed to be valid? How do we deal with matrix violating it elsewhere?

Note that assert shouldn't be used at all, use np.testing.assert_ instead (applies to multiple lines below as well)

In my defense, those other lines were not part of this patch. Do you want me to fix all cases of assert in this file anyway?

Or should this be self.assertIsInstance?

Changed to assert_(isinstance(...)), because that's what other tests do. Switching to self.assertIsInstance can be discussed elsewhere

In my defense, those other lines were not part of this patch. Do you want me to fix all cases of assert in this file anyway?

if you could, that'd be useful

self.assertIsInstance seems fine to use if you prefer that, it has a sane implementation in both unittest and unittest2.

@rgommers: I've patched the other ones in that file

Can someone comment on whether this is actually correct, or if it can just be omitted altogether

__array_prepare__ seems to always do something screwy for builtin array subclasses. It's not clear to me what the contract entails, and whether I'm violating it, or np.matrix and np.ma.MaskedArray are

@mhvk may know

I think MaskedConstant.__array_prepare__ was broken - fixed in #8508. I think a corollary of this is that np.ma.apply_along_axis can basically defer now.

@mhvk may know

I recognize that this works to update the view inplace, but it still makes me a little nervous. It would be more direct and obviously correct to do the transposing afterwards.

This would change the memory layout of the result array. Transposing afterwards would ensure that array is written with contiguous writes, which could result in a marginal performance improvement (at least for this function).

Yep, my goal here was to leave the memory layout of the result array unaffected when compared to the old version.

In particular, the old code would always give a contiguous output, something that may have been relied upon

I think we reserve the right to change the memory layout of arrays returned by NumPy functions unless it is documented as part of the API. Users who truly rely on contiguous output should be using np.ascontiguousarray.

I guess transposing afterwards would fix the above issue. If you don't think a contiguous output is beneficial, then I'll move the transpose

We have done it before (e.g. even indexing I did it). If it is obvious/possible, then "the same as before" should be preferred most of the time I would say.

I'd argue that the correctness is more obvious this way around, because otherwise the inverse transpose order is needed, which is more troublesome to generate

I mostly like transposing afterwards more. I don't care so much about the memory layout.

Actually I think you're right, it is simpler - the transpose order becomes dims_out[:axis] + dims_out[-res.ndim:] + dims_out[axis+1:] (more complicated than that), which better conveys the insertion of the added axes

Should the transpose happen before or after the __array_wrap__?

Ok, I'm sold on your suggestion, but for a different reason: when res.__array_prepare__ is called, the axes of res (a single output) and the final output are paired as if they were broadcasted, which seems desirable for any kind of axis metadata copying

This assumes that transpose always returns a view (and not sometimes a copy). is this acceptable?

(edit: updated mailing list url, which was dead. Not sure it points to the right thing any longer. https://mail.python.org/pipermail/numpy-discussion/2013-June/066822.html might be what I meant)

Actually, only the output transpose does, which only matters if we use __array_prepare__

As @seberg writes in that mailing list discussion, transpose on base numpy arrays always returns a view, never a copy.

But this might indeed break in surprising ways for ndarray subclasses

@shoyer: Would an acceptable compromise be to not call __array_prepare__ if transpose is a copy?

while we're at it, also check that axis is not still negative?

I like the idea of killing this altogether, and delegating parsing axis to moveaxis

Alarmingly, I can't even find a precedent where numpy checks for this anywhere else within shapebase

@eric-wieser - when you add the changelog, could you also move this check up to before axis is possibly changed and check for too-negative as well? That way, the error message gives the actual input rather than something that's already changed. I.e.,

if axis < -nd or axis >= nd:
    raise ...

if axis < 0:
    axis += nd

@mhvk: Done

When I looked at this first, I thought you might consider using the (relatively new) np.moveaxis, i.e.,

inarr_view = moveaxis(inarr, axis, -1)

This will not be any faster (as it just sets up a transpose, like here), but is perhaps clearer. As a side benefit, you can remove above the checks on the validity of axis, as moveaxis does those anyway.

However, looking further down, it is not obvious moveaxis would work there, since res.ndim could in principle be >1-d (for which your code adds support, which I think is very nice!). But one could steal a bit from the moveaxis code and write here

in_permute = [n for n in range(nd) if n != axis] + [axis]
inarr_view = transpose(arr, in_permute)

Oh, good point. moveaxis works with multiple axes too, so should work in both cases.

For restoring axes, you could maybe write something like:

moveaxis(pre_arr, [i + nd - 1 for i in range(res.ndim)],
         [i + axis for i in range(res.ndim)])

I think this is a slight improvement in clarity over building up the permutation axes for transpose directly.

I think moveaxis should learn a shorthand for dest such that this would work:

moveaxis(pre_arr, [i + nd - 1 for i in range(res.ndim)], axis)

Ie, if dest is a scalar, move all the source axes to that location, in the order they were passed in src

@shoyer: Based on @mhvk's revised opinion, are you happy for this to remain as it is? Any further pain-points blocking this being merged?

Really like using the iterator here. Much cleaner!

Here you can just use the shape of inarr_view, i.e.,

pre_shape = inarr_view.shape + res.shape

Nice!

Maybe a bit faster (no list->ndarray->list) and clearer (to me at least):

if res.ndim > 0:
   pre_permute = pre_dims[:axis] + pre_dims[-res.ndim:] + pre_dims[axis+red.ndim:]

(you'd need to rename pre_dims to pre_permute if you want the transpose at the end even for res.ndim=0)

Yeah, that line was inspired by a talk extolling the virtues of std::rotate.

Had that before, but without the if - as you correctly spot, it fails when res.ndim == 0. Not a huge fan of special casing that.

Either way, I reckon moveaxis will do the job here too.

Or how about just

pre_permute = pre_dims[:axis] + pre_dims[nd-res.ndim:] + pre_dims[axis+red.ndim:]

Adding the nd there fixes the error when res.ndim= 0

For speed for the more common case (and required with my comment above):

return pre_arr if res.ndim == 0 else transpose(pre_arr, pre_permute)

same as above

Just on one line?

assert_array_equal(result, expected)

same

Beginning to think this would be a fair bit slower, due to the error checking in moveaxis, and isn't much clearer.

Ooh, how about np.r_[:axis, axis+1:nd, axis]

Reading over these tests one last time, it occurs to me that we should be more comprehensive for testing insertion of new axes.

Right now you only test inserting at the start, but that doesn't really exercise the permutation logic in a comprehensive way. Let's also test inserting in the middle and at the end. Also, diag is axes order invariant, so none of the tests verify that new axes are inserted in the right order.

Maybe outer product with its reverse would be better than diag then?

I think the outer product is still invariant to transposes. Something like x[:, np.newaxis] - x[np.newaxis, :] would work, though.

On re-reading: yes, outer product with its reverse, would be fine

You didn't miss it the first time, I did indeed suggest a plain outer product beforehand, and realized my mistake. This is now done too.

A cleaner way to check this sort of thing is to construct the desired result array and just call assert_array_equal (e.g., np.stack([np.diag(a[:,i]) for i in range(3)], axis=-1).view(cls)). Otherwise it's easy to overlook one part of the equality check (e.g., shape, in this case).

@shoyer: I was hoping I could be done with this, but you are definitely right. I'll fix those things up soon