From e590547df0876b118967c789640c3a2f89f8883f Mon Sep 17 00:00:00 2001 From: Harrison Cook Date: Thu, 7 Nov 2024 09:19:23 +0000 Subject: [PATCH] Fix issue with input as grib --- src/ai_models/outputs/__init__.py | 26 +++++++++++++++++++++++--- 1 file changed, 23 insertions(+), 3 deletions(-) diff --git a/src/ai_models/outputs/__init__.py b/src/ai_models/outputs/__init__.py index 600e1c8..dd8d473 100644 --- a/src/ai_models/outputs/__init__.py +++ b/src/ai_models/outputs/__init__.py @@ -126,15 +126,35 @@ def __init__(self, owner, path, metadata, **kwargs): def write(self, data, *args, check = False, **kwargs): template = kwargs.pop("template") step = kwargs.pop("step") + + if data is None: + return + import xarray as xr + + if isinstance(template, ekd.readers.grib.codes.GribField): + xarray_obj: xr.Dataset = template.to_xarray() + attrs = xarray_obj.attrs - xarray_obj: xr.DataArray = template.to_xarray() - xarray_obj.data = data - xarray_obj = xarray_obj.assign_coords(step = step) + xarray_obj = xarray_obj[list(xarray_obj.data_vars)[0]] + xarray_obj.attrs.update(attrs) + + xarray_obj.data = data + + if 'levtype' in attrs and 'levelist' in attrs: + xarray_obj = xarray_obj.assign_coords({attrs['levtype']: attrs['levelist']}) + + else: + xarray_obj: xr.DataArray = template.to_xarray() + xarray_obj.data = data + if 'pl' in xarray_obj.coords: xarray_obj = xarray_obj.expand_dims('pl') if 'ml' in xarray_obj.coords: xarray_obj = xarray_obj.expand_dims('ml') + + xarray_obj = xarray_obj.assign_coords(step = step) + xarray_obj.attrs.pop('_earthkit', None) self._outputs[step].append(xarray_obj)