Implement real/imag splitting of arrays

[ethanhs]

Picking things up from #1029

It seems the main thing is that tests were needed looking at the discussion. I've added a few to test mutation, reading the views, etc. I don't have a lot of experience testing matrix libraries so let me know if you have other suggestions for tests.

Closes #1029

[jturner314]

Will you please add a "Co-authored-by" line to the commit message for attribution of my changes? Or, if you'd prefer, you could use separate commits (with the correct authors).

Regarding testing, the following cases tend to be where most bugs occur, so, as a general rule, it's good to test them in particular:

• One or more axes which have length 0.
• One or more axes which have length 1.
• Negative strides (due to slicing with a negative step or using the .invert_axis() method).
• Zero-dimensional arrays (Array0).
• Zero-sized elements.
• Non-standard axis orders (e.g. due to .permuted_axes()).
• Arrays where the elements are not contiguous in memory (due to slicing).

Depending on the functionality being tested, it may not be necessary to have tests for all of these edge cases.

Another edge case which isn't applicable here is: owned arrays where all the strides are positive but the first visible element is not the first element in the underlying allocation (due to slicing).

[bluss]

We should probably adopt quickcheck more to be a bit more efficient and principled about testing "known tricky cases". It's a topic for another day, but would be very welcome. I'm looking at updating us to quickcheck 1.0.

[ethanhs]

Will you please add a "Co-authored-by" line to the commit message for attribution of my changes? Or, if you'd prefer, you could use separate commits (with the correct authors).

Yes of course sorry! I normally use a squash on merge flow, so I figured you'd get attribution there anyway. I've redone things with your commits so you get proper attribution.

Regarding testing, the following cases tend to be where most bugs occur...

Thank you for the recommendations, I have a few more tests that I've added locally and will push.

I have found one edge case where things currently panic:

#[test]
fn test_split_re_im_invert_axis() {
    let s_re: Vec<Complex<f64>> = vec![Complex::one(); 12];
    // negative strides
    let a = ArrayView::from_shape((2, 3, 2).strides((6, (-2isize) as usize, 1)), &s_re).unwrap();
    // panics because the stride doubled overflows usize 
    let Complex { re, im } = a.split_re_im();
}

Perhaps when doubling the strides I should use std::num::Wrapping?

[bluss]

@ethanhs Good find. All stride arithmetic needs to be done as isize.

[ethanhs]

Okay pushed the test and fix for that, as well as a few other tests. Let me know if there are other cases I should consider and I'll add tests.

[bluss]

thanks. Should be good. I pushed in an edit to the test just to add more assertions

[bluss]

Ok thank you! Good to go. I'm fine with names either split_complex, split_real_imag, split_re_im, so we can think of that until the next patch release.

[ethanhs]

Yeah I personally named it split_complex when I implemented it myself, as you are splitting the complex parts. I feel like split_re_im is both harder to type and a shortening of saying "split into real and imaginary parts", which is also accurate, but I felt like split_complex is nicer when reading. But I don't have strong opinions so I just left what @jturner314 already did :)

[bluss]

I would vote split_complex here I think. I see the consistency of using re/im seeing that's what num_complex does, but split complex seems easier to understand.

I have (messy) patches for the related but a bit different [T; 4] -> new dimension and have a working name of "expand" right now, in the sense of expanding the element into new axis. But since it works with axes the name isn't really related. Either way, in that patch it's likely that complex will show up again, allowing one to expand Complex into a new axis just like [T; 2] is allowed to.

Expand unfortunately is not allowed with regular owned arrays, because it's not legal to change the alignment/element type of an allocation (we could save it in metadata, but we don't have a location for that right now).

[jturner314]

I would vote split_complex here I think. I see the consistency of using re/im seeing that's what num_complex does, but split complex seems easier to understand.

Fwiw, I chose the name split_re_im for a few reasons:

1. Consistency with the names of the struct members.
2. It makes it clear that the two parts are the real and imaginary components, instead of the magnitude and angle.
3. It makes it clear that the first item in the resulting tuple is the real part and the second item is the imaginary part.

That said, I think split_complex is clear enough.

I have (messy) patches for the related but a bit different [T; 4] -> new dimension and have a working name of "expand" right now, in the sense of expanding the element into new axis. But since it works with axes the name isn't really related. Either way, in that patch it's likely that complex will show up again, allowing one to expand Complex into a new axis just like [T; 2] is allowed to.

That would be useful functionality, and that's a good point regarding the complex case. "expand" is a reasonable term for this. My first thought was something related to "flattening" the types, but that could easily be confused with flattening/merging axes. Another thought would be a name like elems_as_axis. The case of nested arrays [[T; M]; N] may also be of interest, although the user could deal with that by using multiple "expand" calls.

[bluss]

I made some experiments and I think I went overboard with the possibilities, including allowing expand u128 => 16 x u8 or u128 => 4 x u32 (When these are signed, then it's harder to know that the expand should really mean, even if it's easy to implement).

Multiple expand calls might as well be used, at least in the shape of the API where we ask the user to specify where the new axis is going to be, since that's actually arbitrary. Even though for common cases, it would of course be possible to guess if it should be the first or last, but it's technically arbitrary.

[bluss]

For points (2) and (3) that makes sense, but I think convention is quite strong, those should be the almost obvious choices.

[ethanhs]

I agree the convention is pretty strong, I would in fact be quite surprised if split_complex returned anything but re and im given those are the fields on num_complex::Complex type.

Also, after this was merged I was looking at how pytorch does this (See accessing real and imag).

>>> y.real
tensor([ 0.6125, -0.3773, -0.0861])
>>> y.imag
tensor([-0.1681,  1.3487, -0.7981])

>>> y.real.mul_(2)
tensor([ 1.2250, -0.7546, -0.1722])
>>> y
tensor([ 1.2250-0.1681j, -0.7546+1.3487j, -0.1722-0.7981j])
>>> y.real.stride()
(2,)

It might be nice to provide these as utilities on top of whichever split_complex/split_re_im as well.

I'm happy to make a PR with the name change and add these.

[bluss]

All of those conveniences are already there by virtue of y.re, y.im being array views, I think. We don't have mul_(2) but we have y.re * 2.0 instead?

Let's go for the name change to split_complex

[bluss]

oh, is that about more convenient access to real/imag? It was last discussed here #1029 (comment)

[jturner314]

Oh, I just took another look at the implementation, and realized that split_complex returns Complex<(array type)>. For some reason, I had in my head that it was returning a tuple. With a return type of Complex<(array type)>, I don't see any possible confusion, and I agree that split_complex is the better name.

I made some experiments and I think I went overboard with the possibilities, including allowing expand u128 => 16 x u8 or u128 => 4 x u32 (When these are signed, then it's harder to know that the expand should really mean, even if it's easy to implement).

Yeah, I don't think we should provide those conversions directly. Instead, I'd suggest adding a cast_elems (or "cast_ptr" or just "cast"?) method to raw views so that users could convert e.g. RawArrayView<u128, D> to RawArrayView<[u8; 16], D>. (A casting method is useful for other things, too. It would require that the element sizes be the same, but otherwise impose no restrictions. For non-raw views, we could allow more restrictive but guaranteed safe casting via the traits in the bytemuck crate.) The array of arrays could then be used with the "expand" method.

[bluss]

We already have .cast() for raw views fortunately 🙂 (#734)

I guess I'm thinking about if we eventually want some conveniences to make it easy in safe code to translate between different equivalent representations, thinking of u128 => [u32; 4] but also T <=> Wrapping<T>, Cell<T> => T and so on. Yes, these need to keep the allowed direction(s) in mind to be safe.

It's strange that we have gone so far without addressing ergonomics of wrapping computations yet, but maybe easy translation <=> Wrapping elements is a way to go?

[jturner314]

We already have .cast() for raw views fortunately slightly_smiling_face (#734)

Oh, yeah, I forgot about that. :)

I guess I'm thinking about if we eventually want some conveniences to make it easy in safe code to translate between different equivalent representations, thinking of u128 => [u32; 4] but also T <=> Wrapping<T>, Cell<T> => T and so on. Yes, these need to keep the allowed direction(s) in mind to be safe.

Yeah, I think the best solution is a trait like the EqualRepresentation trait in your expand branch. I would really like to see the Rust community decide on a standard trait for that purpose. I've seen some pre-RFC discussions regarding this type of thing (usually called "safe transmute"), but I'm not sure of their status. AFAIK, the most common solution at the moment is the bytemuck crate, but since it doesn't consider the direction of the casts, bool => u8 is not supported.

It's strange that we have gone so far without addressing ergonomics of wrapping computations yet, but maybe easy translation <=> Wrapping elements is a way to go?

Here are some options:

1. Allow easy conversion to Wrapping.
2. Implement the various Wrapping* traits defined in num_traits::ops::wrapping. Unfortunately, due to the way those traits are defined, we could only do so for &Array<A, D> (op) &Array<A, D> -> Array<A, D>.
3. Add wrapping_* methods to ArrayBase with num_traits::ops::wrapping::Wrapping* bounds on the element types. I don't think that would allow for the same flexibility as our existing (non-wrapping) operator implementations.
4. Create our own Wrapping* traits which are more flexible than the ones provided by num-traits, and implement them for ArrayBase and &ArrayBase. For these new traits, we could provide blanket implementations for types which implement the Wrapping* traits from num-traits.

Option 4 would provide the most convenience for the user, but I like the simplicity of option 1.