Skip to content

Commit 09884fc

Browse files
blussbluss
authored
Merge pull request #570 from jturner314/slice-new-axis
Add support for inserting new axes while slicing
2 parents 9b23ba5 + 7506f90 commit 09884fc

File tree

12 files changed

+828
-353
lines changed

12 files changed

+828
-353
lines changed

‎blas-tests/tests/oper.rs‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,8 @@ extern crate num_traits;
66
use ndarray::linalg::general_mat_mul;
77
use ndarray::linalg::general_mat_vec_mul;
88
use ndarray::prelude::*;
9+
use ndarray::{AxisSliceInfo, Ix, Ixs};
910
use ndarray::{Data, LinalgScalar};
10-
use ndarray::{Ix, Ixs, SliceInfo, SliceOrIndex};
1111

1212
use approx::{assert_abs_diff_eq, assert_relative_eq};
1313
use defmac::defmac;
@@ -420,19 +420,19 @@ fn scaled_add_3() {
420420
let mut answer = a.clone();
421421
let cdim = if n == 1 { vec![q] } else { vec![n, q] };
422422
let cslice = if n == 1 {
423-
vec![SliceOrIndex::from(..).step_by(s2)]
423+
vec![AxisSliceInfo::from(..).step_by(s2)]
424424
} else {
425425
vec![
426-
SliceOrIndex::from(..).step_by(s1),
427-
SliceOrIndex::from(..).step_by(s2),
426+
AxisSliceInfo::from(..).step_by(s1),
427+
AxisSliceInfo::from(..).step_by(s2),
428428
]
429429
};
430430

431431
let c = range_mat64(n, q).into_shape(cdim).unwrap();
432432

433433
{
434434
let mut av = a.slice_mut(s![..;s1, ..;s2]);
435-
let c = c.slice(SliceInfo::<_, IxDyn>::new(cslice).unwrap().as_ref());
435+
let c = c.slice(&*cslice);
436436

437437
let mut answerv = answer.slice_mut(s![..;s1, ..;s2]);
438438
answerv += &(beta * &c);

‎src/dimension/dimension_trait.rs‎

Lines changed: 8 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -13,13 +13,14 @@ use alloc::vec::Vec;
1313

1414
use super::axes_of;
1515
use super::conversion::Convert;
16+
use super::ops::DimAdd;
1617
use super::{stride_offset, stride_offset_checked};
1718
use crate::itertools::{enumerate, zip};
1819
use crate::{Axis, DimMax};
1920
use crate::IntoDimension;
2021
use crate::RemoveAxis;
2122
use crate::{ArrayView1, ArrayViewMut1};
22-
use crate::{Dim, Ix, Ix0, Ix1, Ix2, Ix3, Ix4, Ix5, Ix6, IxDyn, IxDynImpl, Ixs, SliceOrIndex};
23+
use crate::{Dim, Ix, Ix0, Ix1, Ix2, Ix3, Ix4, Ix5, Ix6, IxDyn, IxDynImpl, Ixs};
2324

2425
/// Array shape and index trait.
2526
///
@@ -51,26 +52,17 @@ pub trait Dimension:
5152
+ DimMax<IxDyn, Output=IxDyn>
5253
+ DimMax<<Self as Dimension>::Smaller, Output=Self>
5354
+ DimMax<<Self as Dimension>::Larger, Output=<Self as Dimension>::Larger>
55+
+ DimAdd<Self>
56+
+ DimAdd<<Self as Dimension>::Smaller>
57+
+ DimAdd<<Self as Dimension>::Larger>
58+
+ DimAdd<Ix0, Output = Self>
59+
+ DimAdd<Ix1, Output = <Self as Dimension>::Larger>
60+
+ DimAdd<IxDyn, Output = IxDyn>
5461
{
5562
/// For fixed-size dimension representations (e.g. `Ix2`), this should be
5663
/// `Some(ndim)`, and for variable-size dimension representations (e.g.
5764
/// `IxDyn`), this should be `None`.
5865
const NDIM: Option<usize>;
59-
/// `SliceArg` is the type which is used to specify slicing for this
60-
/// dimension.
61-
///
62-
/// For the fixed size dimensions it is a fixed size array of the correct
63-
/// size, which you pass by reference. For the dynamic dimension it is
64-
/// a slice.
65-
///
66-
/// - For `Ix1`: `[SliceOrIndex; 1]`
67-
/// - For `Ix2`: `[SliceOrIndex; 2]`
68-
/// - and so on..
69-
/// - For `IxDyn`: `[SliceOrIndex]`
70-
///
71-
/// The easiest way to create a `&SliceInfo<SliceArg, Do>` is using the
72-
/// [`s![]`](macro.s!.html) macro.
73-
type SliceArg: ?Sized + AsRef<[SliceOrIndex]>;
7466
/// Pattern matching friendly form of the dimension value.
7567
///
7668
/// - For `Ix1`: `usize`,
@@ -399,7 +391,6 @@ macro_rules! impl_insert_axis_array(
399391

400392
impl Dimension for Dim<[Ix; 0]> {
401393
const NDIM: Option<usize> = Some(0);
402-
type SliceArg = [SliceOrIndex; 0];
403394
type Pattern = ();
404395
type Smaller = Self;
405396
type Larger = Ix1;
@@ -443,7 +434,6 @@ impl Dimension for Dim<[Ix; 0]> {
443434

444435
impl Dimension for Dim<[Ix; 1]> {
445436
const NDIM: Option<usize> = Some(1);
446-
type SliceArg = [SliceOrIndex; 1];
447437
type Pattern = Ix;
448438
type Smaller = Ix0;
449439
type Larger = Ix2;
@@ -559,7 +549,6 @@ impl Dimension for Dim<[Ix; 1]> {
559549

560550
impl Dimension for Dim<[Ix; 2]> {
561551
const NDIM: Option<usize> = Some(2);
562-
type SliceArg = [SliceOrIndex; 2];
563552
type Pattern = (Ix, Ix);
564553
type Smaller = Ix1;
565554
type Larger = Ix3;
@@ -716,7 +705,6 @@ impl Dimension for Dim<[Ix; 2]> {
716705

717706
impl Dimension for Dim<[Ix; 3]> {
718707
const NDIM: Option<usize> = Some(3);
719-
type SliceArg = [SliceOrIndex; 3];
720708
type Pattern = (Ix, Ix, Ix);
721709
type Smaller = Ix2;
722710
type Larger = Ix4;
@@ -839,7 +827,6 @@ macro_rules! large_dim {
839827
($n:expr, $name:ident, $pattern:ty, $larger:ty, { $($insert_axis:tt)* }) => (
840828
impl Dimension for Dim<[Ix; $n]> {
841829
const NDIM: Option<usize> = Some($n);
842-
type SliceArg = [SliceOrIndex; $n];
843830
type Pattern = $pattern;
844831
type Smaller = Dim<[Ix; $n - 1]>;
845832
type Larger = $larger;
@@ -890,7 +877,6 @@ large_dim!(6, Ix6, (Ix, Ix, Ix, Ix, Ix, Ix), IxDyn, {
890877
/// and memory wasteful, but it allows an arbitrary and dynamic number of axes.
891878
impl Dimension for IxDyn {
892879
const NDIM: Option<usize> = None;
893-
type SliceArg = [SliceOrIndex];
894880
type Pattern = Self;
895881
type Smaller = Self;
896882
type Larger = Self;

‎src/dimension/mod.rs‎

Lines changed: 55 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,8 @@
77
// except according to those terms.
88

99
use crate::error::{from_kind, ErrorKind, ShapeError};
10-
use crate::{Ix, Ixs, Slice, SliceOrIndex};
10+
use crate::slice::SliceArg;
11+
use crate::{AxisSliceInfo, Ix, Ixs, Slice};
1112
use num_integer::div_floor;
1213

1314
pub use self::axes::{axes_of, Axes, AxisDescription};
@@ -18,6 +19,7 @@ pub use self::dim::*;
1819
pub use self::dimension_trait::Dimension;
1920
pub use self::dynindeximpl::IxDynImpl;
2021
pub use self::ndindex::NdIndex;
22+
pub use self::ops::DimAdd;
2123
pub use self::remove_axis::RemoveAxis;
2224

2325
use crate::shape_builder::Strides;
@@ -35,6 +37,7 @@ pub mod dim;
3537
mod dimension_trait;
3638
mod dynindeximpl;
3739
mod ndindex;
40+
mod ops;
3841
mod remove_axis;
3942

4043
/// Calculate offset from `Ix` stride converting sign properly
@@ -596,20 +599,24 @@ fn slice_min_max(axis_len: usize, slice: Slice) -> Option<(usize, usize)> {
596599
/// Returns `true` iff the slices intersect.
597600
pub fn slices_intersect<D: Dimension>(
598601
dim: &D,
599-
indices1: &D::SliceArg,
600-
indices2: &D::SliceArg,
602+
indices1: &impl SliceArg<D>,
603+
indices2: &impl SliceArg<D>,
601604
) -> bool {
602-
debug_assert_eq!(indices1.as_ref().len(), indices2.as_ref().len());
603-
for (&axis_len, &si1, &si2) in izip!(dim.slice(), indices1.as_ref(), indices2.as_ref()) {
604-
// The slices do not intersect iff any pair of `SliceOrIndex` does not intersect.
605+
debug_assert_eq!(indices1.in_ndim(), indices2.in_ndim());
606+
for (&axis_len, &si1, &si2) in izip!(
607+
dim.slice(),
608+
indices1.as_ref().iter().filter(|si| !si.is_new_axis()),
609+
indices2.as_ref().iter().filter(|si| !si.is_new_axis()),
610+
) {
611+
// The slices do not intersect iff any pair of `AxisSliceInfo` does not intersect.
605612
match (si1, si2) {
606613
(
607-
SliceOrIndex::Slice {
614+
AxisSliceInfo::Slice {
608615
start: start1,
609616
end: end1,
610617
step: step1,
611618
},
612-
SliceOrIndex::Slice {
619+
AxisSliceInfo::Slice {
613620
start: start2,
614621
end: end2,
615622
step: step2,
@@ -630,8 +637,8 @@ pub fn slices_intersect<D: Dimension>(
630637
return false;
631638
}
632639
}
633-
(SliceOrIndex::Slice { start, end, step }, SliceOrIndex::Index(ind))
634-
| (SliceOrIndex::Index(ind), SliceOrIndex::Slice { start, end, step }) => {
640+
(AxisSliceInfo::Slice { start, end, step }, AxisSliceInfo::Index(ind))
641+
| (AxisSliceInfo::Index(ind), AxisSliceInfo::Slice { start, end, step }) => {
635642
let ind = abs_index(axis_len, ind);
636643
let (min, max) = match slice_min_max(axis_len, Slice::new(start, end, step)) {
637644
Some(m) => m,
@@ -641,13 +648,14 @@ pub fn slices_intersect<D: Dimension>(
641648
return false;
642649
}
643650
}
644-
(SliceOrIndex::Index(ind1), SliceOrIndex::Index(ind2)) => {
651+
(AxisSliceInfo::Index(ind1), AxisSliceInfo::Index(ind2)) => {
645652
let ind1 = abs_index(axis_len, ind1);
646653
let ind2 = abs_index(axis_len, ind2);
647654
if ind1 != ind2 {
648655
return false;
649656
}
650657
}
658+
(AxisSliceInfo::NewAxis, _) | (_, AxisSliceInfo::NewAxis) => unreachable!(),
651659
}
652660
}
653661
true
@@ -719,7 +727,7 @@ mod test {
719727
};
720728
use crate::error::{from_kind, ErrorKind};
721729
use crate::slice::Slice;
722-
use crate::{Dim, Dimension, Ix0, Ix1, Ix2, Ix3, IxDyn};
730+
use crate::{Dim, Dimension, Ix0, Ix1, Ix2, Ix3, IxDyn, NewAxis};
723731
use num_integer::gcd;
724732
use quickcheck::{quickcheck, TestResult};
725733

@@ -993,17 +1001,45 @@ mod test {
9931001

9941002
#[test]
9951003
fn slices_intersect_true() {
996-
assert!(slices_intersect(&Dim([4, 5]), s![.., ..], s![.., ..]));
997-
assert!(slices_intersect(&Dim([4, 5]), s![0, ..], s![0, ..]));
998-
assert!(slices_intersect(&Dim([4, 5]), s![..;2, ..], s![..;3, ..]));
999-
assert!(slices_intersect(&Dim([4, 5]), s![.., ..;2], s![.., 1..;3]));
1004+
assert!(slices_intersect(
1005+
&Dim([4, 5]),
1006+
s![NewAxis, .., NewAxis, ..],
1007+
s![.., NewAxis, .., NewAxis]
1008+
));
1009+
assert!(slices_intersect(
1010+
&Dim([4, 5]),
1011+
s![NewAxis, 0, ..],
1012+
s![0, ..]
1013+
));
1014+
assert!(slices_intersect(
1015+
&Dim([4, 5]),
1016+
s![..;2, ..],
1017+
s![..;3, NewAxis, ..]
1018+
));
1019+
assert!(slices_intersect(
1020+
&Dim([4, 5]),
1021+
s![.., ..;2],
1022+
s![.., 1..;3, NewAxis]
1023+
));
10001024
assert!(slices_intersect(&Dim([4, 10]), s![.., ..;9], s![.., 3..;6]));
10011025
}
10021026

10031027
#[test]
10041028
fn slices_intersect_false() {
1005-
assert!(!slices_intersect(&Dim([4, 5]), s![..;2, ..], s![1..;2, ..]));
1006-
assert!(!slices_intersect(&Dim([4, 5]), s![..;2, ..], s![1..;3, ..]));
1007-
assert!(!slices_intersect(&Dim([4, 5]), s![.., ..;9], s![.., 3..;6]));
1029+
assert!(!slices_intersect(
1030+
&Dim([4, 5]),
1031+
s![..;2, ..],
1032+
s![NewAxis, 1..;2, ..]
1033+
));
1034+
assert!(!slices_intersect(
1035+
&Dim([4, 5]),
1036+
s![..;2, NewAxis, ..],
1037+
s![1..;3, ..]
1038+
));
1039+
assert!(!slices_intersect(
1040+
&Dim([4, 5]),
1041+
s![.., ..;9],
1042+
s![.., 3..;6, NewAxis]
1043+
));
10081044
}
10091045
}

‎src/dimension/ops.rs‎

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,90 @@
1+
use crate::imp_prelude::*;
2+
3+
/// Adds the two dimensions at compile time.
4+
pub trait DimAdd<D: Dimension> {
5+
/// The sum of the two dimensions.
6+
type Output: Dimension;
7+
}
8+
9+
macro_rules! impl_dimadd_const_out_const {
10+
($lhs:expr, $rhs:expr) => {
11+
impl DimAdd<Dim<[usize; $rhs]>> for Dim<[usize; $lhs]> {
12+
type Output = Dim<[usize; $lhs + $rhs]>;
13+
}
14+
};
15+
}
16+
17+
macro_rules! impl_dimadd_const_out_dyn {
18+
($lhs:expr, IxDyn) => {
19+
impl DimAdd<IxDyn> for Dim<[usize; $lhs]> {
20+
type Output = IxDyn;
21+
}
22+
};
23+
($lhs:expr, $rhs:expr) => {
24+
impl DimAdd<Dim<[usize; $rhs]>> for Dim<[usize; $lhs]> {
25+
type Output = IxDyn;
26+
}
27+
};
28+
}
29+
30+
impl<D: Dimension> DimAdd<D> for Ix0 {
31+
type Output = D;
32+
}
33+
34+
impl_dimadd_const_out_const!(1, 0);
35+
impl_dimadd_const_out_const!(1, 1);
36+
impl_dimadd_const_out_const!(1, 2);
37+
impl_dimadd_const_out_const!(1, 3);
38+
impl_dimadd_const_out_const!(1, 4);
39+
impl_dimadd_const_out_const!(1, 5);
40+
impl_dimadd_const_out_dyn!(1, 6);
41+
impl_dimadd_const_out_dyn!(1, IxDyn);
42+
43+
impl_dimadd_const_out_const!(2, 0);
44+
impl_dimadd_const_out_const!(2, 1);
45+
impl_dimadd_const_out_const!(2, 2);
46+
impl_dimadd_const_out_const!(2, 3);
47+
impl_dimadd_const_out_const!(2, 4);
48+
impl_dimadd_const_out_dyn!(2, 5);
49+
impl_dimadd_const_out_dyn!(2, 6);
50+
impl_dimadd_const_out_dyn!(2, IxDyn);
51+
52+
impl_dimadd_const_out_const!(3, 0);
53+
impl_dimadd_const_out_const!(3, 1);
54+
impl_dimadd_const_out_const!(3, 2);
55+
impl_dimadd_const_out_const!(3, 3);
56+
impl_dimadd_const_out_dyn!(3, 4);
57+
impl_dimadd_const_out_dyn!(3, 5);
58+
impl_dimadd_const_out_dyn!(3, 6);
59+
impl_dimadd_const_out_dyn!(3, IxDyn);
60+
61+
impl_dimadd_const_out_const!(4, 0);
62+
impl_dimadd_const_out_const!(4, 1);
63+
impl_dimadd_const_out_const!(4, 2);
64+
impl_dimadd_const_out_dyn!(4, 3);
65+
impl_dimadd_const_out_dyn!(4, 4);
66+
impl_dimadd_const_out_dyn!(4, 5);
67+
impl_dimadd_const_out_dyn!(4, 6);
68+
impl_dimadd_const_out_dyn!(4, IxDyn);
69+
70+
impl_dimadd_const_out_const!(5, 0);
71+
impl_dimadd_const_out_const!(5, 1);
72+
impl_dimadd_const_out_dyn!(5, 2);
73+
impl_dimadd_const_out_dyn!(5, 3);
74+
impl_dimadd_const_out_dyn!(5, 4);
75+
impl_dimadd_const_out_dyn!(5, 5);
76+
impl_dimadd_const_out_dyn!(5, 6);
77+
impl_dimadd_const_out_dyn!(5, IxDyn);
78+
79+
impl_dimadd_const_out_const!(6, 0);
80+
impl_dimadd_const_out_dyn!(6, 1);
81+
impl_dimadd_const_out_dyn!(6, 2);
82+
impl_dimadd_const_out_dyn!(6, 3);
83+
impl_dimadd_const_out_dyn!(6, 4);
84+
impl_dimadd_const_out_dyn!(6, 5);
85+
impl_dimadd_const_out_dyn!(6, 6);
86+
impl_dimadd_const_out_dyn!(6, IxDyn);
87+
88+
impl<D: Dimension> DimAdd<D> for IxDyn {
89+
type Output = IxDyn;
90+
}

‎src/doc/ndarray_for_numpy_users/mod.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -532,7 +532,7 @@
532532
//! `a[:] = b` | [`a.assign(&b)`][.assign()] | copy the data from array `b` into array `a`
533533
//! `np.concatenate((a,b), axis=1)` | [`concatenate![Axis(1), a, b]`][concatenate!] or [`concatenate(Axis(1), &[a.view(), b.view()])`][concatenate()] | concatenate arrays `a` and `b` along axis 1
534534
//! `np.stack((a,b), axis=1)` | [`stack![Axis(1), a, b]`][stack!] or [`stack(Axis(1), vec![a.view(), b.view()])`][stack()] | stack arrays `a` and `b` along axis 1
535-
//! `a[:,np.newaxis]` or `np.expand_dims(a, axis=1)` | [`a.insert_axis(Axis(1))`][.insert_axis()] | create an array from `a`, inserting a new axis 1
535+
//! `a[:,np.newaxis]` or `np.expand_dims(a, axis=1)` | [`a.slice(s![.., NewAxis])`][.slice()] or [`a.insert_axis(Axis(1))`][.insert_axis()] | create an view of 1-D array `a`, inserting a new axis 1
536536
//! `a.transpose()` or `a.T` | [`a.t()`][.t()] or [`a.reversed_axes()`][.reversed_axes()] | transpose of array `a` (view for `.t()` or by-move for `.reversed_axes()`)
537537
//! `np.diag(a)` | [`a.diag()`][.diag()] | view the diagonal of `a`
538538
//! `a.flatten()` | [`use std::iter::FromIterator; Array::from_iter(a.iter().cloned())`][::from_iter()] | create a 1-D array by flattening `a`

0 commit comments

Comments
 (0)