block2.py

import gzip
import xarray as xr
import numpy as np

def create_dataset(filenames, datetimes):
    '''Create a Dataset from binary files.

    Arguments:
        - filenames: list of file names to be concatenated along the time
          dimension.
        - datetimes: list of dates corresponding to the time coordinate.

    Returns:
        - ds: the Dataset created from the files and dates.
    '''
    ds_list = []
    for fname in filenames:
        f = gzip.open(f'trmm_data/{fname}', 'rb')
        data = f.read()
        f.close()
        header_size_id = 'header_byte_length='
        data_head = data[:1000]
        header_size_pos = data_head.decode('utf-8').find(header_size_id) + \
                          len(header_size_id)
        header_size = int(data[header_size_pos:data_head.decode('utf-8')
                                               .find(' ', header_size_pos)])
        data = data[header_size:header_size + 480*1440*2]
        data = np.frombuffer(data, dtype=np.dtype('>i2'), count=480*1440)
        data_array = data.reshape((480, 1440)).astype(np.float32) / 100 # mm/h
        np.clip(data_array, 0, np.inf, out=data_array)
        lat = np.arange(60-0.25/2, -60, -0.25)
        lon = np.arange(0+0.25/2, 360, 0.25)
        ds = xr.Dataset({'precipitation': (['lat', 'lon'], data_array)},
                        coords={'lat':lat, 'lon':lon})
        ds_list.append(ds)
    ds = xr.concat(ds_list, 'time')
    ds = ds.assign_coords(time=datetimes)
    return ds