Skip to content

Commit cd08a06

Browse files
jturner314jturner314
committed
Replace view_re_im with split_re_im
1 parent 9cbc975 commit cd08a06

File tree

3 files changed

+155
-158
lines changed

3 files changed

+155
-158
lines changed

‎src/impl_raw_views.rs‎

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
use num_complex::Complex;
12
use std::mem;
23
use std::ptr::NonNull;
34

@@ -149,6 +150,73 @@ where
149150
}
150151
}
151152

153+
impl<T, D> RawArrayView<Complex<T>, D>
154+
where
155+
D: Dimension,
156+
{
157+
/// Splits the view into views of the real and imaginary components of the
158+
/// elements.
159+
pub fn split_re_im(self) -> Complex<RawArrayView<T, D>> {
160+
// Check that the size and alignment of `Complex<T>` are as expected.
161+
// These assertions should always pass, for arbitrary `T`.
162+
assert_eq!(
163+
mem::size_of::<Complex<T>>(),
164+
mem::size_of::<T>().checked_mul(2).unwrap()
165+
);
166+
assert_eq!(mem::align_of::<Complex<T>>(), mem::align_of::<T>());
167+
168+
let dim = self.dim.clone();
169+
170+
// Double the strides. In the zero-sized element case and for axes of
171+
// length <= 1, we leave the strides as-is to avoid possible overflow.
172+
let mut strides = self.strides.clone();
173+
if mem::size_of::<T>() != 0 {
174+
for ax in 0..strides.ndim() {
175+
if dim[ax] > 1 {
176+
strides[ax] *= 2;
177+
}
178+
}
179+
}
180+
181+
let ptr_re: *mut T = self.ptr.as_ptr().cast();
182+
let ptr_im: *mut T = if self.is_empty() {
183+
// In the empty case, we can just reuse the existing pointer since
184+
// it won't be dereferenced anyway. It is not safe to offset by
185+
// one, since the allocation may be empty.
186+
ptr_re
187+
} else {
188+
// In the nonempty case, we can safely offset into the first
189+
// (complex) element.
190+
unsafe { ptr_re.add(1) }
191+
};
192+
193+
// `Complex` is `repr(C)` with only fields `re: T` and `im: T`. So, the
194+
// real components of the elements start at the same pointer, and the
195+
// imaginary components start at the pointer offset by one, with
196+
// exactly double the strides. The new, doubled strides still meet the
197+
// overflow constraints:
198+
//
199+
// - For the zero-sized element case, the strides are unchanged in
200+
// units of bytes and in units of the element type.
201+
//
202+
// - For the nonzero-sized element case:
203+
//
204+
// - In units of bytes, the strides are unchanged. The only exception
205+
// is axes of length <= 1, but those strides are irrelevant anyway.
206+
//
207+
// - Since `Complex<T>` for nonzero `T` is always at least 2 bytes,
208+
// and the original strides did not overflow in units of bytes, we
209+
// know that the new, doubled strides will not overflow in units of
210+
// `T`.
211+
unsafe {
212+
Complex {
213+
re: RawArrayView::new_(ptr_re, dim.clone(), strides.clone()),
214+
im: RawArrayView::new_(ptr_im, dim, strides),
215+
}
216+
}
217+
}
218+
}
219+
152220
impl<A, D> RawArrayViewMut<A, D>
153221
where
154222
D: Dimension,
@@ -300,3 +368,20 @@ where
300368
unsafe { RawArrayViewMut::new(ptr, self.dim, self.strides) }
301369
}
302370
}
371+
372+
impl<T, D> RawArrayViewMut<Complex<T>, D>
373+
where
374+
D: Dimension,
375+
{
376+
/// Splits the view into views of the real and imaginary components of the
377+
/// elements.
378+
pub fn split_re_im(self) -> Complex<RawArrayViewMut<T, D>> {
379+
let Complex { re, im } = self.into_raw_view().split_re_im();
380+
unsafe {
381+
Complex {
382+
re: RawArrayViewMut::new(re.ptr, re.dim, re.strides),
383+
im: RawArrayViewMut::new(im.ptr, im.dim, im.strides),
384+
}
385+
}
386+
}
387+
}

‎src/impl_views/splitting.rs‎

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88

99
use crate::imp_prelude::*;
1010
use crate::slice::MultiSliceArg;
11+
use num_complex::Complex;
1112

1213
/// Methods for read-only array views.
1314
impl<'a, A, D> ArrayView<'a, A, D>
@@ -95,6 +96,37 @@ where
9596
}
9697
}
9798

99+
impl<'a, T, D> ArrayView<'a, Complex<T>, D>
100+
where
101+
D: Dimension,
102+
{
103+
/// Splits the view into views of the real and imaginary components of the
104+
/// elements.
105+
///
106+
/// ```
107+
/// use ndarray::prelude::*;
108+
/// use num_complex::{Complex, Complex64};
109+
///
110+
/// let arr = array![
111+
/// [Complex64::new(1., 2.), Complex64::new(3., 4.)],
112+
/// [Complex64::new(5., 6.), Complex64::new(7., 8.)],
113+
/// [Complex64::new(9., 10.), Complex64::new(11., 12.)],
114+
/// ];
115+
/// let Complex { re, im } = arr.view().split_re_im();
116+
/// assert_eq!(re, array![[1., 3.], [5., 7.], [9., 11.]]);
117+
/// assert_eq!(im, array![[2., 4.], [6., 8.], [10., 12.]]);
118+
/// ```
119+
pub fn split_re_im(self) -> Complex<ArrayView<'a, T, D>> {
120+
unsafe {
121+
let Complex { re, im } = self.into_raw_view().split_re_im();
122+
Complex {
123+
re: re.deref_into_view(),
124+
im: im.deref_into_view(),
125+
}
126+
}
127+
}
128+
}
129+
98130
/// Methods for read-write array views.
99131
impl<'a, A, D> ArrayViewMut<'a, A, D>
100132
where
@@ -135,3 +167,41 @@ where
135167
info.multi_slice_move(self)
136168
}
137169
}
170+
171+
impl<'a, T, D> ArrayViewMut<'a, Complex<T>, D>
172+
where
173+
D: Dimension,
174+
{
175+
/// Splits the view into views of the real and imaginary components of the
176+
/// elements.
177+
///
178+
/// ```
179+
/// use ndarray::prelude::*;
180+
/// use num_complex::{Complex, Complex64};
181+
///
182+
/// let mut arr = array![
183+
/// [Complex64::new(1., 2.), Complex64::new(3., 4.)],
184+
/// [Complex64::new(5., 6.), Complex64::new(7., 8.)],
185+
/// [Complex64::new(9., 10.), Complex64::new(11., 12.)],
186+
/// ];
187+
///
188+
/// let Complex { mut re, mut im } = arr.view_mut().split_re_im();
189+
/// assert_eq!(re, array![[1., 3.], [5., 7.], [9., 11.]]);
190+
/// assert_eq!(im, array![[2., 4.], [6., 8.], [10., 12.]]);
191+
///
192+
/// re[[0, 1]] = 13.;
193+
/// im[[2, 0]] = 14.;
194+
///
195+
/// assert_eq!(arr[[0, 1]], Complex64::new(13., 4.));
196+
/// assert_eq!(arr[[2, 0]], Complex64::new(9., 14.));
197+
/// ```
198+
pub fn split_re_im(self) -> Complex<ArrayViewMut<'a, T, D>> {
199+
unsafe {
200+
let Complex { re, im } = self.into_raw_view_mut().split_re_im();
201+
Complex {
202+
re: re.deref_into_view_mut(),
203+
im: im.deref_into_view_mut(),
204+
}
205+
}
206+
}
207+
}

‎src/numeric/mod.rs‎

Lines changed: 0 additions & 158 deletions
Original file line numberDiff line numberDiff line change
@@ -1,159 +1 @@
1-
use crate::imp_prelude::*;
2-
use num_complex::Complex;
3-
use rawpointer::PointerExt;
4-
use std::mem;
5-
use std::ptr::NonNull;
6-
71
mod impl_numeric;
8-
9-
impl<T, S, D> ArrayBase<S, D>
10-
where
11-
S: Data<Elem = Complex<T>>,
12-
D: Dimension,
13-
{
14-
/// Returns views of the real and imaginary components of the elements.
15-
///
16-
/// ```
17-
/// use ndarray::prelude::*;
18-
/// use num_complex::{Complex, Complex64};
19-
///
20-
/// let arr = array![
21-
/// [Complex64::new(1., 2.), Complex64::new(3., 4.)],
22-
/// [Complex64::new(5., 6.), Complex64::new(7., 8.)],
23-
/// [Complex64::new(9., 10.), Complex64::new(11., 12.)],
24-
/// ];
25-
/// let Complex { re, im } = arr.view_re_im();
26-
/// assert_eq!(re, array![[1., 3.], [5., 7.], [9., 11.]]);
27-
/// assert_eq!(im, array![[2., 4.], [6., 8.], [10., 12.]]);
28-
/// ```
29-
pub fn view_re_im(&self) -> Complex<ArrayView<'_, T, D>> {
30-
debug_assert!(self.pointer_is_inbounds());
31-
32-
let dim = self.dim.clone();
33-
34-
// Double the strides. In the zero-sized element case and for axes of
35-
// length <= 1, we leave the strides as-is to avoid possible overflow.
36-
let mut strides = self.strides.clone();
37-
if mem::size_of::<T>() != 0 {
38-
for ax in 0..strides.ndim() {
39-
if dim[ax] > 1 {
40-
strides[ax] *= 2;
41-
}
42-
}
43-
}
44-
45-
let ptr_re: NonNull<T> = self.ptr.cast();
46-
let ptr_im: NonNull<T> = if self.is_empty() {
47-
// In the empty case, we can just reuse the existing pointer since
48-
// it won't be dereferenced anyway. It is not safe to offset by one
49-
// since the allocation may be empty.
50-
ptr_re
51-
} else {
52-
// In the nonempty case, we can safely offset into the first
53-
// (complex) element.
54-
unsafe { ptr_re.add(1) }
55-
};
56-
57-
// `Complex` is `repr(C)` with only fields `re: T` and `im: T`. So, the
58-
// real components of the elements start at the same pointer, and the
59-
// imaginary components start at the pointer offset by one, with
60-
// exactly double the strides. The new, doubled strides still meet the
61-
// overflow constraints:
62-
//
63-
// - For the zero-sized element case, the strides are unchanged in
64-
// units of bytes and in units of the element type.
65-
//
66-
// - For the nonzero-sized element case:
67-
//
68-
// - In units of bytes, the strides are unchanged.
69-
//
70-
// - Since `Complex<T>` for nonzero `T` is always at least 2 bytes,
71-
// and the original strides did not overflow in units of bytes, we
72-
// know that the new doubled strides will not overflow in units of
73-
// `T`.
74-
unsafe {
75-
Complex {
76-
re: ArrayView::new(ptr_re, dim.clone(), strides.clone()),
77-
im: ArrayView::new(ptr_im, dim, strides),
78-
}
79-
}
80-
}
81-
82-
/// Returns mutable views of the real and imaginary components of the elements.
83-
///
84-
/// ```
85-
/// use ndarray::prelude::*;
86-
/// use num_complex::{Complex, Complex64};
87-
///
88-
/// let mut arr = array![
89-
/// [Complex64::new(1., 2.), Complex64::new(3., 4.)],
90-
/// [Complex64::new(5., 6.), Complex64::new(7., 8.)],
91-
/// [Complex64::new(9., 10.), Complex64::new(11., 12.)],
92-
/// ];
93-
///
94-
/// let Complex { mut re, mut im } = arr.view_mut_re_im();
95-
/// assert_eq!(re, array![[1., 3.], [5., 7.], [9., 11.]]);
96-
/// assert_eq!(im, array![[2., 4.], [6., 8.], [10., 12.]]);
97-
///
98-
/// re[[0, 1]] = 13.;
99-
/// im[[2, 0]] = 14.;
100-
///
101-
/// assert_eq!(arr[[0, 1]], Complex64::new(13., 4.));
102-
/// assert_eq!(arr[[2, 0]], Complex64::new(9., 14.));
103-
/// ```
104-
pub fn view_mut_re_im(&mut self) -> Complex<ArrayViewMut<'_, T, D>>
105-
where
106-
S: DataMut,
107-
{
108-
self.ensure_unique();
109-
110-
let dim = self.dim.clone();
111-
112-
// Double the strides. In the zero-sized element case and for axes of
113-
// length <= 1, we leave the strides as-is to avoid possible overflow.
114-
let mut strides = self.strides.clone();
115-
if mem::size_of::<T>() != 0 {
116-
for ax in 0..strides.ndim() {
117-
if dim[ax] > 1 {
118-
strides[ax] *= 2;
119-
}
120-
}
121-
}
122-
123-
let ptr_re: NonNull<T> = self.ptr.cast();
124-
let ptr_im: NonNull<T> = if self.is_empty() {
125-
// In the empty case, we can just reuse the existing pointer since
126-
// it won't be dereferenced anyway. It is not safe to offset by one
127-
// since the allocation may be empty.
128-
ptr_re
129-
} else {
130-
// In the nonempty case, we can safely offset into the first
131-
// (complex) element.
132-
unsafe { ptr_re.add(1) }
133-
};
134-
135-
// `Complex` is `repr(C)` with only fields `re: T` and `im: T`. So, the
136-
// real components of the elements start at the same pointer, and the
137-
// imaginary components start at the pointer offset by one, with
138-
// exactly double the strides. The new, doubled strides still meet the
139-
// overflow constraints:
140-
//
141-
// - For the zero-sized element case, the strides are unchanged in
142-
// units of bytes and in units of the element type.
143-
//
144-
// - For the nonzero-sized element case:
145-
//
146-
// - In units of bytes, the strides are unchanged.
147-
//
148-
// - Since `Complex<T>` for nonzero `T` is always at least 2 bytes,
149-
// and the original strides did not overflow in units of bytes, we
150-
// know that the new doubled strides will not overflow in units of
151-
// `T`.
152-
unsafe {
153-
Complex {
154-
re: ArrayViewMut::new(ptr_re, dim.clone(), strides.clone()),
155-
im: ArrayViewMut::new(ptr_im, dim, strides),
156-
}
157-
}
158-
}
159-
}

0 commit comments

Comments
 (0)