ModPieceWrapper.cpp

Go to the documentation of this file.

#include "AllClassWrappers.h"
 
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include <pybind11/eigen.h>
#include "pybind11_json.h"
 
#include "MUQ/Modeling/ConstantVector.h"
#include "MUQ/Modeling/UMBridge/UMBridgeModPiece.h"
#include "MUQ/Modeling/UMBridge/UMBridgeModPieceServer.h"
#include "MUQ/Modeling/ModPiece.h"
#include "MUQ/Modeling/ModGraphPiece.h"
#include "MUQ/Modeling/MultiLogisticLikelihood.h"
#include "MUQ/Modeling/OneStepCachePiece.h"
#include "MUQ/Modeling/PyModPiece.h"
#include "MUQ/Modeling/ReplicateOperator.h"
#include "MUQ/Modeling/SplitVector.h"
#include "MUQ/Modeling/WorkGraph.h"
#include "MUQ/Modeling/SumPiece.h"
 
#include <string>
 
#include <functional>
#include <vector>
 
using namespace muq::Modeling;
namespace py = pybind11;
 
class PyModPieceTramp : public PyModPiece {
public:
  /* Inherit the constructors */
  using PyModPiece::PyModPiece;
 
  /* Trampoline (need one for each virtual function) */
  void EvaluateImpl(std::vector<Eigen::VectorXd> const& input) override {
    PYBIND11_OVERLOAD_PURE(void, PyModPiece, EvaluateImpl, input);
  }
 
  void GradientImpl(unsigned int                 const  outputDimWrt,
                    unsigned int                 const  inputDimWrt,
                    std::vector<Eigen::VectorXd> const& input,
                    Eigen::VectorXd              const& sensitivity) override {
    PYBIND11_OVERLOAD(void, PyModPiece, GradientImpl, outputDimWrt, inputDimWrt, input, sensitivity);
  }
 
  void JacobianImpl(unsigned int                 const  outputDimWrt,
                    unsigned int                 const  inputDimWrt,
                    std::vector<Eigen::VectorXd> const& input) override {
    PYBIND11_OVERLOAD(void, PyModPiece, JacobianImpl, outputDimWrt, inputDimWrt, input);
  }
 
  void ApplyJacobianImpl(unsigned int                 const  outputDimWrt,
                         unsigned int                 const  inputDimWrt,
                         std::vector<Eigen::VectorXd> const& input,
                         Eigen::VectorXd              const& vec) override {
    PYBIND11_OVERLOAD(void, PyModPiece, ApplyJacobianImpl, outputDimWrt, inputDimWrt, input, vec);
  }
 
  void ApplyHessianImpl(unsigned int                 const  outputDimWrt,
                        unsigned int                 const  inputDimWrt1,
                        unsigned int                 const  inputDimWrt2,
                        std::vector<Eigen::VectorXd> const& input,
                        Eigen::VectorXd              const& sens,
                        Eigen::VectorXd              const& vec) override {
    PYBIND11_OVERLOAD(void, PyModPiece, ApplyHessianImpl, outputDimWrt, inputDimWrt1, inputDimWrt2, input, sens, vec);
  }
};
 
// The server itself does not do any Python calls; nevertheless, Python's global interpreter
// lock (GIL) remains active while entering the server loop, unless explicitly instructed not to.
// The server spawns threads for requests which may have to evaluate a Python model, in turn locking the GIL.
// We therefore have to unlock the GIL for the server loop, else we get a deadlock once
// a Python model is evaluated by the server.
void serveModPieceWithoutGIL(std::shared_ptr<ModPiece> modPiece, std::string name, std::string host, int port) {
  Py_BEGIN_ALLOW_THREADS
  muq::Modeling::serveModPiece(modPiece, name, host, port);
  Py_END_ALLOW_THREADS
}
 
class Publicist : public PyModPiece {
public:
    // Expose protected functions
    using PyModPiece::EvaluateImpl;
    using PyModPiece::GradientImpl;
    using PyModPiece::JacobianImpl;
    using PyModPiece::ApplyJacobianImpl;
    using PyModPiece::ApplyHessianImpl;
 
    // Expose protected member variables
    using PyModPiece::outputs;
    using PyModPiece::gradient;
    using PyModPiece::jacobian;
    using PyModPiece::jacobianAction;
    using PyModPiece::hessAction;
};
 
void muq::Modeling::PythonBindings::ModPieceWrapper(py::module &m)
{
  // Define some functions from the WorkPiece base class
  py::class_<ModPiece, WorkPiece, std::shared_ptr<ModPiece>> mp(m, "ModPiece");
  mp
    .def("Evaluate", (std::vector<Eigen::VectorXd> const& (ModPiece::*)(std::vector<Eigen::VectorXd> const&)) &ModPiece::Evaluate)
    .def("Evaluate", (std::vector<Eigen::VectorXd> const& (ModPiece::*)()) &ModPiece::Evaluate)
    .def_readonly("inputSizes", &ModPiece::inputSizes)
    .def_readonly("outputSizes", &ModPiece::outputSizes)
    .def("GetRunTime", &ModPiece::GetRunTime)
    .def("ResetCallTime", &ModPiece::ResetCallTime)
    .def("GetNumCalls", &ModPiece::GetNumCalls)
    .def("Gradient", (Eigen::VectorXd const& (ModPiece::*)(unsigned int, unsigned int, std::vector<Eigen::VectorXd> const&, Eigen::VectorXd const&)) &ModPiece::Gradient)
    .def("Jacobian", (Eigen::MatrixXd const& (ModPiece::*)(unsigned int, unsigned int, std::vector<Eigen::VectorXd> const&)) &ModPiece::Jacobian)
    .def("ApplyJacobian", (Eigen::VectorXd const& (ModPiece::*)(unsigned int, unsigned int, std::vector<Eigen::VectorXd> const&, Eigen::VectorXd const&)) &ModPiece::ApplyJacobian)
    .def("GradientByFD", (Eigen::VectorXd (ModPiece::*)(unsigned int, unsigned int, std::vector<Eigen::VectorXd> const&, Eigen::VectorXd const&)) &ModPiece::GradientByFD)
    .def("JacobianByFD", (Eigen::MatrixXd (ModPiece::*)(unsigned int, unsigned int, std::vector<Eigen::VectorXd> const&)) &ModPiece::JacobianByFD)
    .def("ApplyJacobianByFD", (Eigen::VectorXd (ModPiece::*)(unsigned int, unsigned int, std::vector<Eigen::VectorXd> const&, Eigen::VectorXd const&)) &ModPiece::ApplyJacobianByFD)
    .def("ApplyHessian", (Eigen::VectorXd const& (ModPiece::*)(unsigned int, unsigned int, unsigned int, std::vector<Eigen::VectorXd> const&, Eigen::VectorXd const&, Eigen::VectorXd const&)) &ModPiece::ApplyHessian)
    .def("ApplyHessianByFD", (Eigen::VectorXd (ModPiece::*)(unsigned int, unsigned int, unsigned int, std::vector<Eigen::VectorXd> const&, Eigen::VectorXd const&, Eigen::VectorXd const&)) &ModPiece::ApplyHessianByFD)
    .def("EnableCache", &ModPiece::EnableCache)
    .def("DisableCache", &ModPiece::DisableCache)
    .def("CacheStatus", &ModPiece::CacheStatus)
    .def("SetWarnLevel", &ModPiece::SetWarnLevel);
 
  py::class_<PyModPiece, PyModPieceTramp, ModPiece, WorkPiece, std::shared_ptr<PyModPiece> > pymp(m, "PyModPiece");
  pymp
    .def(py::init<Eigen::VectorXi const&, Eigen::VectorXi const&>())
    .def("EvaluateImpl", (void (PyModPiece::*)(std::vector<Eigen::VectorXd> const&)) &Publicist::EvaluateImpl)
    .def_readwrite("outputs", &Publicist::outputs)
    .def("GradientImpl", (void (PyModPiece::*)(unsigned int const, unsigned int const, std::vector<Eigen::VectorXd> const&, Eigen::VectorXd const&)) &Publicist::GradientImpl)
    .def_readwrite("gradient", &Publicist::gradient)
    .def("JacobianImpl", (void (PyModPiece::*)(unsigned int const, unsigned int const, std::vector<Eigen::VectorXd> const&)) &Publicist::JacobianImpl)
    .def_readwrite("jacobian", &Publicist::jacobian)
    .def("ApplyJacobianImpl", (void (PyModPiece::*)(unsigned int const, unsigned int const, std::vector<Eigen::VectorXd> const&, Eigen::VectorXd const&)) &Publicist::ApplyJacobianImpl)
    .def_readwrite("jacobianAction", &Publicist::jacobianAction)
    .def("ApplyHessianImpl", (void (PyModPiece::*)(unsigned int const, unsigned int const, unsigned int const, std::vector<Eigen::VectorXd> const&, Eigen::VectorXd const&, Eigen::VectorXd const&)) &Publicist::ApplyHessianImpl)
    .def_readwrite("hessAction", &Publicist::hessAction);
 
  py::class_<OneStepCachePiece, ModPiece, WorkPiece, std::shared_ptr<OneStepCachePiece>> ocp(m, "OneStepCachePiece");
  ocp
    .def(py::init<std::shared_ptr<ModPiece>>())
    .def("HitRatio", &OneStepCachePiece::HitRatio);
 
  py::class_<UMBridgeModPiece, ModPiece, WorkPiece, std::shared_ptr<UMBridgeModPiece>> hmp(m, "UMBridgeModPiece");
  hmp
    .def(py::init( [](std::string host, std::string name) {return new UMBridgeModPiece(host, name); }))
    .def(py::init( [](std::string host, std::string name, py::dict config) {return new UMBridgeModPiece(host, name, config); }));
 
  m.def("serveModPiece", &serveModPieceWithoutGIL);
 
  py::class_<ConstantVector, ModPiece, WorkPiece, std::shared_ptr<ConstantVector>> cv(m, "ConstantVector");
  cv
    .def(py::init<Eigen::VectorXd const&>())
    .def("SetValue",  &ConstantVector::SetValue);
 
  py::class_<ReplicateOperator, ModPiece, WorkPiece, std::shared_ptr<ReplicateOperator>> ro(m, "ReplicateOperator");
  ro
    .def(py::init<unsigned int, unsigned int>());
 
  py::class_<SumPiece, ModPiece, WorkPiece, std::shared_ptr<SumPiece>>(m, "SumPiece")
    .def(py::init<unsigned int, unsigned int>(), py::arg("dim"),py::arg("numInputs")=2);
 
 
  py::class_<ModGraphPiece, ModPiece, WorkPiece, std::shared_ptr<ModGraphPiece>> mgp(m, "ModGraphPiece");
  mgp
    .def("GetGraph", &ModGraphPiece::GetGraph)
    .def("GetConstantPieces", &ModGraphPiece::GetConstantPieces)
    .def("GradientGraph", &ModGraphPiece::GradientGraph)
    .def("JacobianGraph", &ModGraphPiece::JacobianGraph);
 
  py::class_<MultiLogisticLikelihood, ModPiece, WorkPiece, std::shared_ptr<MultiLogisticLikelihood>> mll(m, "MultiLogisticLikelihood");
  mll
    .def(py::init<unsigned int, Eigen::VectorXi>());
 
  py::class_<SplitVector, ModPiece, WorkPiece, std::shared_ptr<SplitVector>>(m,"SplitVector")
    .def(py::init<Eigen::VectorXi const&, Eigen::VectorXi const&, unsigned int const>())
    .def("StartIndices", &SplitVector::StartIndices);
 
}