save astropy units in netcdf files

[hafarooki]

Is your feature request related to a problem?

A recent PR made it possible to use astropy Quantities seamlessly with Datasets. However, when writing the dataset to a file with to_netcdf, then loading the netcdf file, the quantities become ndarrays.

Describe the solution you'd like

Would be great if this could be fixed!

Describe alternatives you've considered

Pickle files; not ideal, but a valid temporary solution

Additional context

No response

[welcome]

Thanks for opening your first issue here at xarray! Be sure to follow the issue template!
If you have an idea for a solution, we would really welcome a Pull Request with proposed changes.
See the Contributing Guide for more.
It may take us a while to respond here, but we really value your contribution. Contributors like you help make xarray better.
Thank you!

[hafarooki]

As a temporary fix, using this code before saving and after loading worked for me:

    for variable in dataset.variables:
        if isinstance(dataset[variable].data, u.Quantity):
            dataset[variable].attrs['unit'] = dataset[variable].data.unit.to_string()
            
    dataset.to_netcdf('file_name.nc')

    dataset = xr.open_dataset('file_name.nc')
    
    for variable in dataset.variables:
        variable_attrs = dataset[variable].attrs
        if 'unit' in variable_attrs:
            dataset[variable].data *= astropy.units.Unit(variable_attrs['unit'])

[keewis]

yes, converting to array + metadata before serializing to disk is the only thing we can do right now (this is also what pint-xarray is doing).

[hafarooki]

@keewis thanks for the follow up! So would this be something that goes in xarray or a separate library?

[hafarooki]

In case it's useful to anyone else, I made these slightly more general functions to be used before/after saving:

import xarray as xr
import astropy.units as u

def with_decoded_units(original_dataset: xr.Dataset) -> xr.Dataset:
    modified_dataset: xr.Dataset = original_dataset.copy(deep=False)
    
    for variable in original_dataset.variables:
        variable_attrs: dict = modified_dataset[variable].attrs
      
        if 'unit' in variable_attrs:
            assert not isinstance(original_dataset[variable].data, u.Quantity), f'original dataset already has a quantity for {variable}'
            
            modified_dataset[variable].data *= u.Unit(variable_attrs['unit'])
            del variable_attrs['unit']
            
            assert 'unit' in original_dataset[variable].attrs, f'original dataset should still have unit attribute for {variable}'
            assert 'unit' not in modified_dataset[variable].attrs, f'modified dataset should no longer have unit attribute for {variable}'
            assert not isinstance(original_dataset[variable].data, u.Quantity), f'original dataset should still not have a quantity for {variable}'
            assert isinstance(modified_dataset[variable].data, u.Quantity), f'modified dataset should have a quantity for {variable}'
    
    return modified_dataset
   
   
def with_encoded_units(original_dataset: xr.Dataset) -> xr.Dataset:
    modified_dataset: xr.Dataset = original_dataset.copy(deep=False)
    
    for variable in original_dataset.variables:
        if isinstance(original_dataset[variable].data, u.Quantity):
            quantity: u.Quantity = original_dataset[variable].data
            unit: u.UnitBase | u.StructuredUnit | None = quantity.unit
            
            if unit is None:
                continue
            
            # add unit attribute
            modified_dataset[variable].attrs['unit'] = unit.to_string()
            
            # replace Quantity with ndarray
            modified_dataset[variable].data = quantity.value
            
            assert isinstance(original_dataset[variable].data, u.Quantity), f'original dataset should still have a quantity for {variable}'
            assert not isinstance(modified_dataset[variable].data, u.Quantity), f'modified dataset should no longer have a quantity for {variable}'
    
    return modified_dataset