1"""
2Author: Daniel Berke, berke.daniel@gmail.com
3Date: October 27, 2019
4Requirements: h5py>=2.10.0, unyt>=v2.4.0
5Notes: This short example script shows how to save unit information attached
6to a `unyt_array` using `attrs` in HDF5, and recover it upon reading the file.
7It uses the Unyt package (https://github.com/yt-project/unyt) because that's
8what I'm familiar with, but presumably similar options exist for Pint and
9astropy.units.
10"""
11
12import h5py
13import tempfile
14import unyt as u
15
16# Set up a temporary file for this example.
17tf = tempfile.TemporaryFile()
18f = h5py.File(tf, 'a')
19
20# Create some mock data with moderately complicated units (this is the
21# dimensional representation of Joules of energy).
22test_data = [1, 2, 3, 4, 5] * u.kg * ( u.m / u.s ) ** 2
23print(test_data.units)
24# kg*m**2/s**2
25
26# Create a data set to hold the numerical information:
27f.create_dataset('stored data', data=test_data)
28
29# Save the units information as a string in `attrs`.
30f['stored data'].attrs['units'] = str(test_data.units)
31
32# Now recover the data, using the saved units information to reconstruct the
33# original quantities.
34reconstituted_data = u.unyt_array(f['stored data'],
35                                  units=f['stored data'].attrs['units'])
36
37print(reconstituted_data.units)
38# kg*m**2/s**2
39
40assert reconstituted_data.units == test_data.units
41