134 lines
3.2 KiB
C++
134 lines
3.2 KiB
C++
#pragma once
|
|
#include "lasp_mathtypes.h"
|
|
#include <cstring>
|
|
#include <pybind11/numpy.h>
|
|
#include <type_traits>
|
|
|
|
namespace py = pybind11;
|
|
|
|
template<typename T>
|
|
using pyarray = py::array_t<T, py::array::f_style | py::array::forcecast>;
|
|
using dpyarray = pyarray<d>;
|
|
using cpyarray = pyarray<c>;
|
|
|
|
|
|
/**
|
|
* @brief Convert Armadillo column vector to Numpy 1D array
|
|
*
|
|
* @tparam T Type (double, d, c, etc)
|
|
* @param data Input Armadillo array
|
|
*
|
|
* @return Numpy array with copied data
|
|
*/
|
|
template <typename T> py::array_t<T> ColToNpy(const arma::Col<T> &data) {
|
|
|
|
const us Tsize = sizeof(T);
|
|
|
|
const auto nrows = static_cast<ssize_t>(data.n_rows);
|
|
|
|
const us size = nrows;
|
|
|
|
py::array_t<T> result({nrows}, {Tsize});
|
|
|
|
// Copy over memory
|
|
if (size > 0) {
|
|
memcpy(result.mutable_data(0), data.memptr(), sizeof(T) * nrows);
|
|
}
|
|
|
|
return result;
|
|
}
|
|
|
|
/**
|
|
* @brief Convert Armadillo 2D array Numpy 2D array
|
|
*
|
|
* @tparam T Type (double, d, c, etc)
|
|
* @param data Input Armadillo array
|
|
*
|
|
* @return Numpy array with copied data
|
|
*/
|
|
template <typename T> py::array_t<T> MatToNpy(const arma::Mat<T> &data) {
|
|
|
|
const us Tsize = sizeof(T);
|
|
|
|
const auto nrows = static_cast<ssize_t>(data.n_rows);
|
|
const auto ncols = static_cast<ssize_t>(data.n_cols);
|
|
|
|
const us size = nrows;
|
|
|
|
py::array_t<T> result({nrows, ncols}, {Tsize, nrows * Tsize});
|
|
|
|
// Copy over memory
|
|
if (size > 0) {
|
|
memcpy(result.mutable_data(0, 0), data.memptr(), sizeof(T) * nrows * ncols);
|
|
}
|
|
|
|
return result;
|
|
}
|
|
|
|
/**
|
|
* @brief Convert Armadillo Cube to Numpy 3D array
|
|
*
|
|
* @tparam T Type (double, d, c, etc)
|
|
* @param data Input Armadillo Array
|
|
*
|
|
* @return Numpy array with copied data
|
|
*/
|
|
template <typename T> py::array_t<T> CubeToNpy(const arma::Cube<T> &data) {
|
|
|
|
const us Tsize = sizeof(T);
|
|
|
|
const auto nrows = static_cast<ssize_t>(data.n_rows);
|
|
const auto ncols = static_cast<ssize_t>(data.n_cols);
|
|
const auto nslices = static_cast<ssize_t>(data.n_slices);
|
|
|
|
const us size = nrows * ncols * nslices;
|
|
|
|
py::array_t<T> result({nrows, ncols, nslices},
|
|
{Tsize, nrows * Tsize, nrows * ncols * Tsize});
|
|
|
|
// Copy over memory
|
|
if (size > 0) {
|
|
memcpy(result.mutable_data(0, 0, 0), data.memptr(),
|
|
sizeof(T) * nrows * ncols * nslices);
|
|
}
|
|
|
|
return result;
|
|
}
|
|
|
|
/// BACK converters
|
|
|
|
/**
|
|
* @brief Wrap Numpy array to 2D Armadillo Mat
|
|
*
|
|
* @tparam T Type
|
|
* @tparam copy Whether to copy the data. If true, creates a copy
|
|
* @param data The Numpy array to convert
|
|
*
|
|
* @return Mat instance
|
|
*/
|
|
template <typename T, bool copy = true>
|
|
arma::Mat<T> NpyToMat(pyarray<T> data) {
|
|
if (data.ndim() != 2) {
|
|
throw std::runtime_error("Expects a 2D array");
|
|
}
|
|
return arma::Mat<T>(data.mutable_data(0,0), data.shape(0), data.shape(1), copy);
|
|
}
|
|
|
|
/**
|
|
* @brief Wrap Numpy array to 1D Armadillo column vector
|
|
*
|
|
* @tparam T Type
|
|
* @tparam copy Whether to copy the data. If true, creates a copy
|
|
* @param data The Numpy array to convert
|
|
*
|
|
* @return Armadillo column instance
|
|
*/
|
|
template <typename T, bool copy = true>
|
|
arma::Mat<T> NpyToCol(pyarray<T> data) {
|
|
if (data.ndim() != 1) {
|
|
throw std::runtime_error("Expects a 1D array");
|
|
}
|
|
return arma::Col<T>(data.mutable_data(0), data.size(), copy);
|
|
}
|
|
|