diff --git a/weather_mv/loader_pipeline/ee.py b/weather_mv/loader_pipeline/ee.py index 2ad914dd..ecaf5408 100644 --- a/weather_mv/loader_pipeline/ee.py +++ b/weather_mv/loader_pipeline/ee.py @@ -448,7 +448,6 @@ def convert_to_asset(self, queue: Queue, uri: str): with open_dataset(uri, self.open_dataset_kwargs, self.disable_grib_schema_normalization, - band_names_dict=self.band_names_dict, initialization_time_regex=self.initialization_time_regex, forecast_time_regex=self.forecast_time_regex, group_common_hypercubes=self.group_common_hypercubes) as ds_list: @@ -459,15 +458,17 @@ def convert_to_asset(self, queue: Queue, uri: str): attrs = ds.attrs data = list(ds.values()) asset_name = get_ee_safe_name(uri) - channel_names = [da.name for da in data] - start_time, end_time, is_normalized = (attrs.get(key) for key in - ('start_time', 'end_time', 'is_normalized')) + channel_names = [ + self.band_names_dict.get(da.name, da.name) if self.band_names_dict + else da.name for da in data + ] + dtype, crs, transform = (attrs.pop(key) for key in ['dtype', 'crs', 'transform']) - attrs.update({'is_normalized': str(is_normalized)}) # EE properties does not support bool. # Adding job_start_time to properites. attrs["job_start_time"] = job_start_time # Make attrs EE ingestable. attrs = make_attrs_ee_compatible(attrs) + start_time, end_time = (attrs.get(key) for key in ('start_time', 'end_time')) if self.group_common_hypercubes: level, height = (attrs.pop(key) for key in ['level', 'height']) diff --git a/weather_mv/loader_pipeline/sinks.py b/weather_mv/loader_pipeline/sinks.py index 7f004ca5..52611a78 100644 --- a/weather_mv/loader_pipeline/sinks.py +++ b/weather_mv/loader_pipeline/sinks.py @@ -139,9 +139,14 @@ def rearrange_time_list(order_list: t.List, time_list: t.List) -> t.List: return datetime.datetime(*time_list) -def _preprocess_tif(ds: xr.Dataset, filename: str, tif_metadata_for_start_time: str, - tif_metadata_for_end_time: str, uri: str, band_names_dict: t.Dict, - initialization_time_regex: str, forecast_time_regex: str) -> xr.Dataset: +def _preprocess_tif( + ds: xr.Dataset, + tif_metadata_for_start_time: str, + tif_metadata_for_end_time: str, + uri: str, + initialization_time_regex: str, + forecast_time_regex: str +) -> xr.Dataset: """Transforms (y, x) coordinates into (lat, long) and adds bands data in data variables. This also retrieves datetime from tif's metadata and stores it into dataset. @@ -432,7 +437,6 @@ def open_dataset(uri: str, disable_grib_schema_normalization: bool = False, tif_metadata_for_start_time: t.Optional[str] = None, tif_metadata_for_end_time: t.Optional[str] = None, - band_names_dict: t.Optional[t.Dict] = None, initialization_time_regex: t.Optional[str] = None, forecast_time_regex: t.Optional[str] = None, group_common_hypercubes: t.Optional[bool] = False, @@ -482,11 +486,9 @@ def open_dataset(uri: str, xr_dataset = xr_datasets.sel(time=slice(start_date, end_date)) if uri_extension in ['.tif', '.tiff']: xr_dataset = _preprocess_tif(xr_dataset, - local_path, tif_metadata_for_start_time, tif_metadata_for_end_time, uri, - band_names_dict, initialization_time_regex, forecast_time_regex)