Line data Source code
1 : #pragma once
2 :
3 : #include "box.hpp"
4 :
5 : #include <cassert>
6 : #include <chrono>
7 : #include <functional>
8 : #include <memory>
9 : #include <type_traits>
10 :
11 : namespace pa {
12 :
13 : /**
14 : * @class Problem
15 : * @brief Problem description for minimization problems.
16 : *
17 : * @f[ \begin{aligned}
18 : * & \underset{x}{\text{minimize}}
19 : * & & f(x) &&&& f : \mathbb{R}^n \rightarrow \mathbb{R} \\
20 : * & \text{subject to}
21 : * & & \underline{x} \le x \le \overline{x} \\
22 : * &&& \underline{z} \le g(x) \le \overline{z} &&&& g :
23 : * \mathbb{R}^n \rightarrow \mathbb{R}^m
24 : * \end{aligned} @f]
25 : */
26 20 : struct Problem {
27 : unsigned int n; ///< Number of decision variables, dimension of x
28 : unsigned int m; ///< Number of constraints, dimension of g(x) and z
29 : Box C; ///< Constraints of the decision variables, @f$ x \in C @f$
30 : Box D; ///< Other constraints, @f$ g(x) \in D @f$
31 :
32 : /// Signature of the function that evaluates the cost @f$ f(x) @f$
33 : /// @param [in] x
34 : /// Decision variable @f$ x \in \mathbb{R}^n @f$
35 : using f_sig = real_t(crvec x);
36 : /// Signature of the function that evaluates the gradient of the cost
37 : /// function @f$ \nabla f(x) @f$
38 : /// @param [in] x
39 : /// Decision variable @f$ x \in \mathbb{R}^n @f$
40 : /// @param [out] grad_fx
41 : /// Gradient of cost function @f$ \nabla f(x) \in \mathbb{R}^n @f$
42 : using grad_f_sig = void(crvec x, rvec grad_fx);
43 : /// Signature of the function that evaluates the constraints @f$ g(x) @f$
44 : /// @param [in] x
45 : /// Decision variable @f$ x \in \mathbb{R}^n @f$
46 : /// @param [out] gx
47 : /// Value of the constraints @f$ g(x) \in \mathbb{R}^m @f$
48 : using g_sig = void(crvec x, rvec gx);
49 : /// Signature of the function that evaluates the gradient of the constraints
50 : /// times a vector
51 : /// @f$ \nabla g(x)\ y @f$
52 : /// @param [in] x
53 : /// Decision variable @f$ x \in \mathbb{R}^n @f$
54 : /// @param [in] y
55 : /// Vector @f$ y \in \mathbb{R}^m @f$ to multiply the gradient by
56 : /// @param [out] grad_gxy
57 : /// Gradient of the constraints
58 : /// @f$ \nabla g(x)\ y \in \mathbb{R}^n @f$
59 : using grad_g_prod_sig = void(crvec x, crvec y, rvec grad_gxy);
60 : /// Signature of the function that evaluates the gradient of one specific
61 : /// constraints
62 : /// @f$ \nabla g_i(x) @f$
63 : /// @param [in] x
64 : /// Decision variable @f$ x \in \mathbb{R}^n @f$
65 : /// @param [in] i
66 : /// Which constraint @f$ 0 \le i \lt m @f$
67 : /// @param [out] grad_gi
68 : /// Gradient of the constraint
69 : /// @f$ \nabla g_i(x) \mathbb{R}^n @f$
70 : using grad_gi_sig = void(crvec x, unsigned i, rvec grad_gi);
71 : /// Signature of the function that evaluates the Hessian of the Lagrangian
72 : /// multiplied by a vector
73 : /// @f$ \nabla_{xx}^2L(x, y)\ v @f$
74 : /// @param [in] x
75 : /// Decision variable @f$ x \in \mathbb{R}^n @f$
76 : /// @param [in] y
77 : /// Lagrange multipliers @f$ y \in \mathbb{R}^m @f$
78 : /// @param [in] v
79 : /// Vector to multiply by @f$ v \in \mathbb{R}^n @f$
80 : /// @param [out] Hv
81 : /// Hessian-vector product
82 : /// @f$ \nabla_{xx}^2 L(x, y)\ v \in \mathbb{R}^{n} @f$
83 : using hess_L_prod_sig = void(crvec x, crvec y, crvec v, rvec Hv);
84 : /// Signature of the function that evaluates the Hessian of the Lagrangian
85 : /// @f$ \nabla_{xx}^2L(x, y) @f$
86 : /// @param [in] x
87 : /// Decision variable @f$ x \in \mathbb{R}^n @f$
88 : /// @param [in] y
89 : /// Lagrange multipliers @f$ y \in \mathbb{R}^m @f$
90 : /// @param [out] H
91 : /// Hessian @f$ \nabla_{xx}^2 L(x, y) \in \mathbb{R}^{n\times n} @f$
92 : using hess_L_sig = void(crvec x, crvec y, rmat H);
93 :
94 : /// Cost function @f$ f(x) @f$
95 : std::function<f_sig> f;
96 : /// Gradient of the cost function @f$ \nabla f(x) @f$
97 : std::function<grad_f_sig> grad_f;
98 : /// Constraint function @f$ g(x) @f$
99 : std::function<g_sig> g;
100 : /// Gradient of the constraint function times vector @f$ \nabla g(x)\ y @f$
101 : std::function<grad_g_prod_sig> grad_g_prod;
102 : /// Gradient of a specific constraint @f$ \nabla g_i(x) @f$
103 : std::function<grad_gi_sig> grad_gi;
104 : /// Hessian of the Lagrangian function times vector
105 : /// @f$ \nabla_{xx}^2 L(x, y)\ v @f$
106 : std::function<hess_L_prod_sig> hess_L_prod;
107 : /// Hessian of the Lagrangian function @f$ \nabla_{xx}^2 L(x, y) @f$
108 : std::function<hess_L_sig> hess_L;
109 :
110 16 : Problem() = default;
111 : Problem(unsigned int n, unsigned int m)
112 : : n(n), m(m), C{vec::Constant(n, +inf), vec::Constant(n, -inf)},
113 : D{vec::Constant(m, +inf), vec::Constant(m, -inf)} {}
114 8 : Problem(unsigned n, unsigned int m, Box C, Box D, std::function<f_sig> f,
115 : std::function<grad_f_sig> grad_f, std::function<g_sig> g,
116 : std::function<grad_g_prod_sig> grad_g_prod,
117 : std::function<grad_gi_sig> grad_gi,
118 : std::function<hess_L_prod_sig> hess_L_prod,
119 : std::function<hess_L_sig> hess_L)
120 8 : : n(n), m(m), C(std::move(C)), D(std::move(D)), f(std::move(f)),
121 8 : grad_f(std::move(grad_f)), g(std::move(g)),
122 8 : grad_g_prod(std::move(grad_g_prod)), grad_gi(std::move(grad_gi)),
123 8 : hess_L_prod(std::move(hess_L_prod)), hess_L(std::move(hess_L)) {}
124 : };
125 :
126 : class ParamWrapper {
127 : public:
128 : ParamWrapper(unsigned p) : param(vec::Constant(p, NaN)) {}
129 :
130 : vec param;
131 :
132 : virtual ~ParamWrapper() = default;
133 : virtual void wrap(Problem &) = 0;
134 : virtual std::shared_ptr<ParamWrapper> clone() const = 0;
135 : };
136 :
137 : class ProblemWithParam : public Problem {
138 : public:
139 : using Problem::Problem;
140 : explicit ProblemWithParam(const ProblemWithParam &o)
141 : : Problem(o), wrapper(o.wrapper ? o.wrapper->clone() : nullptr) {
142 : wrapper->wrap(*this);
143 : }
144 : ProblemWithParam &operator=(const ProblemWithParam &o) {
145 : static_cast<Problem &>(*this) = static_cast<const Problem &>(o);
146 : if (o.wrapper) {
147 : this->wrapper = o.wrapper->clone();
148 : this->wrapper->wrap(*this);
149 : }
150 : return *this;
151 : }
152 : ProblemWithParam(ProblemWithParam &&) = default;
153 : ProblemWithParam &operator=(ProblemWithParam &&) = default;
154 :
155 : void set_param(crvec p) {
156 : assert(p.size() == wrapper->param.size());
157 : wrapper->param = p;
158 : }
159 : void set_param(vec &&p) {
160 : assert(p.size() == wrapper->param.size());
161 : wrapper->param = std::move(p);
162 : }
163 : vec &get_param() { return wrapper->param; }
164 : const vec &get_param() const { return wrapper->param; }
165 :
166 : std::shared_ptr<ParamWrapper> wrapper;
167 : };
168 :
169 1 : struct EvalCounter {
170 1 : unsigned f{};
171 1 : unsigned grad_f{};
172 1 : unsigned g{};
173 1 : unsigned grad_g_prod{};
174 1 : unsigned grad_gi{};
175 1 : unsigned hess_L_prod{};
176 1 : unsigned hess_L{};
177 :
178 1 : struct EvalTimer {
179 1 : std::chrono::nanoseconds f{};
180 1 : std::chrono::nanoseconds grad_f{};
181 1 : std::chrono::nanoseconds g{};
182 1 : std::chrono::nanoseconds grad_g_prod{};
183 1 : std::chrono::nanoseconds grad_gi{};
184 1 : std::chrono::nanoseconds hess_L_prod{};
185 1 : std::chrono::nanoseconds hess_L{};
186 : } time;
187 :
188 : void reset() { *this = {}; }
189 : };
190 :
191 : inline EvalCounter &operator+=(EvalCounter &a, const EvalCounter &b) {
192 : a.f += b.f;
193 : a.grad_f += b.grad_f;
194 : a.g += b.g;
195 : a.grad_g_prod += b.grad_g_prod;
196 : a.grad_gi += b.grad_gi;
197 : a.hess_L_prod += b.hess_L_prod;
198 : a.hess_L += b.hess_L;
199 : return a;
200 : }
201 :
202 : inline EvalCounter::EvalTimer &operator+=(EvalCounter::EvalTimer &a,
203 : const EvalCounter::EvalTimer &b) {
204 : a.f += b.f;
205 : a.grad_f += b.grad_f;
206 : a.g += b.g;
207 : a.grad_g_prod += b.grad_g_prod;
208 : a.grad_gi += b.grad_gi;
209 : a.hess_L_prod += b.hess_L_prod;
210 : a.hess_L += b.hess_L;
211 : return a;
212 : }
213 :
214 : inline EvalCounter operator+(EvalCounter a, const EvalCounter &b) {
215 : return a += b;
216 : }
217 :
218 : template <class ProblemT>
219 1 : class ProblemWithCounters : public ProblemT {
220 : public:
221 : ProblemWithCounters(ProblemT &&p) : ProblemT(std::move(p)) {
222 : attach_counters(*this);
223 : }
224 2 : ProblemWithCounters(const ProblemT &p) : ProblemT(p) {
225 1 : attach_counters(*this);
226 1 : }
227 :
228 : ProblemWithCounters(const ProblemWithCounters &) = delete;
229 : ProblemWithCounters &operator=(const ProblemWithCounters &) = delete;
230 : ProblemWithCounters(ProblemWithCounters &&) = default;
231 : ProblemWithCounters &operator=(ProblemWithCounters &&) = default;
232 :
233 : public:
234 1 : std::shared_ptr<EvalCounter> evaluations = std::make_shared<EvalCounter>();
235 :
236 : private:
237 : static void attach_counters(ProblemWithCounters &);
238 : };
239 :
240 : template <class ProblemT>
241 1 : void ProblemWithCounters<ProblemT>::attach_counters(
242 : ProblemWithCounters<ProblemT> &wc) {
243 :
244 2412 : const static auto timed = [](auto &time, const auto &f) -> decltype(f()) {
245 : if constexpr (std::is_same_v<decltype(f()), void>) {
246 1738 : auto t0 = std::chrono::steady_clock::now();
247 1738 : f();
248 1738 : auto t1 = std::chrono::steady_clock::now();
249 1738 : time += t1 - t0;
250 1738 : } else {
251 674 : auto t0 = std::chrono::steady_clock::now();
252 674 : auto res = f();
253 674 : auto t1 = std::chrono::steady_clock::now();
254 674 : time += t1 - t0;
255 1348 : return res;
256 674 : }
257 1738 : };
258 :
259 680 : wc.f = [ev{wc.evaluations}, f{std::move(wc.f)}](crvec x) {
260 674 : ++ev->f;
261 2022 : return timed(ev->time.f, [&] { return f(x); });
262 : };
263 527 : wc.grad_f = [ev{wc.evaluations}, grad_f{std::move(wc.grad_f)}](crvec x,
264 : rvec grad) {
265 521 : ++ev->grad_f;
266 1563 : timed(ev->time.grad_f, [&] { grad_f(x, grad); });
267 521 : };
268 702 : wc.g = [ev{wc.evaluations}, g{std::move(wc.g)}](crvec x, rvec gx) {
269 696 : ++ev->g;
270 2088 : timed(ev->time.g, [&] { g(x, gx); });
271 696 : };
272 528 : wc.grad_g_prod = [ev{wc.evaluations},
273 1 : grad_g_prod{std::move(wc.grad_g_prod)}](crvec x, crvec y,
274 : rvec grad) {
275 521 : ++ev->grad_g_prod;
276 1563 : timed(ev->time.grad_g_prod, [&] { grad_g_prod(x, y, grad); });
277 521 : };
278 6 : wc.grad_gi = [ev{wc.evaluations}, grad_gi{std::move(wc.grad_gi)}](
279 : crvec x, unsigned i, rvec grad) {
280 0 : ++ev->grad_g_prod;
281 0 : timed(ev->time.grad_g_prod, [&] { grad_gi(x, i, grad); });
282 0 : };
283 7 : wc.hess_L_prod = [ev{wc.evaluations},
284 1 : hess_L_prod{std::move(wc.hess_L_prod)}](
285 : crvec x, crvec y, crvec v, rvec Hv) {
286 0 : ++ev->hess_L_prod;
287 0 : timed(ev->time.hess_L_prod, [&] { hess_L_prod(x, y, v, Hv); });
288 0 : };
289 7 : wc.hess_L = [ev{wc.evaluations},
290 1 : hess_L{std::move(wc.hess_L)}](crvec x, crvec y, rmat H) {
291 0 : ++ev->hess_L;
292 0 : timed(ev->time.hess_L, [&] { hess_L(x, y, H); });
293 0 : };
294 1 : }
295 :
296 : /// Moves the state constraints in the set C to the set D, resulting in an
297 : /// unconstraint inner problem. The new constraints function g becomes the
298 : /// concatenation of the original g function and the identity function. The
299 : /// new set D is the cartesian product of the original D × C.
300 1 : class ProblemOnlyD : public Problem {
301 : public:
302 : ProblemOnlyD(Problem &&p) : original(std::move(p)) { transform(); }
303 2 : ProblemOnlyD(const Problem &p) : original(p) { transform(); }
304 :
305 : private:
306 : Problem original; // TODO: Keeping this copy around is unnecessary.
307 : vec work;
308 :
309 : void transform();
310 : };
311 :
312 : } // namespace pa
|