IterateNdDataset.hpp

Go to the documentation of this file.

#ifndef RITSUKO_ITERATE_ND_DATASET_HPP
#define RITSUKO_ITERATE_ND_DATASET_HPP
 
#include "H5Cpp.h"
 
#include <vector>
#include <algorithm>
#include <cmath>
 
namespace ritsuko {
 
namespace hdf5 {
 
struct IterateNdDataset {
    IterateNdDataset(std::vector<hsize_t> d, std::vector<hsize_t> b) : 
        data_extent(std::move(d)), 
        block_extent(std::move(b)), 
        ndims(data_extent.size()), 
        starts_internal(ndims), 
        counts_internal(block_extent), 
        dspace(ndims, data_extent.data()) 
    {
        for (auto b : block_extent) {
            total_size *= b;
        }
 
        if (total_size) {
            dspace.selectHyperslab(H5S_SELECT_SET, counts_internal.data(), starts_internal.data());
            mspace.setExtentSimple(ndims, counts_internal.data());
        } else {
            finished_internal = true;
        }
    }
 
    void next() {
        // Attempting a shift from the last dimension as this is the fastest-changing.
        for (size_t i = ndims; i > 0; --i) {
            auto d = i - 1;
            starts_internal[d] += block_extent[d];
 
            // Shift was possible, breaking out.
            if (starts_internal[d] < data_extent[d]) {
                total_size /= counts_internal[d];
                counts_internal[d] = std::min(data_extent[d] - starts_internal[d], block_extent[d]);
                total_size *= counts_internal[d];
                break;
            }
 
            // Next step isn't possible as we've reached the end of the dataset.
            if (d == 0) {
                finished_internal = true;
                return;
            }
 
            // Reached the end of the current dimension extent; set it to zero, 
            // move to the next dimension and increment it.
            starts_internal[d] = 0;
            total_size /= counts_internal[d];
            counts_internal[d] = std::min(data_extent[d], block_extent[d]);
            total_size *= counts_internal[d];
        }
 
        dspace.selectHyperslab(H5S_SELECT_SET, counts_internal.data(), starts_internal.data());
        mspace.setExtentSimple(ndims, counts_internal.data());
    }
 
public:
    bool finished() const {
        return finished_internal;
    }
 
    size_t current_block_size() const {
        return total_size;
    }
 
    const std::vector<hsize_t>& starts () const {
        return starts_internal;
    }
 
    const std::vector<hsize_t>& counts () const {
        return counts_internal;
    }
 
    const H5::DataSpace& file_space() const {
        return dspace;
    }
 
    const H5::DataSpace& memory_space() const {
        return mspace;
    }
 
    const std::vector<hsize_t>& dimensions() const {
        return data_extent;
    }
 
    const std::vector<hsize_t>& block_dimensions() const {
        return block_extent;
    }
 
private:
    std::vector<hsize_t> data_extent, block_extent;
    size_t ndims;
 
    std::vector<hsize_t> starts_internal, counts_internal;
    H5::DataSpace mspace, dspace;
    bool finished_internal = false;
    size_t total_size = 1;
};
 
}
 
}
 
#endif