autodiff.h 14 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363
  1. // Ceres Solver - A fast non-linear least squares minimizer
  2. // Copyright 2023 Google Inc. All rights reserved.
  3. // http://ceres-solver.org/
  4. //
  5. // Redistribution and use in source and binary forms, with or without
  6. // modification, are permitted provided that the following conditions are met:
  7. //
  8. // * Redistributions of source code must retain the above copyright notice,
  9. // this list of conditions and the following disclaimer.
  10. // * Redistributions in binary form must reproduce the above copyright notice,
  11. // this list of conditions and the following disclaimer in the documentation
  12. // and/or other materials provided with the distribution.
  13. // * Neither the name of Google Inc. nor the names of its contributors may be
  14. // used to endorse or promote products derived from this software without
  15. // specific prior written permission.
  16. //
  17. // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  18. // AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
  19. // IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
  20. // ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
  21. // LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
  22. // CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
  23. // SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
  24. // INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
  25. // CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
  26. // ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
  27. // POSSIBILITY OF SUCH DAMAGE.
  28. //
  29. // Author: keir@google.com (Keir Mierle)
  30. //
  31. // Computation of the Jacobian matrix for vector-valued functions of multiple
  32. // variables, using automatic differentiation based on the implementation of
  33. // dual numbers in jet.h. Before reading the rest of this file, it is advisable
  34. // to read jet.h's header comment in detail.
  35. //
  36. // The helper wrapper AutoDifferentiate() computes the jacobian of
  37. // functors with templated operator() taking this form:
  38. //
  39. // struct F {
  40. // template<typename T>
  41. // bool operator()(const T *x, const T *y, ..., T *z) {
  42. // // Compute z[] based on x[], y[], ...
  43. // // return true if computation succeeded, false otherwise.
  44. // }
  45. // };
  46. //
  47. // All inputs and outputs may be vector-valued.
  48. //
  49. // To understand how jets are used to compute the jacobian, a
  50. // picture may help. Consider a vector-valued function, F, returning 3
  51. // dimensions and taking a vector-valued parameter of 4 dimensions:
  52. //
  53. // y x
  54. // [ * ] F [ * ]
  55. // [ * ] <--- [ * ]
  56. // [ * ] [ * ]
  57. // [ * ]
  58. //
  59. // Similar to the 2-parameter example for f described in jet.h, computing the
  60. // jacobian dy/dx is done by substituting a suitable jet object for x and all
  61. // intermediate steps of the computation of F. Since x is has 4 dimensions, use
  62. // a Jet<double, 4>.
  63. //
  64. // Before substituting a jet object for x, the dual components are set
  65. // appropriately for each dimension of x:
  66. //
  67. // y x
  68. // [ * | * * * * ] f [ * | 1 0 0 0 ] x0
  69. // [ * | * * * * ] <--- [ * | 0 1 0 0 ] x1
  70. // [ * | * * * * ] [ * | 0 0 1 0 ] x2
  71. // ---+--- [ * | 0 0 0 1 ] x3
  72. // | ^ ^ ^ ^
  73. // dy/dx | | | +----- infinitesimal for x3
  74. // | | +------- infinitesimal for x2
  75. // | +--------- infinitesimal for x1
  76. // +----------- infinitesimal for x0
  77. //
  78. // The reason to set the internal 4x4 submatrix to the identity is that we wish
  79. // to take the derivative of y separately with respect to each dimension of x.
  80. // Each column of the 4x4 identity is therefore for a single component of the
  81. // independent variable x.
  82. //
  83. // Then the jacobian of the mapping, dy/dx, is the 3x4 sub-matrix of the
  84. // extended y vector, indicated in the above diagram.
  85. //
  86. // Functors with multiple parameters
  87. // ---------------------------------
  88. // In practice, it is often convenient to use a function f of two or more
  89. // vector-valued parameters, for example, x[3] and z[6]. Unfortunately, the jet
  90. // framework is designed for a single-parameter vector-valued input. The wrapper
  91. // in this file addresses this issue adding support for functions with one or
  92. // more parameter vectors.
  93. //
  94. // To support multiple parameters, all the parameter vectors are concatenated
  95. // into one and treated as a single parameter vector, except that since the
  96. // functor expects different inputs, we need to construct the jets as if they
  97. // were part of a single parameter vector. The extended jets are passed
  98. // separately for each parameter.
  99. //
  100. // For example, consider a functor F taking two vector parameters, p[2] and
  101. // q[3], and producing an output y[4]:
  102. //
  103. // struct F {
  104. // template<typename T>
  105. // bool operator()(const T *p, const T *q, T *z) {
  106. // // ...
  107. // }
  108. // };
  109. //
  110. // In this case, the necessary jet type is Jet<double, 5>. Here is a
  111. // visualization of the jet objects in this case:
  112. //
  113. // Dual components for p ----+
  114. // |
  115. // -+-
  116. // y [ * | 1 0 | 0 0 0 ] --- p[0]
  117. // [ * | 0 1 | 0 0 0 ] --- p[1]
  118. // [ * | . . | + + + ] |
  119. // [ * | . . | + + + ] v
  120. // [ * | . . | + + + ] <--- F(p, q)
  121. // [ * | . . | + + + ] ^
  122. // ^^^ ^^^^^ |
  123. // dy/dp dy/dq [ * | 0 0 | 1 0 0 ] --- q[0]
  124. // [ * | 0 0 | 0 1 0 ] --- q[1]
  125. // [ * | 0 0 | 0 0 1 ] --- q[2]
  126. // --+--
  127. // |
  128. // Dual components for q --------------+
  129. //
  130. // where the 4x2 submatrix (marked with ".") and 4x3 submatrix (marked with "+"
  131. // of y in the above diagram are the derivatives of y with respect to p and q
  132. // respectively. This is how autodiff works for functors taking multiple vector
  133. // valued arguments (up to 6).
  134. //
  135. // Jacobian null pointers (nullptr)
  136. // --------------------------------
  137. // In general, the functions below will accept nullptr for all or some of the
  138. // Jacobian parameters, meaning that those Jacobians will not be computed.
  139. #ifndef CERES_PUBLIC_INTERNAL_AUTODIFF_H_
  140. #define CERES_PUBLIC_INTERNAL_AUTODIFF_H_
  141. #include <array>
  142. #include <cstddef>
  143. #include <utility>
  144. #include "ceres/internal/array_selector.h"
  145. #include "ceres/internal/eigen.h"
  146. #include "ceres/internal/fixed_array.h"
  147. #include "ceres/internal/parameter_dims.h"
  148. #include "ceres/internal/variadic_evaluate.h"
  149. #include "ceres/jet.h"
  150. #include "ceres/types.h"
  151. #include "glog/logging.h"
  152. // If the number of parameters exceeds this values, the corresponding jets are
  153. // placed on the heap. This will reduce performance by a factor of 2-5 on
  154. // current compilers.
  155. #ifndef CERES_AUTODIFF_MAX_PARAMETERS_ON_STACK
  156. #define CERES_AUTODIFF_MAX_PARAMETERS_ON_STACK 50
  157. #endif
  158. #ifndef CERES_AUTODIFF_MAX_RESIDUALS_ON_STACK
  159. #define CERES_AUTODIFF_MAX_RESIDUALS_ON_STACK 20
  160. #endif
  161. namespace ceres::internal {
  162. // Extends src by a 1st order perturbation for every dimension and puts it in
  163. // dst. The size of src is N. Since this is also used for perturbations in
  164. // blocked arrays, offset is used to shift which part of the jet the
  165. // perturbation occurs. This is used to set up the extended x augmented by an
  166. // identity matrix. The JetT type should be a Jet type, and T should be a
  167. // numeric type (e.g. double). For example,
  168. //
  169. // 0 1 2 3 4 5 6 7 8
  170. // dst[0] [ * | . . | 1 0 0 | . . . ]
  171. // dst[1] [ * | . . | 0 1 0 | . . . ]
  172. // dst[2] [ * | . . | 0 0 1 | . . . ]
  173. //
  174. // is what would get put in dst if N was 3, offset was 3, and the jet type JetT
  175. // was 8-dimensional.
  176. template <int j, int N, int Offset, typename T, typename JetT>
  177. struct Make1stOrderPerturbation {
  178. public:
  179. inline static void Apply(const T* src, JetT* dst) {
  180. if (j == 0) {
  181. DCHECK(src);
  182. DCHECK(dst);
  183. }
  184. dst[j] = JetT(src[j], j + Offset);
  185. Make1stOrderPerturbation<j + 1, N, Offset, T, JetT>::Apply(src, dst);
  186. }
  187. };
  188. template <int N, int Offset, typename T, typename JetT>
  189. struct Make1stOrderPerturbation<N, N, Offset, T, JetT> {
  190. public:
  191. static void Apply(const T* /* NOT USED */, JetT* /* NOT USED */) {}
  192. };
  193. // Calls Make1stOrderPerturbation for every parameter block.
  194. //
  195. // Example:
  196. // If one having three parameter blocks with dimensions (3, 2, 4), the call
  197. // Make1stOrderPerturbations<integer_sequence<3, 2, 4>::Apply(params, x);
  198. // will result in the following calls to Make1stOrderPerturbation:
  199. // Make1stOrderPerturbation<0, 3, 0>::Apply(params[0], x + 0);
  200. // Make1stOrderPerturbation<0, 2, 3>::Apply(params[1], x + 3);
  201. // Make1stOrderPerturbation<0, 4, 5>::Apply(params[2], x + 5);
  202. template <typename Seq, int ParameterIdx = 0, int Offset = 0>
  203. struct Make1stOrderPerturbations;
  204. template <int N, int... Ns, int ParameterIdx, int Offset>
  205. struct Make1stOrderPerturbations<std::integer_sequence<int, N, Ns...>,
  206. ParameterIdx,
  207. Offset> {
  208. template <typename T, typename JetT>
  209. inline static void Apply(T const* const* parameters, JetT* x) {
  210. Make1stOrderPerturbation<0, N, Offset, T, JetT>::Apply(
  211. parameters[ParameterIdx], x + Offset);
  212. Make1stOrderPerturbations<std::integer_sequence<int, Ns...>,
  213. ParameterIdx + 1,
  214. Offset + N>::Apply(parameters, x);
  215. }
  216. };
  217. // End of 'recursion'. Nothing more to do.
  218. template <int ParameterIdx, int Total>
  219. struct Make1stOrderPerturbations<std::integer_sequence<int>,
  220. ParameterIdx,
  221. Total> {
  222. template <typename T, typename JetT>
  223. static void Apply(T const* const* /* NOT USED */, JetT* /* NOT USED */) {}
  224. };
  225. // Takes the 0th order part of src, assumed to be a Jet type, and puts it in
  226. // dst. This is used to pick out the "vector" part of the extended y.
  227. template <typename JetT, typename T>
  228. inline void Take0thOrderPart(int M, const JetT* src, T dst) {
  229. DCHECK(src);
  230. for (int i = 0; i < M; ++i) {
  231. dst[i] = src[i].a;
  232. }
  233. }
  234. // Takes N 1st order parts, starting at index N0, and puts them in the M x N
  235. // matrix 'dst'. This is used to pick out the "matrix" parts of the extended y.
  236. template <int N0, int N, typename JetT, typename T>
  237. inline void Take1stOrderPart(const int M, const JetT* src, T* dst) {
  238. DCHECK(src);
  239. DCHECK(dst);
  240. for (int i = 0; i < M; ++i) {
  241. Eigen::Map<Eigen::Matrix<T, N, 1>>(dst + N * i, N) =
  242. src[i].v.template segment<N>(N0);
  243. }
  244. }
  245. // Calls Take1stOrderPart for every parameter block.
  246. //
  247. // Example:
  248. // If one having three parameter blocks with dimensions (3, 2, 4), the call
  249. // Take1stOrderParts<integer_sequence<3, 2, 4>::Apply(num_outputs,
  250. // output,
  251. // jacobians);
  252. // will result in the following calls to Take1stOrderPart:
  253. // if (jacobians[0]) {
  254. // Take1stOrderPart<0, 3>(num_outputs, output, jacobians[0]);
  255. // }
  256. // if (jacobians[1]) {
  257. // Take1stOrderPart<3, 2>(num_outputs, output, jacobians[1]);
  258. // }
  259. // if (jacobians[2]) {
  260. // Take1stOrderPart<5, 4>(num_outputs, output, jacobians[2]);
  261. // }
  262. template <typename Seq, int ParameterIdx = 0, int Offset = 0>
  263. struct Take1stOrderParts;
  264. template <int N, int... Ns, int ParameterIdx, int Offset>
  265. struct Take1stOrderParts<std::integer_sequence<int, N, Ns...>,
  266. ParameterIdx,
  267. Offset> {
  268. template <typename JetT, typename T>
  269. inline static void Apply(int num_outputs, JetT* output, T** jacobians) {
  270. if (jacobians[ParameterIdx]) {
  271. Take1stOrderPart<Offset, N>(num_outputs, output, jacobians[ParameterIdx]);
  272. }
  273. Take1stOrderParts<std::integer_sequence<int, Ns...>,
  274. ParameterIdx + 1,
  275. Offset + N>::Apply(num_outputs, output, jacobians);
  276. }
  277. };
  278. // End of 'recursion'. Nothing more to do.
  279. template <int ParameterIdx, int Offset>
  280. struct Take1stOrderParts<std::integer_sequence<int>, ParameterIdx, Offset> {
  281. template <typename T, typename JetT>
  282. static void Apply(int /* NOT USED*/,
  283. JetT* /* NOT USED*/,
  284. T** /* NOT USED */) {}
  285. };
  286. template <int kNumResiduals,
  287. typename ParameterDims,
  288. typename Functor,
  289. typename T>
  290. inline bool AutoDifferentiate(const Functor& functor,
  291. T const* const* parameters,
  292. int dynamic_num_outputs,
  293. T* function_value,
  294. T** jacobians) {
  295. using JetT = Jet<T, ParameterDims::kNumParameters>;
  296. using Parameters = typename ParameterDims::Parameters;
  297. if (kNumResiduals != DYNAMIC) {
  298. DCHECK_EQ(kNumResiduals, dynamic_num_outputs);
  299. }
  300. ArraySelector<JetT,
  301. ParameterDims::kNumParameters,
  302. CERES_AUTODIFF_MAX_PARAMETERS_ON_STACK>
  303. parameters_as_jets(ParameterDims::kNumParameters);
  304. // Pointers to the beginning of each parameter block
  305. std::array<JetT*, ParameterDims::kNumParameterBlocks> unpacked_parameters =
  306. ParameterDims::GetUnpackedParameters(parameters_as_jets.data());
  307. // If the number of residuals is fixed, we use the template argument as the
  308. // number of outputs. Otherwise we use the num_outputs parameter. Note: The
  309. // ?-operator here is compile-time evaluated, therefore num_outputs is also
  310. // a compile-time constant for functors with fixed residuals.
  311. const int num_outputs =
  312. kNumResiduals == DYNAMIC ? dynamic_num_outputs : kNumResiduals;
  313. DCHECK_GT(num_outputs, 0);
  314. ArraySelector<JetT, kNumResiduals, CERES_AUTODIFF_MAX_RESIDUALS_ON_STACK>
  315. residuals_as_jets(num_outputs);
  316. // Invalidate the output Jets, so that we can detect if the user
  317. // did not assign values to all of them.
  318. for (int i = 0; i < num_outputs; ++i) {
  319. residuals_as_jets[i].a = kImpossibleValue;
  320. residuals_as_jets[i].v.setConstant(kImpossibleValue);
  321. }
  322. Make1stOrderPerturbations<Parameters>::Apply(parameters,
  323. parameters_as_jets.data());
  324. if (!VariadicEvaluate<ParameterDims>(
  325. functor, unpacked_parameters.data(), residuals_as_jets.data())) {
  326. return false;
  327. }
  328. Take0thOrderPart(num_outputs, residuals_as_jets.data(), function_value);
  329. Take1stOrderParts<Parameters>::Apply(
  330. num_outputs, residuals_as_jets.data(), jacobians);
  331. return true;
  332. }
  333. } // namespace ceres::internal
  334. #endif // CERES_PUBLIC_INTERNAL_AUTODIFF_H_