Xarray coordinate imrpovements

python/Cargo.toml

@@ -10,8 +10,8 @@ crate-type = ["cdylib"]
 
 [dependencies]
 gribberish = { path = "../" }
-numpy = "0.17.2"
+numpy = "0.19.0"
 
 [dependencies.pyo3]
-version = "0.17.3"
+version = "0.19.1"
 features = ["extension-module"]

python/examples/dump_dataset.py

@@ -0,0 +1,20 @@
+import argparse
+import gribberish
+
+
+def read_file(filename: str) -> bytes:
+    with open(filename, 'rb') as f:
+        raw_data = f.read()
+        return raw_data
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser('Dump a grib 2 file dataset')
+    parser.add_argument('infile', metavar='i', type=str, help='Path to grib 2 file to ')
+    args = parser.parse_args()
+    input_filename = args.infile
+
+    raw = read_file(input_filename)
+    dataset = gribberish.GribDataset(raw)
+    print(dataset.attrs)
+    print(dataset.vars)
+    print(dataset.coords)

python/examples/dump_raster.py

@@ -25,7 +25,6 @@ def create_filename(input_filename: str, message_index: int) -> str:
     parser = argparse.ArgumentParser('Dump a grib 2 file to a png raster')
     parser.add_argument('infile', metavar='i', type=str, help='Path to grib 2 file to ')
     args = parser.parse_args()
-
     input_filename = args.infile
 
     messages = read_gib_messages(input_filename)

python/examples/requirements.txt

@@ -1,5 +1,4 @@
 numpy 
 matplotlib
 xarray
-rio_xarray
 cf_xarray

python/gribberish/gribberish_backend.py

@@ -1,32 +1,16 @@
 import os
-import numpy as np
 import xarray as xr
 
-from .gribberishpy import parse_grib_mapping, parse_grib_message
+from .gribberishpy import parse_grid_dataset
 from xarray.backends import BackendEntrypoint
 
 
 def read_binary_data(filename: str):
+    # TODO: Make this streamed for large files, etc etc
     with open(filename, 'rb') as f:
         return f.read()
 
 
-def extract_variable_data(grib_message):
-    data = grib_message.data().reshape(grib_message.grid_shape)
-    data = np.expand_dims(data, axis=0)
-    crs = grib_message.crs
-    return (
-        ['time', 'lat', 'lon'] if grib_message.is_regular_grid else ['time', 'y', 'x'],
-        data,
-        {
-            'standard_name': grib_message.var_abbrev,
-            'long_name': grib_message.var_name,
-            'units': grib_message.units,
-            'crs': crs,
-        }
-    )
-
-
 class GribberishBackend(BackendEntrypoint):
     '''
     Custom backend for xarray
@@ -44,82 +28,18 @@ def open_dataset(
     ):
         raw_data = read_binary_data(filename_or_obj)
 
-        # Read the message mapping from the metadata that gives the byte offset
-        # for each variables message
-        var_mapping = parse_grib_mapping(raw_data)
-
-        # For now, any unsupported products get dropped 
-        keys = list(var_mapping.keys())
-        for var in keys: 
-            if var.startswith('(') or var == 'unknown': 
-                var_mapping.pop(var, None)
+        dataset = parse_grid_dataset(raw_data, drop_variables=drop_variables)
+        coords = {k: (v['dims'], v['values'], v['attrs']) for (k, v) in dataset['coords'].items()}
+        data_vars = {k: (v['dims'], v['values'], v['attrs']) for (k, v) in dataset['data_vars'].items()}
+        attrs = dataset['attrs']
 
-        # If there are variabels specified to drop, do so now
-        if drop_variables:
-            for var in drop_variables:
-                var_mapping.pop(var, None)
-
-        # Extract each variables metadata
-        data_vars = {var: extract_variable_data(parse_grib_message(
-            raw_data, lookup[1])) for (var, lookup) in var_mapping.items()}
-
-        # Get the coordinate arrays
-        # TODO: This can be optimized
-        first_message = parse_grib_message(
-            raw_data,
-            list(var_mapping.values())[0][1]
-        )
-
-        lat = first_message.latitudes().reshape(first_message.grid_shape)
-        lng = first_message.longitudes().reshape(first_message.grid_shape)
-
-        if first_message.is_regular_grid:
-            lat = (['lat'], lat[:, 0], {
-                'standard_name': 'latitude',
-                'long_name': 'latitude',
-                'units': 'degrees_north',
-                'axis': 'Y'
-            })
-            lon = (['lon'], lng[0, :], {
-                'standard_name': 'longitude',
-                'long_name': 'longitude',
-                'units': 'degrees_east',
-                'axis': 'X'
-            })
-        else:
-            lat = (['y', 'x'], lat, {
-                'standard_name': 'latitude',
-                'long_name': 'latitude',
-                'units': 'degrees_north',
-                'axis': 'Y'
-            })
-            lon = (['y', 'x'], lng, {
-                'standard_name': 'longitude',
-                'long_name': 'longitude',
-                'units': 'degrees_east',
-                'axis': 'X'
-            })
-
-        coords = {
-            'time': (['time'], [first_message.forecast_date], {
-                'standard_name': 'time',
-                'long_name': 'time',
-                'units': 'seconds since 2010-01-01 00:00:00',
-                'axis': 'T'
-            }),
-            'lat': lat,
-            'lon': lon,
-        }
-
-        # Finally put it all together and create the xarray dataset
         return xr.Dataset(
             data_vars=data_vars,
             coords=coords,
-            attrs={
-                'meta': 'created with gribberish',
-            }
+            attrs=attrs
         )
 
+
     def guess_can_open(self, filename_or_obj):
         try:
             _, ext = os.path.splitext(filename_or_obj)