sinkhorn bindings

Author: ManifoldFR

CMakeLists.txt

@@ -17,11 +17,11 @@ message(STATUS "Python includes at ${Python_INCLUDE_DIRS} ${Python_NumPy_INCLUDE
 set(EntropicMFG_HEADERS
     Kernels.h
     Utils.h
-    MultiSinkhorn.h
     KLOperator.h
     CongestionOperator.h
     RectangularGridLaplacian.h
-    MessagePassing.h)
+    MessagePassing.h
+    MultiSinkhorn.h)
 
 set(EntropicMFG_SOURCES
     MultiSinkhorn.cpp
@@ -32,8 +32,7 @@ add_library(entropicmfg SHARED ${EntropicMFG_SOURCES})
 target_include_directories(entropicmfg PUBLIC .)
 target_link_libraries(entropicmfg Eigen3::Eigen)
 target_link_libraries(entropicmfg OpenMP::OpenMP_CXX)
-set_target_properties(entropicmfg PROPERTIES
-    LINKER_LANGUAGE CXX)
+set_target_properties(entropicmfg PROPERTIES LINKER_LANGUAGE CXX)
 target_compile_options(entropicmfg PRIVATE -fPIC)
 
 

CongestionOperator.h

@@ -1,5 +1,6 @@
 #pragma once
 
+#include <omp.h>
 #include <Eigen/Dense>
 #include "KLOperator.h"
 

KLOperator.h

@@ -1,5 +1,6 @@
 #pragma once
 
+#include <memory>
 #include <Eigen/Core>
 
 
@@ -15,4 +16,6 @@ class BaseProximalOperator {
     virtual MatrixXd operator()(const MatrixXd& x) const = 0;
 };
 
+typedef std::shared_ptr<BaseProximalOperator> ProxPtr;
+
 }  // namespace klprox

MessagePassing.cpp

@@ -1,30 +1,32 @@
 #include "MessagePassing.h"
 #include "Kernels.h"
 
+#include <omp.h>
 #include <vector>
 #include <Eigen/Core>
 
 using kernels::KernelPtr;
 using namespace Eigen;
 
 
-namespace algorithms
+namespace messages
 {
 
 
-MatrixXd contract(std::vector<Ref<MatrixXd>>& potentials, size_t idx, KernelPtr ker) {
+MatrixXd contract(std::vector<Ref<MatrixXd>>& potentials, const int idx, KernelPtr ker) {
     size_t nx = potentials[0].rows();
     size_t ny = potentials[0].cols();
+
     MatrixXd A_ = MatrixXd::Ones(nx, ny);
-    for (size_t k=0; k < idx - 1; k++) {
-        A_ = (*ker)(potentials[k].array() * A_.array());
+    for (int k=0; k < idx; k++) {
+        A_ = ker->operator()(potentials[k].array() * A_.array());
     }
 
+
     MatrixXd B_ = MatrixXd::Ones(nx, ny);
-    for (size_t k=potentials.size()-1; k > idx + 1; k--) {
-        B_ = (*ker)(potentials[k].array() * A_.array());
+    for (int k=potentials.size()-1; k > idx; k--) {
+        B_ = ker->operator()(potentials[k].array() * B_.array());
     }
-
     return A_.array() * B_.array();
 }
 
@@ -35,7 +37,7 @@ std::vector<MatrixXd> compute_marginals(std::vector<Ref<MatrixXd>>& potentials,
     std::vector<MatrixXd> result(num_marginals);
 
     #pragma omp parallel for
-    for (size_t i=0; i < num_marginals; i++) {
+    for (int i=0; i < num_marginals; i++) {
         result[i] = potentials[i].array() * contract(potentials, i, ker).array();
     }
 

MessagePassing.h

@@ -21,7 +21,7 @@ namespace messages
  * @return 
  */
 MatrixXd contract(std::vector<Ref<MatrixXd>>& potentials,
-                  const size_t idx, KernelPtr ker);
+                  const int idx, KernelPtr ker);
 
 
 std::vector<MatrixXd> compute_marginals(std::vector<Ref<MatrixXd>>& potentials, KernelPtr ker);

MultiSinkhorn.cpp

@@ -14,7 +14,6 @@ void MultimarginalSinkhorn::iterate(std::vector<Ref<MatrixXd>>& potentials) {
 
     ArrayXXd conv;
 
-
     conv = contract(potentials, 0, kernel);
     potentials[0] = rho_0.array() / conv;
 
@@ -28,13 +27,13 @@ void MultimarginalSinkhorn::iterate(std::vector<Ref<MatrixXd>>& potentials) {
 
 }
 
-void MultimarginalSinkhorn::run_sinkhorn(std::vector<Ref<MatrixXd>>& potentials, int num_iterations) {
+void MultimarginalSinkhorn::run(std::vector<Ref<MatrixXd>>& potentials, int num_iterations) {
     for (int i = 0; i < num_iterations; i++) {
         this->iterate(potentials);
     }
 }
 
-inline std::vector<MatrixXd> MultimarginalSinkhorn::get_marginals(std::vector<Ref<MatrixXd>>& potentials) {
+std::vector<MatrixXd> MultimarginalSinkhorn::get_marginals(std::vector<Ref<MatrixXd>>& potentials) {
     return compute_marginals(potentials, kernel);
 }
 

MultiSinkhorn.h

@@ -19,32 +19,37 @@ using std::shared_ptr;
 namespace sinkhorn {
 
 using klprox::BaseProximalOperator;
+using klprox::ProxPtr;
 
 class MultimarginalSinkhorn {
     private:
     /// Running cost proximal operator
-    shared_ptr<BaseProximalOperator> running;
+    ProxPtr running;
     /// Terminal cost proximal operator
-    shared_ptr<BaseProximalOperator> terminal;
+    ProxPtr terminal;
     /// Initial marginal
     MatrixXd rho_0;
     /// Kernel
     kernels::KernelPtr kernel;
 
 
     public:
+    void setInitialDistribution(Ref<const MatrixXd>& rho_) {
+        rho_0 = rho_;
+    }
     MultimarginalSinkhorn(
-        shared_ptr<BaseProximalOperator> running,
-        shared_ptr<BaseProximalOperator> terminal,
-        kernels::KernelPtr kernel):
-    running(running), terminal(terminal), kernel(kernel) {}
+        ProxPtr running,
+        ProxPtr terminal,
+        kernels::KernelPtr kernel,
+        MatrixXd& rho):
+    running(running), terminal(terminal), kernel(kernel), rho_0(rho) {}
 
     /// Perform one iterate of the multimarginal Sinkhorn algorithm.
     void iterate(std::vector<Ref<MatrixXd>>& potentials);
 
-    void run_sinkhorn(std::vector<Ref<MatrixXd>>& potentials, int num_iterations);
+    void run(std::vector<Ref<MatrixXd>>& potentials, int num_iterations);
 
-    inline std::vector<MatrixXd> get_marginals(std::vector<Ref<MatrixXd>>& potentials);
+    std::vector<MatrixXd> get_marginals(std::vector<Ref<MatrixXd>>& potentials);
 };
 
 }

examples/crowdmodel.cpp

@@ -52,7 +52,7 @@ int main(int argc, char* argv[]) {
 
     shared_ptr<CongestionObstacleProx> terminal(running);
 
-    sinkhorn::MultimarginalSinkhorn sink(running, terminal, kernel_ptr);
+    sinkhorn::MultimarginalSinkhorn sink(running, terminal, kernel_ptr, rho_0);
 
 
     return 1;

python/CMakeLists.txt

@@ -10,4 +10,5 @@ set(SOURCES
 pybind11_add_module(pyentropicmfg ${SOURCES})
 target_include_directories(pyentropicmfg PRIVATE ../)
 target_link_libraries(pyentropicmfg PRIVATE entropicmfg)
+target_link_libraries(pyentropicmfg PRIVATE OpenMP::OpenMP_CXX)
 target_compile_options(pyentropicmfg PRIVATE -fPIC)

python/bindings.cpp

@@ -1,21 +1,24 @@
 #include <pybind11/pybind11.h>
+#include <pybind11/eigen.h>
+
+#include <omp.h>
 
 
 namespace py = pybind11;
 
 
 void bind_prox(py::module &);
-void bind_sinkhorn(py::module &);
 void bind_kernels(py::module &);
+void bind_sinkhorn(py::module &);
 
 
 PYBIND11_MODULE(pyentropicmfg, m) {
     m.doc() = "A library to solve variational mean-field games using an entropy "
         "minimization approach.";
 
     bind_prox(m);
-    bind_sinkhorn(m);
     bind_kernels(m);
+    bind_sinkhorn(m);
 
 }
 

python/examples/congestion.py

@@ -15,9 +15,6 @@
 xg, yg = np.meshgrid(xar, xar)
 extent = [xg.min(), xg.max(), yg.min(), yg.max()]
 
-epsilon = 0.001
-kernel = mfg.kernels.EuclideanKernel(
-    nx, nx, 0, 1, 0, 1, epsilon)
 
 
 # Define domain mask
@@ -39,15 +36,9 @@ def mask_to_img(mask: np.ndarray):
 rho_0 /= rho_0.sum()
 
 fig = plt.figure()
-plt.subplot(1,2,1)
 plt.imshow(mask_img, zorder=2)
 plt.imshow(rho_0, cmap=plt.cm.Blues, zorder=0, origin='lower')
 
-plt.subplot(1,2,2)
-plt.imshow(mask_img, zorder=2)
-plt.imshow(kernel(rho_0), cmap=plt.cm.Blues, zorder=0, origin='lower')
-
-
 ## Problem setup
 
 congest_max = 1.01 * rho_0.max()
@@ -60,12 +51,6 @@ def mask_to_img(mask: np.ndarray):
 exit_img = mask_to_img(exit_mask)
 exit_img[exit_mask.astype(bool), 0] = .8
 
-plt.figure()
-plt.imshow(mask_img, zorder=2)
-plt.imshow(exit_img, zorder=1, origin='lower')
-plt.imshow(rho_0, cmap=plt.cm.Blues, zorder=0, origin='lower')
-plt.show()
-
 boundary_ = np.ma.MaskedArray(1. - exit_mask, mask=mask)
 potential = skfmm.travel_time(boundary_, np.ones_like(boundary_))
 
@@ -74,10 +59,52 @@ def mask_to_img(mask: np.ndarray):
 plt.imshow(exit_img, zorder=1, origin='lower', extent=extent)
 ct = plt.contourf(potential, zorder=1, levels=40, extent=extent)
 plt.title("Potential function $\\Psi$")
-plt.show()
 
-prox = mfg.prox.CongestionObstacleProx(mask, congest_max, potential)
 
+terminal_prox = mfg.prox.CongestionObstacleProx(mask, congest_max, potential)
+running_prox = mfg.prox.CongestionObstacleProx(mask, congest_max, np.zeros_like(potential))
 
+N_t = 31
+dt = 1./ (N_t - 1)
+epsilon = 0.1
+kernel = mfg.kernels.EuclideanKernel(
+    nx, nx, 0, 1, 0, 1, epsilon * dt)
 
 
+sinkhorn = mfg.sinkhorn.MultiSinkhorn(
+    running_prox, terminal_prox,
+    kernel, rho_0)
+
+
+a_s = [
+    np.ones_like(rho_0, order='F') for _ in range(N_t)
+]
+
+print("Running sinkhorn...")
+import time
+t_a = time.time()
+num_iters = 1
+sinkhorn.run(a_s, num_iters)
+print("Elapsed time:", time.time() - t_a)
+
+print("Computing marginals...")
+marginals = sinkhorn.get_marginals(a_s)
+
+skip = 5
+steps_to_plot = list(np.arange(N_t)[::skip])
+
+ncols = 3
+nrows = len(steps_to_plot) // 3
+
+fig, axes = plt.subplots(nrows, ncols)
+axes = axes.ravel()
+
+for i, t in enumerate(steps_to_plot):
+    m = marginals[t]
+    if i < len(axes):
+        ax = axes[i]
+        ax.imshow(mask_img, zorder=2, origin='lower', extent=extent)
+        ax.imshow(m, zorder=1, origin='lower', extent=extent, cmap=plt.cm.Blues)
+        ax.set_title("Time step $t=%d$" % t)
+
+plt.show()

python/kernels.cpp

@@ -34,10 +34,10 @@ void bind_kernels(py::module& m) {
         "kernels", "Gibbs kernels used in the optimal transport formulation of the MFG. "
                    "These represent the 2-marginal of the Wiener measure.");
 
-    py::class_<BaseKernel, PyBaseKernel>(m2, "BaseKernel")
+    py::class_<BaseKernel, PyBaseKernel, KernelPtr>(m2, "BaseKernel")
         .def("__call__", &BaseKernel::operator());
 
-    py::class_<Kernel2D, BaseKernel>(m2, "EuclideanKernel")
+    py::class_<Kernel2D, BaseKernel, std::shared_ptr<Kernel2D>>(m2, "EuclideanKernel")
         .def(py::init<size_t, size_t, double, double,
                       double, double, double>());
 

python/operators.cpp

@@ -2,6 +2,7 @@
 #include <pybind11/eigen.h>
 #include <pybind11/operators.h>
 
+#include "omp.h"
 #include "KLOperator.h"
 #include "CongestionOperator.h"
 
@@ -31,16 +32,16 @@ void bind_prox(py::module& m) {
 
     py::module m2 = m.def_submodule("prox", "Proximal operators");
 
-    py::class_<BaseProximalOperator, PyBaseProximalOperator>(m2, "BaseProximalOperator")
+    py::class_<BaseProximalOperator, PyBaseProximalOperator, ProxPtr>(m2, "BaseProximalOperator")
         .def("__call__", &BaseProximalOperator::operator());
 
-    py::class_<ObstacleProx, BaseProximalOperator>(m2, "ObstacleProx")
+    py::class_<ObstacleProx, BaseProximalOperator, std::shared_ptr<ObstacleProx>>(m2, "ObstacleProx")
         .def(py::init<const Eigen::ArrayXXi&>());
 
-    py::class_<CongestionPotentialProx, BaseProximalOperator>(m2, "CongestionPotentialProx")
+    py::class_<CongestionPotentialProx, BaseProximalOperator, std::shared_ptr<CongestionPotentialProx>>(m2, "CongestionPotentialProx")
         .def(py::init<double, const Eigen::MatrixXd&>());
 
-    py::class_<CongestionObstacleProx, CongestionPotentialProx, ObstacleProx>(m2, "CongestionObstacleProx")
+    py::class_<CongestionObstacleProx, CongestionPotentialProx, ObstacleProx, std::shared_ptr<CongestionObstacleProx>>(m2, "CongestionObstacleProx")
         .def(py::init<const ArrayXXi&, double, const Eigen::MatrixXd&>());
 
 }

python/sinkhorn.cpp

@@ -1,4 +1,5 @@
 #include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
 #include <pybind11/eigen.h>
 
 #include "MultiSinkhorn.h"
@@ -16,7 +17,14 @@ void bind_sinkhorn(py::module& m) {
                     "to update the dual potentials and compute the marginals.");
 
 
-
+    py::class_<MultimarginalSinkhorn>(m2, "MultiSinkhorn")
+        .def(py::init<ProxPtr, ProxPtr,
+                      kernels::KernelPtr,
+                      Eigen::MatrixXd&>())
+        .def("iterate", &MultimarginalSinkhorn::iterate)
+        .def("run", &MultimarginalSinkhorn::run)
+        .def("get_marginals", &MultimarginalSinkhorn::get_marginals);
+    
 
 }