h5multi.h

// -*- C++ -*-

#ifndef h5multi_
#define h5multi_

#include <type_traits>
#include <memory>
#include <string>
#include <map>
#include "H5Cpp.h"
#include "multidim.h"

namespace h5multi {
using std::string;
using std::shared_ptr;
using std::remove_reference;
using namespace H5;
using namespace multidim;

struct HDF5 {
  shared_ptr<H5File> h5;
  void open(const string &name, bool rw = false, bool erase = false) {
    if (rw) {
      if (erase) {
        h5.reset(new H5File(name, H5F_ACC_TRUNC));
      } else {
        h5.reset(new H5File(name, H5F_ACC_RDWR));
      }
    } else {
      h5.reset(new H5File(name, H5F_ACC_RDONLY));
    }
  }
  ~HDF5() { h5->close(); }
  H5::PredType pred_type(int) { return PredType::NATIVE_INT; }
  H5::PredType pred_type(float) { return PredType::NATIVE_FLOAT; }
  H5::PredType pred_type(double) { return PredType::NATIVE_DOUBLE; }

  struct dims_ {
    int dims[8];
    operator int *() { return dims; }
  };
  inline static dims_ dims(int *p, int n) {
    dims_ result;
    for (int i = 0; i < 8; i++) result.dims[i] = i < n ? p[i] : 0;
    return result;
  }

  template <class T>
  void put(T &a, const string &name, int rank = 2) {
    DSetCreatPropList plist;  // setFillValue, etc.
    hsize_t dims[8];
    for (int i = 0; i < 8; i++) dims[i] = a.dims[i];
    DataSpace fspace(a.rank(), dims);
    DataSet dataset =
        h5->createDataSet(name, pred_type(*a.data), fspace, plist);
    hsize_t start[] = {0, 0, 0, 0, 0, 0, 0, 0};
    hsize_t count[8];
    for (int i = 0; i < 8; i++) count[i] = a.dims[i];
    DataSpace mspace(rank, dims);
    fspace.selectHyperslab(H5S_SELECT_SET, count, start);
    mspace.selectHyperslab(H5S_SELECT_SET, count, start);
    dataset.write(a.data, pred_type(*a.data), mspace, fspace);
  }
  template <class T>
  void shape(mdarray<T> &a, const string &name) {
    DataSet dataset = h5->openDataSet(name);
    DataSpace space = dataset.getSpace();
    hsize_t count[] = {0, 0, 0, 0, 0, 0, 0, 0};
    int rank = space.getSimpleExtentDims(count);
    a.resize(rank);
    for (int i = 0; count[i]; i++) a[i] = count[i];
  }
  template <class T>
  void get(mdarray<T> &a, const string &name) {
    DataSet dataset = h5->openDataSet(name);
    DataSpace space = dataset.getSpace();
    hsize_t offset[] = {0, 0, 0, 0, 0, 0, 0, 0};
    hsize_t count[] = {0, 0, 0, 0, 0, 0, 0, 0};
    int rank = space.getSimpleExtentDims(count);
    a.resize_(count);
    space.selectHyperslab(H5S_SELECT_SET, count, offset);
    DataSpace mem(rank, count);
    mem.selectHyperslab(H5S_SELECT_SET, count, offset);
    dataset.read(a.data, pred_type(*a.data), mem, space);
  }
  template <class T>
  void getrow(mdarray<T> &a, int index, const string &name) {
    DataSet dataset = h5->openDataSet(name);
    DataSpace fspace = dataset.getSpace();
    hsize_t start0[] = {0, 0, 0, 0, 0, 0, 0, 0};
    hsize_t dims[] = {0, 0, 0, 0, 0, 0, 0, 0};
    int rank = fspace.getSimpleExtentDims(dims);
    a.resize_(dims + 1);
    hsize_t count[8];
    for (int i = 0; i < 8; i++) count[i] = dims[i];
    count[0] = 1;
    hsize_t start[] = {hsize_t(index), 0, 0, 0, 0, 0, 0, 0};
    DataSpace mspace(rank, count);
    fspace.selectHyperslab(H5S_SELECT_SET, count, start);
    mspace.selectHyperslab(H5S_SELECT_SET, count, start0);
    dataset.read(a.data, pred_type(*a.data), mspace, fspace);
  }
  template <class T>
  void getvlrow(T &a, int index, const string &name) {
    typedef typename remove_reference<decltype(a[0])>::type S;
    DataSet dataset = h5->openDataSet(name);
    DataSpace space = dataset.getSpace();
    hsize_t dims[] = {0, 0, 0, 0};
    int rank = space.getSimpleExtentDims(dims);
    if (rank != 1) THROW("wrong rank");
    hsize_t start0[] = {0, 0};
    hsize_t start[] = {hsize_t(index), 0};
    hsize_t count[] = {1, 0};
    DataSpace fspace(1, dims);
    DataSpace mspace(1, count);
    fspace.selectHyperslab(H5S_SELECT_SET, count, start);
    mspace.selectHyperslab(H5S_SELECT_SET, count, start0);
    hvl_t vl[1];
    DataType ftype(pred_type(S(0)));
    VarLenType dtype(&ftype);
    dataset.read(vl, dtype, mspace, fspace);
    S *data = (S *)vl[0].p;
    int N = vl[0].len;
    a.resize(N);
    for (int i = 0; i < N; i++) a.data[i] = data[i];
    dataset.vlenReclaim(dtype, mspace, DSetMemXferPropList::DEFAULT, vl);
  }
  template <class T>
  void getarow(mdarray<T> &a, int index, const string &name) {
    mdarray<int> dims;
    shape(dims, name);
    if (dims.size() == 1)
      getvlrow(a, index, name);
    else
      getrow(a, index, name);
  }
  template <class T>
  void getdrow(mdarray<T> &a, int index, const string &name) {
    mdarray<int> dims;
    shape(dims, name);
    if (dims.size() == 1) {
      // variable-shape row; take shape from _dims array
      getvlrow(a, index, name);
      try {
        string sname(name);
        sname += "_dims";
        mdarray<int> ndims;
        getarow(ndims, index, sname.c_str());
        int ndims0[8];
        for (int i = 0; i < ndims.size(); i++) ndims0[i] = ndims[i];
        ndims0[ndims.size()] = 0;
        a.reshape_(ndims0);
      } catch (H5::Exception e) {
        // skip resizing if we can't get _dims, assume it's 1D
      }
    } else {
      // fixed-shape row; take shape from array shape
      getrow(a, index, name);
    }
  }
};

inline HDF5 *make_HDF5() { return new HDF5(); }
}

#endif