ufunc.hpp

Go to the documentation of this file.

// Copyright Jim Bosch 2010-2012.
// Distributed under the Boost Software License, Version 1.0.
//    (See accompanying file LICENSE_1_0.txt or copy at
//          http://www.boost.org/LICENSE_1_0.txt)
#ifndef BOOST_NUMPY_UFUNC_HPP_INCLUDED
#define BOOST_NUMPY_UFUNC_HPP_INCLUDED

#include <boost/python.hpp>
#include <boost/numpy/numpy_object_mgr_traits.hpp>
#include <boost/numpy/dtype.hpp>
#include <boost/numpy/ndarray.hpp>

namespace boost 
{
namespace numpy 
{

class multi_iter : public python::object 
{
public:

  BOOST_PYTHON_FORWARD_OBJECT_CONSTRUCTORS(multi_iter, python::object);

  void next();

  bool not_done() const;

  char * get_data(int n) const;

  int const get_nd() const;
    
  Py_intptr_t const * get_shape() const;

  Py_intptr_t const shape(int n) const;
    
};

multi_iter make_multi_iter(python::object const & a1);

multi_iter make_multi_iter(python::object const & a1, python::object const & a2);

multi_iter make_multi_iter(python::object const & a1, python::object const & a2, python::object const & a3);

template <typename TUnaryFunctor, 
          typename TArgument=typename TUnaryFunctor::argument_type,
          typename TResult=typename TUnaryFunctor::result_type>
struct unary_ufunc 
{

  static python::object call(TUnaryFunctor & self, python::object const & input, python::object const & output) 
  {
    dtype in_dtype = dtype::get_builtin<TArgument>();
    dtype out_dtype = dtype::get_builtin<TResult>();
    ndarray in_array = from_object(input, in_dtype, ndarray::ALIGNED);
    ndarray out_array = (output != python::object()) ? 
      from_object(output, out_dtype, ndarray::ALIGNED | ndarray::WRITEABLE)
      : zeros(in_array.get_nd(), in_array.get_shape(), out_dtype);
    multi_iter iter = make_multi_iter(in_array, out_array);
    while (iter.not_done()) 
    {
      TArgument * argument = reinterpret_cast<TArgument*>(iter.get_data(0));
      TResult * result = reinterpret_cast<TResult*>(iter.get_data(1));
      *result = self(*argument);
      iter.next();
    } 
    return out_array.scalarize();
  }

  static python::object make() 
  {
    namespace p = python;
    return p::make_function(call, p::default_call_policies(), (p::arg("input"), p::arg("output")=p::object())); 
  }
};

template <typename TBinaryFunctor, 
          typename TArgument1=typename TBinaryFunctor::first_argument_type,
          typename TArgument2=typename TBinaryFunctor::second_argument_type,
          typename TResult=typename TBinaryFunctor::result_type>
struct binary_ufunc 
{

  static python::object 
  call(TBinaryFunctor & self, python::object const & input1, python::object const & input2, 
       python::object const & output) 
  {
    dtype in1_dtype = dtype::get_builtin<TArgument1>();
    dtype in2_dtype = dtype::get_builtin<TArgument2>();
    dtype out_dtype = dtype::get_builtin<TResult>();
    ndarray in1_array = from_object(input1, in1_dtype, ndarray::ALIGNED);
    ndarray in2_array = from_object(input2, in2_dtype, ndarray::ALIGNED);
    multi_iter iter = make_multi_iter(in1_array, in2_array);
    ndarray out_array = (output != python::object())
      ? from_object(output, out_dtype, ndarray::ALIGNED | ndarray::WRITEABLE)
      : zeros(iter.get_nd(), iter.get_shape(), out_dtype);
    iter = make_multi_iter(in1_array, in2_array, out_array);
    while (iter.not_done()) 
    {
      TArgument1 * argument1 = reinterpret_cast<TArgument1*>(iter.get_data(0));
      TArgument2 * argument2 = reinterpret_cast<TArgument2*>(iter.get_data(1));
      TResult * result = reinterpret_cast<TResult*>(iter.get_data(2));
      *result = self(*argument1, *argument2);
      iter.next();
    } 
    return out_array.scalarize();
  }

  static python::object make() 
  {
    namespace p = python;
    return p::make_function(call, p::default_call_policies(), 
                            (p::arg("input1"), p::arg("input2"), p::arg("output")=p::object())); 
  }

};

} // namespace boost::numpy

namespace python
{
namespace converter 
{

NUMPY_OBJECT_MANAGER_TRAITS(numpy::multi_iter);

} // namespace boost::python::converter
} // namespace boost::python
} // namespace boost

#endif // !BOOST_NUMPY_UFUNC_HPP_INCLUDED