eic-opticks/sysrap/NP_nanobind.h

0001 #pragma once
0002 /**
0003 NP_nanobind.h
0004 ==============
0005
0006
0007 **/
0008
0009 #include "NP.hh"
0010 #include <nanobind/nanobind.h>
0011 #include <nanobind/ndarray.h>
0012 #include <nanobind/stl/string.h>
0013
0014 struct NP_nanobind
0015 {
0016     static NP*  NP_copy_of_numpy_array(          nanobind::ndarray<nanobind::numpy> a) ;
0017     static NP*  NP_copy_of_numpy_array_with_meta(nanobind::ndarray<nanobind::numpy> a, nanobind::str meta) ;
0018
0019     static nanobind::dlpack::dtype      dtype_for_NP(char uifc, size_t ebyte);
0020     static nanobind::capsule*           owner_for_NP(char uifc, size_t ebyte, void* data);
0021
0022     static nanobind::ndarray<nanobind::numpy> numpy_array_view_of_NP(const NP* a);
0023     static nanobind::tuple                    numpy_array_view_of_NP_with_meta(const NP* a );
0024
0025     static nanobind::ndarray<nanobind::numpy> example_numpy_array_view_of_NP(int code);
0026
0027     static nanobind::ndarray<nanobind::numpy> roundtrip_numpy_array_via_NP(nanobind::ndarray<nanobind::numpy> src);
0028 };
0029
0030
0031 /**
0032 NP_copy_of_numpy_array : python -> C++
0033 ---------------------------------------------
0034
0035 Currently just copying. Can that be avoided ?
0036
0037 **/
0038
0039 inline NP* NP_nanobind::NP_copy_of_numpy_array(nanobind::ndarray<nanobind::numpy> _a) // static
0040 {
0041     void* data = _a.data();
0042     size_t ndim = _a.ndim();
0043     nanobind::dlpack::dtype _dtype = _a.dtype();
0044
0045     std::string dtype ;
0046     if(      _dtype == nanobind::dtype<float>() )        dtype = descr_<float>::dtype()    ;
0047     else if( _dtype == nanobind::dtype<double>())        dtype = descr_<double>::dtype()   ;
0048     else if( _dtype == nanobind::dtype<double>())        dtype = descr_<double>::dtype()   ;
0049     else if( _dtype == nanobind::dtype<int>())           dtype = descr_<int>::dtype()      ;
0050     else if( _dtype == nanobind::dtype<long>())          dtype = descr_<long>::dtype()     ;
0051     else if( _dtype == nanobind::dtype<unsigned>())      dtype = descr_<unsigned>::dtype() ;
0052     else if( _dtype == nanobind::dtype<unsigned long>()) dtype = descr_<unsigned long>::dtype() ;
0053
0054     std::vector<NP::INT> shape(ndim);
0055     for(size_t i=0 ; i < ndim ; i++ ) shape[i] = _a.shape(i) ;
0056
0057     NP* a = new NP(dtype.c_str(), shape );
0058     assert( a->uarr_bytes() == _a.nbytes() );
0059     a->read_bytes( (char*)data );
0060
0061     return a ;
0062 }
0063
0064 inline NP* NP_nanobind::NP_copy_of_numpy_array_with_meta(nanobind::ndarray<nanobind::numpy> _a, nanobind::str _meta) // static
0065 {
0066     NP* a = NP_copy_of_numpy_array(_a);
0067     a->meta = nanobind::cast<std::string>(_meta);
0068     return a ;
0069 }
0070
0071
0072
0073 inline nanobind::dlpack::dtype NP_nanobind::dtype_for_NP(char uifc, size_t ebyte) // static
0074 {
0075     nanobind::dlpack::dtype dtype ;
0076     if(      uifc == 'f' && ebyte == 4 ) dtype = nanobind::dtype<float>();
0077     else if( uifc == 'f' && ebyte == 8 ) dtype = nanobind::dtype<double>();
0078     else if( uifc == 'u' && ebyte == 4 ) dtype = nanobind::dtype<unsigned>();
0079     else if( uifc == 'u' && ebyte == 8 ) dtype = nanobind::dtype<unsigned long>();
0080     else if( uifc == 'i' && ebyte == 4 ) dtype = nanobind::dtype<int>();
0081     else if( uifc == 'i' && ebyte == 8 ) dtype = nanobind::dtype<long>();
0082    return dtype ;
0083 }
0084
0085 inline nanobind::capsule* NP_nanobind::owner_for_NP(char uifc, size_t ebyte, void* data) // static
0086 {
0087     nanobind::capsule* owner = nullptr ;
0088     if(      uifc == 'f' && ebyte == 4 ) owner = new nanobind::capsule(data, [](void *p) noexcept { delete[] (float *)p ; });
0089     else if( uifc == 'f' && ebyte == 8 ) owner = new nanobind::capsule(data, [](void *p) noexcept { delete[] (double *)p ; });
0090     else if( uifc == 'u' && ebyte == 4 ) owner = new nanobind::capsule(data, [](void *p) noexcept { delete[] (unsigned *)p ; });
0091     else if( uifc == 'u' && ebyte == 8 ) owner = new nanobind::capsule(data, [](void *p) noexcept { delete[] (unsigned long *)p ; });
0092     else if( uifc == 'i' && ebyte == 4 ) owner = new nanobind::capsule(data, [](void *p) noexcept { delete[] (int *)p ; });
0093     else if( uifc == 'i' && ebyte == 8 ) owner = new nanobind::capsule(data, [](void *p) noexcept { delete[] (long *)p ; });
0094     return owner ;
0095 }
0096
0097
0098 /**
0099 numpy_array_view_of_NP : C++ -> python
0100 --------------------------------------------
0101
0102 No copying, just adopts the same data pointer.
0103
0104
0105
0106
0107 **/
0108
0109 inline nanobind::ndarray<nanobind::numpy> NP_nanobind::numpy_array_view_of_NP(const NP* a) // static
0110 {
0111     void* data = (void*)a->bytes() ;
0112     std::vector<size_t> sh ;
0113     a->get_shape(sh);
0114     size_t ndim = sh.size();
0115     const size_t* shape = sh.data();
0116     nanobind::capsule* owner = owner_for_NP(a->uifc, a->ebyte, data);
0117     const int64_t* strides = nullptr ;
0118     nanobind::dlpack::dtype dtype = dtype_for_NP(a->uifc, a->ebyte);
0119     int device_type = nanobind::device::cpu::value ;
0120     int device_id = 0 ;
0121     char order = 'C' ;
0122     return nanobind::ndarray<nanobind::numpy>(data, ndim, shape, *owner, strides, dtype, device_type, device_id, order);
0123 }
0124
0125 inline nanobind::tuple NP_nanobind::numpy_array_view_of_NP_with_meta(const NP* a )
0126 {
0127     nanobind::ndarray<nanobind::numpy> arr = numpy_array_view_of_NP(a);
0128     return nanobind::make_tuple(arr, a->meta );
0129 }
0130
0131
0132 inline nanobind::ndarray<nanobind::numpy> NP_nanobind::example_numpy_array_view_of_NP(int code) // static
0133 {
0134     NP* a = nullptr ;
0135     switch(code)
0136     {
0137         case 0: a = NP::Make<float>(3,6,4);         break ;
0138         case 1: a = NP::Make<double>(3,6,4);        break ;
0139         case 2: a = NP::Make<int>(3,6,4);           break ;
0140         case 3: a = NP::Make<long>(3,6,4);          break ;
0141         case 4: a = NP::Make<unsigned>(3,6,4);      break ;
0142         case 5: a = NP::Make<unsigned long>(3,6,4); break ;
0143     }
0144     a->fillIndexFlat();
0145     return numpy_array_view_of_NP(a);
0146 }
0147
0148
0149 inline nanobind::ndarray<nanobind::numpy> NP_nanobind::roundtrip_numpy_array_via_NP(nanobind::ndarray<nanobind::numpy> src) // static
0150 {
0151     NP* a_src = NP_copy_of_numpy_array(src);
0152     nanobind::ndarray<nanobind::numpy> a_dst = numpy_array_view_of_NP(a_src) ;
0153     return a_dst ;
0154 }
0155
0156
0157