Add HDF5 support

Project.toml

@@ -6,6 +6,7 @@ version = "0.11.0"
 [deps]
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
+HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
 KM3NeTTestData = "3249b543-581e-4f22-b7da-6c2cdf549b24"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"

docs/src/api.md

@@ -26,6 +26,12 @@ MCTrk
 
 ```
 
+## HDF5
+```@docs
+H5File
+H5CompoundDataset
+```
+
 ## Hardware
 
 ```@docs

docs/src/manual/hdf5.md

+# HDF5 Files
+
+...

src/KM3io.jl

@@ -12,7 +12,14 @@ using DocStringExtensions
 using StaticArrays: FieldVector, SVector
 import UnROOT
 
+using HDF5
+# TODO: this will be fixed when https://github.com/JuliaIO/HDF5.jl/pull/1069
+# is merged. HDF5 dependency should then be set to ^0.16.14 and the following
+# line removed
+HDF5.datatype(::Type{T}) where {T} = HDF5.Datatype(HDF5.hdf5_type_id(T), isstructtype(T))
+
 export ROOTFile
+export H5File, create_dataset
 
 export Direction, Position, UTMPosition, Location, Quaternion, Track, AbstractCalibratedHit
 export Detector, DetectorModule, PMT, Tripod, Hydrophone, center, isbasemodule
@@ -74,6 +81,7 @@ include("hardware.jl")
 include("root/online.jl")
 include("root/offline.jl")
 include("root/root.jl")
+include("hdf5/hdf5.jl")
 include("daq.jl")
 include("acoustics.jl")
 include("calibration.jl")

src/hdf5/hdf5.jl

+struct H5CompoundDatasetCache{T}
+    buffer::Vector{T}
+    size::Int
+end
+
+"""
+
+A flat HDF5 compound dataset which is essentially a vector of structs. It has a
+cache which is filled when elements are pushed to it. The cache is automatically
+written to the target HDF5 path when full.
+
+"""
+struct H5CompoundDataset{T}
+    dset::HDF5.Dataset
+    cache::H5CompoundDatasetCache{T}
+end
+Base.length(c::H5CompoundDatasetCache) = length(c.buffer)
+isfull(c::H5CompoundDatasetCache) = length(c) >= c.size
+
+"""
+
+Forces the cache to be written to the HDF5 file.
+
+"""
+function Base.flush(d::H5CompoundDataset)
+    current_dims, _ = HDF5.get_extent_dims(d.dset)
+    idx = first(current_dims)
+    n = length(d.cache)
+    HDF5.set_extent_dims(d.dset, (idx + n,))
+    d.dset[idx+1:idx+n] = d.cache.buffer
+    empty!(d.cache.buffer)
+end
+
+"""
+
+A wrapper for an HDF5 file used in KM3NeT.
+
+"""
+struct H5File
+    _h5f
+    _datasets::Dict{String, H5CompoundDataset}
+
+    function H5File(fname::AbstractString)
+        h5f = h5open(fname, "cw")
+        new(h5f, Dict{String, H5CompoundDataset}())
+    end
+end
+function Base.close(f::H5File)
+    for dset in values(f._datasets)
+        flush(dset)
+    end
+    close(f._h5f)
+end
+function Base.write(f::H5File, path::AbstractString, data)
+    write_dataset(f._h5f, path, data)
+end
+
+"""
+
+Creates a one-dimensional compound dataset [`H5CompoundDataset`](@ref) of a
+given type which can be extended one-by-one. The cache is used to accumulate
+data and reduce the number of dataset extensions. Each time the cache is full,
+the HDF5 dataset will be extended, the buffer written and cleared.
+
+To force the writing, use [`flush`](@ref)
+"""
+function HDF5.create_dataset(f::H5File, path::AbstractString, ::Type{T}; cache_size=1000) where T
+    dset = HDF5.create_dataset(f._h5f, path, T, ((0,), (-1,)); chunk=(100,))
+    cache = H5CompoundDatasetCache(T[], cache_size)
+    d = H5CompoundDataset(dset, cache)
+    f._datasets[path] = d
+    d
+end
+function Base.push!(d::H5CompoundDataset{T}, element::T) where T
+    push!(d.cache.buffer, element)
+    isfull(d.cache) && flush(d)
+end