"""
Author: Daniel Berke, berke.daniel@gmail.com
Date: October 27, 2019
Requirements: h5py>=2.10.0, unyt>=v2.4.0
Notes: This short example script shows how to save unit information attached
to a `unyt_array` using `attrs` in HDF5, and recover it upon reading the file.
It uses the Unyt package (https://github.com/yt-project/unyt) because that's 
what I'm familiar with, but presumably similar options exist for Pint and
astropy.units. 
"""

import h5py
import tempfile
import unyt as u

# Set up a temporary file for this example.
tf = tempfile.TemporaryFile()
f = h5py.File(tf, 'a')

# Create some mock data with moderately complicated units (this is the 
# dimensional representation of Joules of energy).
test_data = [1, 2, 3, 4, 5] * u.kg * ( u.m / u.s ) ** 2
print(test_data.units)
# kg*m**2/s**2

# Create a data set to hold the numerical information:
f.create_dataset('stored data', data=test_data)

# Save the units information as a string in `attrs`.
f['stored data'].attrs['units'] = str(test_data.units)

# Now recover the data, using the saved units information to reconstruct the
# original quantities.
reconstituted_data = u.unyt_array(f['stored data'],
                                  units=f['stored data'].attrs['units'])

print(reconstituted_data.units)
# kg*m**2/s**2

assert reconstituted_data.units == test_data.units

Create an example of storing units in HDF5
Create an example of storing units in HDF5

This example script demonstrates how to store unit information attached to an array of numbers using HDF5 in such a way that the original data can be recovered upon reading the file. It uses the Unyt package (https://github.com/yt-project/unyt) for handling units, but presumably similar methods would work for people using astropy.units or Pint. It relies on Unyt being able to be able to parse its own string-ified versions of the units attached correctly, but as it should hopefully create such strings consistently it should work for the majority of use cases.