Skip to main content

jlrs_ndarray/
lib.rs

1//! Borrow data from Julia arrays as `ndarray`'s `ArrayView` and `ArrayViewMut`.
2//!
3//! This crate defines a single trait, `NdArray`, that provides methods that return an immutable
4//! or a mutable view of the array data and is implemented by `Array` and `TypedArray` from jlrs.
5//! It's easier to use this trait with `TypedArray`, you'll likely have to provide type
6//! annotations with `Array`.
7
8use jlrs::error::other;
9use jlrs::prelude::*;
10use ndarray::{ArrayView, ArrayViewMut, Dim, IntoDimension, IxDynImpl, ShapeBuilder};
11
12mod private {
13    use jlrs::traits::ValidLayout;
14    use jlrs::value::array::{Array, TypedArray};
15
16    pub trait Sealed {}
17    impl<'frame, 'data> Sealed for Array<'frame, 'data> {}
18    impl<'frame, 'data, T> Sealed for TypedArray<'frame, 'data, T> where T: Copy + ValidLayout {}
19}
20
21/// Trait to borrow Julia arrays with inline data as `ndarray`'s `ArrayView` and `ArrayViewMut`.
22pub trait NdArray<'borrow, T>: private::Sealed {
23    /// Borrow the data in the array as an `ArrayView`. Returns an error if the wrong type is
24    /// provided or the data is not stored inline.
25    fn array_view<'frame: 'borrow, F>(
26        self,
27        frame: &'borrow F,
28    ) -> JlrsResult<ArrayView<'borrow, T, Dim<IxDynImpl>>>
29    where
30        F: Frame<'frame>,
31        T: ValidLayout + Copy;
32
33    /// Mutably borrow the data in the array as an `ArrayViewMut`. Returns an error if the wrong
34    /// type is provided or the data is not stored inline.
35    fn array_view_mut<'frame: 'borrow, F>(
36        self,
37        frame: &'borrow mut F,
38    ) -> JlrsResult<ArrayViewMut<'borrow, T, Dim<IxDynImpl>>>
39    where
40        F: Frame<'frame>,
41        T: ValidLayout + Copy;
42}
43
44impl<'frame: 'borrow, 'data: 'borrow, 'borrow, T: ValidLayout + Copy> NdArray<'borrow, T>
45    for Array<'frame, 'data>
46{
47    fn array_view<'fr: 'borrow, F>(
48        self,
49        frame: &'borrow F,
50    ) -> JlrsResult<ArrayView<'borrow, T, Dim<IxDynImpl>>>
51    where
52        F: Frame<'fr>,
53        T: ValidLayout + Copy,
54    {
55        let data = self.inline_data::<T, _>(frame)?;
56        let shape = data.dimensions().as_slice().into_dimension().f();
57        match ArrayView::from_shape(shape, data.into_slice()) {
58            Ok(arr) => Ok(arr),
59            Err(e) => other(e)?,
60        }
61    }
62
63    fn array_view_mut<'fr: 'borrow, F>(
64        self,
65        frame: &'borrow mut F,
66    ) -> JlrsResult<ArrayViewMut<'borrow, T, Dim<IxDynImpl>>>
67    where
68        F: Frame<'fr>,
69        T: ValidLayout + Copy,
70    {
71        let data = self.inline_data_mut::<T, _>(frame)?;
72        let shape = data.dimensions().as_slice().into_dimension().f();
73        let raw = data.into_mut_slice();
74        match ArrayViewMut::from_shape(shape, raw) {
75            Ok(arr) => Ok(arr),
76            Err(e) => other(e)?,
77        }
78    }
79}
80
81impl<'frame: 'borrow, 'data: 'borrow, 'borrow, T: ValidLayout + Copy> NdArray<'borrow, T>
82    for TypedArray<'frame, 'data, T>
83{
84    fn array_view<'fr: 'borrow, F>(
85        self,
86        frame: &'borrow F,
87    ) -> JlrsResult<ArrayView<'borrow, T, Dim<IxDynImpl>>>
88    where
89        F: Frame<'fr>,
90        T: ValidLayout,
91    {
92        let data = self.inline_data(frame)?;
93        let shape = data.dimensions().as_slice().into_dimension().f();
94        match ArrayView::from_shape(shape, data.into_slice()) {
95            Ok(arr) => Ok(arr),
96            Err(e) => other(e)?,
97        }
98    }
99
100    fn array_view_mut<'fr: 'borrow, F>(
101        self,
102        frame: &'borrow mut F,
103    ) -> JlrsResult<ArrayViewMut<'borrow, T, Dim<IxDynImpl>>>
104    where
105        F: Frame<'fr>,
106        T: ValidLayout,
107    {
108        let data = self.inline_data_mut(frame)?;
109        let shape = data.dimensions().as_slice().into_dimension().f();
110        let raw = data.into_mut_slice();
111        match ArrayViewMut::from_shape(shape, raw) {
112            Ok(arr) => Ok(arr),
113            Err(e) => other(e)?,
114        }
115    }
116}
117
118#[cfg(test)]
119mod tests {
120    use super::NdArray;
121    use jlrs::prelude::*;
122    use ndarray::{ArrayView, ArrayViewMut, IxDyn};
123
124    use std::cell::RefCell;
125
126    thread_local! {
127        pub static JULIA: RefCell<Julia> = RefCell::new(unsafe { Julia::init(32).unwrap() });
128    }
129
130    #[test]
131    fn array_view() {
132        JULIA.with(|j| {
133            let mut julia = j.borrow_mut();
134
135            julia
136                .dynamic_frame(|_global, frame| {
137                    let mut data = vec![1usize, 2, 3, 4, 5, 6];
138                    let slice = &mut data.as_mut_slice();
139                    let borrowed = Value::borrow_array(frame, slice, (3, 2))?;
140
141                    let jl_array = borrowed.cast::<Array>()?;
142                    let x = jl_array.inline_data::<usize, _>(frame)?[(1, 0)];
143
144                    let array: ArrayView<usize, _> = jl_array.array_view(frame)?;
145                    assert_eq!(array[IxDyn(&[1, 0])], x);
146
147                    Ok(())
148                })
149                .unwrap();
150        });
151    }
152
153    #[test]
154    fn array_view_wrong_type() {
155        JULIA.with(|j| {
156            let mut julia = j.borrow_mut();
157
158            julia
159                .dynamic_frame(|_global, frame| {
160                    let mut data = vec![1usize, 2, 3, 4, 5, 6];
161                    let slice = &mut data.as_mut_slice();
162                    let borrowed = Value::borrow_array(frame, slice, (3, 2))?;
163
164                    let jl_array = borrowed.cast::<Array>()?;
165                    let view: Result<ArrayView<isize, _>, _> = jl_array.array_view(frame);
166                    assert!(view.is_err());
167                    Ok(())
168                })
169                .unwrap();
170        });
171    }
172
173    #[test]
174    fn array_view_mut() {
175        JULIA.with(|j| {
176            let mut julia = j.borrow_mut();
177
178            julia
179                .dynamic_frame(|_global, frame| {
180                    let mut data = vec![1usize, 2, 3, 4, 5, 6];
181                    let slice = &mut data.as_mut_slice();
182                    let borrowed = Value::borrow_array(frame, slice, (3, 2))?;
183
184                    let jl_array = borrowed.cast::<Array>()?;
185                    let mut inline = jl_array.inline_data_mut::<usize, _>(frame)?;
186                    let x = inline[(1, 0)];
187
188                    inline[(1, 0)] = x + 1;
189
190                    let mut array: ArrayViewMut<usize, _> = jl_array.array_view_mut(frame)?;
191                    assert_eq!(array[IxDyn(&[1, 0])], x + 1);
192                    array[IxDyn(&[1, 0])] -= 1;
193
194                    let inline = jl_array.inline_data_mut::<usize, _>(frame)?;
195                    assert_eq!(inline[(1, 0)], x);
196                    Ok(())
197                })
198                .unwrap();
199        });
200    }
201
202    #[test]
203    fn array_view_mut_wrong_type() {
204        JULIA.with(|j| {
205            let mut julia = j.borrow_mut();
206
207            julia
208                .dynamic_frame(|_global, frame| {
209                    let mut data = vec![1usize, 2, 3, 4, 5, 6];
210                    let slice = &mut data.as_mut_slice();
211                    let borrowed = Value::borrow_array(frame, slice, (3, 2))?;
212
213                    let jl_array = borrowed.cast::<Array>()?;
214                    let view: Result<ArrayViewMut<isize, _>, _> = jl_array.array_view_mut(frame);
215                    assert!(view.is_err());
216                    Ok(())
217                })
218                .unwrap();
219        });
220    }
221
222    #[test]
223    fn typed_array_view() {
224        JULIA.with(|j| {
225            let mut julia = j.borrow_mut();
226
227            julia
228                .dynamic_frame(|_global, frame| {
229                    let mut data = vec![1usize, 2, 3, 4, 5, 6];
230                    let slice = &mut data.as_mut_slice();
231                    let borrowed = Value::borrow_array(frame, slice, (3, 2))?;
232
233                    let jl_array = borrowed.cast::<TypedArray<usize>>()?;
234                    let x = jl_array.inline_data(frame)?[(1, 0)];
235
236                    let array: ArrayView<usize, _> = jl_array.array_view(frame)?;
237                    assert_eq!(array[IxDyn(&[1, 0])], x);
238
239                    Ok(())
240                })
241                .unwrap();
242        });
243    }
244
245    #[test]
246    fn typed_array_view_mut() {
247        JULIA.with(|j| {
248            let mut julia = j.borrow_mut();
249
250            julia
251                .dynamic_frame(|_global, frame| {
252                    let mut data = vec![1usize, 2, 3, 4, 5, 6];
253                    let slice = &mut data.as_mut_slice();
254                    let borrowed = Value::borrow_array(frame, slice, (3, 2))?;
255
256                    let jl_array = borrowed.cast::<TypedArray<usize>>()?;
257                    let mut inline = jl_array.inline_data_mut(frame)?;
258                    let x = inline[(1, 0)];
259
260                    inline[(1, 0)] = x + 1;
261
262                    let mut array: ArrayViewMut<usize, _> = jl_array.array_view_mut(frame)?;
263                    assert_eq!(array[IxDyn(&[1, 0])], x + 1);
264                    array[IxDyn(&[1, 0])] -= 1;
265
266                    let inline = jl_array.inline_data_mut(frame)?;
267                    assert_eq!(inline[(1, 0)], x);
268                    Ok(())
269                })
270                .unwrap();
271        });
272    }
273
274    #[test]
275    fn example() {
276        JULIA.with(|j| {
277            let mut julia = j.borrow_mut();
278
279            julia
280                .dynamic_frame(|_global, frame| {
281                    let mut data = vec![1usize, 2, 3, 4, 5, 6];
282                    let slice = &mut data.as_mut_slice();
283                    let borrowed = Value::borrow_array(frame, slice, (3, 2))?;
284
285                    let _array = borrowed.cast::<TypedArray<usize>>()?.array_view(frame)?;
286
287                    Ok(())
288                })
289                .unwrap();
290        });
291    }
292}