Line data Source code
1 : #include "CartesianIndices.h"
2 : #include "TypeCasts.hpp"
3 : #include <iostream>
4 : #include <iterator>
5 : #include <numeric>
6 :
7 : namespace elsa
8 : {
9 : template <typename Fn>
10 : auto map(std::vector<IndexRange> v, Fn fn)
11 : {
12 : IndexVector_t result(v.size());
13 : std::transform(v.begin(), v.end(), result.begin(), std::move(fn));
14 : return result;
15 : }
16 :
17 : template <std::size_t N, typename Fn>
18 : auto map(std::array<IndexRange, N> v, index_t size, Fn fn)
19 : {
20 : IndexVector_t result(size);
21 : std::transform(v.begin(), v.begin() + size, result.begin(), std::move(fn));
22 : return result;
23 : }
24 :
25 0 : CartesianIndices::CartesianIndices(std::vector<IndexRange> ranges) : size_(ranges.size()) {}
26 :
27 : CartesianIndices::CartesianIndices(const IndexVector_t& to)
28 : : size_(to.size()), first_(IndexVector_t::Zero(size_)), last_(to)
29 1665 : {
30 1665 : }
31 :
32 : CartesianIndices::CartesianIndices(const IndexVector_t& from, const IndexVector_t& to)
33 : : size_(to.size()), first_(from), last_(to)
34 7944 : {
35 7944 : if (from.size() != to.size()) {
36 0 : throw InvalidArgumentError("CartesianIndices: vectors must be of same size");
37 0 : }
38 7944 : }
39 :
40 20 : auto CartesianIndices::dims() const -> index_t { return size_; }
41 :
42 491 : auto CartesianIndices::size() const -> index_t { return (last_ - first_).prod(); }
43 :
44 9 : auto CartesianIndices::range(index_t i) const -> IndexRange { return {first_[i], last_[i]}; }
45 :
46 24096 : auto CartesianIndices::first() -> IndexVector_t& { return first_; }
47 :
48 66 : auto CartesianIndices::first() const -> const IndexVector_t& { return first_; }
49 :
50 16065 : auto CartesianIndices::last() -> IndexVector_t& { return last_; }
51 :
52 40 : auto CartesianIndices::last() const -> const IndexVector_t& { return last_; }
53 :
54 8043 : auto CartesianIndices::begin() -> iterator { return iterator(first(), first(), last()); }
55 :
56 26 : auto CartesianIndices::begin() const -> iterator { return {first(), first(), last()}; }
57 :
58 8029 : auto CartesianIndices::end() -> iterator { return {as_sentinel_t, first(), last()}; }
59 :
60 11 : auto CartesianIndices::end() const -> iterator { return {as_sentinel_t, first(), last()}; }
61 :
62 : /****** Iterator implementation ******/
63 : CartesianIndices::iterator::iterator(as_sentinel, const IndexVector_t& begin,
64 : const IndexVector_t& end)
65 : : cur_(end), begins_(end), ends_(end)
66 8038 : {
67 8038 : cur_.tail(cur_.size() - 1) = begin.tail(cur_.size() - 1);
68 8038 : }
69 :
70 : CartesianIndices::iterator::iterator(const IndexVector_t& vec)
71 : : cur_(vec), begins_(vec), ends_(vec)
72 0 : {
73 0 : }
74 :
75 : CartesianIndices::iterator::iterator(const IndexVector_t& cur, const IndexVector_t& begin,
76 : const IndexVector_t& end)
77 : : cur_(cur), begins_(begin), ends_(end)
78 8068 : {
79 8068 : }
80 :
81 200898 : auto CartesianIndices::iterator::operator*() const -> const IndexVector_t& { return cur_; }
82 :
83 : auto CartesianIndices::iterator::operator++() -> iterator&
84 201390 : {
85 201390 : inc();
86 201390 : return *this;
87 201390 : }
88 :
89 : auto CartesianIndices::iterator::operator++(int) -> iterator
90 1 : {
91 1 : auto copy = *this;
92 1 : ++(*this);
93 1 : return copy;
94 1 : }
95 :
96 : auto CartesianIndices::iterator::operator--() -> iterator&
97 2 : {
98 2 : prev();
99 2 : return *this;
100 2 : }
101 :
102 : auto CartesianIndices::iterator::operator--(int) -> iterator
103 1 : {
104 1 : auto copy = *this;
105 1 : --(*this);
106 1 : return copy;
107 1 : }
108 :
109 : auto CartesianIndices::iterator::operator+=(difference_type n) -> iterator&
110 34 : {
111 34 : advance(n);
112 34 : return *this;
113 34 : }
114 :
115 : auto CartesianIndices::iterator::operator-=(difference_type n) -> iterator&
116 14 : {
117 14 : advance(-n);
118 14 : return *this;
119 14 : }
120 :
121 : auto operator+(const CartesianIndices::iterator& iter,
122 : CartesianIndices::iterator::difference_type n) -> CartesianIndices::iterator
123 12 : {
124 12 : auto copy = iter;
125 12 : copy += n;
126 12 : return copy;
127 12 : }
128 :
129 : auto operator+(CartesianIndices::iterator::difference_type n,
130 : const CartesianIndices::iterator& iter) -> CartesianIndices::iterator
131 6 : {
132 6 : return iter + n;
133 6 : }
134 :
135 : auto operator-(const CartesianIndices::iterator& iter,
136 : CartesianIndices::iterator::difference_type n) -> CartesianIndices::iterator
137 6 : {
138 6 : auto copy = iter;
139 6 : copy -= n;
140 6 : return copy;
141 6 : }
142 :
143 : auto operator-(const CartesianIndices::iterator& lhs, const CartesianIndices::iterator& rhs)
144 : -> CartesianIndices::iterator::difference_type
145 35 : {
146 35 : return rhs.distance_to(lhs);
147 35 : }
148 :
149 : auto CartesianIndices::iterator::operator[](difference_type n) const -> value_type
150 5 : {
151 : // TODO: Make this more efficient
152 5 : auto copy = *this;
153 5 : copy += n;
154 5 : return *copy;
155 5 : }
156 :
157 : auto operator<(const CartesianIndices::iterator& lhs, const CartesianIndices::iterator& rhs)
158 : -> bool
159 16 : {
160 : // Not sure why I have to instantiate them this way, else it doesn't work...
161 16 : using Array = Eigen::Array<index_t, Eigen::Dynamic, 1>;
162 16 : const Array a1 = lhs.cur_.array();
163 16 : const Array a2 = rhs.cur_.array();
164 :
165 16 : return (a1 < a2).any();
166 16 : }
167 :
168 : auto operator>(const CartesianIndices::iterator& lhs, const CartesianIndices::iterator& rhs)
169 : -> bool
170 8 : {
171 8 : return rhs < lhs;
172 8 : }
173 :
174 : auto operator<=(const CartesianIndices::iterator& lhs, const CartesianIndices::iterator& rhs)
175 : -> bool
176 4 : {
177 4 : return !(lhs > rhs);
178 4 : }
179 :
180 : auto operator>=(const CartesianIndices::iterator& lhs, const CartesianIndices::iterator& rhs)
181 : -> bool
182 4 : {
183 4 : return !(lhs < rhs);
184 4 : }
185 :
186 0 : auto CartesianIndices::iterator::at_end() -> bool { return cur_[0] == ends_[0]; }
187 :
188 0 : auto CartesianIndices::iterator::at_end() const -> bool { return cur_[0] == ends_[0]; }
189 :
190 : auto CartesianIndices::iterator::inc_recursive(index_t N) -> void
191 244706 : {
192 244706 : auto& iter = cur_[N];
193 :
194 244706 : if (++iter == ends_[N] && N > 0) {
195 43461 : iter = begins_[N];
196 43461 : inc_recursive(N - 1);
197 43461 : }
198 244706 : }
199 :
200 201412 : auto CartesianIndices::iterator::inc() -> void { inc_recursive(cur_.size() - 1); }
201 :
202 : auto CartesianIndices::iterator::prev_recursive(index_t N) -> void
203 2 : {
204 2 : auto& iter = cur_[N];
205 :
206 2 : if (iter == begins_[N] && N > 0) {
207 0 : iter = ends_[N];
208 0 : inc_recursive(N - 1);
209 0 : }
210 2 : --iter;
211 2 : }
212 :
213 2 : auto CartesianIndices::iterator::prev() -> void { prev_recursive(cur_.size() - 1); }
214 :
215 : auto CartesianIndices::iterator::distance_to_recusrive(const iterator& other,
216 : difference_type N) const
217 : -> difference_type
218 70 : {
219 70 : if (N == 0) {
220 35 : return other.cur_[N] - cur_[N];
221 35 : } else {
222 35 : const auto d = distance_to_recusrive(other, N - 1);
223 35 : const auto scale = ends_[N] - begins_[N];
224 35 : const auto increment = other.cur_[N] - cur_[N];
225 :
226 35 : return difference_type{(d * scale) + increment};
227 35 : }
228 70 : }
229 :
230 : auto CartesianIndices::iterator::distance_to(const iterator& other) const -> difference_type
231 35 : {
232 35 : auto N = cur_.size() - 1;
233 35 : return distance_to_recusrive(other, N);
234 35 : }
235 :
236 : auto CartesianIndices::iterator::advance_recursive(difference_type n, index_t N) -> void
237 76 : {
238 76 : if (n == 0) {
239 18 : return;
240 18 : }
241 :
242 58 : auto& iter = cur_[N];
243 58 : const auto size = ends_[N] - begins_[N];
244 58 : const auto first = begins_[N];
245 :
246 58 : auto const idx = as<difference_type>(iter - first);
247 58 : n += idx;
248 :
249 58 : auto div = size ? n / size : 0;
250 58 : auto mod = size ? n % size : 0;
251 :
252 58 : if (N != 0) {
253 28 : if (mod < 0) {
254 0 : mod += size;
255 0 : div--;
256 0 : }
257 28 : advance_recursive(div, N - 1);
258 30 : } else {
259 30 : if (div > 0) {
260 7 : mod = size;
261 7 : }
262 30 : }
263 :
264 58 : iter = first + mod;
265 58 : }
266 :
267 : auto CartesianIndices::iterator::advance(difference_type n) -> void
268 48 : {
269 48 : advance_recursive(n, cur_.size() - 1);
270 48 : }
271 :
272 : CartesianIndices neighbours_in_slice(const IndexVector_t& pos, const IndexVector_t& dist)
273 2 : {
274 2 : Eigen::IOFormat format(4, 0, ", ", "", "", "", "[", "]");
275 2 : IndexVector_t from = pos;
276 2 : from.tail(pos.size() - 1).array() -= dist.array();
277 :
278 2 : IndexVector_t to = pos;
279 2 : to.tail(pos.size() - 1).array() += dist.array() + 1;
280 2 : to[0] += 1; // Be sure this is incremented, so we actually iterate over it
281 :
282 2 : return {from, to};
283 2 : }
284 :
285 : CartesianIndices neighbours_in_slice(const IndexVector_t& pos, index_t dist)
286 2 : {
287 2 : return neighbours_in_slice(pos, IndexVector_t::Constant(pos.size() - 1, dist));
288 2 : }
289 :
290 : CartesianIndices neighbours_in_slice(const IndexVector_t& pos, const IndexVector_t& dist,
291 : const IndexVector_t& lower, const IndexVector_t& upper)
292 7935 : {
293 7935 : IndexVector_t from = pos;
294 7935 : from.array() -= dist.array();
295 7935 : from = from.cwiseMax(lower);
296 :
297 7935 : IndexVector_t to = pos;
298 7935 : to.array() += dist.array() + 1;
299 7935 : to = to.cwiseMin(upper);
300 :
301 : // FIXME: sometimes this happens when padding is added to an aabb
302 7935 : from = (to.array() == from.array()).select(from.array() - 1, from);
303 :
304 7935 : return {from, to};
305 7935 : }
306 :
307 : CartesianIndices neighbours_in_slice(const IndexVector_t& pos, index_t dist, index_t leadingDim,
308 : const IndexVector_t& lower, const IndexVector_t& upper)
309 2 : {
310 2 : IndexVector_t distvec = IndexVector_t::Constant(pos.size(), dist);
311 2 : distvec[leadingDim] = 0;
312 :
313 2 : return neighbours_in_slice(pos, distvec, lower, upper);
314 2 : }
315 : } // namespace elsa
|