ceph/src/boost/libs/python/example/numpy/gaussian.cpp

   1 // Copyright Jim Bosch 2010-2012.
   2 // Distributed under the Boost Software License, Version 1.0.
   3 //    (See accompanying file LICENSE_1_0.txt or copy at
   4 //          http://www.boost.org/LICENSE_1_0.txt)
   5
   6 #include <boost/python/numpy.hpp>
   7
   8 #include <cmath>
   9 #include <memory>
  10
  11 #ifndef M_PI
  12 #include <boost/math/constants/constants.hpp>
  13 const double M_PI = boost::math::constants::pi<double>();
  14 #endif
  15
  16 namespace bp = boost::python;
  17 namespace bn = boost::python::numpy;
  18
  19 /**
  20  *  A 2x2 matrix class, purely for demonstration purposes.
  21  *
  22  *  Instead of wrapping this class with Boost.Python, we'll convert it to/from numpy.ndarray.
  23  */
  24 class matrix2 {
  25 public:
  26
  27     double & operator()(int i, int j) {
  28         return _data[i*2 + j];
  29     }
  30
  31     double const & operator()(int i, int j) const {
  32         return _data[i*2 + j];
  33     }
  34
  35     double const * data() const { return _data; }
  36
  37 private:
  38     double _data[4];
  39 };
  40
  41 /**
  42  *  A 2-element vector class, purely for demonstration purposes.
  43  *
  44  *  Instead of wrapping this class with Boost.Python, we'll convert it to/from numpy.ndarray.
  45  */
  46 class vector2 {
  47 public:
  48
  49     double & operator[](int i) {
  50         return _data[i];
  51     }
  52
  53     double const & operator[](int i) const {
  54         return _data[i];
  55     }
  56
  57     double const * data() const { return _data; }
  58
  59     vector2 operator+(vector2 const & other) const {
  60         vector2 r;
  61         r[0] = _data[0] + other[0];
  62         r[1] = _data[1] + other[1];
  63         return  r;
  64     }
  65
  66     vector2 operator-(vector2 const & other) const {
  67         vector2 r;
  68         r[0] = _data[0] - other[0];
  69         r[1] = _data[1] - other[1];
  70         return  r;
  71     }
  72
  73 private:
  74     double _data[2];
  75 };
  76
  77 /**
  78  *  Matrix-vector multiplication.
  79  */
  80 vector2 operator*(matrix2 const & m, vector2 const & v) {
  81     vector2 r;
  82     r[0] = m(0, 0) * v[0] + m(0, 1) * v[1];
  83     r[1] = m(1, 0) * v[0] + m(1, 1) * v[1];
  84     return r;
  85 }
  86
  87 /**
  88  *  Vector inner product.
  89  */
  90 double dot(vector2 const & v1, vector2 const & v2) {
  91     return v1[0] * v2[0] + v1[1] * v2[1];
  92 }
  93
  94 /**
  95  *  This class represents a simple 2-d Gaussian (Normal) distribution, defined by a
  96  *  mean vector 'mu' and a covariance matrix 'sigma'.
  97  */
  98 class bivariate_gaussian {
  99 public:
 100
 101     vector2 const & get_mu() const { return _mu; }
 102
 103     matrix2 const & get_sigma() const { return _sigma; }
 104
 105     /**
 106      *  Evaluate the density of the distribution at a point defined by a two-element vector.
 107      */
 108     double operator()(vector2 const & p) const {
 109         vector2 u = _cholesky * (p - _mu);
 110         return 0.5 * _cholesky(0, 0) * _cholesky(1, 1) * std::exp(-0.5 * dot(u, u)) / M_PI;
 111     }
 112
 113     /**
 114      *  Evaluate the density of the distribution at an (x, y) point.
 115      */
 116     double operator()(double x, double y) const {
 117         vector2 p;
 118         p[0] = x;
 119         p[1] = y;
 120         return operator()(p);
 121     }
 122
 123     /**
 124      *  Construct from a mean vector and covariance matrix.
 125      */
 126     bivariate_gaussian(vector2 const & mu, matrix2 const & sigma)
 127         : _mu(mu), _sigma(sigma), _cholesky(compute_inverse_cholesky(sigma))
 128     {}
 129
 130 private:
 131
 132     /**
 133      *  This evaluates the inverse of the Cholesky factorization of a 2x2 matrix;
 134      *  it's just a shortcut in evaluating the density.
 135      */
 136     static matrix2 compute_inverse_cholesky(matrix2 const & m) {
 137         matrix2 l;
 138         // First do cholesky factorization: l l^t = m
 139         l(0, 0) = std::sqrt(m(0, 0));
 140         l(0, 1) = m(0, 1) / l(0, 0);
 141         l(1, 1) = std::sqrt(m(1, 1) - l(0,1) * l(0,1));
 142         // Now do forward-substitution (in-place) to invert:
 143         l(0, 0) = 1.0 / l(0, 0);
 144         l(1, 0) = l(0, 1) = -l(0, 1) / l(1, 1);
 145         l(1, 1) = 1.0 / l(1, 1);
 146         return l;
 147     }
 148
 149     vector2 _mu;
 150     matrix2 _sigma;
 151     matrix2 _cholesky;
 152
 153 };
 154
 155 /*
 156  *  We have a two options for wrapping get_mu and get_sigma into NumPy-returning Python methods:
 157  *   - we could deep-copy the data, making totally new NumPy arrays;
 158  *   - we could make NumPy arrays that point into the existing memory.
 159  *  The latter is often preferable, especially if the arrays are large, but it's dangerous unless
 160  *  the reference counting is correct: the returned NumPy array needs to hold a reference that
 161  *  keeps the memory it points to from being deallocated as long as it is alive.  This is what the
 162  *  "owner" argument to from_data does - the NumPy array holds a reference to the owner, keeping it
 163  *  from being destroyed.
 164  *
 165  *  Note that this mechanism isn't completely safe for data members that can have their internal
 166  *  storage reallocated.  A std::vector, for instance, can be invalidated when it is resized,
 167  *  so holding a Python reference to a C++ class that holds a std::vector may not be a guarantee
 168  *  that the memory in the std::vector will remain valid.
 169  */
 170
 171 /**
 172  *  These two functions are custom wrappers for get_mu and get_sigma, providing the shallow-copy
 173  *  conversion with reference counting described above.
 174  *
 175  *  It's also worth noting that these return NumPy arrays that cannot be modified in Python;
 176  *  the const overloads of vector::data() and matrix::data() return const references,
 177  *  and passing a const pointer to from_data causes NumPy's 'writeable' flag to be set to false.
 178  */
 179 static bn::ndarray py_get_mu(bp::object const & self) {
 180     vector2 const & mu = bp::extract<bivariate_gaussian const &>(self)().get_mu();
 181     return bn::from_data(
 182         mu.data(),
 183         bn::dtype::get_builtin<double>(),
 184         bp::make_tuple(2),
 185         bp::make_tuple(sizeof(double)),
 186         self
 187     );
 188 }
 189 static bn::ndarray py_get_sigma(bp::object const & self) {
 190     matrix2 const & sigma = bp::extract<bivariate_gaussian const &>(self)().get_sigma();
 191     return bn::from_data(
 192         sigma.data(),
 193         bn::dtype::get_builtin<double>(),
 194         bp::make_tuple(2, 2),
 195         bp::make_tuple(2 * sizeof(double), sizeof(double)),
 196         self
 197     );
 198 }
 199
 200 /**
 201  *  To allow the constructor to work, we need to define some from-Python converters from NumPy arrays
 202  *  to the matrix/vector types.  The rvalue-from-python functionality is not well-documented in Boost.Python
 203  *  itself; you can learn more from boost/python/converter/rvalue_from_python_data.hpp.
 204  */
 205
 206 /**
 207  *  We start with two functions that just copy a NumPy array into matrix/vector objects.  These will be used
 208  *  in the templated converted below.  The first just uses the operator[] overloads provided by
 209  *  bp::object.
 210  */
 211 static void copy_ndarray_to_mv2(bn::ndarray const & array, vector2 & vec) {
 212     vec[0] = bp::extract<double>(array[0]);
 213     vec[1] = bp::extract<double>(array[1]);
 214 }
 215
 216 /**
 217  *  Here, we'll take the alternate approach of using the strides to access the array's memory directly.
 218  *  This can be much faster for large arrays.
 219  */
 220 static void copy_ndarray_to_mv2(bn::ndarray const & array, matrix2 & mat) {
 221     // Unfortunately, get_strides() can't be inlined, so it's best to call it once up-front.
 222     Py_intptr_t const * strides = array.get_strides();
 223     for (int i = 0; i < 2; ++i) {
 224         for (int j = 0; j < 2; ++j) {
 225             mat(i, j) = *reinterpret_cast<double const *>(array.get_data() + i * strides[0] + j * strides[1]);
 226         }
 227     }
 228 }
 229
 230 /**
 231  *  Here's the actual converter.  Because we've separated the differences into the above functions,
 232  *  we can write a single template class that works for both matrix2 and vector2.
 233  */
 234 template <typename T, int N>
 235 struct mv2_from_python {
 236
 237     /**
 238      *  Register the converter.
 239      */
 240     mv2_from_python() {
 241         bp::converter::registry::push_back(
 242             &convertible,
 243             &construct,
 244             bp::type_id< T >()
 245         );
 246     }
 247
 248     /**
 249      *  Test to see if we can convert this to the desired type; if not return zero.
 250      *  If we can convert, returned pointer can be used by construct().
 251      */
 252     static void * convertible(PyObject * p) {
 253         try {
 254             bp::object obj(bp::handle<>(bp::borrowed(p)));
 255             std::auto_ptr<bn::ndarray> array(
 256                 new bn::ndarray(
 257                     bn::from_object(obj, bn::dtype::get_builtin<double>(), N, N, bn::ndarray::V_CONTIGUOUS)
 258                 )
 259             );
 260             if (array->shape(0) != 2) return 0;
 261             if (N == 2 && array->shape(1) != 2) return 0;
 262             return array.release();
 263         } catch (bp::error_already_set & err) {
 264             bp::handle_exception();
 265             return 0;
 266         }
 267     }
 268
 269     /**
 270      *  Finish the conversion by initializing the C++ object into memory prepared by Boost.Python.
 271      */
 272     static void construct(PyObject * obj, bp::converter::rvalue_from_python_stage1_data * data) {
 273         // Extract the array we passed out of the convertible() member function.
 274         std::auto_ptr<bn::ndarray> array(reinterpret_cast<bn::ndarray*>(data->convertible));
 275         // Find the memory block Boost.Python has prepared for the result.
 276         typedef bp::converter::rvalue_from_python_storage<T> storage_t;
 277         storage_t * storage = reinterpret_cast<storage_t*>(data);
 278         // Use placement new to initialize the result.
 279         T * m_or_v = new (storage->storage.bytes) T();
 280         // Fill the result with the values from the NumPy array.
 281         copy_ndarray_to_mv2(*array, *m_or_v);
 282         // Finish up.
 283         data->convertible = storage->storage.bytes;
 284     }
 285
 286 };
 287
 288
 289 BOOST_PYTHON_MODULE(gaussian) {
 290     bn::initialize();
 291
 292     // Register the from-python converters
 293     mv2_from_python< vector2, 1 >();
 294     mv2_from_python< matrix2, 2 >();
 295
 296     typedef double (bivariate_gaussian::*call_vector)(vector2 const &) const;
 297
 298     bp::class_<bivariate_gaussian>("bivariate_gaussian", bp::init<bivariate_gaussian const &>())
 299
 300         // Declare the constructor (wouldn't work without the from-python converters).
 301         .def(bp::init< vector2 const &, matrix2 const & >())
 302
 303         // Use our custom reference-counting getters
 304         .add_property("mu", &py_get_mu)
 305         .add_property("sigma", &py_get_sigma)
 306
 307         // First overload accepts a two-element array argument
 308         .def("__call__", (call_vector)&bivariate_gaussian::operator())
 309
 310         // This overload works like a binary NumPy universal function: you can pass
 311         // in scalars or arrays, and the C++ function will automatically be called
 312         // on each element of an array argument.
 313         .def("__call__", bn::binary_ufunc<bivariate_gaussian,double,double,double>::make())
 314         ;
 315 }