conversion_utils.hpp

/** @file python/conversion_utils.hpp
 *  @brief A set of functions to convert std::vectors to and from python objects
 *
 *  @author Thomas A. R. Purcell (tpurcell)
 *  @bug No known bugs.
 */
#ifndef UTILS_CONVERSION
#define UTILS_CONVERSION

#include <algorithm>
#include <array>
#include <vector>
#include <boost/python.hpp>
#include <boost/python/numpy.hpp>

#include <iostream>

namespace python_conv_utils
{
    namespace py = boost::python;
    namespace np = boost::python::numpy;

    /**
     * @brief Convert a python list object to a std::vector
     *
     * @param lst python list to convert to a std::vector
     * @tparam T type of objects inside the list
     * @return A vector of the python list
     */
    template<typename T>
    std::vector<T> from_list(py::list lst)
    {
        std::vector<T> vec(py::len(lst));
        for(int ll = 0; ll < vec.size(); ++ll)
            vec[ll] = py::extract<T>(lst[ll]);
        return vec;
    }

    /**
     * @brief Convert a numpy array object to a std::vector
     *
     * @param lst numpy array to convert to a std::vector
     * @tparam T type of objects inside the numpy array
     * @return A vector of the numpy array
     */
    template<typename T>
    std::vector<T> from_ndarray(np::ndarray arr)
    {
        if(arr.get_dtype() != np::dtype::get_builtin<T>())
            throw std::logic_error("arr is not of the correct type.");

        std::vector<T> vec(arr.shape(0));
        for(int ii = 0; ii < arr.shape(0); ++ii)
            vec[ii] = py::extract<T>(arr[ii]);

        return vec;
    }

    /**
     * @brief Convert a python list object to a std::array
     *
     * @param lst python list to convert to a std::vector
     * @tparam T type of objects inside the list
     * @tparam N the size of the array
     * @return An array from the python list of the python list
     */
    template<typename T, size_t N>
    std::array<T, N> arr_from_list(py::list lst)
    {
        if(N != py::len(lst))
            throw std::logic_error("N parameter is not the same size as the list");

        std::array<T, N> arr;
        for(int ll = 0; ll < N; ++ll)
            arr[ll] = py::extract<T>(lst[ll]);
        return arr;
    }

    /**
     * @brief Convert a numpy array object to a std::vector
     *
     * @param lst numpy array to convert to a std::array
     * @tparam T type of objects inside the numpy array
     * @tparam N the size of the array
     * @return An array of the numpy array
     */
    template<typename T, size_t N>
    std::array<T, N> arr_from_ndarray(np::ndarray ndarr)
    {
        if(ndarr.get_dtype() != np::dtype::get_builtin<T>())
            throw std::logic_error("arr is not of the correct type.");

        if(N != ndarr)
            throw std::logic_error("N parameter is not the same size as the ndarray");

        std::array<T, N> arr;
        for(int ii = 0; ii < N; ++ii)
            arr[ii] = py::extract<T>(ndarr[ii]);

        return arr;
    }

    /**
     * @brief Convert a python list into a vector of std::shared_ptrs to the object
     *
     * @param lst list to convert
     * @tparam T_ptr the shared_ptr type after the conversion
     * @tparam T_base The base type of the objects in the list
     * @return The vector of the shared_ptrs that represent the objects in the list
     */
    template<typename T_ptr, typename T_base>
    std::vector<std::shared_ptr<T_ptr>> shared_ptr_vec_from_list(py::list lst)
    {
        std::vector<std::shared_ptr<T_ptr>> vec(py::len(lst));
        for(int ll = 0; ll < vec.size(); ++ll)
            vec[ll] = std::make_shared<T_base>(py::extract<T_base>(lst[ll]));
        return vec;
    }

    /**
     * @brief Converts a std::vector to a python list
     *
     * @param vec vector to convert
     * @tparam T Type of the objects inside the vector
     * @return The python list
     */
    template<typename T>
    py::list to_list(std::vector<T> vec)
    {
        py::list lst;
        for(auto& item : vec)
            lst.append<T>(item);
        return lst;
    }

    /**
     * @brief Converts a std::vector to a numpy array
     *
     * @param vec vector to convert
     * @tparam T Type of the objects inside the vector
     * @return The numpy array
     */
    template<typename T>
    np::ndarray to_ndarray(std::vector<T> vec)
    {
        np::ndarray arr = np::zeros(py::make_tuple(vec.size()), np::dtype::get_builtin<T>());
        std::copy_n(vec.data(), vec.size(), reinterpret_cast<T*>(arr.get_data()));
        return arr;
    }

    /**
     * @brief Converts a std::vector to a python list
     *
     * @param vec vector to convert
     * @tparam T Type of the objects inside the vector
     * @return The python list
     */
    template<typename T, int N>
    py::list to_list(std::array<T, N> arr)
    {
        py::list lst;
        for(auto& item : arr)
            lst.append<T>(item);
        return lst;
    }

    /**
     * @brief Converts a std::vector to a numpy array
     *
     * @param vec vector to convert
     * @tparam T Type of the objects inside the vector
     * @return The numpy array
     */
    template<typename T, int N>
    np::ndarray to_ndarray(std::array<T, N> arr)
    {
        np::ndarray np_arr = np::zeros(py::make_tuple(arr.size()), np::dtype::get_builtin<T>());
        std::copy_n(arr.data(), arr.size(), reinterpret_cast<T*>(np_arr.get_data()));
        return np_arr;
    }

    template<typename key, typename val>
    py::dict to_dict(std::map<key, val> map)
    {
        py::dict dct;
        for(auto& iter : map)
            dct[iter.first] = iter.second;

        return dct;
    }
}

#endif