Mssl compilation fixes

cpp/examples/forward_backward/inpainting.cc

@@ -7,7 +7,7 @@
 #include <ctime>
 
 #include "sopt/imaging_forward_backward.h"
-#include "sopt/l1_g_proximal.h"
+#include "sopt/l1_non_diff_function.h"
 #include "sopt/logging.h"
 #include "sopt/maths.h"
 #include "sopt/relative_variation.h"
@@ -111,7 +111,7 @@ int main(int argc, char const **argv) {
     .Psi(psi);
 
   // Once the properties are set, inject it into the ImagingForwardBackward object
-  fb.g_proximal(gp);
+  fb.g_function(gp);
 
   SOPT_HIGH_LOG("Starting Forward Backward");
   // Alternatively, forward-backward can be called with a tuple (x, residual) as argument

cpp/examples/forward_backward/inpainting_credible_interval.cc

@@ -8,7 +8,7 @@
 
 #include "sopt/credible_region.h"
 #include "sopt/imaging_forward_backward.h"
-#include "sopt/l1_g_proximal.h"
+#include "sopt/l1_non_diff_function.h"
 #include "sopt/logging.h"
 #include "sopt/maths.h"
 #include "sopt/relative_variation.h"
@@ -112,7 +112,7 @@ int main(int argc, char const **argv) {
     .Psi(psi);
 
   // Once the properties are set, inject it into the ImagingForwardBackward object
-  fb.g_proximal(gp);
+  fb.g_function(gp);
 
   SOPT_HIGH_LOG("Starting Forward Backward");
   // Alternatively, forward-backward can be called with a tuple (x, residual) as argument

cpp/examples/forward_backward/inpainting_joint_map.cc

@@ -7,7 +7,7 @@
 #include <ctime>
 
 #include "sopt/imaging_forward_backward.h"
-#include "sopt/l1_g_proximal.h"
+#include "sopt/l1_non_diff_function.h"
 #include "sopt/joint_map.h"
 #include "sopt/logging.h"
 #include "sopt/maths.h"
@@ -111,7 +111,7 @@ int main(int argc, char const **argv) {
     .Psi(psi);
 
   // Once the properties are set, inject it into the ImagingForwardBackward object
-  fb->g_proximal(gp);
+  fb->g_function(gp);
 
   SOPT_HIGH_LOG("Starting Forward Backward");
   // Alternatively, forward-backward can be called with a tuple (x, residual) as argument

cpp/sopt/CMakeLists.txt

@@ -3,15 +3,16 @@ set(headers
   bisection_method.h chained_operators.h credible_region.h
   imaging_padmm.h logging.h
   forward_backward.h imaging_forward_backward.h
-  g_proximal.h l1_g_proximal.h joint_map.h
+  non_differentiable_func.h l1_non_diff_function.h real_indicator.h joint_map.h
+  differentiable_func.h
   imaging_primal_dual.h primal_dual.h
   maths.h proximal.h relative_variation.h sdmm.h
   wavelets.h conjugate_gradient.h l1_proximal.h padmm.h proximal_expression.h
   reweighted.h types.h wrapper.h exception.h linear_transform.h positive_quadrant.h
   tv_primal_dual.h gradient_operator.h
   l2_forward_backward.h l2_primal_dual.h
   real_type.h sampling.h power_method.h objective_functions.h ${PROJECT_BINARY_DIR}/include/sopt/config.h
-)
+  )
 
 set(wavelet_headers
   wavelets/direct.h wavelets/indirect.h wavelets/innards.impl.h wavelets/sara.h
@@ -31,7 +32,9 @@ if(SOPT_MPI)
 endif()
 
 if(onnxrt)
-  list(APPEND headers ort_session.h tf_g_proximal.h)
+
+  list(APPEND headers ort_session.h onnx_differentiable_func.h tf_non_diff_function.h)
+
 endif()
 
 add_library(sopt SHARED ${sources})
@@ -52,6 +55,7 @@ endif()
 if(EIGEN3_INCLUDE_DIR)
   target_include_directories(sopt SYSTEM PUBLIC ${EIGEN3_INCLUDE_DIR})
 endif()
+message(STATUS "Eigen3 include dir ${EIGEN3_INCLUDE_DIR}")
 
 if(SOPT_OPENMP)
   target_link_libraries(sopt OpenMP::OpenMP_CXX)
@@ -62,9 +66,16 @@ if(SOPT_MPI)
   target_include_directories(sopt SYSTEM PUBLIC ${MPI_CXX_INCLUDE_PATH})
 endif()
 
+message(STATUS "onnxrt enabled? ${onnxrt}")
 if(onnxrt)
+  if(${onnxruntime_FOUND})
+    message(STATUS "ONNX Source dir ${onnxruntime_SOURCE_DIR}")
+    message(STATUS "ONNX Include dir ${onnxruntime_INCLUDE_DIR}")
+  endif()
   target_link_libraries(sopt ${onnxruntime_LIBRARIES})
   target_include_directories(sopt SYSTEM PUBLIC ${onnxruntime_INCLUDE_DIR})
+else()
+  message(STATUS "ONNXrt not found. ONNXrt support disabled.")
 endif()
 
 target_link_libraries(sopt ${CONAN_LIBS})

cpp/sopt/config.in.h

@@ -24,6 +24,7 @@
 
 #include <string>
 #include <tuple>
+#include <cstdint>
 
 namespace sopt {
 

cpp/sopt/differentiable_func.h

@@ -0,0 +1,43 @@
+#ifndef DIFFERENTIABLE_FUNC_H
+#define DIFFERENTIABLE_FUNC_H
+
+#include "sopt/forward_backward.h"
+
+// Abstract base class providing the interface for differentiable functions f(x)
+// with a defined gradient.
+template <typename SCALAR> class DifferentiableFunc
+{
+
+public:
+  using FB = sopt::algorithm::ForwardBackward<SCALAR>;
+  using Real = typename FB::Real;
+  using t_Vector = typename FB::t_Vector;
+  using t_Gradient = typename FB::t_Gradient;
+  using t_LinearTransform = typename FB::t_LinearTransform;
+
+    // A function that prints a log message
+  virtual void log_message() const = 0;
+
+  // Return a function representing the proximal operator for this function.
+  // Function must be of type t_Proximal, that is
+  // void proximal_operator(Vector, real, Vector)
+  virtual t_Gradient gradient() const
+  {
+    return [this](t_Vector &output, const t_Vector &image, const t_Vector &residual,
+                  const t_LinearTransform &Phi) -> void { this->gradient(output, image, residual, Phi); };
+  }
+
+  // Calculate the gradient directly
+  virtual void gradient(t_Vector &output, const t_Vector &image, const t_Vector &residual,
+                        const t_LinearTransform &Phi) const = 0;
+
+  // Calculate the function directly
+  virtual Real function(t_Vector const &image, t_Vector const &y, t_LinearTransform const &Phi) const = 0;
+
+  // Transforms input image to a different basis.
+  // Return linear_transform_identity() if transform not necessary.
+  //virtual const t_LinearTransform &Phi() const = 0;
+
+};
+
+#endif

cpp/sopt/forward_backward.h

@@ -13,8 +13,14 @@
 
 namespace sopt::algorithm {
 
-//! \brief Forward Backward Splitting
-//! \details \f$\min_{x} f(\Phi x - y) + g(z)\f$. \f$y\f$ is a target vector.
+/*! \brief Forward Backward Splitting 
+  \An optimisation method to solve the problem \f$y = \Phi(x) + N(\sigma)\f
+  \f$\min_{x} f(x, y, \Phi) + g(x)\f$. \f$y\f$ is a target vector, while x is the current solution.
+  \f$f$ is a differentiable function. It is necessary to supply the gradient. 
+  \f$f$ is represented using a DifferentiableFunc object, which supplies the function and its gradient.
+  \f$g$ is a non-differentiable function. It is necessary to supply a proximal operator (or similar stepping function e.g. tensor-flow denoiser).
+  \f$g$ is represented using a NonDifferentiableFunc object, which supplies the function and its proximal operator.
+*/
 template <typename SCALAR>
 class ForwardBackward {
  public:
@@ -33,7 +39,8 @@ class ForwardBackward {
   //! Type of the proximal operator
   using t_Proximal = ProximalFunction<Scalar>;
   //! Type of the gradient
-  using t_Gradient = typename std::function<void(t_Vector &, const t_Vector &)>;
+  // The first argument is the output vector, the others are inputs
+  using t_Gradient = std::function<void(t_Vector &gradient, const t_Vector &image, const t_Vector &residual, const t_LinearTransform& Phi)>;
 
   //! Values indicating how the algorithm ran
   struct Diagnostic {
@@ -56,8 +63,8 @@ class ForwardBackward {
   };
 
   //! Setups ForwardBackward
-  //! \param[in] f_gradient: gradient of the \f$f\f$ function.
-  //! \param[in] g_proximal: proximal operator of the \f$g\f$ function
+  //! \param[in] f_function: the differentiable function \f$f\f$ with a gradient
+  //! \param[in] g_function: the non-differentiable function \f$g\f$ with a proximal operator
   template <typename DERIVED>
   ForwardBackward(t_Gradient const &f_gradient, t_Proximal const &g_proximal,
                   Eigen::MatrixBase<DERIVED> const &target)
@@ -66,7 +73,7 @@ class ForwardBackward {
         beta_(1),
         nu_(1),
         is_converged_(),
-	fista_(true),
+	      fista_(true),
         Phi_(linear_transform_identity<Scalar>()),
         f_gradient_(f_gradient),
         g_proximal_(g_proximal),
@@ -107,9 +114,9 @@ class ForwardBackward {
   //! Second proximal
   SOPT_MACRO(g_proximal, t_Proximal);
 #undef SOPT_MACRO
-  //! \brief Simplifies calling the proximal of f.
-  void f_gradient(t_Vector &out, t_Vector const &x) const { f_gradient()(out, x); }
-  //! \brief Simplifies calling the proximal of f.
+  //! \brief Simplifies calling the gradient function
+  void f_gradient(t_Vector &out, t_Vector const &x, t_Vector const &res, t_LinearTransform const &Phi) const { f_gradient()(out, x, res, Phi); }
+  //! \brief Simplifies calling the proximal function
   void g_proximal(t_Vector &out, Real gamma, t_Vector const &x) const {
     g_proximal()(out, gamma, x);
   }
@@ -231,15 +238,21 @@ class ForwardBackward {
 };
 
 template <typename SCALAR>
-void ForwardBackward<SCALAR>::iteration_step(t_Vector &out, t_Vector &residual, t_Vector &p,
-                                             t_Vector &z, const t_real lambda) const {
-  p = out;
-  f_gradient(z, residual);
-  const t_Vector input = out - beta() / nu() * (Phi().adjoint() * z);
+void ForwardBackward<SCALAR>::iteration_step(t_Vector &image_new, t_Vector &residual, t_Vector &image_current,
+                                             t_Vector &gradient_current, const t_real FISTA_ratio) const {
+  image_current = image_new;
+  SOPT_LOW_LOG("Calculate gradient");
+  f_gradient(gradient_current, image_current, residual, Phi());  // takes residual and calculates the grad = 1/sig^2 residual
+  SOPT_LOW_LOG("Take a gradient step");
+  image_new = image_current - beta() / nu() * gradient_current;  // step to new image using gradient
+  SOPT_LOW_LOG("Calculate the weight");
   const Real weight = gamma() * beta();
-  g_proximal(out, weight, input);
-  p = out + lambda * (out - p);
-  residual = (Phi() * p) - target();
+  SOPT_LOW_LOG("Apply proximal operator");
+  g_proximal(image_new, weight, image_new);  // apply proximal operator to new image
+  SOPT_LOW_LOG("FISTA acceleration step");
+  image_current = image_new + FISTA_ratio * (image_new - image_current);  // FISTA acceleration step  
+  SOPT_LOW_LOG("Calculate the residual");
+  residual = (Phi() * image_current) - target();  // calculates the residual for the NEXT iteration.
 }
 
 template <typename SCALAR>
@@ -253,30 +266,33 @@ typename ForwardBackward<SCALAR>::Diagnostic ForwardBackward<SCALAR>::operator()
   }
   sanity_check(x_guess, res_guess);
 
-  t_Vector p = t_Vector::Zero(x_guess.size());
-  t_Vector z = t_Vector::Zero(target().size());
+  const uint image_size = x_guess.size();
+
+  t_Vector image_current = x_guess;
   t_Vector residual = res_guess;
+  t_Vector gradient_current = t_Vector::Zero(image_size);
   out = x_guess;
 
   t_uint niters(0);
   bool converged = false;
   Real theta = 1.0;
   Real theta_new = 1.0;
-  Real lambda = 0.0;
+  Real FISTA_ratio = 0.0;
   for (; (not converged) && (niters < itermax()); ++niters) {
-    SOPT_LOW_LOG("    - [FB] Iteration {}/{}", niters, itermax());
+    SOPT_HIGH_LOG("    - [FB] Iteration {}/{}", niters, itermax());
     if (fista()) {
       theta_new = (1 + std::sqrt(1 + 4 * theta * theta)) / 2.;
-      lambda = (theta - 1) / (theta_new);
+      FISTA_ratio = (theta - 1) / (theta_new);
       theta = theta_new;
     }
-    iteration_step(out, residual, p, z, lambda);
-    SOPT_LOW_LOG("      - [FB] Sum of residuals: {}", residual.array().abs().sum());
+    SOPT_HIGH_LOG("      - Call iteration step");
+    iteration_step(out, residual, image_current, gradient_current, FISTA_ratio);
+    SOPT_HIGH_LOG("      - [FB] Sum of residuals: {}", residual.array().abs().sum());
     converged = is_converged(out, residual);
   }
 
   if (converged) {
-    SOPT_MEDIUM_LOG("    - [FB] converged in {} of {} iterations", niters, itermax());
+    SOPT_HIGH_LOG("    - [FB] converged in {} of {} iterations", niters, itermax());
   } else if (static_cast<bool>(is_converged())) {
     // not meaningful if not convergence function
     SOPT_ERROR("    - [FB] did not converge within {} iterations", itermax());