Line data Source code
1 1 : // Distributed under the MIT License.
2 : // See LICENSE.txt for details.
3 :
4 : /// \file
5 : /// Declares function RootFinder::toms748
6 :
7 : #pragma once
8 :
9 : #include <cstdint>
10 : #include <functional>
11 : #include <limits>
12 : #include <optional>
13 : #include <ostream>
14 : #include <type_traits>
15 : #include <utility>
16 :
17 : #include "DataStructures/DataVector.hpp"
18 : #include "Utilities/ErrorHandling/Assert.hpp"
19 : #include "Utilities/ErrorHandling/CaptureForError.hpp"
20 : #include "Utilities/ErrorHandling/Error.hpp"
21 : #include "Utilities/Expected.hpp"
22 : #include "Utilities/Simd/Simd.hpp"
23 :
24 : namespace RootFinder {
25 : /// \ingroup NumericalAlgorithmsGroup
26 : /// Why a `RootFinder::toms748` root find failed
27 1 : enum class Toms748Failure : uint8_t {
28 : /// `f(lower_bound)` and `f(upper_bound)` are nonzero and have the same sign
29 : RootNotBracketed,
30 : /// `lower_bound > upper_bound`
31 : LowerBoundLargerThanUpperBound,
32 : /// The tolerance wasn't met after the maximum number of iterations
33 : MaxIterations
34 : };
35 :
36 : /*!
37 : * \ingroup NumericalAlgorithmsGroup
38 : * \brief The error of a failed `RootFinder::toms748` root find.
39 : *
40 : * Holds the arguments of the root find. Stream it to get an error message.
41 : */
42 : template <typename T>
43 1 : struct Toms748Error {
44 0 : Toms748Failure failure{};
45 0 : T lower_bound{};
46 0 : T upper_bound{};
47 0 : T f_at_lower_bound{};
48 0 : T f_at_upper_bound{};
49 0 : simd::scalar_type_t<T> absolute_tolerance{};
50 0 : simd::scalar_type_t<T> relative_tolerance{};
51 : };
52 :
53 : template <typename T>
54 0 : std::ostream& operator<<(std::ostream& os, const Toms748Error<T>& error) {
55 : switch (error.failure) {
56 : case Toms748Failure::RootNotBracketed:
57 : return os << "Root not bracketed: f(" << error.lower_bound
58 : << ") = " << error.f_at_lower_bound << ", f("
59 : << error.upper_bound << ") = " << error.f_at_upper_bound;
60 : case Toms748Failure::LowerBoundLargerThanUpperBound:
61 : return os << "Lower bound is larger than upper bound: "
62 : << error.lower_bound << " > " << error.upper_bound;
63 : case Toms748Failure::MaxIterations:
64 : return os << "toms748 reached max iterations without converging.\n"
65 : << "Absolute tolerance: " << error.absolute_tolerance
66 : << "\nRelative tolerance: " << error.relative_tolerance;
67 : default:
68 : ERROR("Unknown Toms748Failure");
69 : }
70 : }
71 :
72 : namespace toms748_detail {
73 : // Original implementation of TOMS748 is from Boost:
74 : // (C) Copyright John Maddock 2006.
75 : // Use, modification and distribution are subject to the
76 : // Boost Software License, Version 1.0. (See accompanying file
77 : // LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)
78 : //
79 : // Significant changes made to pretty much all of it to support SIMD by Nils
80 : // Deppe. Changes are copyrighted by SXS Collaboration under the MIT License.
81 :
82 : template <typename T>
83 : T safe_div(const T& num, const T& denom, const T& r) {
84 : if constexpr (std::is_floating_point_v<T>) {
85 : using std::abs;
86 : if (abs(denom) < static_cast<T>(1)) {
87 : if (abs(denom * std::numeric_limits<T>::max()) <= abs(num)) {
88 : return r;
89 : }
90 : }
91 : return num / denom;
92 : } else {
93 : // return num / denom without overflow, return r if overflow would occur.
94 : //
95 : // To do this we need to handle the following cases:
96 : // 1. fabs(denom) < 1 AND fabs(denom * max) <= fabs(num):
97 : // return r
98 : // 2. return num / denom
99 : //
100 : // We do this by creating 2 masks.
101 : // 1. `mask0` selects where fabs(denom) < 1. This is where we _may_ have
102 : // issues with division by zero or overflow.
103 : // 2. `mask` selects where fabs(denom) < 1 AND fabs(denom * max)<=fabs(num)
104 : // * Note: the edge case of fabs(num)==max could be problematic, but if
105 : // you're dealing with numbers like that you are likely in
106 : // trouble anyway.
107 : // The second mask is where we must avoid division by zero and instead
108 : // return `r`.
109 : const auto mask0 = fabs(denom) < static_cast<T>(1);
110 : // Note: if denom >= 1.0 you get an FPE because of overflow from
111 : // `max() * (1. + value)`, which is why `mask0` is necessary.
112 : const auto mask = fabs(simd::select(mask0, denom, static_cast<T>(1.0)) *
113 : std::numeric_limits<T>::max()) <= fabs(num);
114 : const auto new_denom = simd::select(mask, static_cast<T>(1), denom);
115 : return simd::select(mask, r, num / new_denom);
116 : }
117 : }
118 :
119 : template <typename T>
120 : T secant_interpolate(const T& a, const T& b, const T& fa, const T& fb,
121 : const simd::mask_type_t<T>& incomplete_mask) {
122 : //
123 : // Performs standard secant interpolation of [a,b] given
124 : // function evaluations f(a) and f(b). Performs a bisection
125 : // if secant interpolation would leave us very close to either
126 : // a or b. Rationale: we only call this function when at least
127 : // one other form of interpolation has already failed, so we know
128 : // that the function is unlikely to be smooth with a root very
129 : // close to a or b.
130 : //
131 : // For lanes where `incomplete_mask` is false (already converged), the
132 : // denominator `fb - fa` may be stale or zero. We substitute 1 to avoid
133 : // undefined behavior (e.g. division by zero); the returned value for those
134 : // lanes is meaningless and should not be used by the caller. See
135 : // `cubic_interpolate` for the same pattern.
136 : //
137 : const T tol_batch = std::numeric_limits<T>::epsilon() * static_cast<T>(5);
138 : // WARNING: There are several different ways to implement the interpolation
139 : // that all have different rounding properties. Unfortunately this means that
140 : // tolerances even at 1e-14 can be difficult to achieve generically.
141 : //
142 : // `g` below is:
143 : // const T g = (fa / (fb - fa));
144 : //
145 : // const T c = simd::fma((fa / (fb - fa)), (a - b), a); // fails
146 : //
147 : // const T c = a * (static_cast<T>(1) - g) + g * b; // works
148 : //
149 : // const T c = simd::fma(a, (static_cast<T>(1) - g), g * b); // works
150 : //
151 : // Original Boost code:
152 : // const T c = a - (fa / (fb - fa)) * (b - a); // works
153 :
154 : const T c =
155 : a - (fa / simd::select(incomplete_mask, fb - fa, static_cast<T>(1))) *
156 : (b - a);
157 : return simd::select((c <= simd::fma(fabs(a), tol_batch, a)) or
158 : (c >= simd::fnma(fabs(b), tol_batch, b)),
159 : static_cast<T>(0.5) * (a + b), c);
160 : }
161 :
162 : template <bool AssumeFinite, typename T>
163 : T quadratic_interpolate(const T& a, const T& b, const T& d, const T& fa,
164 : const T& fb, const T& fd,
165 : const simd::mask_type_t<T>& incomplete_mask,
166 : const unsigned count) {
167 : // Performs quadratic interpolation to determine the next point,
168 : // takes count Newton steps to find the location of the
169 : // quadratic polynomial.
170 : //
171 : // Point d must lie outside of the interval [a,b], it is the third
172 : // best approximation to the root, after a and b.
173 : //
174 : // Note: this does not guarantee to find a root
175 : // inside [a, b], so we fall back to a secant step should
176 : // the result be out of range.
177 : //
178 : // Start by obtaining the coefficients of the quadratic polynomial:
179 : const T B = safe_div(fb - fa, b - a, std::numeric_limits<T>::max());
180 : T A = safe_div(fd - fb, d - b, std::numeric_limits<T>::max());
181 : A = safe_div(A - B, d - a, static_cast<T>(0));
182 :
183 : const auto secant_failure_mask = A == static_cast<T>(0) and incomplete_mask;
184 : T result_secant{};
185 : if (UNLIKELY(simd::any(secant_failure_mask))) {
186 : // failure to determine coefficients, try a secant step:
187 : result_secant = secant_interpolate(a, b, fa, fb, secant_failure_mask);
188 : if (UNLIKELY(simd::all(secant_failure_mask or (not incomplete_mask)))) {
189 : return result_secant;
190 : }
191 : }
192 :
193 : // Determine the starting point of the Newton steps:
194 : //
195 : // Note: unlike Boost, we assume A*fa doesn't overflow. This speeds up the
196 : // code quite a bit.
197 : T c = AssumeFinite
198 : ? simd::select(A * fa > static_cast<T>(0), a, b)
199 : : simd::select(simd::sign(A) * simd::sign(fa) > static_cast<T>(0),
200 : a, b);
201 :
202 : // Take the Newton steps:
203 : const T two_A = static_cast<T>(2) * A;
204 : const T half_a_plus_b = 0.5 * (a + b);
205 : const T one_minus_a = static_cast<T>(1) - a;
206 : const T B_minus_A_times_b = B - A * b;
207 : for (unsigned i = 1; i <= count; ++i) {
208 : c -= safe_div(simd::fma(simd::fma(A, c, B_minus_A_times_b), c - a, fa),
209 : simd::fma(two_A, c - half_a_plus_b, B), one_minus_a + c);
210 : }
211 : if (const auto mask = ((c <= a) or (c >= b)) and incomplete_mask;
212 : simd::any(mask)) {
213 : // Failure, try a secant step:
214 : c = simd::select(mask, secant_interpolate(a, b, fa, fb, mask), c);
215 : }
216 : return simd::select(secant_failure_mask, result_secant, c);
217 : }
218 :
219 : template <bool AssumeFinite, typename T>
220 : T cubic_interpolate(const T& a, const T& b, const T& d, const T& e, const T& fa,
221 : const T& fb, const T& fd, const T& fe,
222 : const simd::mask_type_t<T>& incomplete_mask) {
223 : // Uses inverse cubic interpolation of f(x) at points
224 : // [a,b,d,e] to obtain an approximate root of f(x).
225 : // Points d and e lie outside the interval [a,b]
226 : // and are the third and forth best approximations
227 : // to the root that we have found so far.
228 : //
229 : // Note: this does not guarantee to find a root
230 : // inside [a, b], so we fall back to quadratic
231 : // interpolation in case of an erroneous result.
232 : //
233 : // This commented chunk is the original Boost implementation translated into
234 : // simd. The actual code below is a heavily optimized version.
235 : //
236 : // const T q11 = (d - e) * fd / (fe - fd);
237 : // const T q21 = (b - d) * fb / (fd - fb);
238 : // const T q31 = (a - b) * fa / (fb - fa);
239 : // const T d21 = (b - d) * fd / (fd - fb);
240 : // const T d31 = (a - b) * fb / (fb - fa);
241 : //
242 : // const T q22 = (d21 - q11) * fb / (fe - fb);
243 : // const T q32 = (d31 - q21) * fa / (fd - fa);
244 : // const T d32 = (d31 - q21) * fd / (fd - fa);
245 : // const T q33 = (d32 - q22) * fa / (fe - fa);
246 : //
247 : // T c = q31 + q32 + q33 + a;
248 :
249 : // The optimized implementation here is L1-cache bound. That is, we aren't
250 : // able to completely saturate the FP units because we are waiting on the L1
251 : // cache. While not ideal, that's okay and just part of the algorithm.
252 : const T denom_fb_fa = fb - fa;
253 : const T denom_fd_fb = fd - fb;
254 : const T denom_fd_fa = fd - fa;
255 : const T denom_fe_fd = fe - fd;
256 : const T denom_fe_fb = fe - fb;
257 : const T denom =
258 : denom_fe_fb * denom_fe_fd * denom_fd_fa * denom_fd_fb * (fe - fa);
259 :
260 : // Avoid division by zero with mask.
261 : const T fa_by_denom = fa / simd::select(incomplete_mask, denom_fb_fa * denom,
262 : static_cast<T>(1));
263 :
264 : const T d31 = (a - b);
265 : const T q21 = (b - d);
266 : const T q32 = simd::fms(denom_fd_fb, d31, denom_fb_fa * q21) * denom_fe_fb *
267 : denom_fe_fd;
268 : const T q22 = simd::fms(denom_fe_fd, q21, (d - e) * denom_fd_fb) * fd *
269 : denom_fb_fa * denom_fd_fa;
270 :
271 : // Note: the reduction in rounding error that comes from the improvement by
272 : // Stoer & Bulirsch to Neville's algorithm is adding `a` at the very end as
273 : // we do below. Alternative ways of evaluating polynomials do not delay this
274 : // inclusion of `a`, and so then when the correction to `a` is small,
275 : // floating point errors decrease the accuracy of the result.
276 : T c = simd::fma(
277 : fa_by_denom,
278 : simd::fma(fb, simd::fms(q32, (fe + denom_fd_fa), q22), d31 * denom), a);
279 :
280 : if (const auto mask = ((c <= a) or (c >= b)) and incomplete_mask;
281 : simd::any(mask)) {
282 : // Out of bounds step, fall back to quadratic interpolation:
283 : //
284 : // Note: we only apply quadratic interpolation at points where cubic
285 : // failed and that aren't already at a root.
286 : c = simd::select(
287 : mask, quadratic_interpolate<AssumeFinite>(a, b, d, fa, fb, fd, mask, 3),
288 : c);
289 : }
290 :
291 : return c;
292 : }
293 :
294 : template <bool AssumeFinite, typename F, typename T>
295 : void bracket(F f, T& a, T& b, T c, T& fa, T& fb, T& d, T& fd,
296 : const simd::mask_type_t<T>& incomplete_mask) {
297 : // Given a point c inside the existing enclosing interval
298 : // [a, b] sets a = c if f(c) == 0, otherwise finds the new
299 : // enclosing interval: either [a, c] or [c, b] and sets
300 : // d and fd to the point that has just been removed from
301 : // the interval. In other words d is the third best guess
302 : // to the root.
303 : //
304 : // Note: `bracket` will only modify slots marked as `true` in
305 : // `incomplete_mask`
306 : const T tol_batch = std::numeric_limits<T>::epsilon() * static_cast<T>(2);
307 :
308 : // If the interval [a,b] is very small, or if c is too close
309 : // to one end of the interval then we need to adjust the
310 : // location of c accordingly. This is:
311 : //
312 : // if ((b - a) < 2 * tol * a) {
313 : // c = a + (b - a) / 2;
314 : // } else if (c <= a + fabs(a) * tol) {
315 : // c = a + fabs(a) * tol;
316 : // } else if (c >= b - fabs(b) * tol) {
317 : // c = b - fabs(b) * tol;
318 : // }
319 : const T a_filt = simd::fma(fabs(a), tol_batch, a);
320 : const T b_filt = simd::fnma(fabs(b), tol_batch, b);
321 : const T b_minus_a = b - a;
322 : c = simd::select(
323 : (static_cast<T>(2) * tol_batch * a > b_minus_a) and incomplete_mask,
324 : simd::fma(b_minus_a, static_cast<T>(0.5), a),
325 : simd::select(c <= a_filt, a_filt, simd::select(c >= b_filt, b_filt, c)));
326 :
327 : // Invoke f(c):
328 : T fc = f(c);
329 :
330 : // if we have a zero then we have an exact solution to the root:
331 : const auto fc_is_zero_mask = (fc == static_cast<T>(0));
332 : if (const auto mask = fc_is_zero_mask and incomplete_mask;
333 : UNLIKELY(simd::any(mask))) {
334 : a = simd::select(mask, c, a);
335 : fa = simd::select(mask, static_cast<T>(0), fa);
336 : d = simd::select(mask, static_cast<T>(0), d);
337 : fd = simd::select(mask, static_cast<T>(0), fd);
338 : if (UNLIKELY(simd::all(mask or not incomplete_mask))) {
339 : return;
340 : }
341 : }
342 :
343 : // Non-zero fc, update the interval:
344 : //
345 : // Note: unlike Boost, we assume fa*fc doesn't overflow. This speeds up the
346 : // code quite a bit.
347 : //
348 : // Boost code is:
349 : // if (boost::math::sign(fa) * boost::math::sign(fc) < 0) {...} else {...}
350 : using simd::sign;
351 : const auto sign_mask = AssumeFinite
352 : ? (fa * fc < static_cast<T>(0))
353 : : (sign(fa) * sign(fc) < static_cast<T>(0));
354 : const auto mask_if =
355 : (sign_mask and (not fc_is_zero_mask)) and incomplete_mask;
356 : d = simd::select(mask_if, b, d);
357 : fd = simd::select(mask_if, fb, fd);
358 : b = simd::select(mask_if, c, b);
359 : fb = simd::select(mask_if, fc, fb);
360 :
361 : const auto mask_else =
362 : ((not sign_mask) and (not fc_is_zero_mask)) and incomplete_mask;
363 : d = simd::select(mask_else, a, d);
364 : fd = simd::select(mask_else, fa, fd);
365 : a = simd::select(mask_else, c, a);
366 : fa = simd::select(mask_else, fc, fa);
367 : }
368 :
369 : template <bool AssumeFinite, class F, class T, class Tol>
370 : std::pair<T, T> toms748_solve(F f, const T& ax, const T& bx, const T& fax,
371 : const T& fbx, Tol tol,
372 : const simd::mask_type_t<T>& ignore_filter,
373 : size_t& max_iter,
374 : std::optional<Toms748Failure>& failure) {
375 : // Main entry point and logic for Toms Algorithm 748
376 : // root finder.
377 : if (UNLIKELY(simd::any(ax > bx))) {
378 : failure = Toms748Failure::LowerBoundLargerThanUpperBound;
379 : return std::pair{ax, bx};
380 : }
381 :
382 : // Sanity check - are we allowed to iterate at all?
383 : if (UNLIKELY(max_iter == 0)) {
384 : return std::pair{ax, bx};
385 : }
386 :
387 : size_t count = max_iter;
388 : // mu is a parameter in the algorithm that must be between (0, 1).
389 : static const T mu = 0.5f;
390 :
391 : // initialise a, b and fa, fb:
392 : T a = ax;
393 : T b = bx;
394 : T fa = fax;
395 : T fb = fbx;
396 : CAPTURE_FOR_ERROR(ax);
397 : CAPTURE_FOR_ERROR(bx);
398 : CAPTURE_FOR_ERROR(fax);
399 : CAPTURE_FOR_ERROR(fbx);
400 :
401 : const auto fa_is_zero_mask = (fa == static_cast<T>(0));
402 : const auto fb_is_zero_mask = (fb == static_cast<T>(0));
403 : auto completion_mask =
404 : tol(a, b) or fa_is_zero_mask or fb_is_zero_mask or ignore_filter;
405 : auto incomplete_mask = not completion_mask;
406 : if (UNLIKELY(simd::all(completion_mask))) {
407 : max_iter = 0;
408 : return std::pair{simd::select(fb_is_zero_mask, b, a),
409 : simd::select(fa_is_zero_mask, a, b)};
410 : }
411 :
412 : // `toms748_bracket` already checked that [a, b] brackets the root
413 : // dummy value for fd, e and fe:
414 : T fe(static_cast<T>(1e5F));
415 : T e(static_cast<T>(1e5F));
416 : T fd(static_cast<T>(1e5F));
417 :
418 : T c(0.0);
419 : T d(0.0);
420 :
421 : // We can't use signaling_NaN() because we do a max comparison at the end to
422 : // 0, and signaling_NaN when compared to 0 raises an FPE.
423 : const T nan(std::numeric_limits<T>::quiet_NaN());
424 : auto completed_a = simd::select(completion_mask, a, nan);
425 : auto completed_b = simd::select(completion_mask, b, nan);
426 : auto completed_fa = simd::select(completion_mask, fa, nan);
427 : auto completed_fb = simd::select(completion_mask, fb, nan);
428 : const auto update_completed = [&fa, &fb, &completion_mask, &incomplete_mask,
429 : &completed_a, &completed_b, &a, &b,
430 : &completed_fa, &completed_fb, &tol]() {
431 : const auto new_completed =
432 : (fa == static_cast<T>(0) or tol(a, b)) and (not completion_mask);
433 : completed_a = simd::select(new_completed, a, completed_a);
434 : completed_b = simd::select(new_completed, b, completed_b);
435 : completed_fa = simd::select(new_completed, fa, completed_fa);
436 : completed_fb = simd::select(new_completed, fb, completed_fb);
437 : completion_mask = new_completed or completion_mask;
438 : incomplete_mask = not completion_mask;
439 : // returns true if _all_ simd registers have been completed
440 : return simd::all(completion_mask);
441 : };
442 :
443 : if (simd::any(fa != static_cast<T>(0))) {
444 : // On the first step we take a secant step:
445 : c = toms748_detail::secant_interpolate(a, b, fa, fb, incomplete_mask);
446 : toms748_detail::bracket<AssumeFinite>(f, a, b, c, fa, fb, d, fd,
447 : incomplete_mask);
448 : --count;
449 :
450 : // Note: The Boost fa!=0 check is handled with the completion_mask.
451 : if (count and not update_completed()) {
452 : // On the second step we take a quadratic interpolation:
453 : c = toms748_detail::quadratic_interpolate<AssumeFinite>(
454 : a, b, d, fa, fb, fd, incomplete_mask, 2);
455 : e = d;
456 : fe = fd;
457 : toms748_detail::bracket<AssumeFinite>(f, a, b, c, fa, fb, d, fd,
458 : incomplete_mask);
459 : --count;
460 : update_completed();
461 : }
462 : }
463 :
464 : T u(std::numeric_limits<T>::signaling_NaN());
465 : T fu(std::numeric_limits<T>::signaling_NaN());
466 : T a0(std::numeric_limits<T>::signaling_NaN());
467 : T b0(std::numeric_limits<T>::signaling_NaN());
468 :
469 : // Note: The Boost fa!=0 check is handled with the completion_mask.
470 : while (count and not simd::all(completion_mask)) {
471 : // save our brackets:
472 : a0 = a;
473 : b0 = b;
474 : // Starting with the third step taken
475 : // we can use either quadratic or cubic interpolation.
476 : // Cubic interpolation requires that all four function values
477 : // fa, fb, fd, and fe are distinct, should that not be the case
478 : // then variable prof will get set to true, and we'll end up
479 : // taking a quadratic step instead.
480 : static const T min_diff = std::numeric_limits<T>::min() * 32;
481 : bool prof =
482 : simd::any(((fabs(fa - fb) < min_diff) or (fabs(fa - fd) < min_diff) or
483 : (fabs(fa - fe) < min_diff) or (fabs(fb - fd) < min_diff) or
484 : (fabs(fb - fe) < min_diff) or (fabs(fd - fe) < min_diff)) and
485 : incomplete_mask);
486 : if (prof) {
487 : c = toms748_detail::quadratic_interpolate<AssumeFinite>(
488 : a, b, d, fa, fb, fd, incomplete_mask, 2);
489 : } else {
490 : c = toms748_detail::cubic_interpolate<AssumeFinite>(
491 : a, b, d, e, fa, fb, fd, fe, incomplete_mask);
492 : }
493 : // re-bracket, and check for termination:
494 : e = d;
495 : fe = fd;
496 : toms748_detail::bracket<AssumeFinite>(f, a, b, c, fa, fb, d, fd,
497 : incomplete_mask);
498 : if ((0 == --count) or update_completed()) {
499 : break;
500 : }
501 : // Now another interpolated step:
502 : prof =
503 : simd::any(((fabs(fa - fb) < min_diff) or (fabs(fa - fd) < min_diff) or
504 : (fabs(fa - fe) < min_diff) or (fabs(fb - fd) < min_diff) or
505 : (fabs(fb - fe) < min_diff) or (fabs(fd - fe) < min_diff)) and
506 : incomplete_mask);
507 : if (prof) {
508 : c = toms748_detail::quadratic_interpolate<AssumeFinite>(
509 : a, b, d, fa, fb, fd, incomplete_mask, 3);
510 : } else {
511 : c = toms748_detail::cubic_interpolate<AssumeFinite>(
512 : a, b, d, e, fa, fb, fd, fe, incomplete_mask);
513 : }
514 : // Bracket again, and check termination condition, update e:
515 : toms748_detail::bracket<AssumeFinite>(f, a, b, c, fa, fb, d, fd,
516 : incomplete_mask);
517 : if ((0 == --count) or update_completed()) {
518 : break;
519 : }
520 :
521 : // Now we take a double-length secant step:
522 : const auto fabs_fa_less_fabs_fb_mask =
523 : (fabs(fa) < fabs(fb)) and incomplete_mask;
524 : u = simd::select(fabs_fa_less_fabs_fb_mask, a, b);
525 : fu = simd::select(fabs_fa_less_fabs_fb_mask, fa, fb);
526 : const T b_minus_a = b - a;
527 : // Assumes that bounds a & b are not so close that fa == fb. Boost makes
528 : // this assumption too. If this is violated then the algorithm doesn't
529 : // work since the function at least appears constant.
530 :
531 : // Safeguard complete lanes: `fb - fa` may be stale or zero for lanes
532 : // that have already converged. Use 1 as a harmless dummy denominator;
533 : // the result for complete lanes is discarded downstream.
534 : c = simd::fnma(
535 : static_cast<T>(2) *
536 : (fu / simd::select(incomplete_mask, fb - fa, static_cast<T>(1))),
537 : b_minus_a, u);
538 : c = simd::select(static_cast<T>(2) * fabs(c - u) > b_minus_a,
539 : simd::fma(static_cast<T>(0.5), b_minus_a, a), c);
540 :
541 : // Bracket again, and check termination condition:
542 : e = d;
543 : fe = fd;
544 : toms748_detail::bracket<AssumeFinite>(f, a, b, c, fa, fb, d, fd,
545 : incomplete_mask);
546 : if ((0 == --count) or update_completed()) {
547 : break;
548 : }
549 :
550 : // And finally... check to see if an additional bisection step is
551 : // to be taken, we do this if we're not converging fast enough:
552 : const auto bisection_mask = (b - a) >= mu * (b0 - a0) and incomplete_mask;
553 : if (LIKELY(simd::none(bisection_mask))) {
554 : continue;
555 : }
556 : // bracket again on a bisection:
557 : //
558 : // Note: the mask ensures we only ever modify the slots that the mask has
559 : // identified as needing to be modified.
560 : e = simd::select(bisection_mask, d, e);
561 : fe = simd::select(bisection_mask, fd, fe);
562 : toms748_detail::bracket<AssumeFinite>(
563 : f, a, b, simd::fma((b - a), static_cast<T>(0.5), a), fa, fb, d, fd,
564 : bisection_mask);
565 : --count;
566 : if (update_completed()) {
567 : break;
568 : }
569 : } // while loop
570 :
571 : max_iter -= count;
572 : completed_b =
573 : simd::select(completed_fa == static_cast<T>(0), completed_a, completed_b);
574 : completed_a =
575 : simd::select(completed_fb == static_cast<T>(0), completed_b, completed_a);
576 : return std::pair{completed_a, completed_b};
577 : }
578 :
579 : // The root find of `toms748`. Returns the final bracket of the root and sets
580 : // `failure` if the root find failed.
581 : template <bool AssumeFinite, typename Function, typename T>
582 : std::pair<T, T> toms748_bracket(std::optional<Toms748Failure>& failure,
583 : const Function& f, const T lower_bound,
584 : const T upper_bound, const T f_at_lower_bound,
585 : const T f_at_upper_bound,
586 : const simd::scalar_type_t<T> absolute_tolerance,
587 : const simd::scalar_type_t<T> relative_tolerance,
588 : const size_t max_iterations,
589 : const simd::mask_type_t<T> ignore_filter) {
590 : ASSERT(relative_tolerance >
591 : std::numeric_limits<simd::scalar_type_t<T>>::epsilon(),
592 : "The relative tolerance is too small. Got "
593 : << relative_tolerance << " but must be at least "
594 : << std::numeric_limits<simd::scalar_type_t<T>>::epsilon());
595 : // Note: unlike Boost, we can assume fa*fb doesn't overflow or underflow when
596 : // possible. This speeds up the code quite a bit.
597 : if (UNLIKELY(simd::any(
598 : AssumeFinite
599 : ? (f_at_lower_bound * f_at_upper_bound > static_cast<T>(0))
600 : : (simd::sign(f_at_lower_bound) * simd::sign(f_at_upper_bound) >
601 : static_cast<T>(0))))) {
602 : failure = Toms748Failure::RootNotBracketed;
603 : return std::pair{lower_bound, upper_bound};
604 : }
605 :
606 : std::size_t max_iters = max_iterations;
607 :
608 : // This solver requires tol to be passed as a termination condition. This
609 : // termination condition is equivalent to the convergence criteria used by the
610 : // GSL
611 : const auto tol = [absolute_tolerance, relative_tolerance](const T& lhs,
612 : const T& rhs) {
613 : return simd::abs(lhs - rhs) <=
614 : simd::fma(T(relative_tolerance),
615 : simd::min(simd::abs(lhs), simd::abs(rhs)),
616 : T(absolute_tolerance));
617 : };
618 : const auto result = toms748_solve<AssumeFinite>(
619 : f, lower_bound, upper_bound, f_at_lower_bound, f_at_upper_bound, tol,
620 : ignore_filter, max_iters, failure);
621 : if (not failure.has_value() and max_iters >= max_iterations) {
622 : failure = Toms748Failure::MaxIterations;
623 : }
624 : return result;
625 : }
626 : } // namespace toms748_detail
627 :
628 : /*!
629 : * \ingroup NumericalAlgorithmsGroup
630 : * \brief Finds the root of the function `f` with the TOMS_748 method.
631 : *
632 : * `f` is a unary invokable that takes a `double` which is the current value at
633 : * which to evaluate `f`. An example is below.
634 : *
635 : * \snippet Test_TOMS748.cpp double_root_find
636 : *
637 : * The TOMS_748 algorithm searches for a root in the interval [`lower_bound`,
638 : * `upper_bound`], and fails if this interval does not bracket a root,
639 : * i.e. if `f(lower_bound) * f(upper_bound) > 0`.
640 : *
641 : * The arguments `f_at_lower_bound` and `f_at_upper_bound` are optional, and
642 : * are the function values at `lower_bound` and `upper_bound`. These function
643 : * values are often known because the user typically checks if a root is
644 : * bracketed before calling `toms748`; passing the function values here saves
645 : * two function evaluations.
646 : *
647 : * \note if `AssumeFinite` is true than the code assumes all numbers are
648 : * finite and that `a > 0` is equivalent to `sign(a) > 0`. This reduces
649 : * runtime but will cause bugs if the numbers aren't finite. It also assumes
650 : * that products like `fa * fb` are also finite.
651 : *
652 : * \requires Function `f` is invokable with a `double`
653 : *
654 : * \returns the root, or a `Toms748Error` if the interval doesn't bracket a root
655 : * or the requested tolerance is not met after `max_iterations` iterations.
656 : * Call `value()` on the result to get the root and error if the root find
657 : * failed.
658 : */
659 : template <bool AssumeFinite = false, typename Function, typename T>
660 1 : Expected<T, Toms748Error<T>> toms748(
661 : const Function& f, const T lower_bound, const T upper_bound,
662 : const T f_at_lower_bound, const T f_at_upper_bound,
663 : const simd::scalar_type_t<T> absolute_tolerance,
664 : const simd::scalar_type_t<T> relative_tolerance,
665 : const size_t max_iterations = 100,
666 : const simd::mask_type_t<T> ignore_filter =
667 : static_cast<simd::mask_type_t<T>>(0)) {
668 : std::optional<Toms748Failure> failure{};
669 : const auto result = toms748_detail::toms748_bracket<AssumeFinite>(
670 : failure, f, lower_bound, upper_bound, f_at_lower_bound, f_at_upper_bound,
671 : absolute_tolerance, relative_tolerance, max_iterations, ignore_filter);
672 : if (UNLIKELY(failure.has_value())) {
673 : return Unexpected{Toms748Error<T>{*failure, lower_bound, upper_bound,
674 : f_at_lower_bound, f_at_upper_bound,
675 : absolute_tolerance, relative_tolerance}};
676 : }
677 : return simd::fma(static_cast<T>(0.5), (result.second - result.first),
678 : result.first);
679 : }
680 :
681 : /*!
682 : * \ingroup NumericalAlgorithmsGroup
683 : * \brief Finds the root of the function `f` with the TOMS_748 method, where
684 : * function values are not supplied at the lower and upper bounds.
685 : *
686 : * \note if `AssumeFinite` is true than the code assumes all numbers are
687 : * finite and that `a > 0` is equivalent to `sign(a) > 0`. This reduces
688 : * runtime but will cause bugs if the numbers aren't finite. It also assumes
689 : * that products like `fa * fb` are also finite.
690 : */
691 : template <bool AssumeFinite = false, typename Function, typename T>
692 1 : Expected<T, Toms748Error<T>> toms748(
693 : const Function& f, const T lower_bound, const T upper_bound,
694 : const simd::scalar_type_t<T> absolute_tolerance,
695 : const simd::scalar_type_t<T> relative_tolerance,
696 : const size_t max_iterations = 100,
697 : const simd::mask_type_t<T> ignore_filter =
698 : static_cast<simd::mask_type_t<T>>(0)) {
699 : return toms748<AssumeFinite>(
700 : f, lower_bound, upper_bound, f(lower_bound), f(upper_bound),
701 : absolute_tolerance, relative_tolerance, max_iterations, ignore_filter);
702 : }
703 :
704 : /*!
705 : * \ingroup NumericalAlgorithmsGroup
706 : * \brief Finds the root of the function `f` with the TOMS_748 method on each
707 : * element in a `DataVector`.
708 : *
709 : * `f` is a binary invokable that takes a `double` as its first argument and a
710 : * `size_t` as its second. The `double` is the current value at which to
711 : * evaluate `f`, and the `size_t` is the current index into the `DataVector`s.
712 : * Below is an example of how to root find different functions by indexing into
713 : * a lambda-captured `DataVector` using the `size_t` passed to `f`.
714 : *
715 : * \snippet Test_TOMS748.cpp datavector_root_find
716 : *
717 : * For each index `i` into the DataVector, the TOMS_748 algorithm searches for a
718 : * root in the interval [`lower_bound[i]`, `upper_bound[i]`], and fails if
719 : * this interval does not bracket a root,
720 : * i.e. if `f(lower_bound[i], i) * f(upper_bound[i], i) > 0`.
721 : *
722 : * See the [Boost](http://www.boost.org/) documentation for more details.
723 : *
724 : * \requires Function `f` be callable with a `double` and a `size_t`
725 : *
726 : * \note if `AssumeFinite` is true than the code assumes all numbers are
727 : * finite and that `a > 0` is equivalent to `sign(a) > 0`. This reduces
728 : * runtime but will cause bugs if the numbers aren't finite. It also assumes
729 : * that products like `fa * fb` are also finite.
730 : *
731 : * \returns the roots, or the `Toms748Error` of the first index whose root find
732 : * failed.
733 : */
734 : template <bool UseSimd = true, bool AssumeFinite = false, typename Function>
735 1 : Expected<DataVector, Toms748Error<double>> toms748(
736 : const Function& f, const DataVector& lower_bound,
737 : const DataVector& upper_bound, const double absolute_tolerance,
738 : const double relative_tolerance, const size_t max_iterations = 100) {
739 : DataVector result_vector{lower_bound.size()};
740 : size_t first_scalar_index = 0;
741 : if constexpr (UseSimd) {
742 : constexpr size_t simd_width{simd::size<
743 : std::decay_t<decltype(simd::load_unaligned(lower_bound.data()))>>()};
744 : const size_t vectorized_size =
745 : lower_bound.size() - lower_bound.size() % simd_width;
746 : for (; first_scalar_index < vectorized_size;
747 : first_scalar_index += simd_width) {
748 : const size_t i = first_scalar_index;
749 : const auto root = toms748<AssumeFinite>(
750 : [&f, i](const auto x) { return f(x, i); },
751 : simd::load_unaligned(&lower_bound[i]),
752 : simd::load_unaligned(&upper_bound[i]), absolute_tolerance,
753 : relative_tolerance, max_iterations);
754 : if (UNLIKELY(not root.has_value())) {
755 : // Root find these points one by one below to get the error of the
756 : // point that failed
757 : break;
758 : }
759 : simd::store_unaligned(&result_vector[i], root.value());
760 : }
761 : }
762 : for (size_t i = first_scalar_index; i < lower_bound.size(); ++i) {
763 : const auto root = toms748<AssumeFinite>(
764 : [&f, i](const double x) { return f(x, i); }, lower_bound[i],
765 : upper_bound[i], absolute_tolerance, relative_tolerance, max_iterations);
766 : if (UNLIKELY(not root.has_value())) {
767 : return Unexpected{root.error()};
768 : }
769 : result_vector[i] = root.value();
770 : }
771 : return result_vector;
772 : }
773 :
774 : /*!
775 : * \ingroup NumericalAlgorithmsGroup
776 : * \brief Finds the root of the function `f` with the TOMS_748 method on each
777 : * element in a `DataVector`, where function values are supplied at the lower
778 : * and upper bounds.
779 : *
780 : * Supplying function values is an optimization that saves two
781 : * function calls per point. The function values are often available
782 : * because one often checks if the root is bracketed before calling `toms748`.
783 : *
784 : * \note if `AssumeFinite` is true than the code assumes all numbers are
785 : * finite and that `a > 0` is equivalent to `sign(a) > 0`. This reduces
786 : * runtime but will cause bugs if the numbers aren't finite. It also assumes
787 : * that products like `fa * fb` are also finite.
788 : */
789 : template <bool UseSimd = true, bool AssumeFinite = false, typename Function>
790 1 : Expected<DataVector, Toms748Error<double>> toms748(
791 : const Function& f, const DataVector& lower_bound,
792 : const DataVector& upper_bound, const DataVector& f_at_lower_bound,
793 : const DataVector& f_at_upper_bound, const double absolute_tolerance,
794 : const double relative_tolerance, const size_t max_iterations = 100) {
795 : DataVector result_vector{lower_bound.size()};
796 : size_t first_scalar_index = 0;
797 : if constexpr (UseSimd) {
798 : constexpr size_t simd_width{simd::size<
799 : std::decay_t<decltype(simd::load_unaligned(lower_bound.data()))>>()};
800 : const size_t vectorized_size =
801 : lower_bound.size() - lower_bound.size() % simd_width;
802 : for (; first_scalar_index < vectorized_size;
803 : first_scalar_index += simd_width) {
804 : const size_t i = first_scalar_index;
805 : const auto root = toms748<AssumeFinite>(
806 : [&f, i](const auto x) { return f(x, i); },
807 : simd::load_unaligned(&lower_bound[i]),
808 : simd::load_unaligned(&upper_bound[i]),
809 : simd::load_unaligned(&f_at_lower_bound[i]),
810 : simd::load_unaligned(&f_at_upper_bound[i]), absolute_tolerance,
811 : relative_tolerance, max_iterations);
812 : if (UNLIKELY(not root.has_value())) {
813 : // Root find these points one by one below to get the error of the
814 : // point that failed
815 : break;
816 : }
817 : simd::store_unaligned(&result_vector[i], root.value());
818 : }
819 : }
820 : for (size_t i = first_scalar_index; i < lower_bound.size(); ++i) {
821 : const auto root = toms748<AssumeFinite>(
822 : [&f, i](const double x) { return f(x, i); }, lower_bound[i],
823 : upper_bound[i], f_at_lower_bound[i], f_at_upper_bound[i],
824 : absolute_tolerance, relative_tolerance, max_iterations);
825 : if (UNLIKELY(not root.has_value())) {
826 : return Unexpected{root.error()};
827 : }
828 : result_vector[i] = root.value();
829 : }
830 : return result_vector;
831 : }
832 : } // namespace RootFinder
|
