diff --git a/docs/usage.rst b/docs/usage.rst index 3414a7f8..acaf020b 100644 --- a/docs/usage.rst +++ b/docs/usage.rst @@ -1057,7 +1057,7 @@ HDF5 Files The :mod:`unyt` library provides a hook for writing data both to a new HDF5 file and an existing file and then subsequently reading that data back in to restore the array. This works via the :meth:`unyt_array.write_hdf5 ` and :meth:`unyt_array.from_hdf5 ` methods. The simplest way to use these functions is to write data to a file that does not exist yet: - >>> from unyt import cm, unyt_array + >>> from unyt import cm >>> import os >>> data = [1, 2, 3]*cm >>> data.write_hdf5('my_data.h5') diff --git a/unyt/array.py b/unyt/array.py index b92b2e06..a58f90fd 100644 --- a/unyt/array.py +++ b/unyt/array.py @@ -127,8 +127,13 @@ from unyt.equivalencies import equivalence_registry from unyt._on_demand_imports import _astropy, _pint from unyt._pint_conversions import convert_pint_units +from unyt._unit_lookup_table import default_unit_symbol_lut from unyt.unit_object import _check_em_conversion, _em_conversion, Unit -from unyt.unit_registry import _sanitize_unit_system, UnitRegistry +from unyt.unit_registry import ( + _sanitize_unit_system, + UnitRegistry, + default_unit_registry, +) NULL_UNIT = Unit() POWER_SIGN_MAPPING = {multiply: 1, divide: -1} @@ -1319,7 +1324,11 @@ def write_hdf5(self, filename, dataset_name=None, info=None, group_name=None): info = {} info["units"] = str(self.units) - info["unit_registry"] = np.void(pickle.dumps(self.units.registry.lut)) + lut = {} + for k, v in self.units.registry.lut.items(): + if k not in default_unit_registry.lut: + lut[k] = v + info["unit_registry"] = np.void(pickle.dumps(lut)) if dataset_name is None: dataset_name = "array_data" @@ -1382,7 +1391,9 @@ def from_hdf5(cls, filename, dataset_name=None, group_name=None): dataset = g[dataset_name] data = dataset[:] units = dataset.attrs.get("units", "") - unit_lut = pickle.loads(dataset.attrs["unit_registry"].tostring()) + unit_lut = default_unit_symbol_lut.copy() + unit_lut_load = pickle.loads(dataset.attrs["unit_registry"].tostring()) + unit_lut.update(unit_lut_load) f.close() registry = UnitRegistry(lut=unit_lut, add_default_symbols=False) return cls(data, units, registry=registry)