cres 0.9.1

Cell resampling for collider events
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
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165

use crate::traits::Distance;
use crate::vptree::{NearestNeighbourIter, VPTree};

use noisy_float::prelude::*;
use rayon::prelude::*;

/// Algorithm for nearest-neighbour search
pub trait NeighbourSearchAlgo {
    /// Data structure to hold information for nearest-neighbour searches
    type Output<D>;

    /// Initialise nearest neighbour search
    ///
    /// The arguments are the number of points and a function
    /// returning the distance given the indices of two points
    fn new_with_dist<D: Distance<usize> + Send + Sync>(
        npoints: usize,
        d: D,
        max_dist: N64,
    ) -> Self::Output<D>;
}

/// Nearest neighbour search for indexed points
pub trait NeighbourSearch {
    /// Iterator over nearest neighbours
    ///
    /// This has to implement `Iterator<Item = (usize, N64)>`, where
    /// the first tuple element is the index of the nearest neighbour
    /// and the second one the distance.  At the moment it is
    /// unfortunately impossible to enforce this constraint at the
    /// trait level.
    type Iter;

    /// Return nearest neighbours in order for the point with the given index
    fn nearest(self, point: &usize) -> Self::Iter;
}

/// Nearest-neighbour search using a vantage point tree
#[derive(Copy, Clone, Debug, Default, Eq, PartialEq, Ord, PartialOrd, Hash)]
pub struct TreeSearch {}

impl<'a, D> NeighbourSearch for &'a VPTree<usize, D>
where
    D: Distance<usize>,
{
    type Iter = NearestNeighbourIter<'a, usize, D>;

    fn nearest(self, point: &usize) -> Self::Iter {
        self.nearest_in(point)
    }
}

impl NeighbourSearchAlgo for TreeSearch {
    type Output<D> = VPTree<usize, D>;

    fn new_with_dist<D: Distance<usize> + Send + Sync>(
        npoints: usize,
        d: D,
        max_dist: N64,
    ) -> VPTree<usize, D> {
        let range = (0..npoints).into_par_iter();
        VPTree::from_par_iter_with_dist(range, d).with_max_dist(max_dist)
    }
}

/// Naive nearest neighbour search
#[derive(Copy, Clone, Debug, Default, Eq, PartialEq, Ord, PartialOrd, Hash)]
pub struct NaiveNeighbourSearch {}

/// Data required for naive nearest neighbour search
#[derive(Clone, PartialEq, Eq, Ord, PartialOrd, Hash, Debug, Default)]
pub struct NaiveSearchData<D> {
    dist: D,
    cached_dist: Vec<(usize, N64)>,
    max_dist: N64,
}

impl<D> NaiveSearchData<D> {
    /// Convert into the underlying distance
    pub fn into_dist(self) -> D {
        self.dist
    }

    /// Access the distance function
    pub fn dist(&self) -> &D {
        &self.dist
    }

    /// Get maximum allowed distance for nearest neighbours
    pub fn max_dist(&self) -> N64 {
        self.max_dist
    }
}

impl<D: Distance<usize> + Send + Sync> NeighbourSearch for &NaiveSearchData<D> {
    type Iter = NaiveNeighbourIter;

    fn nearest(self, point: &usize) -> Self::Iter {
        let max_dist = self.max_dist;
        let mut dist = self.cached_dist.clone();
        dist.par_iter_mut().for_each(|(id, dist)| {
            *dist = self.dist.distance(id, point);
        });
        NaiveNeighbourIter::new(dist, *point, max_dist)
    }
}

impl NeighbourSearchAlgo for NaiveNeighbourSearch {
    type Output<D> = NaiveSearchData<D>;

    fn new_with_dist<D: Distance<usize> + Send + Sync>(
        npoints: usize,
        dist: D,
        max_dist: N64,
    ) -> Self::Output<D> {
        NaiveSearchData {
            dist,
            cached_dist: Vec::from_iter((0..npoints).map(|id| (id, n64(0.)))),
            max_dist,
        }
    }
}

/// Iterator over nearest neighbours using a naive search algorithm
#[derive(PartialEq, Eq, Ord, PartialOrd, Hash, Debug, Default)]
pub struct NaiveNeighbourIter {
    dist: Vec<(usize, N64)>,
    candidates: Vec<usize>,
    max_dist: N64,
}

impl NaiveNeighbourIter {
    fn new(dist: Vec<(usize, N64)>, seed: usize, max_dist: N64) -> Self {
        let mut candidates = Vec::from_iter(0..dist.len());
        candidates.swap_remove(seed);
        Self {
            dist,
            candidates,
            max_dist,
        }
    }
}

impl Iterator for NaiveNeighbourIter {
    type Item = (usize, N64);

    fn next(&mut self) -> Option<Self::Item> {
        let nearest = self
            .candidates
            .par_iter()
            .enumerate()
            .min_by_key(|(_pos, &idx)| self.dist[idx].1);
        if let Some((pos, &idx)) = nearest {
            let dist = self.dist[idx].1;
            if dist <= self.max_dist {
                self.candidates.swap_remove(pos);
                Some((idx, self.dist[idx].1))
            } else {
                None
            }
        } else {
            None
        }
    }
}