arma_npy.h

Go to the documentation of this file.

#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>;
 
 
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;
}
 
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;
}
 
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;
}
 
 
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);
}
 
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);
}