npy.h

Go to the documentation of this file.

 
 
#ifndef _NPY_H_
#define _NPY_H_
 
#include <algorithm>
#include <cassert>
#include <cstdint>
#include <filesystem>
#include <fstream>
#include <map>
#include <sstream>
#include <stdexcept>
#include <string>
#include <vector>
 
#define NPY_VERSION_MAJOR 2
#define NPY_VERSION_MINOR 0
#define NPY_VERSION_PATCH 0
#define NPY_VERSION_STRING "2.0.0"
 
const int STATIC_HEADER_LENGTH = 10;
 
namespace npy {
enum class endian_t : char {
  NATIVE,
  BIG,
  LITTLE
};
 
inline endian_t native_endian() {
  union {
    std::uint32_t i;
    char c[4];
  } endian_test = {0x01020304};
 
  return endian_test.c[0] == 1 ? endian_t::BIG : endian_t::LITTLE;
};
 
enum class data_type_t : char {
  INT8,
  UINT8,
  INT16,
  UINT16,
  INT32,
  UINT32,
  INT64,
  UINT64,
  FLOAT32,
  FLOAT64,
  COMPLEX64,
  COMPLEX128,
  UNICODE_STRING
};
 
const std::string &to_dtype(data_type_t dtype,
                            endian_t endian = endian_t::NATIVE);
 
const std::pair<data_type_t, endian_t> &from_dtype(const std::string &dtype);
 
std::ostream &operator<<(std::ostream &os, const endian_t &obj);
std::ostream &operator<<(std::ostream &os, const data_type_t &obj);
 
struct header_info {
  explicit header_info(const std::string &dictionary);
 
  header_info(data_type_t dtype, npy::endian_t endianness, bool fortran_order,
              const std::vector<size_t> &shape);
 
  data_type_t dtype;
 
  npy::endian_t endianness;
 
  bool fortran_order;
 
  std::vector<size_t> shape;
 
  std::size_t max_element_length;
};
 
template <typename CHAR>
void write_npy_header(std::basic_ostream<CHAR> &output,
                      const std::string &dtype, bool fortran_order,
                      const std::vector<size_t> &shape) {
  std::ostringstream buff;
  buff << "{'descr': '" << dtype;
  buff << "', 'fortran_order': " << (fortran_order ? "True" : "False");
  buff << ", 'shape': (";
  for (auto dim = shape.begin(); dim < shape.end(); ++dim) {
    buff << *dim;
    if (dim < shape.end() - 1) {
      buff << ", ";
    }
  }
 
  if (shape.size() == 1) {
    buff << ",";
  }
 
  buff << "), }";
  std::string dictionary = buff.str();
  auto dict_length = dictionary.size() + 1;
  std::string end = "\n";
  auto header_length = dict_length + STATIC_HEADER_LENGTH;
  if (header_length % 64 != 0) {
    header_length = ((header_length / 64) + 1) * 64;
    dict_length = header_length - STATIC_HEADER_LENGTH;
    end = std::string(dict_length - dictionary.length(), ' ');
    end.back() = '\n';
  }
 
  const char header[STATIC_HEADER_LENGTH] = {
      static_cast<char>(0x93),        'N', 'U', 'M', 'P', 'Y', 0x01, 0x00,
      static_cast<char>(dict_length), 0x00};
  output.write(header, STATIC_HEADER_LENGTH);
  output.write(reinterpret_cast<const CHAR *>(dictionary.data()),
               dictionary.length());
  output.write(reinterpret_cast<const CHAR *>(end.data()), end.length());
}
 
template <typename T, typename CHAR>
void write_values(std::basic_ostream<CHAR> &output, const T *data_ptr,
                  size_t num_elements, endian_t endianness);
 
template <typename T, typename CHAR>
void save(std::basic_ostream<CHAR> &output, const T &tensor,
          endian_t endianness = npy::endian_t::NATIVE) {
  std::vector<size_t> shape;
  for (size_t d = 0; d < tensor.ndim(); ++d) {
    shape.push_back(tensor.shape(d));
  }
 
  write_npy_header(output, tensor.dtype(endianness), tensor.fortran_order(),
                   shape);
  tensor.save(output, endianness);
};
 
template <typename T, template <typename> class TENSOR, typename CHAR>
void save(std::basic_ostream<CHAR> &output, const TENSOR<T> &tensor,
          endian_t endianness = npy::endian_t::NATIVE) {
  save<TENSOR<T>, CHAR>(output, tensor, endianness);
}
 
template <typename T>
void save(const std::string &path, T &tensor,
          endian_t endianness = npy::endian_t::NATIVE) {
  std::ofstream output(path, std::ios::out | std::ios::binary);
  if (!output.is_open()) {
    throw std::invalid_argument("path");
  }
 
  save<T>(output, tensor, endianness);
};
 
template <typename T, template <typename> class TENSOR>
void save(const std::string &path, T &tensor,
          endian_t endianness = npy::endian_t::NATIVE) {
  save<TENSOR<T>>(path, tensor, endianness);
};
 
template <typename CHAR>
header_info read_npy_header(std::basic_istream<CHAR> &input) {
  std::uint8_t header[STATIC_HEADER_LENGTH];
  input.read(reinterpret_cast<CHAR *>(header), STATIC_HEADER_LENGTH);
  assert(header[0] == 0x93);
  assert(header[1] == 'N');
  assert(header[2] == 'U');
  assert(header[3] == 'M');
  assert(header[4] == 'P');
  assert(header[5] == 'Y');
  size_t dict_length = 0;
  if (header[6] == 0x01 && header[7] == 0x00) {
    dict_length = header[8] | (header[9] << 8);
  } else if (header[6] == 0x02 && header[7] == 0x00) {
    std::uint8_t extra[2];
    input.read(reinterpret_cast<CHAR *>(extra), 2);
    dict_length =
        header[8] | (header[9] << 8) | (extra[0] << 16) | (extra[1] << 24);
  }
 
  std::vector<CHAR> buffer(dict_length);
  input.read(buffer.data(), dict_length);
  std::string dictionary(buffer.begin(), buffer.end());
  return header_info(dictionary);
}
 
template <typename T, typename CHAR>
void read_values(std::basic_istream<CHAR> &input, T *data_ptr,
                 size_t num_elements, const header_info &info);
 
template <typename T, typename CHAR> T load(std::basic_istream<CHAR> &input) {
  header_info info = read_npy_header(input);
  return T::load(input, info);
}
 
template <typename T> T load(const std::string &path) {
  std::ifstream input(path, std::ios::in | std::ios::binary);
  if (!input.is_open()) {
    throw std::invalid_argument("path");
  }
 
  return load<T>(input);
}
 
template <typename T, template <typename> class TENSOR>
TENSOR<T> load(const std::string &path) {
  return load<TENSOR<T>>(path);
}
 
template <typename CHAR> header_info peek(std::basic_istream<CHAR> &input) {
  return read_npy_header(input);
}
 
header_info peek(const std::string &path);
 
enum class compression_method_t : std::uint16_t {
  STORED = 0,
  DEFLATED = 8
};
 
struct file_entry {
  std::string filename;
  std::uint32_t crc32;
  std::uint64_t compressed_size;
  std::uint64_t uncompressed_size;
  std::uint16_t compression_method;
  std::uint64_t offset;
 
  bool check(const file_entry &other) const;
};
 
class npzstringwriter {
public:
  npzstringwriter(
      compression_method_t compression = compression_method_t::STORED,
      endian_t endianness = npy::endian_t::NATIVE);
 
  ~npzstringwriter();
 
  std::string str() const;
 
  void close();
 
  template <typename T>
  void write(const std::string &filename, const T &tensor) {
    if (m_closed) {
      throw std::runtime_error("Stream is closed");
    }
 
    std::ostringstream output;
    save<T>(output, tensor, m_endianness);
 
    std::string suffix = ".npy";
    std::string name = filename;
    if (name.size() < 4 ||
        !std::equal(suffix.rbegin(), suffix.rend(), name.rbegin())) {
      name += ".npy";
    }
 
    write_file(name, output.str());
  }
 
private:
  void write_file(const std::string &filename, std::string &&bytes);
 
  bool m_closed;
  std::ostringstream m_output;
  compression_method_t m_compression_method;
  endian_t m_endianness;
  std::vector<file_entry> m_entries;
};
 
class npzfilewriter {
public:
  npzfilewriter(const std::string &path,
                compression_method_t compression = compression_method_t::STORED,
                endian_t endianness = npy::endian_t::NATIVE);
 
  npzfilewriter(const char *path,
                compression_method_t compression = compression_method_t::STORED,
                endian_t endianness = npy::endian_t::NATIVE);
 
  npzfilewriter(const std::filesystem::path &path,
                compression_method_t compression = compression_method_t::STORED,
                endian_t endianness = npy::endian_t::NATIVE);
 
  ~npzfilewriter();
 
  bool is_open() const;
 
  void close();
 
  template <typename T>
  void write(const std::string &filename, const T &tensor) {
    if (m_closed) {
      throw std::runtime_error("Stream is closed");
    }
 
    std::ostringstream output;
    save<T>(output, tensor, m_endianness);
 
    std::string suffix = ".npy";
    std::string name = filename;
    if (name.size() < 4 ||
        !std::equal(suffix.rbegin(), suffix.rend(), name.rbegin())) {
      name += ".npy";
    }
 
    write_file(name, output.str());
  }
 
private:
  void write_file(const std::string &filename, std::string &&bytes);
 
  bool m_closed;
  std::ofstream m_output;
  compression_method_t m_compression_method;
  endian_t m_endianness;
  std::vector<file_entry> m_entries;
};
 
class npzstringreader {
public:
  npzstringreader(const std::string &bytes);
 
  npzstringreader(std::string &&bytes);
 
  const std::vector<std::string> &keys() const;
 
  bool contains(const std::string &filename);
 
  header_info peek(const std::string &filename);
 
  template <typename T> T read(const std::string &filename) {
    std::istringstream stream(read_file(filename));
    return load<T>(stream);
  }
 
  template <typename T, template <typename> class TENSOR>
  TENSOR<T> read(const std::string &filename) {
    return read<TENSOR<T>>(filename);
  }
 
private:
  std::string read_file(const std::string &filename);
 
  void read_entries();
 
  std::istringstream m_input;
  std::map<std::string, file_entry> m_entries;
  std::vector<std::string> m_keys;
};
 
class npzfilereader {
public:
  npzfilereader(const std::string &path);
 
  npzfilereader(const char *path);
 
  npzfilereader(const std::filesystem::path &path);
 
  bool is_open() const;
 
  void close();
 
  const std::vector<std::string> &keys() const;
 
  bool contains(const std::string &filename);
 
  header_info peek(const std::string &filename);
 
  template <typename T> T read(const std::string &filename) {
    std::istringstream stream(read_file(filename));
    return load<T>(stream);
  }
 
  template <typename T, template <typename> class TENSOR>
  TENSOR<T> read(const std::string &filename) {
    read<TENSOR<T>>(filename);
  }
 
private:
  std::string read_file(const std::string &filename);
 
  void read_entries();
 
  std::ifstream m_input;
  std::map<std::string, file_entry> m_entries;
  std::vector<std::string> m_keys;
};
 
template <typename T> class tensor {
public:
  typedef T value_type;
  typedef value_type &reference;
  typedef const value_type &const_reference;
  typedef value_type *pointer;
  typedef const value_type *const_pointer;
 
  tensor(const std::vector<size_t> &shape) : tensor(shape, false) {}
 
  tensor(const std::vector<size_t> &shape, bool fortran_order)
      : m_shape(shape),
        m_ravel_strides(tensor<T>::get_ravel_strides(shape, fortran_order)),
        m_fortran_order(fortran_order), m_dtype(tensor<T>::get_dtype()),
        m_values(tensor<T>::get_size(shape)) {}
 
  tensor(const tensor<T> &other)
      : m_shape(other.m_shape), m_ravel_strides(other.m_ravel_strides),
        m_fortran_order(other.m_fortran_order), m_dtype(other.m_dtype),
        m_values(other.m_values) {}
 
  tensor(tensor<T> &&other)
      : m_shape(std::move(other.m_shape)),
        m_ravel_strides(std::move(other.m_ravel_strides)),
        m_fortran_order(other.m_fortran_order), m_dtype(other.m_dtype),
        m_values(std::move(other.m_values)) {}
 
  static tensor<T> from_file(const std::string &path) {
    return npy::load<tensor<T>>(path);
  }
 
  static tensor<T> load(std::basic_istream<char> &input,
                        const header_info &info) {
    tensor<T> result(info.shape, info.fortran_order);
    if (info.dtype != result.dtype()) {
      throw std::runtime_error("requested dtype does not match stream's dtype");
    }
 
    read_values(input, result.m_values.data(), result.m_values.size(), info);
 
    return result;
  }
 
  void save(std::basic_ostream<char> &output, endian_t endianness) const {
    write_values(output, m_values.data(), m_values.size(), endianness);
  }
 
  template <typename... Indices> const T &operator()(Indices... index) const {
    return m_values[ravel(std::vector<std::int32_t>({index...}))];
  }
 
  const T &operator()(const std::vector<std::size_t> &multi_index) const {
    return m_values[ravel(multi_index)];
  }
 
  template <typename... Indices> T &operator()(Indices... index) {
    return m_values[ravel(std::vector<std::int32_t>({index...}))];
  }
 
  T &operator()(const std::vector<std::size_t> &multi_index) {
    return m_values[ravel(multi_index)];
  }
 
  typename std::vector<T>::iterator begin() { return m_values.begin(); }
 
  typename std::vector<T>::const_iterator begin() const {
    return m_values.begin();
  }
 
  typename std::vector<T>::iterator end() { return m_values.end(); }
 
  typename std::vector<T>::const_iterator end() const { return m_values.end(); }
 
  void set(const std::vector<std::int32_t> &multi_index, const T &value) {
    m_values[ravel(multi_index)] = value;
  }
 
  const T &get(const std::vector<std::int32_t> &multi_index) const {
    return m_values[ravel(multi_index)];
  }
 
  std::string dtype(endian_t endianness) const {
    return to_dtype(m_dtype, endianness);
  }
 
  data_type_t dtype() const { return m_dtype; };
 
  const std::vector<T> &values() const { return m_values; }
 
  void copy_from(const T *source, size_t nitems) {
    if (nitems != size()) {
      throw std::invalid_argument("nitems");
    }
 
    std::copy(source, source + nitems, m_values.begin());
  }
 
  void copy_from(const std::vector<T> &source) {
    if (source.size() != size()) {
      throw std::invalid_argument("source.size");
    }
 
    std::copy(source.begin(), source.end(), m_values.begin());
  }
 
  void move_from(std::vector<T> &&source) {
    if (source.size() != size()) {
      throw std::invalid_argument("source.size");
    }
 
    m_values = std::move(source);
  }
 
  T *data() { return m_values.data(); }
 
  const T *data() const { return m_values.data(); }
 
  size_t size() const { return m_values.size(); }
 
  const std::vector<size_t> &shape() const { return m_shape; }
 
  size_t shape(int index) const { return m_shape[index]; }
 
  size_t ndim() const { return m_shape.size(); }
 
  bool fortran_order() const { return m_fortran_order; }
 
  tensor<T> &operator=(const tensor<T> &other) {
    m_shape = other.m_shape;
    m_ravel_strides = other.m_ravel_strides;
    m_fortran_order = other.m_fortran_order;
    m_dtype = other.m_dtype;
    m_values = other.m_values;
    return *this;
  }
 
  tensor<T> &operator=(tensor<T> &&other) {
    m_shape = std::move(other.m_shape);
    m_ravel_strides = std::move(other.m_ravel_strides);
    m_fortran_order = other.m_fortran_order;
    m_dtype = other.m_dtype;
    m_values = std::move(other.m_values);
    return *this;
  }
 
  void save(const std::string &path,
            endian_t endianness = npy::endian_t::NATIVE) {
    npy::save(path, *this, endianness);
  }
 
  template <class INDEX_IT, class SHAPE_IT>
  size_t ravel(INDEX_IT index, SHAPE_IT shape) const {
    std::size_t ravel = 0;
    for (auto stride = m_ravel_strides.begin(); stride < m_ravel_strides.end();
         ++index, ++shape, ++stride) {
      if (*index >= *shape) {
        throw std::invalid_argument("multi_index");
      }
 
      ravel += *index * *stride;
    }
 
    return ravel;
  }
 
  size_t ravel(const std::vector<std::int32_t> &multi_index) const {
    if (multi_index.size() != m_shape.size()) {
      throw std::invalid_argument("multi_index");
    }
 
    std::vector<std::size_t> abs_multi_index(multi_index.size());
    std::transform(multi_index.begin(), multi_index.end(), m_shape.begin(),
                   abs_multi_index.begin(),
                   [](std::int32_t index, std::size_t shape) -> std::size_t {
                     if (index < 0) {
                       return static_cast<std::size_t>(shape + index);
                     }
 
                     return static_cast<std::size_t>(index);
                   });
 
    return ravel(abs_multi_index);
  }
 
  size_t ravel(const std::vector<std::size_t> &abs_multi_index) const {
    if (m_fortran_order) {
      return ravel(abs_multi_index.rbegin(), m_shape.rbegin());
    }
 
    return ravel(abs_multi_index.begin(), m_shape.begin());
  }
 
private:
  std::vector<size_t> m_shape;
  std::vector<size_t> m_ravel_strides;
  bool m_fortran_order;
  data_type_t m_dtype;
  std::vector<T> m_values;
 
  static data_type_t get_dtype();
 
  static size_t get_size(const std::vector<size_t> &shape) {
    size_t size = 1;
    for (auto &dim : shape) {
      size *= dim;
    }
 
    return size;
  }
 
  static std::vector<size_t> get_ravel_strides(const std::vector<size_t> &shape,
                                               bool fortran_order) {
    std::vector<size_t> ravel_strides(shape.size());
    size_t stride = 1;
    auto ravel = ravel_strides.rbegin();
    if (fortran_order) {
      for (auto max_index = shape.begin(); max_index < shape.end();
           ++max_index, ++ravel) {
        *ravel = stride;
        stride *= *max_index;
      }
    } else {
      for (auto max_index = shape.rbegin(); max_index < shape.rend();
           ++max_index, ++ravel) {
        *ravel = stride;
        stride *= *max_index;
      }
    }
 
    return ravel_strides;
  }
};
 
template <>
inline std::string tensor<std::wstring>::dtype(endian_t endianness) const {
  std::size_t max_length = 0;
  for (const auto &element : m_values) {
    if (element.size() > max_length) {
      max_length = element.size();
    }
  }
 
  if (endianness == npy::endian_t::NATIVE) {
    endianness = native_endian();
  }
 
  if (endianness == npy::endian_t::LITTLE) {
    return "<U" + std::to_string(max_length);
  }
 
  return ">U" + std::to_string(max_length);
}
 
} // namespace npy
 
#endif