metatensor 0.3.0

Self-describing sparse tensor data format for atomistic machine learning and beyond
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94

#[cfg(not(feature = "bench"))]
compile_error!("the bench feature is required for bencharks, use `cargo bench --features bench`");

#[cfg(feature = "bench")]
mod benchmarks {
    use criterion::{BatchSize, BenchmarkId, Criterion};
    pub use criterion::{criterion_group, criterion_main};
    use metatensor::{Labels, LabelsBuilder, MtsArray, TensorBlock, TensorMap};

    fn tensor_map(n_blocks: usize, n_samples: usize, n_properties: usize) -> TensorMap {
        let mut keys = LabelsBuilder::new(vec!["key_1", "key_2"]);
        for i in 0..n_blocks {
            keys.add(&[i, 0]);
        }
        let keys = keys.finish();

        let mut blocks = Vec::new();
        for _ in 0..n_blocks {
            let mut samples = LabelsBuilder::new(vec!["samples"]);
            for i in 0..n_samples {
                samples.add(&[i]);
            }
            let samples = samples.finish();

            let mut components_builder = LabelsBuilder::new(vec!["components"]);
            components_builder.add(&[0]);
            let components = [components_builder.finish()];

            let mut properties = LabelsBuilder::new(vec!["properties"]);
            for i in 0..n_properties {
                properties.add(&[i]);
            }
            let properties = properties.finish();

            let shape = vec![samples.count(), components[0].count(), properties.count()];
            let data = ndarray::Array::from_elem(shape, 1.0);

            blocks.push(TensorBlock::new(data, &samples, &components, &properties).unwrap());
        }

        TensorMap::new(keys, blocks).unwrap()
    }

    fn make_fill_value(scalar: f64) -> MtsArray {
        MtsArray::from(ndarray::Array::from_elem(vec![], scalar))
    }

    pub fn keys_to_samples(c: &mut Criterion) {
        let mut group = c.benchmark_group("TensorMap::keys_to_samples");
        // reduce sample size for faster benchmark execution during dev/test
        group.sample_size(10);

        for n_blocks in [1, 10, 100] {
            let tensor = tensor_map(n_blocks, 1000, 1000);
            let keys_to_move = Labels::empty(vec!["key_1"]);
            group.bench_function(BenchmarkId::new("n_blocks", n_blocks), |b| {
                b.iter_batched(
                    || tensor.try_clone().unwrap(),
                    |tensor| std::hint::black_box(tensor.keys_to_samples(&keys_to_move, make_fill_value(0.0), true)),
                    BatchSize::LargeInput,
                );
            });
        }
        group.finish();
    }

    pub fn keys_to_properties(c: &mut Criterion) {
        let mut group = c.benchmark_group("TensorMap::keys_to_properties");
        group.sample_size(10);

        for n_blocks in [1, 10, 100] {
            let tensor = tensor_map(n_blocks, 1000, 1000);
            let keys_to_move = Labels::empty(vec!["key_1"]);
            group.bench_function(BenchmarkId::new("n_blocks", n_blocks), |b| {
                b.iter_batched(
                    || tensor.try_clone().unwrap(),
                    |tensor| std::hint::black_box(tensor.keys_to_properties(&keys_to_move, make_fill_value(0.0), true)),
                    BatchSize::LargeInput,
                );
            });
        }
        group.finish();
    }
}


#[cfg(feature = "bench")]
benchmarks::criterion_group!(benches, benchmarks::keys_to_samples, benchmarks::keys_to_properties);

#[cfg(feature = "bench")]
benchmarks::criterion_main!(benches);

#[cfg(not(feature = "bench"))]
fn main() {}