Add support for boolean tensors

include/npy/npy.h

@@ -85,10 +85,46 @@ enum class data_type_t : char {
   COMPLEX64,
   /// 128-bit complex number (std::complex<double>)
   COMPLEX128,
+  /// Boolean value
+  BOOL,
   /// Unicode string (std::wstring)
   UNICODE_STRING
 };
 
+/// @brief Boolean datatype which uses 1-byte storage
+/// @details std::vector<bool> uses bitfields to store boolean values, which
+/// makes it inefficient for reading and copying data from numpy, which stores
+/// boolean values as bytes. This struct mimics this scheme while remaining
+/// implicitly convertable to and from boolean values and expressions.
+struct boolean {
+  /// @brief The storage of the boolean, either 1 (true) or 0 (false)
+  std::uint8_t value;
+
+  /// @brief Default constructor, sets value to false.
+  boolean() { value = 0; }
+
+  /// @brief Templated constructor, allowing conversion from any type that can
+  /// be interpreted as a boolean.
+  /// @param v The value to convert into a boolean
+  /// @tparam T Must be implicitly convertable to bool
+  template <typename T> boolean(const T &v) { value = v > 0 ? 1 : 0; }
+
+  /// @brief Constructor from boolean values.
+  /// @param b The bool to store
+  boolean(bool b) { value = b ? 1 : 0; }
+
+  /// @brief Implicit cast operator to bool
+  operator bool() const { return value != 0; }
+
+  /// @brief Assignment operator
+  /// @param v The value to convert into a boolean
+  /// @tparam T Must be implicitly convertable to bool
+  template <typename T> boolean &operator=(const T &v) {
+    value = v ? 1 : 0;
+    return *this;
+  }
+};
+
 /// @brief Convert a data type and endianness to a NPY dtype string.
 /// @param dtype the data type
 /// @param endian the endianness. Defaults to the current endianness of the
@@ -377,8 +413,8 @@ class npzstringwriter {
       compression_method_t compression = compression_method_t::STORED,
       endian_t endianness = npy::endian_t::NATIVE);
 
-  /// @brief Destructor. This will call @ref npy::npzstringwriter::close, if it has
-  /// not been called already.
+  /// @brief Destructor. This will call @ref npy::npzstringwriter::close, if it
+  /// has not been called already.
   ~npzstringwriter();
 
   /// @brief Returns the contents of the string stream as a string.
@@ -452,8 +488,8 @@ class npzfilewriter {
                 compression_method_t compression = compression_method_t::STORED,
                 endian_t endianness = npy::endian_t::NATIVE);
 
-  /// @brief Destructor. This will call @ref npy::npzfilewriter::close, if it has
-  /// not been called already.
+  /// @brief Destructor. This will call @ref npy::npzfilewriter::close, if it
+  /// has not been called already.
   ~npzfilewriter();
 
   /// @brief Returns whether the NPZ file is open.

src/dtype.cpp

@@ -9,13 +9,13 @@
 #define GETC(x) static_cast<char>((x).get())
 
 namespace {
-std::array<std::string, 13> BIG_ENDIAN_DTYPES = {"|i1", "|u1", ">i2", ">u2",
-                                                 ">i4", ">u4", ">i8", ">u8",
-                                                 ">f4", ">f8", ">c8", ">c16"};
+std::array<std::string, 13> BIG_ENDIAN_DTYPES = {
+    "|i1", "|u1", ">i2", ">u2", ">i4",  ">u4", ">i8",
+    ">u8", ">f4", ">f8", ">c8", ">c16", "|b1"};
 
 std::array<std::string, 13> LITTLE_ENDIAN_DTYPES = {
-    "|i1", "|u1", "<i2", "<u2", "<i4", "<u4",
-    "<i8", "<u8", "<f4", "<f8", "<c8", "<c16"};
+    "|i1", "|u1", "<i2", "<u2", "<i4",  "<u4", "<i8",
+    "<u8", "<f4", "<f8", "<c8", "<c16", "|b1"};
 
 std::map<std::string, std::pair<npy::data_type_t, npy::endian_t>> DTYPE_MAP = {
     {"|u1", {npy::data_type_t::UINT8, npy::endian_t::NATIVE}},
@@ -39,7 +39,8 @@ std::map<std::string, std::pair<npy::data_type_t, npy::endian_t>> DTYPE_MAP = {
     {"<c8", {npy::data_type_t::COMPLEX64, npy::endian_t::LITTLE}},
     {">c8", {npy::data_type_t::COMPLEX64, npy::endian_t::BIG}},
     {"<c16", {npy::data_type_t::COMPLEX128, npy::endian_t::LITTLE}},
-    {">c16", {npy::data_type_t::COMPLEX128, npy::endian_t::BIG}}};
+    {">c16", {npy::data_type_t::COMPLEX128, npy::endian_t::BIG}},
+    {"|b1", {npy::data_type_t::BOOL, npy::endian_t::NATIVE}}};
 } // namespace
 
 namespace npy {
@@ -99,6 +100,19 @@ void read_values<>(std::basic_istream<char> &input, int8_t *data_ptr,
   input.read(start, num_elements);
 }
 
+template <>
+void write_values<>(std::basic_ostream<char> &output, const boolean *data_ptr,
+                    size_t num_elements, endian_t) {
+  output.write(reinterpret_cast<const char *>(data_ptr), num_elements);
+}
+
+template <>
+void read_values<>(std::basic_istream<char> &input, boolean *data_ptr,
+                   size_t num_elements, const header_info &) {
+  char *start = reinterpret_cast<char *>(data_ptr);
+  input.read(start, num_elements);
+}
+
 template <>
 void write_values<>(std::basic_ostream<char> &output,
                     const uint_least16_t *data_ptr, size_t num_elements,

src/tensor.cpp

@@ -54,4 +54,8 @@ template <> data_type_t tensor<std::wstring>::get_dtype() {
   return data_type_t::UNICODE_STRING;
 }
 
+template <> data_type_t tensor<boolean>::get_dtype() {
+  return data_type_t::BOOL;
+}
+
 } // namespace npy

test/libnpy_tests.h

@@ -190,6 +190,17 @@ inline npy::tensor<std::wstring> test_tensor(const std::vector<size_t> &shape) {
   return tensor;
 }
 
+template <>
+inline npy::tensor<npy::boolean> test_tensor(const std::vector<size_t> &shape) {
+  npy::tensor<npy::boolean> tensor(shape);
+  int i = 0;
+  for (auto it = tensor.begin(); it != tensor.end(); ++it, ++i) {
+    *it = (i % 2) == 1; // Alternating false, true, false, true, ...
+  }
+
+  return tensor;
+}
+
 template <typename T> npy::tensor<T> test_fortran_tensor() {
   std::vector<int> values = {0,  10, 20, 30, 40, 5,  15, 25, 35, 45, 1,  11, 21,
                              31, 41, 6,  16, 26, 36, 46, 2,  12, 22, 32, 42, 7,
@@ -220,6 +231,21 @@ template <> inline npy::tensor<std::wstring> test_fortran_tensor() {
   return tensor;
 }
 
+template <> inline npy::tensor<npy::boolean> test_fortran_tensor() {
+  std::vector<int> values = {0,  10, 20, 30, 40, 5,  15, 25, 35, 45, 1,  11, 21,
+                             31, 41, 6,  16, 26, 36, 46, 2,  12, 22, 32, 42, 7,
+                             17, 27, 37, 47, 3,  13, 23, 33, 43, 8,  18, 28, 38,
+                             48, 4,  14, 24, 34, 44, 9,  19, 29, 39, 49};
+  npy::tensor<npy::boolean> tensor({5, 2, 5}, true);
+  auto dst = tensor.begin();
+  auto src = values.begin();
+  for (; dst != tensor.end(); ++src, ++dst) {
+    *dst = (*src % 2) == 1;
+  }
+
+  return tensor;
+}
+
 template <typename T>
 std::string npy_stream(npy::endian_t endianness = npy::endian_t::NATIVE) {
   std::ostringstream actual_stream;

test/npy_read.cpp

@@ -22,6 +22,7 @@ int test_npy_read() {
   test_read<std::complex<float>>(result, "complex64");
   test_read<std::complex<double>>(result, "complex128");
   test_read<std::wstring>(result, "unicode");
+  test_read<npy::boolean>(result, "bool");
 
   return result;
 }

test/npy_write.cpp

@@ -78,5 +78,9 @@ int test_npy_write() {
   actual = test::npy_stream<std::wstring>(npy::endian_t::LITTLE);
   test::assert_equal(expected, actual, result, "npy_write_unicode");
 
+  expected = test::read_asset("bool.npy");
+  actual = test::npy_stream<npy::boolean>();
+  test::assert_equal(expected, actual, result, "npy_write_bool");
+
   return result;
 };