dtc_super_resolution.rcm_emulator.ml_module.inference_helpers

Functions needed for using the trained model on new data.

Functions

load_model(config)

Load a trained model from the specified path.

predict(model, input_patches, metadata, device)

Run inference on the model using the provided DataLoader.

run_postprocessing(monthly_predictions, ...)

Post-process the model's predictions to match the original input shape.

run_postprocessing_loc(predictions, ...)

Post-process the model's predictions to match the original input shape for a specific location.
load_model(config: dict) → Module

Load a trained model from the specified path.

Parameters:

config (dict) – A dictionary containing the model configuration, including the model type and weights path.

Returns:

The loaded PyTorch model.

Return type:

nn.Module

predict(model: Module, input_patches: ndarray, metadata: list, device: device) → ndarray

Run inference on the model using the provided DataLoader.

Parameters:

  • model (nn.Module) – The trained PyTorch model.

  • input_patches (np.ndarray) – A list of numpy arrays containing preprocessed input patches for inference.

  • metadata (list) – A list of dictionaries containing metadata for each input patch.

  • device (torch.device) – The device to run the inference on (CPU or GPU).

Returns:

Model predictions with shape (time, tiles, height, width).

Return type:

np.ndarray

run_postprocessing(monthly_predictions: ndarray, metadata: list, tiles: list, rcm_ds: Dataset, era5_xy: Dataset) → ndarray

Post-process the model’s predictions to match the original input shape.

Parameters:

  • monthly_predictions (np.ndarray) – The model’s predictions.

  • metadata (list) – A list of dictionaries containing metadata for each input patch.

  • tiles (list) – A list of dictionaries containing metadata for each tile.

  • rcm_ds (xr.Dataset) – The RCM dataset.

  • era5_xy (xr.Dataset) – The ERA5 dataset reprojected to the RCM grid.

Returns:

The post-processed predictions.

Return type:

np.ndarray

run_postprocessing_loc(predictions: ndarray, tile_meta: list, rcm_ds: Dataset, era5_xy: Dataset, input_patches: ndarray, target_patches: ndarray) → ndarray

Post-process the model’s predictions to match the original input shape for a specific location.

Parameters:

  • predictions (np.ndarray) – The model’s predictions.

  • tile_meta (list) – A list of dictionaries containing metadata for each input patch.

  • rcm_ds (xr.Dataset) – The RCM dataset.

  • era5_xy (xr.Dataset) – The ERA5 dataset reprojected to the RCM grid.

  • input_patches (np.ndarray) – The input patches extracted from the ERA5 dataset.

  • target_patches (np.ndarray) – The target patches extracted from the RCM dataset.

Returns:

The post-processed predictions for the specific location.

Return type:

np.ndarray