ReadIO.hpp

Go to the documentation of this file.

#pragma once
 
#include <any>
#include <array>
#include <cassert>
#include <cstdint>
#include <iostream>
#include <memory>
#include <numeric>
#include <stdexcept>
#include <string>
#include <type_traits>
#include <typeindex>
#include <variant>
#include <vector>
 
#include "BaseIO.hpp"
#include "Types.hpp"
 
using StorageObjectType = AQNWB::Types::StorageObjectType;
using SizeType = AQNWB::Types::SizeType;

namespace AQNWB::IO
{

class DataBlockGeneric
{
public:
  std::any data;
  SizeArray shape;

  std::type_index typeIndex = typeid(void);

  IO::BaseDataType baseDataType;

  DataBlockGeneric() = default;

  DataBlockGeneric(const std::any& inData,
                   const SizeArray& inShape,
                   const std::type_index& inTypeIndex,
                   const IO::BaseDataType inBaseDataType)
      : data(inData)
      , shape(inShape)
      , typeIndex(inTypeIndex)
      , baseDataType(inBaseDataType)
  {
  }

  inline BaseDataType getBaseDataType() const { return baseDataType; }

  BaseDataType::BaseDataVectorVariant as_variant() const
  {
    try {
      switch (baseDataType.type) {
        case BaseDataType::T_U8:
          return std::any_cast<std::vector<uint8_t>>(data);
        case BaseDataType::T_U16:
          return std::any_cast<std::vector<uint16_t>>(data);
        case BaseDataType::T_U32:
          return std::any_cast<std::vector<uint32_t>>(data);
        case BaseDataType::T_U64:
          return std::any_cast<std::vector<uint64_t>>(data);
        case BaseDataType::T_I8:
          return std::any_cast<std::vector<int8_t>>(data);
        case BaseDataType::T_I16:
          return std::any_cast<std::vector<int16_t>>(data);
        case BaseDataType::T_I32:
          return std::any_cast<std::vector<int32_t>>(data);
        case BaseDataType::T_I64:
          return std::any_cast<std::vector<int64_t>>(data);
        case BaseDataType::T_F32:
          return std::any_cast<std::vector<float>>(data);
        case BaseDataType::T_F64:
          return std::any_cast<std::vector<double>>(data);
        case BaseDataType::T_STR:
          return std::any_cast<std::vector<std::string>>(data);
        default:
          return std::monostate {};
      }
    } catch (const std::bad_any_cast&) {
      // If the actual type stored in `data` does not match the expected type
      // based on `baseDataType`, a `std::bad_any_cast` exception will be
      // thrown. In this case, we catch the exception and return
      // `std::monostate` to indicate that the conversion failed.
      return std::monostate {};
    }
  }
};

template<typename DTYPE, std::size_t NDIMS>
class ConstMultiArrayView
{
public:
  using size_type = std::size_t;

  ConstMultiArrayView(const DTYPE* data,
                      const std::array<size_type, NDIMS>& shape,
                      const std::array<size_type, NDIMS>& strides)
      : m_data(data)
      , m_shape(shape)
      , m_strides(strides)
  {
  }

  template<size_type N = NDIMS, typename std::enable_if_t<N == 1, int> = 0>
  const DTYPE& operator[](size_type index) const
  {
    assert(index < m_shape[0]);
    return m_data[index * m_strides[0]];
  }

  template<size_type N = NDIMS, typename std::enable_if_t<(N > 1), int> = 0>
  ConstMultiArrayView<DTYPE, NDIMS - 1> operator[](size_type index) const
  {
    assert(index < m_shape[0]);
    std::array<size_type, NDIMS - 1> sub_shape {};
    std::array<size_type, NDIMS - 1> sub_strides {};
    for (size_type i = 1; i < NDIMS; ++i) {
      sub_shape[i - 1] = m_shape[i];
      sub_strides[i - 1] = m_strides[i];
    }
    return ConstMultiArrayView<DTYPE, NDIMS - 1>(
        m_data + index * m_strides[0], sub_shape, sub_strides);
  }

  template<size_type N = NDIMS, typename std::enable_if_t<N == 1, int> = 0>
  const DTYPE* begin() const
  {
    assert(m_shape[0] == 0 || m_data != nullptr);
    return m_data;
  }

  template<size_type N = NDIMS, typename std::enable_if_t<N == 1, int> = 0>
  const DTYPE* end() const
  {
    assert(m_shape[0] == 0 || m_data != nullptr);
    if (m_shape[0] == 0) {
      return m_data;
    }
    return m_data + (m_shape[0] * m_strides[0]);
  }

  const std::array<size_type, NDIMS>& shape() const { return m_shape; }
 
private:
  const DTYPE* m_data;  
  std::array<size_type, NDIMS> m_shape;  
  std::array<size_type, NDIMS> m_strides;  
};

template<typename DTYPE>
class DataBlock
{
public:
  std::vector<DTYPE> data;
  SizeArray shape;
  const std::type_index typeIndex = typeid(DTYPE);

  DataBlock(const std::vector<DTYPE>& inData, const SizeArray& inShape)
      : data(inData)
      , shape(inShape)
  {
  }

  template<std::size_t NDIMS>
  inline ConstMultiArrayView<DTYPE, NDIMS> as_multi_array() const
  {
    if (shape.size() != NDIMS) {
      throw std::invalid_argument(
          "Shape size does not match the number of dimensions.");
    }
 
    // Calculate the total number of elements expected
    SizeType expected_size = std::accumulate(
        shape.begin(), shape.end(), SizeType {1}, std::multiplies<SizeType> {});
 
    if (data.size() != expected_size) {
      throw std::invalid_argument("Data size does not match the shape.");
    }
 
    std::array<std::size_t, NDIMS> shape_array {};
    for (std::size_t i = 0; i < NDIMS; ++i) {
      shape_array[i] = static_cast<std::size_t>(shape[i]);
    }
 
    std::array<std::size_t, NDIMS> strides {};
    std::size_t stride = 1;
    for (std::size_t i = NDIMS; i-- > 0;) {
      strides[i] = stride;
      stride *= shape_array[i];
    }
 
    return ConstMultiArrayView<DTYPE, NDIMS>(data.data(), shape_array, strides);
  }

  inline static DataBlock<DTYPE> fromGeneric(
      const DataBlockGeneric& genericData)
  {
    auto result = DataBlock<DTYPE>(
        std::any_cast<std::vector<DTYPE>>(genericData.data), genericData.shape);
    return result;
  }

  inline BaseDataType getBaseDataType() const
  {
    return BaseDataType::fromTypeId(typeIndex);
  }
};  // class DataBlock

template<StorageObjectType T>
struct isAllowedStorageObjectType : std::false_type
{
};

template<>
struct isAllowedStorageObjectType<StorageObjectType::Dataset> : std::true_type
{
};

template<>
struct isAllowedStorageObjectType<StorageObjectType::Attribute> : std::true_type
{
};

template<StorageObjectType OTYPE, typename VTYPE = std::any>
class ReadDataWrapper
{
  // Embedded traits for compile time checking of allowed OTYPE for the class
  // and methods
private:
  template<StorageObjectType U>
  struct isDataset
      : std::integral_constant<bool, (U == StorageObjectType::Dataset)>
  {
  };

  static_assert(isAllowedStorageObjectType<OTYPE>::value,
                "StorageObjectType not allowed for ReadDataWrapper");
 
  // Actual definition of the class
public:
  ReadDataWrapper(const std::shared_ptr<IO::BaseIO> io, const std::string& path)
      : m_io(io)
      , m_path(path)
  {
  }

  static inline StorageObjectType getStorageObjectType() { return OTYPE; }

  template<typename T>
  static constexpr bool isType()
  {
    return std::is_same_v<VTYPE, T>;
  }

  ReadDataWrapper(const ReadDataWrapper&) = delete;

  ReadDataWrapper& operator=(const ReadDataWrapper&) = delete;

  virtual ~ReadDataWrapper() {}

  inline const std::string& getPath() const { return m_path; }

  inline std::shared_ptr<IO::BaseIO> getIO() const { return m_io; }

  inline SizeArray getShape() const
  {
    return m_io->getStorageObjectShape(m_path);
  }

  inline SizeType getNumDimensions() const { return this->getShape().size(); }

  inline IO::BaseDataType getDataType() const
  {
    return m_io->getStorageObjectDataType(m_path);
  }

  inline SizeArray getChunking() const
  {
    return m_io->getStorageObjectChunking(m_path);
  }

  template<StorageObjectType U = OTYPE,
           typename std::enable_if<isDataset<U>::value, int>::type = 0>
  inline IO::LinkArrayDataSetConfig toLinkArrayDataSetConfig() const
  {
    return IO::LinkArrayDataSetConfig(m_path);
  }

  inline bool exists() const
  {
    switch (OTYPE) {
      case StorageObjectType::Dataset: {
        return m_io->objectExists(m_path);
      }
      case StorageObjectType::Attribute: {
        return m_io->attributeExists(m_path);
      }
      default: {
        throw std::runtime_error("Unsupported StorageObjectType");
      }
    }
  }

  inline DataBlockGeneric valuesGeneric() const
  {
    switch (OTYPE) {
      case StorageObjectType::Dataset: {
        return m_io->readDataset(m_path);
      }
      case StorageObjectType::Attribute: {
        return m_io->readAttribute(m_path);
      }
      default: {
        throw std::runtime_error("Unsupported StorageObjectType");
      }
    }
  }

  template<StorageObjectType U = OTYPE,
           typename std::enable_if<isDataset<U>::value, int>::type = 0>
  inline DataBlockGeneric valuesGeneric(const SizeArray& start,
                                        const SizeArray& count = {},
                                        const SizeArray& stride = {},
                                        const SizeArray& block = {}) const
  {
    // The function is only enabled for datasets so we don't need to check
    // for attributes here.
    return m_io->readDataset(m_path, start, count, stride, block);
  }

  template<typename T = VTYPE>
  inline DataBlock<VTYPE> values() const
  {
    return DataBlock<T>::fromGeneric(this->valuesGeneric());
  }

  template<typename T = VTYPE,
           StorageObjectType U = OTYPE,
           typename std::enable_if<isDataset<U>::value, int>::type = 0>
  inline DataBlock<VTYPE> values(const SizeArray& start,
                                 const SizeArray& count = {},
                                 const SizeArray& stride = {},
                                 const SizeArray& block = {}) const
  {
    // The function is only enabled for datasets so we don't need to check
    // for attributes here.
    return DataBlock<VTYPE>::fromGeneric(
        this->valuesGeneric(start, count, stride, block));
  }
 
protected:
  const std::shared_ptr<IO::BaseIO> m_io;
  std::string m_path;
};  // ReadDataWrapper
 
}  // namespace AQNWB::IO