include/panoc-alm/util/problem.hpp

Go to the documentation of this file.

#pragma once
 
#include "box.hpp"
 
#include <cassert>
#include <chrono>
#include <functional>
#include <memory>
#include <type_traits>
 
namespace pa {
 
struct Problem {
    unsigned int n; 
    unsigned int m; 
    Box C;          
    Box D;          
 
    using f_sig = real_t(crvec x);
    using grad_f_sig = void(crvec x, rvec grad_fx);
    using g_sig = void(crvec x, rvec gx);
    using grad_g_prod_sig = void(crvec x, crvec y, rvec grad_gxy);
    using grad_gi_sig = void(crvec x, unsigned i, rvec grad_gi);
    using hess_L_prod_sig = void(crvec x, crvec y, crvec v, rvec Hv);
    using hess_L_sig = void(crvec x, crvec y, rmat H);
 
    std::function<f_sig> f;
    std::function<grad_f_sig> grad_f;
    std::function<g_sig> g;
    std::function<grad_g_prod_sig> grad_g_prod;
    std::function<grad_gi_sig> grad_gi;
    std::function<hess_L_prod_sig> hess_L_prod;
    std::function<hess_L_sig> hess_L;
 
    Problem() = default;
    Problem(unsigned int n, unsigned int m)
        : n(n), m(m), C{vec::Constant(n, +inf), vec::Constant(n, -inf)},
          D{vec::Constant(m, +inf), vec::Constant(m, -inf)} {}
    Problem(unsigned n, unsigned int m, Box C, Box D, std::function<f_sig> f,
            std::function<grad_f_sig> grad_f, std::function<g_sig> g,
            std::function<grad_g_prod_sig> grad_g_prod,
            std::function<grad_gi_sig> grad_gi,
            std::function<hess_L_prod_sig> hess_L_prod,
            std::function<hess_L_sig> hess_L)
        : n(n), m(m), C(std::move(C)), D(std::move(D)), f(std::move(f)),
          grad_f(std::move(grad_f)), g(std::move(g)),
          grad_g_prod(std::move(grad_g_prod)), grad_gi(std::move(grad_gi)),
          hess_L_prod(std::move(hess_L_prod)), hess_L(std::move(hess_L)) {}
};
 
class ParamWrapper {
  public:
    ParamWrapper(unsigned p) : param(vec::Constant(p, NaN)) {}
 
    vec param;
 
    virtual ~ParamWrapper()                             = default;
    virtual void wrap(Problem &)                        = 0;
    virtual std::shared_ptr<ParamWrapper> clone() const = 0;
};
 
class ProblemWithParam : public Problem {
  public:
    using Problem::Problem;
    explicit ProblemWithParam(const ProblemWithParam &o)
        : Problem(o), wrapper(o.wrapper ? o.wrapper->clone() : nullptr) {
        wrapper->wrap(*this);
    }
    ProblemWithParam &operator=(const ProblemWithParam &o) {
        static_cast<Problem &>(*this) = static_cast<const Problem &>(o);
        if (o.wrapper) {
            this->wrapper = o.wrapper->clone();
            this->wrapper->wrap(*this);
        }
        return *this;
    }
    ProblemWithParam(ProblemWithParam &&) = default;
    ProblemWithParam &operator=(ProblemWithParam &&) = default;
 
    void set_param(crvec p) {
        assert(p.size() == wrapper->param.size());
        wrapper->param = p;
    }
    void set_param(vec &&p) {
        assert(p.size() == wrapper->param.size());
        wrapper->param = std::move(p);
    }
    vec &get_param() { return wrapper->param; }
    const vec &get_param() const { return wrapper->param; }
 
    std::shared_ptr<ParamWrapper> wrapper;
};
 
struct EvalCounter {
    unsigned f{};
    unsigned grad_f{};
    unsigned g{};
    unsigned grad_g_prod{};
    unsigned grad_gi{};
    unsigned hess_L_prod{};
    unsigned hess_L{};
 
    struct EvalTimer {
        std::chrono::nanoseconds f{};
        std::chrono::nanoseconds grad_f{};
        std::chrono::nanoseconds g{};
        std::chrono::nanoseconds grad_g_prod{};
        std::chrono::nanoseconds grad_gi{};
        std::chrono::nanoseconds hess_L_prod{};
        std::chrono::nanoseconds hess_L{};
    } time;
 
    void reset() { *this = {}; }
};
 
inline EvalCounter &operator+=(EvalCounter &a, const EvalCounter &b) {
    a.f += b.f;
    a.grad_f += b.grad_f;
    a.g += b.g;
    a.grad_g_prod += b.grad_g_prod;
    a.grad_gi += b.grad_gi;
    a.hess_L_prod += b.hess_L_prod;
    a.hess_L += b.hess_L;
    return a;
}
 
inline EvalCounter::EvalTimer &operator+=(EvalCounter::EvalTimer &a,
                                          const EvalCounter::EvalTimer &b) {
    a.f += b.f;
    a.grad_f += b.grad_f;
    a.g += b.g;
    a.grad_g_prod += b.grad_g_prod;
    a.grad_gi += b.grad_gi;
    a.hess_L_prod += b.hess_L_prod;
    a.hess_L += b.hess_L;
    return a;
}
 
inline EvalCounter operator+(EvalCounter a, const EvalCounter &b) {
    return a += b;
}
 
template <class ProblemT>
class ProblemWithCounters : public ProblemT {
  public:
    ProblemWithCounters(ProblemT &&p) : ProblemT(std::move(p)) {
        attach_counters(*this);
    }
    ProblemWithCounters(const ProblemT &p) : ProblemT(p) {
        attach_counters(*this);
    }
 
    ProblemWithCounters(const ProblemWithCounters &) = delete;
    ProblemWithCounters &operator=(const ProblemWithCounters &) = delete;
    ProblemWithCounters(ProblemWithCounters &&)                 = default;
    ProblemWithCounters &operator=(ProblemWithCounters &&) = default;
 
  public:
    std::shared_ptr<EvalCounter> evaluations = std::make_shared<EvalCounter>();
 
  private:
    static void attach_counters(ProblemWithCounters &);
};
 
template <class ProblemT>
void ProblemWithCounters<ProblemT>::attach_counters(
    ProblemWithCounters<ProblemT> &wc) {
 
    const static auto timed = [](auto &time, const auto &f) -> decltype(f()) {
        if constexpr (std::is_same_v<decltype(f()), void>) {
            auto t0 = std::chrono::steady_clock::now();
            f();
            auto t1 = std::chrono::steady_clock::now();
            time += t1 - t0;
        } else {
            auto t0  = std::chrono::steady_clock::now();
            auto res = f();
            auto t1  = std::chrono::steady_clock::now();
            time += t1 - t0;
            return res;
        }
    };
 
    wc.f = [ev{wc.evaluations}, f{std::move(wc.f)}](crvec x) {
        ++ev->f;
        return timed(ev->time.f, [&] { return f(x); });
    };
    wc.grad_f = [ev{wc.evaluations}, grad_f{std::move(wc.grad_f)}](crvec x,
                                                                   rvec grad) {
        ++ev->grad_f;
        timed(ev->time.grad_f, [&] { grad_f(x, grad); });
    };
    wc.g = [ev{wc.evaluations}, g{std::move(wc.g)}](crvec x, rvec gx) {
        ++ev->g;
        timed(ev->time.g, [&] { g(x, gx); });
    };
    wc.grad_g_prod = [ev{wc.evaluations},
                      grad_g_prod{std::move(wc.grad_g_prod)}](crvec x, crvec y,
                                                              rvec grad) {
        ++ev->grad_g_prod;
        timed(ev->time.grad_g_prod, [&] { grad_g_prod(x, y, grad); });
    };
    wc.grad_gi = [ev{wc.evaluations}, grad_gi{std::move(wc.grad_gi)}](
                     crvec x, unsigned i, rvec grad) {
        ++ev->grad_g_prod;
        timed(ev->time.grad_g_prod, [&] { grad_gi(x, i, grad); });
    };
    wc.hess_L_prod = [ev{wc.evaluations},
                      hess_L_prod{std::move(wc.hess_L_prod)}](
                         crvec x, crvec y, crvec v, rvec Hv) {
        ++ev->hess_L_prod;
        timed(ev->time.hess_L_prod, [&] { hess_L_prod(x, y, v, Hv); });
    };
    wc.hess_L = [ev{wc.evaluations},
                 hess_L{std::move(wc.hess_L)}](crvec x, crvec y, rmat H) {
        ++ev->hess_L;
        timed(ev->time.hess_L, [&] { hess_L(x, y, H); });
    };
}
 
class ProblemOnlyD : public Problem {
  public:
    ProblemOnlyD(Problem &&p) : original(std::move(p)) { transform(); }
    ProblemOnlyD(const Problem &p) : original(p) { transform(); }
 
  private:
    Problem original; // TODO: Keeping this copy around is unnecessary.
    vec work;
 
    void transform();
};
 
} // namespace pa