File size: 4,396 Bytes
7180154 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 | # Copyright 2024 DeepMind Technologies Limited.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS-IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Wrappers for Predictors which allow them to work with data cleaned of NaNs.
The Predictor which is wrapped sees inputs and targets without NaNs, and makes
NaNless predictions.
"""
from typing import Optional, Tuple
from . import predictor_base as base
import numpy as np
import xarray
class NaNCleaner(base.Predictor):
"""A predictor wrapper than removes NaNs from ingested data.
The Predictor which is wrapped sees inputs and targets without NaNs.
"""
def __init__(
self,
predictor: base.Predictor,
var_to_clean: str,
fill_value: xarray.Dataset,
reintroduce_nans: bool = False,
):
"""Initializes the NaNCleaner."""
self._predictor = predictor
self._fill_value = fill_value[var_to_clean]
self._var_to_clean = var_to_clean
self._reintroduce_nans = reintroduce_nans
def _clean(self, dataset: xarray.Dataset) -> xarray.Dataset:
"""Cleans the dataset of NaNs."""
data_array = dataset[self._var_to_clean]
dataset = dataset.assign(
{self._var_to_clean: data_array.fillna(self._fill_value)}
)
return dataset
def _maybe_reintroduce_nans(
self, stale_inputs: xarray.Dataset, predictions: xarray.Dataset
) -> xarray.Dataset:
# NaN positions don't change between input frames, if they do then
# we should be more careful about re-introducing them.
if self._var_to_clean in predictions.keys():
nan_mask = np.isnan(stale_inputs[self._var_to_clean]).any(dim='time')
with_nan_values = predictions[self._var_to_clean].where(~nan_mask, np.nan)
predictions = predictions.assign({self._var_to_clean: with_nan_values})
return predictions
def __call__(
self,
inputs: xarray.Dataset,
targets_template: xarray.Dataset,
forcings: Optional[xarray.Dataset] = None,
**kwargs,
) -> xarray.Dataset:
if self._reintroduce_nans:
# Copy inputs before cleaning so that we can reintroduce NaNs later.
original_inputs = inputs.copy()
if self._var_to_clean in inputs.keys():
inputs = self._clean(inputs)
if forcings and self._var_to_clean in forcings.keys():
forcings = self._clean(forcings)
predictions = self._predictor(
inputs, targets_template, forcings, **kwargs
)
if self._reintroduce_nans:
predictions = self._maybe_reintroduce_nans(original_inputs, predictions)
return predictions
def loss(
self,
inputs: xarray.Dataset,
targets: xarray.Dataset,
forcings: Optional[xarray.Dataset] = None,
**kwargs,
) -> base.LossAndDiagnostics:
if self._var_to_clean in inputs.keys():
inputs = self._clean(inputs)
if self._var_to_clean in targets.keys():
targets = self._clean(targets)
if forcings and self._var_to_clean in forcings.keys():
forcings = self._clean(forcings)
return self._predictor.loss(
inputs, targets, forcings, **kwargs
)
def loss_and_predictions(
self,
inputs: xarray.Dataset,
targets: xarray.Dataset,
forcings: Optional[xarray.Dataset] = None,
**kwargs,
) -> Tuple[base.LossAndDiagnostics, xarray.Dataset]:
if self._reintroduce_nans:
# Copy inputs before cleaning so that we can reintroduce NaNs later.
original_inputs = inputs.copy()
if self._var_to_clean in inputs.keys():
inputs = self._clean(inputs)
if self._var_to_clean in targets.keys():
targets = self._clean(targets)
if forcings and self._var_to_clean in forcings.keys():
forcings = self._clean(forcings)
loss, predictions = self._predictor.loss_and_predictions(
inputs, targets, forcings, **kwargs
)
if self._reintroduce_nans:
predictions = self._maybe_reintroduce_nans(original_inputs, predictions)
return loss, predictions
|