eic-opticks/sysrap/np.h

0001 #pragma once
0002 /**
0003 np.h
0004 =====
0005
0006 https://github.com/simoncblyth/np/
0007
0008 Extract from NP.hh NPU.hh minimal-ish code to write a NumPy file
0009
0010 **/
0011 #include <cassert>
0012 #include <string>
0013 #include <sstream>
0014 #include <fstream>
0015 #include <vector>
0016 #include <cstdint>
0017
0018 struct np
0019 {
0020     static constexpr char* MAGIC = (char*)"\x93NUMPY" ;
0021     static constexpr bool FORTRAN_ORDER = false ;
0022     static char constexpr ENDIAN_LITTLE = '<' ;
0023     static char constexpr ENDIAN_BIG = '>' ;
0024
0025     static std::string _make_header(const std::vector<int>& shape, const char* descr="<f8" );
0026     static std::string _make_header(const std::string& dict);
0027     static std::string _make_preamble( int major=1, int minor=0 );
0028     static std::string _little_endian_short_string( uint16_t dlen ) ;
0029     static std::string _make_tuple(const std::vector<int>& shape, bool json );
0030     static std::string _make_dict(const std::vector<int>& shape, const char* descr );
0031
0032     template<typename T>
0033     static void Write(const char* path, const std::vector<int>& shape, const T* data, const char* descr="<f8" );
0034
0035     template<typename T>
0036     static void Write(const char* dir, const char* name, const std::vector<int>& shape, const T* data, const char* descr="<f8" );
0037
0038     static std::string FormPath( const char* dir, const char* name );
0039     static void WriteString( const char* dir, const char* name, const char* txt );
0040
0041
0042 };
0043
0044 // NPU::_make_header
0045 inline std::string np::_make_header(const std::vector<int>& shape, const char* descr )
0046 {
0047     std::string dict = _make_dict( shape, descr );
0048     std::string header = _make_header( dict );
0049     return header ;
0050 }
0051
0052 // NPU::_make_header
0053 inline std::string np::_make_header(const std::string& dict)
0054 {
0055     uint16_t dlen = dict.size() ;
0056     uint16_t padding = 16 - ((10 + dlen ) % 16 ) - 1 ;
0057     padding += 3*16 ; // adhoc extra padding for bit-perfect matching to NumPy (for test array anyhow)
0058     uint16_t hlen = dlen + padding + 1 ;
0059
0060     assert( (hlen + 10) % 16 == 0 );
0061     std::stringstream ss ;
0062     ss << _make_preamble() ;
0063     ss << _little_endian_short_string( hlen ) ;
0064     ss << dict ;
0065
0066     for(int i=0 ; i < padding ; i++ ) ss << " " ;
0067     ss << "\n" ;
0068
0069     return ss.str();
0070 }
0071
0072 // NPU::_make_preamble
0073 inline std::string np::_make_preamble( int major, int minor )
0074 {
0075     std::string preamble(MAGIC) ;
0076     preamble.push_back((char)major);
0077     preamble.push_back((char)minor);
0078     return preamble ;
0079 }
0080
0081 // NPU::_little_endian_short_string
0082 inline std::string np::_little_endian_short_string( uint16_t dlen )
0083 {
0084     union u16c2_t {
0085         uint16_t u16 ;
0086         char     c[2] ;
0087     };
0088
0089     u16c2_t len ;
0090     len.u16 = dlen ;
0091
0092     unsigned one = 1u ;
0093     char e = *(char *)&one == 1 ? ENDIAN_LITTLE : ENDIAN_BIG ;
0094     std::string hlen(2, ' ') ;
0095     hlen[0] = e == ENDIAN_LITTLE ? len.c[0] : len.c[1] ;
0096     hlen[1] = e == ENDIAN_LITTLE ? len.c[1] : len.c[0] ;
0097     return hlen ;
0098 }
0099
0100 // NPU::_make_tuple
0101 inline std::string np::_make_tuple( const std::vector<int>& shape, bool json )
0102 {
0103     int ndim = shape.size() ;
0104     std::stringstream ss ;
0105     ss <<  ( json ? "[" : "(" ) ;
0106
0107     if( ndim == 1)
0108     {
0109         ss << shape[0] << "," ;
0110     }
0111     else
0112     {
0113         for(int i=0 ; i < ndim ; i++ ) ss << shape[i] << ( i == ndim - 1 ? "" : ", " )  ;
0114     }
0115     ss << ( json ?  "] " : "), " ) ;    // hmm assuming shape comes last in json
0116     return ss.str();
0117 }
0118
0119 // NPU::_make_dict
0120 inline std::string np::_make_dict(const std::vector<int>& shape, const char* descr )
0121 {
0122     std::stringstream ss ;
0123     ss << "{" ;
0124     ss << "'descr': '" << descr << "', " ;
0125     ss << "'fortran_order': " << ( FORTRAN_ORDER ? "True" : "False" ) << ", " ;
0126     ss << "'shape': " ;
0127     bool json = false ;
0128     ss << _make_tuple( shape, json ) ;
0129     ss << "}" ;
0130     return ss.str();
0131 }
0132
0133
0134
0135 // adapt from NP::save
0136 template<typename T>
0137 inline void np::Write(const char* path, const std::vector<int>& shape, const T* data, const char* descr )
0138 {
0139     size_t nv = 1 ;
0140     for(int i=0 ; i < int(shape.size()) ; i++) nv *= shape[i] ;
0141     size_t arr_bytes = sizeof(T)*nv ;
0142
0143     std::string hdr = _make_header(shape, descr);
0144     std::ofstream fpa(path, std::ios::out|std::ios::binary);
0145     fpa << hdr ;
0146     fpa.write( (char*)data, arr_bytes );
0147 }
0148
0149 template<typename T>
0150 inline void np::Write(const char* dir, const char* name, const std::vector<int>& shape, const T* data, const char* descr )
0151 {
0152     std::string path = FormPath(dir, name);
0153     Write<T>( path.c_str(), shape, data, descr );
0154 }
0155
0156 inline std::string np::FormPath( const char* dir, const char* name )
0157 {
0158     std::stringstream ss ;
0159     if(dir) ss << dir << "/" ;
0160     if(name) ss << name ;
0161     std::string path = ss.str();
0162     return path ;
0163 }
0164
0165 inline void np::WriteString( const char* dir, const char* name, const char* txt )
0166 {
0167     std::string path = FormPath(dir, name);
0168     if(txt == nullptr) return ;
0169     std::ofstream fp(path.c_str(), std::ios::out);
0170     fp << txt ;
0171     fp.close();
0172 }
0173
0174