Add TSWriter

.github/workflows/publish_pypi.yml

@@ -74,7 +74,7 @@ jobs:
           CIBW_REPAIR_WHEEL_COMMAND_WINDOWS: "delvewheel repair -w {dest_dir} {wheel}"
           CIBW_ARCHS: ${{ matrix.cibw_archs }}
           CIBW_BEFORE_TEST_LINUX: yum -y install maven java
-          CIBW_TEST_REQUIRES: bfio requests numpy ome_zarr
+          CIBW_TEST_REQUIRES: bfio requests numpy==1.24.0 ome_zarr
           CIBW_TEST_COMMAND: python -W default -m unittest discover -s {project}/tests -v
 
       - name: Install Dependencies
@@ -128,7 +128,7 @@ jobs:
           CIBW_ENVIRONMENT_MACOS: REPAIR_LIBRARY_PATH="/tmp/bfiocpp_bld/local_install/lib:/tmp/bfiocpp_bld/local_install/lib64" ON_GITHUB="TRUE" BFIOCPP_DEP_DIR="/tmp/bfiocpp_bld/local_install" CMAKE_ARGS="-DTENSORSTORE_USE_SYSTEM_JPEG=ON"
           CIBW_REPAIR_WHEEL_COMMAND_MACOS: DYLD_LIBRARY_PATH=$REPAIR_LIBRARY_PATH delocate-listdeps {wheel} && DYLD_LIBRARY_PATH=$REPAIR_LIBRARY_PATH delocate-wheel --require-archs {delocate_archs} -w {dest_dir} {wheel}
           CIBW_ARCHS: ${{ matrix.cibw_archs }}
-          CIBW_TEST_REQUIRES: bfio requests numpy ome_zarr
+          CIBW_TEST_REQUIRES: bfio requests numpy==1.24.0 ome_zarr
           CIBW_TEST_COMMAND: python -W default -m unittest discover -s {project}/tests -v
 
       - name: Install Dependencies

.github/workflows/wheel_build.yml

@@ -72,7 +72,7 @@ jobs:
           CIBW_REPAIR_WHEEL_COMMAND_WINDOWS: "delvewheel repair -w {dest_dir} {wheel}"
           CIBW_ARCHS: ${{ matrix.cibw_archs }}
           CIBW_BEFORE_TEST_LINUX: yum -y install maven java
-          CIBW_TEST_REQUIRES: bfio requests numpy ome_zarr
+          CIBW_TEST_REQUIRES: bfio requests numpy==1.24.0 ome_zarr
           CIBW_TEST_COMMAND: python -W default -m unittest discover -s {project}/tests -v
 
       - name: Upload Artifact
@@ -124,7 +124,7 @@ jobs:
           CIBW_ENVIRONMENT_MACOS: REPAIR_LIBRARY_PATH="/tmp/bfiocpp_bld/local_install/lib:/tmp/bfiocpp_bld/local_install/lib64" ON_GITHUB="TRUE" BFIOCPP_DEP_DIR="/tmp/bfiocpp_bld/local_install" CMAKE_ARGS="-DTENSORSTORE_USE_SYSTEM_JPEG=ON"
           CIBW_REPAIR_WHEEL_COMMAND_MACOS: DYLD_LIBRARY_PATH=$REPAIR_LIBRARY_PATH delocate-listdeps {wheel} && DYLD_LIBRARY_PATH=$REPAIR_LIBRARY_PATH delocate-wheel --require-archs {delocate_archs} -w {dest_dir} {wheel}
           CIBW_ARCHS: ${{ matrix.cibw_archs }}
-          CIBW_TEST_REQUIRES: bfio requests numpy ome_zarr
+          CIBW_TEST_REQUIRES: bfio requests numpy==1.24.0 ome_zarr zarr
           CIBW_TEST_COMMAND: python -W default -m unittest discover -s {project}/tests -v
 
       - name: Upload Artifact

CMakeLists.txt

@@ -24,7 +24,8 @@ set(SOURCE
           src/cpp/ts_driver/ometiff/driver.cc
           src/cpp/interface/interface.cpp
           src/cpp/reader/tsreader.cpp
-          src/cpp/reader/utilities.cpp
+          src/cpp/utilities/utilities.cpp
+          src/cpp/writer/tswriter.cpp
 )
 
 include(FetchContent)
@@ -68,4 +69,4 @@ if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU")
 endif()
 
 target_link_libraries(libbfiocpp PRIVATE tensorstore::tensorstore tensorstore::all_drivers)
-target_link_libraries(libbfiocpp PRIVATE ${Build_LIBRARIES})  
+target_link_libraries(libbfiocpp PRIVATE ${Build_LIBRARIES})  

setup.py

@@ -104,6 +104,6 @@ def build_extension(self, ext):
     zip_safe=False,
     python_requires=">=3.8",
     install_requires=[
-        "numpy",
+        "numpy<2.0.0",
     ],
 )

src/cpp/interface/interface.cpp

@@ -1,10 +1,12 @@
 #include <pybind11/pybind11.h>
 #include <pybind11/stl_bind.h>
+#include <pybind11/stl.h>
 #include <pybind11/numpy.h>
 #include <tuple>
 #include "../reader/tsreader.h"
 #include "../reader/sequence.h"
-#include "../reader/utilities.h"
+#include "../utilities/utilities.h"
+#include "../writer/tswriter.h"
 
 namespace py = pybind11;
 using bfiocpp::Seq;
@@ -116,4 +118,10 @@ PYBIND11_MODULE(libbfiocpp, m) {
         .export_values();
 
     m.def("get_ome_xml", &bfiocpp::GetOmeXml);
-}
+
+
+    // Writer class
+    py::class_<bfiocpp::TsWriterCPP, std::shared_ptr<bfiocpp::TsWriterCPP>>(m, "TsWriterCPP") 
+    .def(py::init<const std::string&, const std::vector<std::int64_t>&, const std::vector<std::int64_t>&, const std::string&>())
+    .def("write", &bfiocpp::TsWriterCPP::write_image);
+}

src/cpp/reader/tsreader.cpp

@@ -7,7 +7,7 @@
 #include "tensorstore/open.h"
 
 #include "tsreader.h"
-#include "utilities.h"
+#include "../utilities/utilities.h"
 #include "type_info.h"
 
 

src/cpp/reader/utilities.cpp → src/cpp/utilities/utilities.cpp

@@ -4,6 +4,7 @@
 #include "utilities.h"
 #include <cassert>
 #include <tiffio.h>
+#include <thread>
 
 namespace bfiocpp {
 tensorstore::Spec GetOmeTiffSpecToRead(const std::string& filename){
@@ -29,7 +30,6 @@ tensorstore::Spec GetZarrSpecToRead(const std::string& filename){
 }
 
 
-
 uint16_t GetDataTypeCode (std::string_view type_name){
 
   if (type_name == std::string_view{"uint8"}) {return 1;}
@@ -108,4 +108,53 @@ std::string GetOmeXml(const std::string& file_path){
     return OmeXmlInfo;
 }
 
+tensorstore::Spec GetZarrSpecToWrite(   const std::string& filename, 
+                                        const std::vector<std::int64_t>& image_shape, 
+                                        const std::vector<std::int64_t>& chunk_shape,
+                                        const std::string& dtype){
+    return tensorstore::Spec::FromJson({{"driver", "zarr"},
+                            {"kvstore", {{"driver", "file"},
+                                         {"path", filename}}
+                            },
+                            {"context", {
+                              {"cache_pool", {{"total_bytes_limit", 1000000000}}},
+                              {"data_copy_concurrency", {{"limit", std::thread::hardware_concurrency()}}},
+                              {"file_io_concurrency", {{"limit", std::thread::hardware_concurrency()}}},
+                            }},
+                            {"metadata", {
+                                          {"zarr_format", 2},
+                                          {"shape", image_shape},
+                                          {"chunks", chunk_shape},
+                                          {"dtype", dtype},
+                                          },
+                            }}).value();
+                                        }
+
+// Function to get the TensorStore DataType based on a string identifier
+tensorstore::DataType GetTensorStoreDataType(const std::string& type_str) {
+    if (type_str == "uint8") {
+        return tensorstore::dtype_v<std::uint8_t>;
+    } else if (type_str == "uint16") {
+        return tensorstore::dtype_v<std::uint16_t>;
+    } else if (type_str == "uint32") {
+        return tensorstore::dtype_v<std::uint32_t>;
+    } else if (type_str == "uint64") {
+        return tensorstore::dtype_v<std::uint64_t>;
+    } else if (type_str == "int8") {
+        return tensorstore::dtype_v<std::int8_t>;
+    } else if (type_str == "int16") {
+        return tensorstore::dtype_v<std::int16_t>;
+    } else if (type_str == "int32") {
+        return tensorstore::dtype_v<std::int32_t>;
+    } else if (type_str == "int64") {
+        return tensorstore::dtype_v<std::int64_t>;
+    } else if (type_str == "float") {
+        return tensorstore::dtype_v<float>;
+    } else if (type_str == "double" || type_str == "float64") { // handle float64 from numpy
+        return tensorstore::dtype_v<double>;
+    } else {
+        throw std::invalid_argument("Unknown data type string: " + type_str);
+    }
+}
+
 } // ns bfiocpp

src/cpp/reader/utilities.h → src/cpp/utilities/utilities.h

@@ -18,4 +18,9 @@ uint16_t GetDataTypeCode (std::string_view type_name);
 std::string GetUTCString();
 std::string GetOmeXml(const std::string& file_path);
 std::tuple<std::optional<int>, std::optional<int>, std::optional<int>>ParseMultiscaleMetadata(const std::string& axes_list, int len);
+tensorstore::DataType GetTensorStoreDataType(const std::string& type_str);
+tensorstore::Spec GetZarrSpecToWrite(const std::string& filename, 
+                                    const std::vector<std::int64_t>& image_shape, 
+                                    const std::vector<std::int64_t>& chunk_shape,
+                                    const std::string& dtype);
 } // ns bfiocpp

src/cpp/writer/tswriter.cpp

@@ -0,0 +1,142 @@
+#include "tswriter.h"
+
+#include "../utilities/utilities.h"
+
+#include <variant>
+#include <string>
+
+using ::tensorstore::internal_zarr::ChooseBaseDType;
+
+namespace bfiocpp {
+
+TsWriterCPP::TsWriterCPP(
+    const std::string& fname, 
+    const std::vector<std::int64_t>& image_shape, 
+    const std::vector<std::int64_t>& chunk_shape,
+    const std::string& dtype_str
+  ): _filename(fname), _image_shape(image_shape), _chunk_shape(chunk_shape) {
+
+  _dtype_code = GetDataTypeCode(dtype_str); 
+
+  std::string dtype_str_converted = (dtype_str == "float64") ? "double" : dtype_str; // change float64 numpy type to double
+
+  auto dtype = GetTensorStoreDataType(dtype_str_converted);
+
+  auto dtype_base = ChooseBaseDType(dtype).value().encoded_dtype;
+
+  auto spec = GetZarrSpecToWrite(_filename, image_shape, chunk_shape, dtype_base);
+
+  TENSORSTORE_CHECK_OK_AND_ASSIGN(_source, tensorstore::Open(
+                            spec,
+                            tensorstore::OpenMode::create |
+                            tensorstore::OpenMode::delete_existing,
+                            tensorstore::ReadWriteMode::write).result());
+}
+
+
+void TsWriterCPP::write_image(py::array& py_image) {
+
+  // use switch instead of template to avoid creating functions for each datatype
+  switch(_dtype_code)
+  {
+    case (1): {
+        auto data_array = tensorstore::Array(py_image.mutable_unchecked<std::uint8_t, 1>().data(0), _image_shape, tensorstore::c_order);
+
+        // Write data array to TensorStore
+        auto write_result = tensorstore::Write(tensorstore::UnownedToShared(data_array), _source).result();
+
+        if (!write_result.ok()) {
+            std::cerr << "Error writing image: " << write_result.status() << std::endl;
+        }
+
+        break;
+    }
+    case (2): {
+        auto data_array = tensorstore::Array(py_image.mutable_unchecked<std::uint16_t, 1>().data(0), _image_shape, tensorstore::c_order);
+
+        auto write_result = tensorstore::Write(tensorstore::UnownedToShared(data_array), _source).result();
+        if (!write_result.ok()) {
+            std::cerr << "Error writing image: " << write_result.status() << std::endl;
+        }
+        break;  
+    }  
+    case (4): {
+        auto data_array = tensorstore::Array(py_image.mutable_unchecked<std::uint32_t, 1>().data(0), _image_shape, tensorstore::c_order);
+
+        auto write_result = tensorstore::Write(tensorstore::UnownedToShared(data_array), _source).result();
+        if (!write_result.ok()) {
+            std::cerr << "Error writing image: " << write_result.status() << std::endl;
+        }
+        break;
+    }
+    case (8): {
+        auto data_array = tensorstore::Array(py_image.mutable_unchecked<std::uint64_t, 1>().data(0), _image_shape, tensorstore::c_order);
+
+        auto write_result = tensorstore::Write(tensorstore::UnownedToShared(data_array), _source).result();
+        if (!write_result.ok()) {
+            std::cerr << "Error writing image: " << write_result.status() << std::endl;
+        }
+        break;
+    }
+    case (16): {
+        auto data_array = tensorstore::Array(py_image.mutable_unchecked<std::int8_t, 1>().data(0), _image_shape, tensorstore::c_order);
+
+        auto write_result = tensorstore::Write(tensorstore::UnownedToShared(data_array), _source).result();
+        if (!write_result.ok()) {
+            std::cerr << "Error writing image: " << write_result.status() << std::endl;
+        }
+        break;
+    }
+    case (32): {
+        auto data_array = tensorstore::Array(py_image.mutable_unchecked<std::int16_t, 1>().data(0), _image_shape, tensorstore::c_order);
+
+        auto write_result = tensorstore::Write(tensorstore::UnownedToShared(data_array), _source).result();
+        if (!write_result.ok()) {
+            std::cerr << "Error writing image: " << write_result.status() << std::endl;
+        }
+        break;   
+    } 
+    case (64): {
+        auto data_array = tensorstore::Array(py_image.mutable_unchecked<std::int32_t, 1>().data(0), _image_shape, tensorstore::c_order);
+
+        auto write_result = tensorstore::Write(tensorstore::UnownedToShared(data_array), _source).result();
+        if (!write_result.ok()) {
+            std::cerr << "Error writing image: " << write_result.status() << std::endl;
+        }
+        break;
+    }
+    case (128): {
+        auto data_array = tensorstore::Array(py_image.mutable_unchecked<std::int64_t, 1>().data(0), _image_shape, tensorstore::c_order);
+
+        // Write data array to TensorStore
+        auto write_result = tensorstore::Write(tensorstore::UnownedToShared(data_array), _source).result();
+        if (!write_result.ok()) {
+            std::cerr << "Error writing image: " << write_result.status() << std::endl;
+        }
+        break;
+    }
+    case (256): {
+        auto data_array = tensorstore::Array(py_image.mutable_unchecked<float, 1>().data(0), _image_shape, tensorstore::c_order);
+
+        auto write_result = tensorstore::Write(tensorstore::UnownedToShared(data_array), _source).result();
+        if (!write_result.ok()) {
+            std::cerr << "Error writing image: " << write_result.status() << std::endl;
+        }
+        break;
+    }
+    case (512): {
+        auto data_array = tensorstore::Array(py_image.mutable_unchecked<double, 1>().data(0), _image_shape, tensorstore::c_order);
+
+        auto write_result = tensorstore::Write(tensorstore::UnownedToShared(data_array), _source).result();
+        if (!write_result.ok()) {
+            std::cerr << "Error writing image: " << write_result.status() << std::endl;
+        }
+        break;
+    }
+    default: {
+        // should not be reached
+        std::cerr << "Error writing image: unsupported data type" << std::endl;
+    }
+  }
+  }
+} 

src/cpp/writer/tswriter.h

@@ -0,0 +1,46 @@
+#pragma once
+
+#include <string>
+#include <memory>
+#include <vector>
+#include <variant>
+#include <iostream>
+#include <tuple>
+#include <optional>
+#include <unordered_map>
+#include "../reader/sequence.h"
+
+#include "tensorstore/tensorstore.h"
+#include "tensorstore/context.h"
+#include "tensorstore/array.h"
+#include "tensorstore/driver/zarr/dtype.h"
+#include "tensorstore/index_space/dim_expression.h"
+#include "tensorstore/kvstore/kvstore.h"
+#include "tensorstore/open.h"
+
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+#include <pybind11/numpy.h>
+
+namespace py = pybind11;
+
+namespace bfiocpp{
+
+class TsWriterCPP{
+public:
+    TsWriterCPP(const std::string& fname, const std::vector<std::int64_t>& image_shape, const std::vector<std::int64_t>& chunk_shape, const std::string& dtype);
+
+    void write_image(py::array& py_image);
+
+private:
+    std::string _filename;
+
+    std::vector<std::int64_t> _image_shape, _chunk_shape;
+
+    uint16_t _dtype_code;
+
+    tensorstore::TensorStore<void, -1, tensorstore::ReadWriteMode::dynamic> _source;
+
+};
+}
+

src/python/bfiocpp/__init__.py

@@ -1,4 +1,5 @@
 from .tsreader import TSReader, Seq, FileType, get_ome_xml  # NOQA: F401
+from .tswriter import TSWriter  # NOQA: F401
 from . import _version
 
 __version__ = _version.get_versions()["version"]