Eigen inference but slow

src/Chemistry/SurrogateModels/AI_Python_functions/keras_AI_surrogate.py

@@ -17,4 +17,12 @@ def training_step(model, x, y, x_val, y_val, batch_size, epochs):
 
 def prediction_step(model, x, batch_size):
     prediction = model.predict(x, batch_size)
-    return np.array(prediction, dtype=np.float64)
+    print("Prediction from Python", flush=True)
+    print(prediction, flush=True)
+    return np.array(prediction, dtype=np.float64)
+
+def get_weights(model):
+    weights = model.get_weights()
+    # Convert all weight and bias arrays to float64 (double precision)
+    weights_as_double = [w.astype(np.float64) for w in weights]
+    return weights_as_double

src/Chemistry/SurrogateModels/AI_functions.cpp

@@ -4,6 +4,7 @@
 #include <vector>
 #include <Python.h>
 #include <numpy/arrayobject.h>
+#include <Eigen/Dense>
 #include "AI_functions.hpp"
 
 
@@ -16,6 +17,7 @@ namespace poet {
  * @brief Loads the Python interpreter and functions  
  * @param functions_file_path Path to the Python file where the AI surrogate
  * functions are defined
+ * @return 0 if function was succesful
  */
 int Python_Keras_setup(std::string functions_file_path) {
   // Initialize Python functions
@@ -34,6 +36,7 @@ int Python_Keras_setup(std::string functions_file_path) {
  * @brief Loads the user-supplied Keras model  
  * @param model_file_path Path to a .keras file that the user must supply as
  * a variable "model_file_path" in the R input script
+ * @return 0 if function was succesful
  */
 int Python_Keras_load_model(std::string model_file_path) {
   // Initialize Keras model
@@ -50,6 +53,7 @@ int Python_Keras_load_model(std::string model_file_path) {
  * @brief Converts the std::vector 2D matrix representation of a POET Field object to a numpy array
  * for use in the Python AI surrogate functions
  * @param field 2D-Matrix with the content of a Field object
+ * @return Numpy representation of the input vector
  */
 PyObject* vector_to_numpy_array(const std::vector<std::vector<double>>& field) {
   // import numpy and numpy array API
@@ -71,46 +75,44 @@ PyObject* vector_to_numpy_array(const std::vector<std::vector<double>>& field) {
 }
 
 /**
- * @brief Converts a Pyton matrix object to a std::vector vector 2D matrix  
+ * @brief Converts a Pyton matrix object to a std::vector vector
  * @param py_matrix Pyobject that must be a 2D matrix
+ * @result Vector that can be used similar to the return value of the Field object
+ * Field.AsVector() method.
  */
-std::vector<std::vector<double>> numpy_array_to_vector(PyObject* pyArray) {
-  std::vector<std::vector<double>> result;
-  if (!PyArray_Check(pyArray)) {
+std::vector<double> numpy_array_to_vector(PyObject* py_array) {
+  std::vector<double> result;
+  if (!PyArray_Check(py_array)) {
     std::cerr << "The model's predictions are not a numpy array." << std::endl;
     return result;
   }
   // Cast generic PyObject to PyArrayObject
-  PyArrayObject* npArray = reinterpret_cast<PyArrayObject*>(pyArray);
+  PyArrayObject* np_array = reinterpret_cast<PyArrayObject*>(py_array);
   
   // Get shape
-  int numDims = PyArray_NDIM(npArray);
-  npy_intp* shape = PyArray_SHAPE(npArray);
-
+  int numDims = PyArray_NDIM(np_array);
+  npy_intp* shape = PyArray_SHAPE(np_array);
   if (numDims != 2) {
     std::cerr << "The model's predictions are not a 2D matrix." << std::endl;
     return result;
   }
 
   // Copy data into std::vector format
-  double* data = static_cast<double*>(PyArray_DATA(npArray));
-  result.resize(shape[0]);
-  for (npy_intp i = 0; i < shape[0]; ++i) {
-    result[i].resize(shape[1]);
-    for (npy_intp j = 0; j < shape[1]; ++j) {
-      result[i][j] = data[i * shape[1] + j];
-    }
-  }
+  double* data = static_cast<double*>(PyArray_DATA(np_array));
+  npy_intp size = PyArray_SIZE(np_array);
+  result.resize(size);
+  std::copy(data, data + size, result.begin());
   return result;
 }
 
-
 /**
- * @brief Converts the 2D-Matrix represantation of a POET Field object to a numpy array
- * for use in the Python AI surrogate functions
- * @param field 2D-Matrix with the content of a Field object
+ * @brief Uses the Python Keras functions to calculate predictions from a neural network.
+ * @param x 2D-Matrix with the content of a Field object
+ * @param batch_size size for mini-batches that are used in the Keras model.predict() method
+ * @return Predictions that the neural network made from the input values x. The predictions are
+ * represented as a vector similar to the representation from the Field.AsVector() method
  */
-std::vector<std::vector<double>> Python_keras_predict(std::vector<std::vector<double>> x, int batch_size) {  
+std::vector<double> Python_keras_predict(std::vector<std::vector<double>> x, int batch_size) {  
   // Prepare data for python
   PyObject* py_df_x = vector_to_numpy_array(x);
   
@@ -127,7 +129,7 @@ std::vector<std::vector<double>> Python_keras_predict(std::vector<std::vector<do
   PyObject *py_predictions = PyObject_CallObject(py_inference_function, args);
 
   // Check py_rv and return as 2D vector
-  std::vector<std::vector<double>> predictions = numpy_array_to_vector(py_predictions); 
+  std::vector<double> predictions = numpy_array_to_vector(py_predictions); 
   // Clean up
   PyErr_Print(); // Ensure that python errors make it to stdout 
   Py_DECREF(py_df_x);
@@ -138,6 +140,207 @@ std::vector<std::vector<double>> Python_keras_predict(std::vector<std::vector<do
   Py_DECREF(args);
   Py_DECREF(py_predictions);
   return predictions;
-} 
+}
 
+/**
+ * @brief Uses Eigen for fast inference with the weights and biases of a neural network 
+ * @param input_batch Batch of input data that must fit the size of the neural networks input layer
+ * @param model Struct of aligned Eigen vectors that hold the neural networks weights and biases.
+ * Only supports simple fully connected feed forward networks.
+ * @return The batch of predictions made with the neural network weights and biases and the data
+ * in input_batch
+ */
+Eigen::MatrixXd eigen_inference_batched(const Eigen::Ref<Eigen::MatrixXd>& input_batch, const EigenModel& model) {
+    Eigen::MatrixXd current_layer = input_batch;
+    std::cout << "input_batch ROWS " << current_layer.rows() << std::endl;
+    std::cout << "input_batch COLS " << current_layer.cols() << std::endl;
+    // Process all hidden layers
+    for (size_t layer = 0; layer < model.weight_matrices.size() - 1; ++layer) {
+
+        std::cout << "WEIGHTS LAYER " << layer << std::endl;
+        std::cout << "WEIGHTS ROWS " << model.weight_matrices[layer].rows() << std::endl;
+        std::cout << "WEIGHTS COLS " << model.weight_matrices[layer].cols() << std::endl;
+        std::cout << "BIASES SIZE " << model.biases[layer].size() << std::endl;
+
+        current_layer = (model.weight_matrices[layer] * current_layer);
+        std::cout << "MULTIPLIED" << std::endl;
+        current_layer = current_layer.colwise() + model.biases[layer];
+        current_layer = current_layer.array().max(0.0);
+    }
+
+    // Process output layer (without ReLU)
+    size_t output_layer = model.weight_matrices.size() - 1;
+    return (model.weight_matrices[output_layer] * current_layer).colwise() + model.biases[output_layer];
+}
+/**
+ * @brief Converts the weights and biases from a Python Keras model to Eigen matrices
+ * @return The struct containing the model weights and biases as aligned Eigen matrices 
+ */
+#include <iostream>
+
+EigenModel Python_Keras_get_weights_as_Eigen() {
+    EigenModel eigen_model;
+
+    PyObject* py_main_module = PyImport_AddModule("__main__");
+    if (!py_main_module) throw std::runtime_error("Failed to import Python main module");
+
+    PyObject* py_global_dict = PyModule_GetDict(py_main_module);
+    if (!py_global_dict) throw std::runtime_error("Failed to get global dictionary");
+
+    PyObject* py_keras_model = PyDict_GetItemString(py_global_dict, "model");
+    if (!py_keras_model) throw std::runtime_error("Failed to get Keras model from Python dictionary");
+
+    PyObject* py_get_weights_function = PyDict_GetItemString(py_global_dict, "get_weights");
+    if (!py_get_weights_function) throw std::runtime_error("Failed to get Keras get_weights function");
+
+    PyObject* args = Py_BuildValue("(O)", py_keras_model);
+    if (!args) throw std::runtime_error("Failed to create argument tuple");
+
+    // Call Python function
+    PyObject* py_weights_list = PyObject_CallObject(py_get_weights_function, args);
+    Py_DECREF(args); // Cleanup args
+
+    if (!py_weights_list) {
+        PyErr_Print(); // Print Python error
+        throw std::runtime_error("Failed to get weights from Keras model");
+    }
+
+    Py_ssize_t num_layers = PyList_Size(py_weights_list);
+    for (Py_ssize_t i = 0; i < num_layers; i += 2) {
+        // Get weight matrix
+        PyObject* weight_array = PyList_GetItem(py_weights_list, i);
+        if (!weight_array) throw std::runtime_error("Failed to get weight array");
+        if (!PyArray_Check(weight_array)) throw std::runtime_error("Weight array is not a NumPy array");
+
+        PyArrayObject* weight_np = reinterpret_cast<PyArrayObject*>(weight_array);
+        if (PyArray_NDIM(weight_np) != 2) throw std::runtime_error("Weight array is not 2-dimensional");
+
+        // Check data type
+        int dtype = PyArray_TYPE(weight_np);
+        if (dtype == NPY_FLOAT32) {
+          std::cout << "DTYPE: SINGLE" << std::endl;
+        }
+        if (dtype == NPY_DOUBLE) {
+          std::cout << "DTYPE: DOUBLE" << std::endl;
+        }
+        if (dtype != NPY_FLOAT32 && dtype != NPY_DOUBLE) {
+            throw std::runtime_error("Unsupported data type for weight array");
+        }
+
+        // Cast the data correctly depending on the type
+        void* weight_data = PyArray_DATA(weight_np);
+        if (!weight_data) throw std::runtime_error("Failed to get weight data pointer");
+
+        int num_rows = PyArray_DIM(weight_np, 0);  // NumPy's row count
+        int num_cols = PyArray_DIM(weight_np, 1);  // NumPy's column count
+
+        Eigen::MatrixXd weight_matrix; // Use double-precision matrix
+
+        if (dtype == NPY_FLOAT32) {
+            // Handle float32 array from NumPy
+            float* weight_data_float = static_cast<float*>(weight_data);
+            weight_matrix = Eigen::Map<Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>(
+                weight_data_float, num_rows, num_cols).cast<double>().transpose();
+        } else if (dtype == NPY_DOUBLE) {
+            // Handle double (float64) array from NumPy
+            double* weight_data_double = static_cast<double*>(weight_data);
+            weight_matrix = Eigen::Map<Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>(
+                weight_data_double, num_rows, num_cols).transpose();
+        }
+
+        // Get bias vector
+        PyObject* bias_array = PyList_GetItem(py_weights_list, i + 1);
+        if (!bias_array) throw std::runtime_error("Failed to get bias array");
+        if (!PyArray_Check(bias_array)) throw std::runtime_error("Bias array is not a NumPy array");
+
+        PyArrayObject* bias_np = reinterpret_cast<PyArrayObject*>(bias_array);
+        if (PyArray_NDIM(bias_np) != 1) throw std::runtime_error("Bias array is not 1-dimensional");
+
+        // Check bias data type and cast accordingly
+        int bias_dtype = PyArray_TYPE(bias_np);
+        void* bias_data = PyArray_DATA(bias_np);
+        if (!bias_data) throw std::runtime_error("Failed to get bias data pointer");
+
+        int bias_size = PyArray_DIM(bias_np, 0);
+        Eigen::VectorXd bias_vector; // Use double-precision vector
+
+        if (bias_dtype == NPY_FLOAT32) {
+            float* bias_data_float = static_cast<float*>(bias_data);
+            bias_vector = Eigen::Map<Eigen::Matrix<float, Eigen::Dynamic, 1>>(bias_data_float, bias_size).cast<double>();
+        } else if (bias_dtype == NPY_DOUBLE) {
+            double* bias_data_double = static_cast<double*>(bias_data);
+            bias_vector = Eigen::Map<Eigen::VectorXd>(bias_data_double, bias_size);
+        }
+
+        // Add to EigenModel
+        eigen_model.weight_matrices.push_back(weight_matrix);
+        eigen_model.biases.push_back(bias_vector);
+    }
+
+    // Clean up
+    Py_DECREF(py_weights_list); // Clean up Python list object
+    Py_DECREF(py_main_module); // Clean up Python module object
+
+    return eigen_model;
+}
+
+
+
+
+
+/**
+ * @brief Gets the model weights and biases of the Keras model and uses Eigen for fast inference
+ * @param x 2D-Matrix with the content of a Field object
+ * @param batch_size size for mini-batches that are used in the Keras model.predict() method
+ * @return Predictions that the neural network made from the input values x. The predictions are
+ * represented as a vector similar to the representation from the Field.AsVector() method
+ */
+std::vector<double> Eigen_predict(const EigenModel& model, std::vector<std::vector<double>> x, int batch_size) {
+  // Convert input data to Eigen matrix
+  const int num_samples = x[0].size();
+  const int num_features = x.size();
+  
+  std::cout << "num_samples " << num_samples << std::endl;
+  std::cout << "num_features " << num_features << std::endl;
+
+  Eigen::MatrixXd full_input_matrix(num_features, num_samples);
+  for (int i = 0; i < num_samples; ++i) {
+    for (int j = 0; j < num_features; ++j) {
+      full_input_matrix(j, i) = x[j][i];
+      if (i < 6) {std::cout << full_input_matrix.coeff(j, i) << "  ";}
+    }
+    if (i < 6) {std::cout << std::endl;}
+  }
+    
+
+  std::cout << "Eigen rows " << full_input_matrix.rows() << std::endl;
+  std::cout << "Eigen cols " << full_input_matrix.cols() << std::endl;
+
+    std::vector<double> result;
+    result.reserve(num_samples * model.biases.back().size());
+    
+    int num_batches = std::ceil(static_cast<double>(num_samples) / batch_size);
+    for (int batch = 0; batch < num_batches; ++batch) {
+        int start_idx = batch * batch_size;
+        int end_idx = std::min((batch + 1) * batch_size, num_samples);
+        int current_batch_size = end_idx - start_idx;
+        std::cout << "CREATE INPUT BATCH" << std::endl;
+        // Extract the current batch
+        Eigen::MatrixXd input_batch(num_features, current_batch_size); 
+        input_batch = full_input_matrix.block(0, start_idx, num_features, current_batch_size);
+        // Perform inference on the batch
+        Eigen::MatrixXd output_batch(num_features, current_batch_size);
+        output_batch = eigen_inference_batched(input_batch, model);
+        for (int i = 0; i < num_samples; ++i) {
+          for (int j = 0; j < num_features; ++j) {
+            if (i < 6) {std::cout << output_batch.coeff(j, i) << "  ";}
+          }
+          if (i < 6) {std::cout << std::endl;}
+        }
+        
+        //result.insert(result.end(), output_batch.data(), output_batch.data() + output_batch.size());
+    }
+    
+    return result;
+}
 }

src/Chemistry/SurrogateModels/AI_functions.hpp

@@ -17,6 +17,12 @@
 #include <Python.h>
 #include <string>
 
+// PhreeqC definition of pi clashes with Eigen macros so we have to temporarily undef it 
+#pragma push_macro("pi")
+#undef pi
+#include <Eigen/Dense> 
+#pragma pop_macro("pi")
+
 namespace poet {
 
 int Python_Keras_setup(std::string functions_file_path);
@@ -25,11 +31,27 @@ int Python_Keras_load_model(std::string model_file_path);
 
 PyObject* vector_to_numpy_array(const std::vector<std::vector<double>>& field);
 
-std::vector<std::vector<double>> numpy_array_to_vector(PyObject* py_matrix);
+std::vector<double> numpy_array_to_vector(PyObject* py_matrix);
 
-std::vector<std::vector<double>> Python_keras_predict(std::vector<std::vector<double>> x, int batch_size);  
+std::vector<double> Python_keras_predict(std::vector<std::vector<double>> x, int batch_size);  
+
+// Define an aligned allocator for std::vector
+template<typename T>
+using aligned_vector = std::vector<T, Eigen::aligned_allocator<T>>;
+// Define a structure to hold the weights in Eigen matrices
+struct EigenModel {
+    aligned_vector<Eigen::MatrixXd> weight_matrices;
+    aligned_vector<Eigen::VectorXd> biases;
+};
+
+EigenModel Python_Keras_get_weights_as_Eigen();
+
+Eigen::MatrixXd Eigen_batched_inference(const Eigen::Ref<Eigen::MatrixXd>& input_batch, const EigenModel& model);
+
+std::vector<double> Eigen_predict(const EigenModel& model, std::vector<std::vector<double>> x, int batch_size);  
 
 //int Python_keras_train(Field &x, Field &y, Field &x_val, Field &y_val, int batch_size, std::string pid)
 
+
 } // namespace poet
 #endif // AI_FUNCTIONS_HPP

src/poet.cpp

@@ -275,6 +275,10 @@ static Rcpp::List RunMasterLoop(RInsidePOET &R, const RuntimeParameters &params,
   /* Iteration Count is dynamic, retrieving value from R (is only needed by
    * master for the following loop) */
   uint32_t maxiter = params.timesteps.size();
+  
+  /* Stores the weights and biases of the AI surrogate model
+   * for use in an inference function with Eigen */ 
+  EigenModel Eigen_model;
 
   if (params.print_progress) {
     chem.setProgressBarPrintout(true);
@@ -314,12 +318,25 @@ static Rcpp::List RunMasterLoop(RInsidePOET &R, const RuntimeParameters &params,
       MSG("AI Preprocessing");
       R.parseEval("predictors_scaled <- preprocess(predictors)");
 
-      
       // Predict
       MSG("AI: Predict");
-      R["predictions_scaled"] = Rcpp::wrap(Python_keras_predict(R["predictors_scaled"], params.batch_size));
-      R.parseEval("setNames(predictions_scaled, TMP_POPS)");
+      R["TMP"] = Python_keras_predict(R["predictors_scaled"], params.batch_size);
+      R.parseEval("predictions_scaled <- matrix(TMP, nrow=nrow(predictors), byrow = TRUE)");
+      R.parseEval("predictions_scaled <- setNames(data.frame(predictions_scaled), colnames(predictors))");
       
+      MSG("Prediction from Python sent to R")
+      R.parseEval("print(head(predictions_scaled))");
+
+      EigenModel Eigen_model = Python_Keras_get_weights_as_Eigen();
+      MSG("EIGEN WEIGHTS SUCESFULLY IMPORTED");
+      R["TMP"] = Eigen_predict(Eigen_model, R["predictors_scaled"], params.batch_size);
+      R.parseEval("predictions_scaled <- matrix(TMP, nrow=nrow(predictors), byrow = TRUE)");
+      MSG("EIGEN SUCESFULLY PREDICTED");
+      R.parseEval("predictions_scaled <- setNames(data.frame(predictions_scaled), colnames(predictors))");
+
+      MSG("Prediction from Eigen sent to R")
+      R.parseEval("print(head(predictions_scaled))");
+
       // after this comes old R code!
       // Apply postprocessing
       MSG("AI Postprocesing");

src/poet.hpp.in

@@ -67,4 +67,5 @@ struct RuntimeParameters {
   std::uint32_t interp_bucket_entries;
 
   bool use_ai_surrogate = false;
+  int batch_size;
 };