metatensor/
tensor.rs

1use std::ffi::{CStr, CString};
2use std::iter::FusedIterator;
3
4use crate::block::TensorBlockRefMut;
5use crate::c_api::mts_tensormap_t;
6
7use crate::errors::{check_status, check_ptr};
8use crate::{Error, LabelValue, Labels, MtsArray, TensorBlock, TensorBlockRef};
9
10/// [`TensorMap`] is the main user-facing struct of this library, and can
11/// store any kind of data used in atomistic machine learning.
12///
13/// A tensor map contains a list of `TensorBlock`s, each one associated with a
14/// key in the form of a single `Labels` entry.
15///
16/// It provides functions to merge blocks together by moving some of these keys
17/// to the samples or properties labels of the blocks, transforming the sparse
18/// representation of the data to a dense one.
19pub struct TensorMap {
20    pub(crate) ptr: *mut mts_tensormap_t,
21    /// cache for the keys labels
22    keys: Labels,
23}
24
25// SAFETY: Send is fine since we can free a TensorMap from any thread
26unsafe impl Send for TensorMap {}
27// SAFETY: Sync is fine since there is no internal mutability in TensorMap
28unsafe impl Sync for TensorMap {}
29
30impl std::fmt::Debug for TensorMap {
31    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
32        use crate::labels::pretty_print_labels;
33        writeln!(f, "Tensormap @ {:p} {{", self.ptr)?;
34
35        write!(f, "    keys: ")?;
36        pretty_print_labels(self.keys(), "    ", f)?;
37        writeln!(f, "}}")
38    }
39}
40
41impl std::ops::Drop for TensorMap {
42    #[allow(unused_must_use)]
43    fn drop(&mut self) {
44        unsafe {
45            crate::c_api::mts_tensormap_free(self.ptr);
46        }
47    }
48}
49
50impl TensorMap {
51    /// Create a new `TensorMap` with the given keys and blocks.
52    ///
53    /// The number of keys must match the number of blocks, and all the blocks
54    /// must contain the same kind of data (same labels names, same gradients
55    /// defined on all blocks).
56    #[allow(clippy::needless_pass_by_value)]
57    #[inline]
58    pub fn new(keys: Labels, mut blocks: Vec<TensorBlock>) -> Result<TensorMap, Error> {
59        let ptr = unsafe {
60            crate::c_api::mts_tensormap(
61                keys.as_mts_labels_t(),
62                // this cast is fine because TensorBlock is `repr(transparent)`
63                // to a `*mut mts_block_t` (through `TensorBlockRefMut`, and
64                // `TensorBlockRef`).
65                blocks.as_mut_ptr().cast::<*mut crate::c_api::mts_block_t>(),
66                blocks.len()
67            )
68        };
69
70        for block in blocks {
71            // we give ownership of the blocks to the new tensormap, so we
72            // should not free them again from Rust
73            std::mem::forget(block);
74        }
75
76        check_ptr(ptr)?;
77
78        return Ok(unsafe { TensorMap::from_raw(ptr) });
79    }
80
81    /// Create a new `TensorMap` from a raw pointer.
82    ///
83    /// This function takes ownership of the pointer, and will call
84    /// `mts_tensormap_free` on it when the `TensorMap` goes out of scope.
85    ///
86    /// # Safety
87    ///
88    /// The pointer must be non-null and created by
89    /// [`crate::c_api::mts_tensormap`] or [`TensorMap::into_raw`].
90    pub unsafe fn from_raw(ptr: *mut mts_tensormap_t) -> TensorMap {
91        assert!(!ptr.is_null());
92
93        let keys_ptr = unsafe {
94            crate::c_api::mts_tensormap_keys(ptr)
95        };
96        assert!(!keys_ptr.is_null(), "failed to get the keys");
97        let keys = unsafe { Labels::from_raw(keys_ptr) };
98
99        return TensorMap {
100            ptr,
101            keys
102        };
103    }
104
105    /// Extract the underlying raw pointer.
106    ///
107    /// The pointer should be passed back to [`TensorMap::from_raw`] or
108    /// [`crate::c_api::mts_tensormap_free`] to release the memory corresponding
109    /// to this `TensorMap`.
110    pub fn into_raw(mut tensor: TensorMap) -> *mut mts_tensormap_t {
111        return std::mem::replace(&mut tensor.ptr, std::ptr::null_mut());
112    }
113
114    /// Get the underlying raw pointer.
115    ///
116    /// After a call, this `TensorMap` is still managing the corresponding
117    /// memory. To fully release the pointer, use [`TensorMap::into_raw`].
118    pub fn as_ptr(&self) -> *const mts_tensormap_t {
119        self.ptr
120    }
121
122    /// Get the underlying (mutable) raw pointer
123    ///
124    /// After a call, this `TensorMap` is still managing the corresponding
125    /// memory. To fully release the pointer, use [`TensorMap::into_raw`].
126    pub fn as_mut_ptr(&mut self) -> *mut mts_tensormap_t {
127        self.ptr
128    }
129
130    /// Clone this `TensorMap`, cloning all the data and metadata contained inside.
131    ///
132    /// This can fail if the external data held inside an `mts_array_t` can not
133    /// be cloned.
134    #[inline]
135    pub fn try_clone(&self) -> Result<TensorMap, Error> {
136        let ptr = unsafe {
137            crate::c_api::mts_tensormap_copy(self.ptr)
138        };
139        crate::errors::check_ptr(ptr)?;
140
141        return Ok(unsafe { TensorMap::from_raw(ptr) });
142    }
143
144    /// Load a `TensorMap` from the file at `path`
145    ///
146    /// This is a convenience function calling [`crate::io::load`]
147    pub fn load(path: impl AsRef<std::path::Path>) -> Result<TensorMap, Error> {
148        return crate::io::load(path);
149    }
150
151    /// Load a `TensorMap` from an in-memory buffer
152    ///
153    /// This is a convenience function calling [`crate::io::load_buffer`]
154    pub fn load_buffer(buffer: &[u8]) -> Result<TensorMap, Error> {
155        return crate::io::load_buffer(buffer);
156    }
157
158    /// Save the given tensor to the file at `path`
159    ///
160    /// This is a convenience function calling [`crate::io::save`]
161    pub fn save(&self, path: impl AsRef<std::path::Path>) -> Result<(), Error> {
162        return crate::io::save(path, self);
163    }
164
165    /// Save the given tensor to an in-memory buffer
166    ///
167    /// This is a convenience function calling [`crate::io::save_buffer`]
168    pub fn save_buffer(&self, buffer: &mut Vec<u8>) -> Result<(), Error> {
169        return crate::io::save_buffer(self, buffer);
170    }
171
172    /// Get the device on which the values of this `TensorMap` are stored.
173    #[inline]
174    pub fn device(&self) -> Result<dlpk::sys::DLDevice, Error> {
175        let mut device = dlpk::sys::DLDevice::cpu();
176        unsafe {
177            check_status(crate::c_api::mts_tensormap_device(
178                self.ptr,
179                &mut device,
180            ))?;
181        }
182        return Ok(device);
183    }
184
185    /// Get the data type of the values of this `TensorMap`.
186    #[inline]
187    pub fn dtype(&self) -> Result<dlpk::sys::DLDataType, Error> {
188        let mut dtype = dlpk::sys::DLDataType {
189            code: dlpk::sys::DLDataTypeCode::kDLFloat,
190            bits: 0,
191            lanes: 0,
192        };
193        unsafe {
194            check_status(crate::c_api::mts_tensormap_dtype(
195                self.ptr,
196                &mut dtype,
197            ))?;
198        }
199        return Ok(dtype);
200    }
201
202    /// Get the keys defined in this `TensorMap`
203    #[inline]
204    pub fn keys(&self) -> &Labels {
205        &self.keys
206    }
207
208    /// Get a reference to the block at the given `index` in this `TensorMap`
209    ///
210    /// # Panics
211    ///
212    /// If the index is out of bounds
213    #[inline]
214    pub fn block_by_id(&self, index: usize) -> TensorBlockRef<'_> {
215
216        let mut block = std::ptr::null_mut();
217        unsafe {
218            check_status(crate::c_api::mts_tensormap_block_by_id(
219                self.ptr,
220                &mut block,
221                index,
222            )).expect("failed to get a block");
223        }
224
225        return unsafe { TensorBlockRef::from_raw(block) }
226    }
227
228    /// Get a mutable reference to the block at the given `index` in this `TensorMap`
229    ///
230    /// # Panics
231    ///
232    /// If the index is out of bounds
233    #[inline]
234    pub fn block_mut_by_id(&mut self, index: usize) -> TensorBlockRefMut<'_> {
235        return unsafe { TensorMap::raw_block_mut_by_id(self.ptr, index) };
236    }
237
238    /// Implementation of `block_mut_by_id` which does not borrow the
239    /// `mts_tensormap_t` pointer.
240    ///
241    /// This is used to provide references to multiple blocks at the same time
242    /// in the iterators.
243    ///
244    /// # Safety
245    ///
246    /// This should be called with a valid `mts_tensormap_t`, and the lifetime
247    /// `'a` should be properly constrained to the lifetime of the owner of
248    /// `ptr`.
249    #[inline]
250    unsafe fn raw_block_mut_by_id<'a>(ptr: *mut mts_tensormap_t, index: usize) -> TensorBlockRefMut<'a> {
251        unsafe {
252            let mut block = std::ptr::null_mut();
253
254            check_status(
255                crate::c_api::mts_tensormap_block_by_id(
256                ptr,
257                &mut block,
258                index,
259            )).expect("failed to get a block");
260
261            return TensorBlockRefMut::from_raw(block);
262        }
263    }
264
265    /// Get a reference to the block matching the given selection.
266    #[inline]
267    pub fn block(&self, selection: &Labels) -> Result<TensorBlockRef<'_>, Error> {
268        let matching = self.keys.select(selection)?;
269        if matching.len() != 1 {
270            let selection_str = selection.names()
271                .iter()
272                .zip(&selection[0])
273                .map(|(name, value)| format!("{} = {}", name, value))
274                .collect::<Vec<_>>()
275                .join(", ");
276
277            if matching.is_empty() {
278                return Err(Error {
279                    code: None,
280                    message: format!(
281                        "no blocks matched the selection ({})",
282                        selection_str
283                    ),
284                });
285            } else {
286                return Err(Error {
287                    code: None,
288                    message: format!(
289                        "{} blocks matched the selection ({}), expected only one",
290                        matching.len(),
291                        selection_str
292                    ),
293                });
294            }
295        }
296
297        return Ok(self.block_by_id(matching[0]));
298    }
299
300    /// Get a reference to every blocks in this `TensorMap`
301    #[inline]
302    pub fn blocks(&self) -> Vec<TensorBlockRef<'_>> {
303        let mut blocks = Vec::new();
304        for i in 0..self.keys().count() {
305            blocks.push(self.block_by_id(i));
306        }
307        return blocks;
308    }
309
310    /// Get a mutable reference to every blocks in this `TensorMap`
311    #[inline]
312    pub fn blocks_mut(&mut self) -> Vec<TensorBlockRefMut<'_>> {
313        let mut blocks = Vec::new();
314        for i in 0..self.keys().count() {
315            blocks.push(unsafe { TensorMap::raw_block_mut_by_id(self.ptr, i) });
316        }
317        return blocks;
318    }
319
320    /// Merge blocks with the same value for selected keys dimensions along the
321    /// samples axis.
322    ///
323    /// The dimensions (names) of `keys_to_move` will be moved from the keys to
324    /// the sample labels, and blocks with the same remaining keys dimensions
325    /// will be merged together along the sample axis.
326    ///
327    /// `keys_to_move` must be empty (`keys_to_move.count() == 0`), and the new
328    /// sample labels will contain entries corresponding to the merged blocks'
329    /// keys.
330    ///
331    /// The new sample labels will contain all of the merged blocks sample
332    /// labels. The order of the samples is controlled by `sort_samples`. If
333    /// `sort_samples` is true, samples are re-ordered to keep them
334    /// lexicographically sorted. Otherwise they are kept in the order in which
335    /// they appear in the blocks.
336    #[inline]
337    pub fn keys_to_samples(&self, keys_to_move: &Labels, fill_value: MtsArray, sort_samples: bool) -> Result<TensorMap, Error> {
338        let ptr = unsafe {
339            crate::c_api::mts_tensormap_keys_to_samples(
340                self.ptr,
341                keys_to_move.as_mts_labels_t(),
342                fill_value.into_raw(),
343                sort_samples,
344            )
345        };
346
347        check_ptr(ptr)?;
348        return Ok(unsafe { TensorMap::from_raw(ptr) });
349    }
350
351    /// Merge blocks with the same value for selected keys dimensions along the
352    /// property axis.
353    ///
354    /// The dimensions (names) of `keys_to_move` will be moved from the keys to
355    /// the property labels, and blocks with the same remaining keys dimensions
356    /// will be merged together along the property axis.
357    ///
358    /// If `keys_to_move` does not contain any entries (`keys_to_move.count()
359    /// == 0`), then the new property labels will contain entries corresponding
360    /// to the merged blocks only. For example, merging a block with key `a=0`
361    /// and properties `p=1, 2` with a block with key `a=2` and properties `p=1,
362    /// 3` will produce a block with properties `a, p = (0, 1), (0, 2), (2, 1),
363    /// (2, 3)`.
364    ///
365    /// If `keys_to_move` contains entries, then the property labels must be the
366    /// same for all the merged blocks. In that case, the merged property labels
367    /// will contain each of the entries of `keys_to_move` and then the current
368    /// property labels. For example, using `a=2, 3` in `keys_to_move`, and
369    /// blocks with properties `p=1, 2` will result in `a, p = (2, 1), (2, 2),
370    /// (3, 1), (3, 2)`.
371    ///
372    /// The new sample labels will contain all of the merged blocks sample
373    /// labels. The order of the samples is controlled by `sort_samples`. If
374    /// `sort_samples` is true, samples are re-ordered to keep them
375    /// lexicographically sorted. Otherwise they are kept in the order in which
376    /// they appear in the blocks.
377    #[inline]
378    pub fn keys_to_properties(&self, keys_to_move: &Labels, fill_value: MtsArray, sort_samples: bool) -> Result<TensorMap, Error> {
379        let ptr = unsafe {
380            crate::c_api::mts_tensormap_keys_to_properties(
381                self.ptr,
382                keys_to_move.as_mts_labels_t(),
383                fill_value.into_raw(),
384                sort_samples,
385            )
386        };
387
388        check_ptr(ptr)?;
389        return Ok(unsafe { TensorMap::from_raw(ptr) });
390    }
391
392    /// Move the given dimensions from the component labels to the property
393    /// labels for each block in this `TensorMap`.
394    #[inline]
395    pub fn components_to_properties(&self, dimensions: &[&str]) -> Result<TensorMap, Error> {
396        let dimensions_c = dimensions.iter()
397            .map(|&v| CString::new(v).expect("unexpected NULL byte"))
398            .collect::<Vec<_>>();
399
400        let dimensions_ptr = dimensions_c.iter()
401            .map(|v| v.as_ptr())
402            .collect::<Vec<_>>();
403
404
405        let ptr = unsafe {
406            crate::c_api::mts_tensormap_components_to_properties(
407                self.ptr,
408                dimensions_ptr.as_ptr(),
409                dimensions.len(),
410            )
411        };
412
413        check_ptr(ptr)?;
414        return Ok(unsafe { TensorMap::from_raw(ptr) });
415    }
416
417    /// Get an iterator over the keys and associated blocks
418    #[inline]
419    pub fn iter(&self) -> TensorMapIter<'_> {
420        return TensorMapIter {
421            inner: self.keys().into_iter().zip(self.blocks())
422        };
423    }
424
425    /// Get an iterator over the keys and associated blocks, with read-write
426    /// access to the blocks
427    #[inline]
428    pub fn iter_mut(&mut self) -> TensorMapIterMut<'_> {
429        // we can not use `self.blocks_mut()` here, since it would
430        // double-borrow self
431        let mut blocks = Vec::new();
432        for i in 0..self.keys().count() {
433            blocks.push(unsafe { TensorMap::raw_block_mut_by_id(self.ptr, i) });
434        }
435
436        return TensorMapIterMut {
437            inner: self.keys().into_iter().zip(blocks)
438        };
439    }
440
441    /// Get a parallel iterator over the keys and associated blocks
442    #[cfg(feature = "rayon")]
443    #[inline]
444    pub fn par_iter(&self) -> TensorMapParIter<'_> {
445        use rayon::prelude::*;
446        TensorMapParIter {
447            inner: self.keys().par_iter().zip_eq(self.blocks().into_par_iter())
448        }
449    }
450
451    /// Get a parallel iterator over the keys and associated blocks, with
452    /// read-write access to the blocks
453    #[cfg(feature = "rayon")]
454    #[inline]
455    pub fn par_iter_mut(&mut self) -> TensorMapParIterMut<'_> {
456        use rayon::prelude::*;
457
458        // we can not use `self.blocks_mut()` here, since it would
459        // double-borrow self
460        let mut blocks = Vec::new();
461        for i in 0..self.keys().count() {
462            blocks.push(unsafe { TensorMap::raw_block_mut_by_id(self.ptr, i) });
463        }
464
465        TensorMapParIterMut {
466            inner: self.keys().par_iter().zip_eq(blocks)
467        }
468    }
469
470    /// Set or update the info (i.e. global metadata) `value` associated with
471    /// `key` for this `TensorMap`.
472    pub fn set_info(&mut self, key: &str, value: &str) {
473        let mut key = key.to_owned().into_bytes();
474        key.push(b'\0');
475
476        let mut value = value.to_owned().into_bytes();
477        value.push(b'\0');
478
479        unsafe {
480            check_status(crate::c_api::mts_tensormap_set_info(
481                self.ptr, key.as_ptr().cast(), value.as_ptr().cast()
482            )).expect("failed to set info");
483        }
484    }
485
486    /// Get the info (i.e. global metadata) with the given `key` for this
487    /// `TensorMap`.
488    pub fn get_info(&self, key: &str) -> Option<&str> {
489        let mut key = key.to_owned().into_bytes();
490        key.push(b'\0');
491
492        let mut value = std::ptr::null();
493
494        unsafe {
495            check_status(crate::c_api::mts_tensormap_get_info(
496                self.ptr, key.as_ptr().cast(), &mut value
497            )).expect("failed to set info");
498        }
499
500        if value.is_null() {
501            return None;
502        }
503
504        let c_str = unsafe { CStr::from_ptr(value) };
505        return Some(c_str.to_str().expect("invalid UTF-8 string"));
506    }
507
508    /// Get an iterator over all the key/value info pairs stored in this
509    /// `TensorMap`.
510    pub fn info(&self) -> TensorMapInfoIter<'_> {
511        let mut keys = std::ptr::null();
512        let mut count = 0;
513        unsafe {
514            check_status(crate::c_api::mts_tensormap_info_keys(
515                self.ptr,
516                &mut keys,
517                &mut count,
518            )).expect("failed to get info keys");
519        };
520
521        let keys = unsafe {
522            std::slice::from_raw_parts(keys, count)
523        };
524        let keys = keys.iter()
525            .map(|&k| {
526                let c_str = unsafe { CStr::from_ptr(k) };
527                c_str.to_str().expect("invalid UTF-8 string")
528            })
529            .collect::<Vec<_>>();
530
531        TensorMapInfoIter {
532            keys: keys,
533            tensor: self,
534            index: 0,
535            count,
536        }
537    }
538}
539
540/******************************************************************************/
541
542/// Iterator over key/block pairs in a [`TensorMap`]
543pub struct TensorMapIter<'a> {
544    inner: std::iter::Zip<crate::labels::LabelsIter<'a>, std::vec::IntoIter<TensorBlockRef<'a>>>
545}
546
547impl<'a> Iterator for TensorMapIter<'a> {
548    type Item = (&'a [LabelValue], TensorBlockRef<'a>);
549
550    #[inline]
551    fn next(&mut self) -> Option<Self::Item> {
552        self.inner.next()
553    }
554
555    fn size_hint(&self) -> (usize, Option<usize>) {
556        self.inner.size_hint()
557    }
558}
559
560impl ExactSizeIterator for TensorMapIter<'_> {
561    #[inline]
562    fn len(&self) -> usize {
563        self.inner.len()
564    }
565}
566
567impl FusedIterator for TensorMapIter<'_> {}
568
569impl<'a> IntoIterator for &'a TensorMap {
570    type Item = (&'a [LabelValue], TensorBlockRef<'a>);
571
572    type IntoIter = TensorMapIter<'a>;
573
574    fn into_iter(self) -> Self::IntoIter {
575        self.iter()
576    }
577}
578
579/******************************************************************************/
580
581/// Iterator over key/block pairs in a [`TensorMap`], with mutable access to the
582/// blocks
583pub struct TensorMapIterMut<'a> {
584    inner: std::iter::Zip<crate::labels::LabelsIter<'a>, std::vec::IntoIter<TensorBlockRefMut<'a>>>
585}
586
587impl<'a> Iterator for TensorMapIterMut<'a> {
588    type Item = (&'a [LabelValue], TensorBlockRefMut<'a>);
589
590    #[inline]
591    fn next(&mut self) -> Option<Self::Item> {
592        self.inner.next()
593    }
594
595    fn size_hint(&self) -> (usize, Option<usize>) {
596        self.inner.size_hint()
597    }
598}
599
600impl ExactSizeIterator for TensorMapIterMut<'_> {
601    #[inline]
602    fn len(&self) -> usize {
603        self.inner.len()
604    }
605}
606
607impl FusedIterator for TensorMapIterMut<'_> {}
608
609impl<'a> IntoIterator for &'a mut TensorMap {
610    type Item = (&'a [LabelValue], TensorBlockRefMut<'a>);
611
612    type IntoIter = TensorMapIterMut<'a>;
613
614    fn into_iter(self) -> Self::IntoIter {
615        self.iter_mut()
616    }
617}
618
619
620/******************************************************************************/
621
622/// Parallel iterator over key/block pairs in a [`TensorMap`]
623#[cfg(feature = "rayon")]
624pub struct TensorMapParIter<'a> {
625    inner: rayon::iter::ZipEq<crate::labels::LabelsParIter<'a>, rayon::vec::IntoIter<TensorBlockRef<'a>>>,
626}
627
628#[cfg(feature = "rayon")]
629impl<'a> rayon::iter::ParallelIterator for TensorMapParIter<'a> {
630    type Item = (&'a [LabelValue], TensorBlockRef<'a>);
631
632    #[inline]
633    fn drive_unindexed<C>(self, consumer: C) -> C::Result
634    where
635        C: rayon::iter::plumbing::UnindexedConsumer<Self::Item> {
636        self.inner.drive_unindexed(consumer)
637    }
638}
639
640#[cfg(feature = "rayon")]
641impl rayon::iter::IndexedParallelIterator for TensorMapParIter<'_> {
642    #[inline]
643    fn len(&self) -> usize {
644        self.inner.len()
645    }
646
647    #[inline]
648    fn drive<C: rayon::iter::plumbing::Consumer<Self::Item>>(self, consumer: C) -> C::Result {
649        self.inner.drive(consumer)
650    }
651
652    #[inline]
653    fn with_producer<CB: rayon::iter::plumbing::ProducerCallback<Self::Item>>(self, callback: CB) -> CB::Output {
654        self.inner.with_producer(callback)
655    }
656}
657
658/******************************************************************************/
659
660/// Parallel iterator over key/block pairs in a [`TensorMap`], with mutable
661/// access to the blocks
662#[cfg(feature = "rayon")]
663pub struct TensorMapParIterMut<'a> {
664    inner: rayon::iter::ZipEq<crate::labels::LabelsParIter<'a>, rayon::vec::IntoIter<TensorBlockRefMut<'a>>>,
665}
666
667#[cfg(feature = "rayon")]
668impl<'a> rayon::iter::ParallelIterator for TensorMapParIterMut<'a> {
669    type Item = (&'a [LabelValue], TensorBlockRefMut<'a>);
670
671    #[inline]
672    fn drive_unindexed<C>(self, consumer: C) -> C::Result
673    where
674        C: rayon::iter::plumbing::UnindexedConsumer<Self::Item> {
675        self.inner.drive_unindexed(consumer)
676    }
677}
678
679#[cfg(feature = "rayon")]
680impl rayon::iter::IndexedParallelIterator for TensorMapParIterMut<'_> {
681    #[inline]
682    fn len(&self) -> usize {
683        self.inner.len()
684    }
685
686    #[inline]
687    fn drive<C: rayon::iter::plumbing::Consumer<Self::Item>>(self, consumer: C) -> C::Result {
688        self.inner.drive(consumer)
689    }
690
691    #[inline]
692    fn with_producer<CB: rayon::iter::plumbing::ProducerCallback<Self::Item>>(self, callback: CB) -> CB::Output {
693        self.inner.with_producer(callback)
694    }
695}
696
697/******************************************************************************/
698
699/// Iterator over info key/value pairs in a `TensorMap`
700pub struct TensorMapInfoIter<'a> {
701    keys: Vec<&'a str>,
702    tensor: &'a TensorMap,
703    index: usize,
704    count: usize,
705}
706
707impl<'a> Iterator for TensorMapInfoIter<'a> {
708    type Item = (&'a str, &'a str);
709
710    #[inline]
711    fn next(&mut self) -> Option<Self::Item> {
712        if self.index >= self.count {
713            return None;
714        }
715        let key = self.keys[self.index];
716        let value = self.tensor.get_info(key).expect("missing info");
717        self.index += 1;
718        return Some((key, value));
719    }
720
721    fn size_hint(&self) -> (usize, Option<usize>) {
722        (self.count, Some(self.count))
723    }
724}
725
726impl ExactSizeIterator for TensorMapInfoIter<'_> {
727    #[inline]
728    fn len(&self) -> usize {
729        self.count
730    }
731}
732
733impl FusedIterator for TensorMapInfoIter<'_> {}
734
735
736/******************************************************************************/
737
738#[cfg(test)]
739#[allow(clippy::float_cmp)]
740mod tests {
741    use crate::{Labels, TensorBlock, TensorMap};
742
743    fn test_tensor() -> TensorMap {
744        let block_1 = TensorBlock::new(
745            ndarray::Array::from_elem(vec![2, 3], 1.0),
746            &Labels::new(["samples"], [[0], [1]]),
747            &[],
748            &Labels::new(["properties"], [[-2], [0], [1]]),
749        ).unwrap();
750
751        let block_2 = TensorBlock::new(
752            ndarray::Array::from_elem(vec![1, 1], 3.0),
753            &Labels::new(["samples"], [[1]]),
754            &[],
755            &Labels::new(["properties"], [[1]]),
756        ).unwrap();
757
758        let block_3 = TensorBlock::new(
759            ndarray::Array::from_elem(vec![3, 2], -4.0),
760            &Labels::new(["samples"], [[0], [1], [3]]),
761            &[],
762            &Labels::new(["properties"], [[-2], [1]]),
763        ).unwrap();
764
765        return TensorMap::new(
766            Labels::new(["key", "other"], [[1, 0], [3, 1], [-4, 0]]),
767            vec![block_1, block_2, block_3],
768        ).unwrap();
769    }
770
771    #[test]
772    fn block_access() {
773        let mut tensor = test_tensor();
774
775        let block = tensor.block_by_id(1);
776        assert_eq!(block.values().shape().unwrap(), [1, 1]);
777
778        let block = tensor.block_mut_by_id(2);
779        assert_eq!(block.values().shape().unwrap(), [3, 2]);
780
781        let selection = Labels::new(["key"], [[1]]);
782
783        let block = tensor.block(&selection).unwrap();
784        {
785            let values = block.values().to_ndarray_lock::<f64>().read().unwrap();
786            assert_eq!(values.shape(), [2, 3]);
787        }
788
789        let blocks = tensor.blocks();
790        assert_eq!(blocks[0].values().shape().unwrap(), [2, 3]);
791        assert_eq!(blocks[1].values().shape().unwrap(), [1, 1]);
792        assert_eq!(blocks[2].values().shape().unwrap(), [3, 2]);
793
794        let blocks = tensor.blocks_mut();
795        assert_eq!(blocks[0].values().shape().unwrap(), [2, 3]);
796        assert_eq!(blocks[1].values().shape().unwrap(), [1, 1]);
797        assert_eq!(blocks[2].values().shape().unwrap(), [3, 2]);
798    }
799
800    #[test]
801    fn iter() {
802        let mut tensor = test_tensor();
803
804        // iterate over keys & blocks
805        for (key, block) in &tensor {
806            let values = block.values().to_ndarray_lock::<f64>().read().unwrap();
807            assert_eq!(values[[0, 0]], f64::from(key[0].i32()));
808        }
809
810        // iterate over keys & blocks mutably
811        for (key, mut block) in &mut tensor {
812            let array = block.values_mut().get_ndarray_mut::<f64>();
813            *array *= 2.0;
814            assert_eq!(array[[0, 0]], 2.0 * f64::from(key[0].i32()));
815        }
816    }
817
818    #[cfg(feature = "rayon")]
819    #[test]
820    fn par_iter() {
821        use rayon::iter::ParallelIterator;
822
823        let mut tensor = test_tensor();
824
825        // iterate over keys & blocks
826        tensor.par_iter().for_each(|(key, block)| {
827            let values = block.values().to_ndarray_lock::<f64>().read().unwrap();
828            assert_eq!(values[[0, 0]], f64::from(key[0].i32()));
829        });
830
831        // iterate over keys & blocks mutably
832        tensor.par_iter_mut().for_each(|(key, mut block)| {
833            let array = block.values_mut().get_ndarray_mut::<f64>();
834            *array *= 2.0;
835            assert_eq!(array[[0, 0]], 2.0 * f64::from(key[0].i32()));
836        });
837    }
838
839    #[test]
840    fn info() {
841        let mut tensor = test_tensor();
842        tensor.set_info("creator", "unit test");
843        tensor.set_info("version", "1.0");
844
845        assert_eq!(tensor.get_info("creator").unwrap(), "unit test");
846        assert_eq!(tensor.get_info("version").unwrap(), "1.0");
847        assert!(tensor.get_info("missing").is_none());
848
849        let mut info_iter = tensor.info();
850        let (key, value) = info_iter.next().unwrap();
851        assert_eq!(key, "creator");
852        assert_eq!(value, "unit test");
853        let (key, value) = info_iter.next().unwrap();
854        assert_eq!(key, "version");
855        assert_eq!(value, "1.0");
856        assert!(info_iter.next().is_none());
857    }
858
859    #[test]
860    fn device_and_dtype() {
861        let tensor = test_tensor();
862
863        let device = tensor.device().unwrap();
864        assert_eq!(device.device_type, dlpk::sys::DLDeviceType::kDLCPU);
865
866        let dtype = tensor.dtype().unwrap();
867        assert_eq!(dtype.code, dlpk::sys::DLDataTypeCode::kDLFloat);
868        assert_eq!(dtype.bits, 64);
869    }
870
871    #[test]
872    fn tensor_map_into_raw() {
873        let tensor = test_tensor();
874        let raw = TensorMap::into_raw(tensor);
875
876        let recovered = unsafe { TensorMap::from_raw(raw) };
877        assert_eq!(
878            recovered.keys(),
879            &Labels::new(["key", "other"], [[1, 0], [3, 1], [-4, 0]])
880        );
881    }
882}
metatensor/tensor.rs

metatensor/
tensor.rs
