molframe-spatial 0.6.2

Deterministic spatial indices and neighbour search for molframe
Documentation 
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
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129

use super::*;
use crate::queries::lattice_fixture::lattice;
use molframe_core::ExecutionContext;

fn context() -> ExecutionContext {
    ExecutionContext::default()
}

fn pair_indices(pairs: &[NeighborPair]) -> Vec<(u32, u32)> {
    pairs.iter().map(|pair| (pair.first, pair.second)).collect()
}

#[test]
fn parallel_pairs_match_serial_for_every_worker_count() {
    let positions = lattice(16); // 4096 atoms -> cell-list backend
    let all = AtomSelection::All(4096);
    let options = SpatialSearchOptions::with_backend(SpatialBackend::CellList);
    let serial = pairs_within_with_options(&positions, &all, &all, 1.5, options, None, &context())
        .expect("serial cell-list search is valid");
    for workers in [1, 2, 3, 4, 8, 16] {
        let worker_context = ExecutionContext::builder()
            .worker_budget(workers)
            .build()
            .expect("worker context is valid");
        let parallel =
            pairs_within_with_options(&positions, &all, &all, 1.5, options, None, &worker_context)
                .expect("parallel cell-list search is valid");
        assert_eq!(
            pair_indices(&serial),
            pair_indices(&parallel),
            "worker count {workers} changed the pair set"
        );
        // Distances must match too, not only the index pairs.
        assert_eq!(
            serial
                .iter()
                .map(|p| p.distance_squared.to_bits())
                .collect::<Vec<_>>(),
            parallel
                .iter()
                .map(|p| p.distance_squared.to_bits())
                .collect::<Vec<_>>(),
            "worker count {workers} changed a squared distance"
        );
    }
}

fn budgeted(workers: usize, bytes: usize) -> ExecutionContext {
    ExecutionContext::builder()
        .worker_budget(workers)
        .memory_budget(molframe_core::MemoryBudget::new(bytes).expect("a positive budget"))
        .scratch_policy(molframe_core::ScratchPolicy::new(0))
        .build()
        .expect("a valid context")
}

#[test]
fn a_bounded_budget_streams_the_pair_blocks_and_matches_the_unbounded_run() {
    let positions = lattice(16);
    let all = AtomSelection::All(4096);
    let options = SpatialSearchOptions::with_backend(SpatialBackend::CellList);
    let reference = pairs_within_with_options(
        &positions,
        &all,
        &all,
        1.5,
        options,
        None,
        &budgeted(1, 512 * 1024 * 1024),
    )
    .expect("serial run");
    for workers in [2, 4, 8] {
        let context = budgeted(workers, 64 * 1024 * 1024);
        let streamed =
            pairs_within_with_options(&positions, &all, &all, 1.5, options, None, &context)
                .expect("a bounded budget suffices");
        assert_eq!(
            pair_indices(&streamed),
            pair_indices(&reference),
            "{workers} workers"
        );
        assert_eq!(context.reserved_bytes(), 0);
    }
}

#[test]
fn a_budget_below_one_pair_block_is_refused() {
    let positions = lattice(16);
    let all = AtomSelection::All(4096);
    let options = SpatialSearchOptions::with_backend(SpatialBackend::CellList);
    let tiny = budgeted(4, 1024);
    assert!(pairs_within_with_options(&positions, &all, &all, 1.5, options, None, &tiny).is_err());
}

#[test]
fn sparse_index_validation_borrows_the_existing_allocation() {
    let selection = AtomSelection::Sparse(vec![1, 5, 9]);
    let indices = super::checked_indices(&selection, 10).expect("valid indices");
    let AtomSelection::Sparse(original) = &selection else {
        panic!("sparse input")
    };
    assert!(matches!(indices, std::borrow::Cow::Borrowed(_)));
    assert_eq!(indices.as_ptr(), original.as_ptr());
    assert_eq!(
        super::checked_indices(&selection, 9),
        Err(SpatialError::AtomOutOfBounds(9))
    );
}

#[test]
fn merged_ranges_match_brute_force_when_atoms_span_several_ranges() {
    // 22^3 = 10648 atoms exceeds one merge range of first-atom indices.
    let positions = lattice(22);
    let all = AtomSelection::All(10_648);
    let cell = SpatialSearchOptions::with_backend(SpatialBackend::CellList);
    let brute = SpatialSearchOptions::with_backend(SpatialBackend::BruteForce);
    let expected = pairs_within_with_options(&positions, &all, &all, 1.2, brute, None, &context())
        .expect("brute-force search is valid");
    for workers in [1, 2, 5, 16] {
        let worker_context = ExecutionContext::builder()
            .worker_budget(workers)
            .build()
            .expect("worker context is valid");
        let found =
            pairs_within_with_options(&positions, &all, &all, 1.2, cell, None, &worker_context)
                .expect("cell-list search is valid");
        assert_eq!(found, expected, "worker count {workers}");
    }
}