diff --git a/py4cast/datasets/base.py b/py4cast/datasets/base.py index 950e619f..7b2a889a 100644 --- a/py4cast/datasets/base.py +++ b/py4cast/datasets/base.py @@ -44,11 +44,13 @@ class Item: inputs has shape (timestep, lat, lon, features) outputs has shape (timestep, lat, lon, features) forcing has shape (timestep, lat, lon, features) + valdity_times has shape (timestep) """ inputs: NamedTensor | None forcing: NamedTensor | None outputs: NamedTensor + validity_times: list[dt.datetime] def unsqueeze_(self, dim_name: str, dim_index: int): """ @@ -116,20 +118,21 @@ def __str__(self) -> str: """ table = [] for attr in (f.name for f in fields(self)): - nt: NamedTensor = getattr(self, attr) - if nt is not None: - for feature_name in nt.feature_names: - tensor = nt[feature_name] - table.append( - [ - attr, - nt.names, - list(nt[feature_name].shape), - feature_name, - tensor.min(), - tensor.max(), - ] - ) + if attr != "validity_times": + nt: NamedTensor = getattr(self, attr) + if nt is not None: + for feature_name in nt.feature_names: + tensor = nt[feature_name] + table.append( + [ + attr, + nt.names, + list(nt[feature_name].shape), + feature_name, + tensor.min(), + tensor.max(), + ] + ) headers = [ "Type", "Dimension Names", @@ -175,16 +178,19 @@ def collate_fn(items: List[Item]) -> ItemBatch: """ # Here we postpone that for each batch the same dimension should be present. batch_of_items = {} - # Iterate over inputs, outputs and forcing fields for field_name in (f.name for f in fields(Item)): - batched_tensor = collate_tensor_fn( - [getattr(item, field_name).tensor for item in items] - ).type(torch.float32) - - batch_of_items[field_name] = NamedTensor.expand_to_batch_like( - batched_tensor, getattr(items[0], field_name) - ) + if field_name == "validity_times": + batched_valid_times = [getattr(item, field_name) for item in items] + batch_of_items[field_name] = batched_valid_times + else: + batched_tensor = collate_tensor_fn( + [getattr(item, field_name).tensor for item in items] + ).type(torch.float32) + + batch_of_items[field_name] = NamedTensor.expand_to_batch_like( + batched_tensor, getattr(items[0], field_name) + ) return ItemBatch(**batch_of_items) @@ -517,6 +523,7 @@ def load(self, no_standardize: bool = False) -> Item: inputs=inputs, outputs=outputs, forcing=forcing, + validity_times=self.output_timestamps.validity_times, ) def plot(self, item: Item, step: int, save_path: Path = None) -> None: diff --git a/py4cast/lightning.py b/py4cast/lightning.py index 078bfabc..5cc313c3 100644 --- a/py4cast/lightning.py +++ b/py4cast/lightning.py @@ -318,6 +318,11 @@ def setup(self, stage=None): self.psd_plot_metric = MetricPSDK(self.save_path, pred_step=max_pred_step) self.acc_metric = MetricACC(self.dataset_info) self.configure_loggers() + self.list_metrics = [ + self.acc_metric, + self.psd_plot_metric, + self.rmse_psd_plot_metric, + ] def configure_loggers(self): layout = { @@ -516,6 +521,8 @@ def _common_step( self.original_shape = None + ds = self.training_strategy == "downscaling_only" + if self.model.model_type == ModelType.GRAPH: # Stack original shape to reshape later self.original_shape = batch.inputs.tensor.shape @@ -531,11 +538,24 @@ def _common_step( # to check at inference time if the feature names are the same # also useful to build NamedTensor outputs with same feature and dim names # If model type is graph, flat the lon/lat dim before saving the dims - if batch_idx == 0 and phase == "train": + if batch_idx == 0: self.input_feature_names = batch.inputs.feature_names self.output_feature_names = batch.outputs.feature_names self.output_dim_names = batch.outputs.names self.output_dtype = batch.outputs.tensor.dtype + # save the indices of the features common to the forcings and outputs in downscaling_only mode + # useful for reconstructing the prediction from the residual + if ds: + forcing_feature_names = batch.forcing.feature_names + common_features_idx = [] + for output_feature_name in self.output_feature_names: + for i, feature_forcing_name in enumerate(forcing_feature_names): + if ( + output_feature_name.split("_")[1:] + == feature_forcing_name.split("_")[1:] + ): + common_features_idx.append(i) + self.common_features_idx = common_features_idx prev_states = batch.inputs prediction_list = [] @@ -575,8 +595,6 @@ def _common_step( else: y = self.model(x) - ds = self.training_strategy == "downscaling_only" - # select the last timestep last_prev_state = prev_states.select_tensor_dim("timestep", -1).clone() if self.mask_on_nan: @@ -590,6 +608,17 @@ def _common_step( + y * step_diff_std + step_diff_mean ) + elif ds: + # update the coarse forcings with our netwwork output + last_coarse_step = batch.forcing.select_tensor_dim( + "timestep", -1 + ).clone() + if self.mask_on_nan: + last_coarse_step = torch.nan_to_num(last_coarse_step, nan=0) + # only add common features + predicted_state = ( + last_coarse_step[:, :, :, self.common_features_idx] + y + ) else: predicted_state = last_prev_state * (1 - ds) + y @@ -919,10 +948,9 @@ def on_validation_epoch_end(self): if self.logging_enabled: # Get dict of metrics' results - dict_metrics = dict() - dict_metrics.update(self.psd_plot_metric.compute()) - dict_metrics.update(self.rmse_psd_plot_metric.compute()) - dict_metrics.update(self.acc_metric.compute()) + dict_metrics = {} + for metric in self.list_metrics: + dict_metrics.update(metric.compute()) for name, elmnt in dict_metrics.items(): if isinstance(elmnt, matplotlib.figure.Figure): # Tensorboard logger @@ -1040,9 +1068,8 @@ def on_test_epoch_end(self): """ if self.logging_enabled: dict_metrics = {} - dict_metrics.update(self.psd_plot_metric.compute(prefix="test")) - dict_metrics.update(self.rmse_psd_plot_metric.compute(prefix="test")) - dict_metrics.update(self.acc_metric.compute(prefix="test")) + for metric in self.list_metrics: + dict_metrics.update(metric.compute(prefix="test")) for name, elmnt in dict_metrics.items(): if isinstance(elmnt, matplotlib.figure.Figure): diff --git a/tests/test_datasets.py b/tests/test_datasets.py index 06e4b929..ee806da0 100644 --- a/tests/test_datasets.py +++ b/tests/test_datasets.py @@ -33,7 +33,11 @@ def test_item(): names=["lat", "lon", "features"], feature_names=[f"forcing_{i}" for i in range(2)], ) - item = Item(inputs=inputs, outputs=outputs, forcing=forcing) + + validity_times = [datetime.datetime(year=2023, month=1, day=1, hour=18)] + item = Item( + inputs=inputs, outputs=outputs, forcing=forcing, validity_times=validity_times + ) print(item) # test collate_fn @@ -59,7 +63,12 @@ def test_item(): feature_names=[f"feature_{i}" for i in range(4)], ) - item = Item(inputs=inputs, outputs=outputs, forcing=forcing) + item = Item( + inputs=inputs, + outputs=outputs, + forcing=forcing, + validity_times=validity_times, + ) # Input and Output must have the same feature names with pytest.raises(ValueError): @@ -74,7 +83,12 @@ def test_item(): feature_names=[f"f_{i}" for i in range(5)], ) - item = Item(inputs=inputs, outputs=outputs, forcing=forcing) + item = Item( + inputs=inputs, + outputs=outputs, + forcing=forcing, + validity_times=validity_times, + ) # Input and Output must have the same dim names with pytest.raises(ValueError): @@ -89,7 +103,12 @@ def test_item(): feature_names=[f"feature_{i}" for i in range(5)], ) - item = Item(inputs=inputs, outputs=outputs, forcing=forcing) + item = Item( + inputs=inputs, + outputs=outputs, + forcing=forcing, + validity_times=validity_times, + ) def test_date_forcing():