Support indexing by fixed size arrays: a[[i, j]] etc.

examples/convo.rs

@@ -0,0 +1,68 @@
+#![allow(unused)]
+#[macro_use(s)]
+extern crate ndarray;
+extern crate num;
+
+use num::Float;
+
+use ndarray::{
+    ArrayView,
+    ArrayViewMut,
+    OwnedArray,
+    Ix,
+};
+
+type Ix2 = (Ix, Ix);
+const SOBEL_X: [[f32; 3]; 3] = [[-1., 0., 1.], [-2., 0., 2.], [-1., 0., 1.]];
+const SOBEL_Y: [[f32; 3]; 3] = [[ 1., 2., 1.], [ 0., 0., 0.], [-1., -2., -1.]];
+const SHARPEN: [[f32; 3]; 3] = [[0., -1., 0.], [ -1., 5., -1.], [0., -1., 0.]];
+
+type Kernel3x3<A> = [[A; 3]; 3];
+
+#[inline(never)]
+fn conv_3x3<F>(a: &ArrayView<F, Ix2>, out: &mut ArrayViewMut<F, Ix2>, kernel: &Kernel3x3<F>) 
+    where F: Float,
+{
+    let (n, m) = a.dim();
+    let (np, mp) = out.dim();
+    if n < 3 || m < 3 {
+        return;
+    }
+    assert!(np >= n && mp >= m);
+    // i, j offset by -1 so that we can use unsigned indices
+    unsafe {
+        for i in 0..n - 2 {
+            for j in 0..m - 2 {
+                let mut conv = F::zero();
+                for k in 0..3 {
+                    for l in 0..3 {
+                        conv = conv + *a.uget((i + k, j + l)) * kernel[k][l];
+                        //conv += a[[i + k, j + l]] * x_kernel[k][l];
+                    }
+                }
+                *out.uget_mut((i + 1, j + 1)) = conv;
+            }
+        }
+    }
+}
+
+fn main() {
+    let n = 16;
+    let mut a = OwnedArray::zeros((n, n));
+    // make a circle
+    let c = (8., 8.);
+    for ((i, j), elt) in a.indexed_iter_mut() {
+        {
+            let s = ((i as f32) - c.0).powi(2) + (j as f32 - c.1).powi(2);
+            if s.sqrt() > 3. && s.sqrt() < 6. {
+                *elt = 1.;
+            }
+        }
+    }
+    println!("{:?}", a);
+    let mut res = OwnedArray::zeros(a.dim());
+    for _ in 0..1000 {
+        conv_3x3(&a.view(), &mut res.view_mut(), &SOBEL_X);
+    }
+    println!("{:?}", res);
+}

examples/life.rs

@@ -63,10 +63,11 @@ fn iterate(z: &mut Board, scratch: &mut Board) {
 }
 
 fn turn_on_corners(z: &mut Board) {
-    z.slice_mut(s![1..2, 1..2]).assign_scalar(&1);
-    z.slice_mut(s![1..2, -2..-1]).assign_scalar(&1);
-    z.slice_mut(s![-2..-1, 1..2]).assign_scalar(&1);
-    z.slice_mut(s![-2..-1, -2..-1]).assign_scalar(&1);
+    let (n, m) = z.dim();
+    z[[1    , 1    ]] = 1;
+    z[[1    , m - 2]] = 1;
+    z[[n - 2, 1    ]] = 1;
+    z[[n - 2, m - 2]] = 1;
 }
 
 fn render(a: &Board) {

src/arraytraits.rs

@@ -19,36 +19,44 @@ use super::{
     Data,
     DataMut,
     DataOwned,
+    NdIndex,
 };
 
+#[cold]
+#[inline(never)]
+fn array_out_of_bounds() -> ! {
+    panic!("ndarray: index out of bounds");
+}
+
 /// Access the element at **index**.
 ///
 /// **Panics** if index is out of bounds.
-impl<S, D> Index<D> for ArrayBase<S, D>
+impl<S, D, I> Index<I> for ArrayBase<S, D>
     where D: Dimension,
+          I: NdIndex<Dim=D>,
           S: Data,
 {
     type Output = S::Elem;
     #[inline]
-    fn index(&self, index: D) -> &S::Elem {
-        self.get(index).expect("Array::index: out of bounds")
+    fn index(&self, index: I) -> &S::Elem {
+        self.get(index).unwrap_or_else(|| array_out_of_bounds())
     }
 }
 
 /// Access the element at **index** mutably.
 ///
 /// **Panics** if index is out of bounds.
-impl<S, D> IndexMut<D> for ArrayBase<S, D>
+impl<S, D, I> IndexMut<I> for ArrayBase<S, D>
     where D: Dimension,
+          I: NdIndex<Dim=D>,
           S: DataMut,
 {
     #[inline]
-    fn index_mut(&mut self, index: D) -> &mut S::Elem {
-        self.get_mut(index).expect("Array::index_mut: out of bounds")
+    fn index_mut(&mut self, index: I) -> &mut S::Elem {
+        self.get_mut(index).unwrap_or_else(|| array_out_of_bounds())
     }
 }
 
-
 impl<S, S2, D> PartialEq<ArrayBase<S2, D>> for ArrayBase<S, D>
     where D: Dimension,
           S: Data,

src/dimension.rs

@@ -254,6 +254,7 @@ unsafe impl Dimension for Ix {
     }
 
     /// Return stride offset for this dimension and index.
+    #[inline]
     fn stride_offset_checked(&self, stride: &Ix, index: &Ix) -> Option<isize>
     {
         if *index < *self {
@@ -309,6 +310,7 @@ unsafe impl Dimension for (Ix, Ix) {
     }
 
     /// Return stride offset for this dimension and index.
+    #[inline]
     fn stride_offset_checked(&self, strides: &(Ix, Ix), index: &(Ix, Ix)) -> Option<isize>
     {
         let (m, n) = *self;
@@ -459,3 +461,87 @@ impl RemoveAxis for Vec<Ix>
     }
 }
 
+/// A tuple or fixed size array that can be used to index an array.
+///
+/// ```
+/// use ndarray::arr2;
+///
+/// let mut a = arr2(&[[0, 1], [0, 0]]);
+/// a[[1, 1]] = 1;
+/// assert_eq!(a[[0, 1]], 1);
+/// assert_eq!(a[[1, 1]], 1);
+/// ```
+///
+/// **Note** the blanket implementation that's not visible in rustdoc:
+/// `impl<D> NdIndex for D where D: Dimension { ... }`
+pub unsafe trait NdIndex {
+    type Dim: Dimension;
+    #[doc(hidden)]
+    fn index_checked(&self, dim: &Self::Dim, strides: &Self::Dim) -> Option<isize>;
+}
+
+unsafe impl<D> NdIndex for D where D: Dimension {
+    type Dim = D;
+    fn index_checked(&self, dim: &Self::Dim, strides: &Self::Dim) -> Option<isize> {
+        dim.stride_offset_checked(strides, self)
+    }
+}
+
+unsafe impl NdIndex for [Ix; 0] {
+    type Dim = ();
+    #[inline]
+    fn index_checked(&self, dim: &Self::Dim, strides: &Self::Dim) -> Option<isize> {
+        dim.stride_offset_checked(strides, &())
+    }
+}
+
+unsafe impl NdIndex for [Ix; 1] {
+    type Dim = Ix;
+    #[inline]
+    fn index_checked(&self, dim: &Self::Dim, strides: &Self::Dim) -> Option<isize> {
+        dim.stride_offset_checked(strides, &self[0])
+    }
+}
+
+unsafe impl NdIndex for [Ix; 2] {
+    type Dim = (Ix, Ix);
+    #[inline]
+    fn index_checked(&self, dim: &Self::Dim, strides: &Self::Dim) -> Option<isize> {
+        let index = (self[0], self[1]);
+        dim.stride_offset_checked(strides, &index)
+    }
+}
+
+unsafe impl NdIndex for [Ix; 3] {
+    type Dim = (Ix, Ix, Ix);
+    #[inline]
+    fn index_checked(&self, dim: &Self::Dim, strides: &Self::Dim) -> Option<isize> {
+        let index = (self[0], self[1], self[2]);
+        dim.stride_offset_checked(strides, &index)
+    }
+}
+
+unsafe impl NdIndex for [Ix; 4] {
+    type Dim = (Ix, Ix, Ix, Ix);
+    #[inline]
+    fn index_checked(&self, dim: &Self::Dim, strides: &Self::Dim) -> Option<isize> {
+        let index = (self[0], self[1], self[2], self[3]);
+        dim.stride_offset_checked(strides, &index)
+    }
+}
+
+unsafe impl<'a> NdIndex for &'a [Ix] {
+    type Dim = Vec<Ix>;
+    fn index_checked(&self, dim: &Self::Dim, strides: &Self::Dim) -> Option<isize> {
+        let mut offset = 0;
+        for ((&d, &i), &s) in zipsl(zipsl(&dim[..], &self[..]),
+                                    strides.slice())
+        {
+            if i >= d {
+                return None;
+            }
+            offset += stride_offset(i, s);
+        }
+        Some(offset)
+    }
+}

src/lib.rs

@@ -81,6 +81,7 @@ use std::slice::{self, Iter, IterMut};
 use it::ZipSlices;
 
 pub use dimension::{Dimension, RemoveAxis};
+pub use dimension::NdIndex;
 pub use indexes::Indexes;
 pub use shape_error::ShapeError;
 pub use si::{Si, S};
@@ -343,10 +344,13 @@ pub unsafe trait Data {
 /// Array’s writable inner representation.
 pub unsafe trait DataMut : Data {
     fn slice_mut(&mut self) -> &mut [Self::Elem];
+    #[inline]
     fn ensure_unique<D>(&mut ArrayBase<Self, D>)
         where Self: Sized, D: Dimension
     {
     }
+    #[inline]
+    fn is_unique(&mut self) -> bool { true }
 }
 
 /// Clone an Array’s storage.
@@ -384,6 +388,10 @@ unsafe impl<A> DataMut for Rc<Vec<A>> where A: Clone {
             self_.ptr = rvec.as_mut_ptr().offset(our_off);
         }
     }
+
+    fn is_unique(&mut self) -> bool {
+        Rc::get_mut(self).is_some()
+    }
 }
 
 unsafe impl<A> DataClone for Rc<Vec<A>> {
@@ -874,12 +882,15 @@ impl<A, S, D> ArrayBase<S, D> where S: Data<Elem=A>, D: Dimension
     /// assert!(
     ///     a.get((0, 1)) == Some(&2.) &&
     ///     a.get((0, 2)) == None &&
-    ///     a[(0, 1)] == 2.
+    ///     a[(0, 1)] == 2. &&
+    ///     a[[0, 1]] == 2.
     /// );
     /// ```
-    pub fn get(&self, index: D) -> Option<&A> {
+    pub fn get<I>(&self, index: I) -> Option<&A>
+        where I: NdIndex<Dim=D>,
+    {
         let ptr = self.ptr;
-        self.dim.stride_offset_checked(&self.strides, &index)
+        index.index_checked(&self.dim, &self.strides)
             .map(move |offset| unsafe {
                 &*ptr.offset(offset)
             })
@@ -893,12 +904,13 @@ impl<A, S, D> ArrayBase<S, D> where S: Data<Elem=A>, D: Dimension
 
     /// Return a mutable reference to the element at `index`, or return `None`
     /// if the index is out of bounds.
-    pub fn get_mut(&mut self, index: D) -> Option<&mut A>
+    pub fn get_mut<I>(&mut self, index: I) -> Option<&mut A>
         where S: DataMut,
+              I: NdIndex<Dim=D>,
     {
         self.ensure_unique();
         let ptr = self.ptr;
-        self.dim.stride_offset_checked(&self.strides, &index)
+        index.index_checked(&self.dim, &self.strides)
             .map(move |offset| unsafe {
                 &mut *ptr.offset(offset)
             })
@@ -941,7 +953,7 @@ impl<A, S, D> ArrayBase<S, D> where S: Data<Elem=A>, D: Dimension
     pub unsafe fn uget_mut(&mut self, index: D) -> &mut A
         where S: DataMut
     {
-        //debug_assert!(Rc::get_mut(&mut self.data).is_some());
+        debug_assert!(self.data.is_unique());
         debug_assert!(self.dim.stride_offset_checked(&self.strides, &index).is_some());
         let off = Dimension::stride_offset(&index, &self.strides);
         &mut *self.ptr.offset(off)

src/linalg.rs

@@ -124,7 +124,7 @@ pub fn least_squares<A: ComplexField>(a: &Mat<A>, b: &Col<A>) -> Col<A>
         let (m, _) = L.dim();
         for i in 1..m {
             for j in 0..i {
-                let elt = &mut L[(i, j)];
+                let elt = &mut L[[i, j]];
                 *elt = elt.conjugate();
             }
         }
@@ -177,7 +177,7 @@ pub fn cholesky<A: ComplexField>(a: Mat<A>) -> Mat<A>
             }
 
             // L_ij = [ A_ij - Sum(k = 1 .. j) L_ik L_jk ] / L_jj
-            L[(i, j)] = (L[(i, j)] - lik_ljk_sum) / L[(j, j)];
+            L[[i, j]] = (L[[i, j]] - lik_ljk_sum) / L[[j, j]];
         }
 
         // Diagonal where i == j
@@ -188,12 +188,12 @@ pub fn cholesky<A: ComplexField>(a: Mat<A>) -> Mat<A>
         for &ljk in L.row_iter(j).take(j as usize) {
             ljk_sum = ljk_sum + ljk * ljk.conjugate();
         }
-        L[(j, j)] = (L[(j, j)] - ljk_sum).sqrt_real();
+        L[[j, j]] = (L[[j, j]] - ljk_sum).sqrt_real();
 
         // After the diagonal
         // L_ij = 0 for j > i
         for j in i + 1..n {
-            L[(i, j)] = z;
+            L[[i, j]] = z;
         }
     }
     L